[PEFT] feat: reduce LoRA sequence-parallel activation memory

[cuichenx]

What changed

• Add opt-in LoRA(sequence_parallel_input_regather=True).
• For eligible non-expert QKV/FC1 LoRA-A projections, retain the sequence-local input instead of the full gathered input.
• Have Transformer Engine return the local LayerNorm-output shard and use Megatron-Core's existing sequence-parallel column-linear path.
• Fall back conservatively for unsupported execution modes, including TP=1, with focused unit and distributed coverage.
• Add PEFT+SP input re-gather to the memory-tuning skill, including enablement, constraints, measured tradeoffs, diagnosis, and verification.

Why

The existing LoRA sequence-parallel path gathers the full LayerNorm output before LoRA-A. Because LoRA-A is trainable, autograd retains that full gathered activation until backward, materially increasing peak memory.

With this option, the full LoRA-A input is gathered temporarily in forward and released after the LoRA-A computation. Megatron-Core retains only the local S / TP shard, issues an asynchronous all-gather during backward for the LoRA-A weight gradient, overlaps that communication with dgrad compute, and uses its existing reduce-scatter/wgrad scheduling. It adds no custom autograd and does not modify the Megatron-Core submodule.

Scope and compatibility

• Disabled by default; existing behavior is unchanged.
• Has no effect when sequence parallelism is disabled.
• Limited to eligible non-expert QKV/FC1 column-parallel adapters.
• Expert layers, row-parallel adapters, TP=1, checkpoint/full-recompute conflicts, CUDA graphs, CPU activation offload, and other unsupported surfaces use the existing path.
• This is a memory optimization. Some throughput degradation from the backward re-gather is expected.

Validation

• uv run --no-sync pre-commit run --all-files
• Memory-tuning skill card parses successfully as YAML.
• Focused PEFT unit tests in the cw-dfw Transformer Engine container: 7 passed
• Real MCore two-rank CUDA regression test: passed on both ranks
  • exact output and input-gradient parity
  • exact LoRA-A/LoRA-B fused main_grad parity
  • exact two-microbatch gradient-accumulation parity
• TP=2 BF16 and FP32 parity checks: exact output, input-gradient, LoRA-A gradient, and LoRA-B gradient parity

Real-model H100 training measurements using real checkpoints and SQuAD:

Model/config	Baseline step (ms)	Input-regather step (ms)	Throughput change	Baseline peak	Input-regather peak
Qwen3-8B, TP2, seq 8192	917.7	984.0	-6.74%	47.545 GB	42.814 GB
Qwen3-30B-A3B, TP4/EP4	1069.9	1101.7	-2.89%	29.890 GB	28.321 GB
GPT-OSS-120B, TP2/EP8	1217.2	1221.3	-0.34%	52.185 GB	51.371 GB

All measured training runs had finite losses with zero skipped or NaN iterations.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[claude]

Light Code Review

Clean implementation with thorough fallback guards and good test coverage.

Observations

1. Unused parameter x in _sequence_parallel_recompute_eligibility (src/megatron/bridge/peft/utils.py): The method signature accepts x: torch.Tensor but never reads it. If it exists for future use or API consistency with other eligibility helpers, a leading underscore (_x) would signal intent and silence linters.

2. linear_in.sequence_parallel mutated every forward call (utils.py:1182): Setting a module attribute during forward() works in practice since TP ranks execute synchronously, but it couples forward-pass behavior to mutable module state. A brief comment noting why this is safe (e.g., "TP ranks are synchronous; no concurrent forward on this module") would help future readers.

3. Doc numbering in performance-guide.md: The new block is numbered > 3. under item 1 ("Activation recompute"), but the previous sub-item is > 2. and the next top-level item is 2. Activation offloading to host memory. The numbering reads correctly, just noting it could confuse readers who skim for "item 3" at the top level.

No critical bugs or logic errors found. The eligibility gates are conservative and well-tested. The TE integration (toggling return_layernorm_output_gathered) is correct -- the adapter gets the local shard when the optimization is active and the gathered tensor otherwise.

Suggested test cases

No perf tests impacted.

[svcnvidia-nemo-ci]

/nvskills-ci

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cuichenx]

/ok to test 323fc4c

[svcnvidia-nemo-ci]

/nvskills-ci

[cuichenx]

/ok to test 4026628

[cuichenx]

/ok to test eacb624

[cuichenx]

/ok to test 337243e

[svcnvidia-nemo-ci]

/nvskills-ci