fix(peft)!: make expert LoRA exact under suppressed tensor-parallel communication

[tigist-far]

What does this PR do ?

LoRA adapters on MoE expert layers silently compute the wrong numbers whenever expert_tensor_parallel_size > 1, and one variant is wrong at plain TP > 1 when moe_shared_expert_overlap is enabled. Nothing errors and nothing warns: training runs to completion and the adapter is wrong. In one path the adapter trains nothing at all, because its output is detached from autograd. Every shipped recipe pins expert_tensor_parallel_size=1, so stock configurations are unaffected; custom MoE topologies are not.

This PR makes all four affected paths exact, and makes the four configurations for which no correct implementation exists yet fail loudly instead of silently.

How it breaks

A LoRA adapter is a branch added to a base linear's output, so it must deliver its result in the same parallel state as the base: replicated, a shard, or a partial awaiting a sum. At expert_tensor_parallel_size > 1 Megatron-Core sets explicit_expert_comm and deletes the expert linear's own collectives, because the MoE token dispatcher performs that sum later in combine. The base's output state therefore changes from replicated to partial. The adapters were never told: they kept gathering a full-width delta that the dispatcher then counted once per rank, and internally they consumed their own un-summed partial as though it were complete.

Base linear	Base's output state	What the adapter must emit
column-parallel	shard	shard (already correct)
row-parallel, ordinary TP	replicated	full width (already correct)
row-parallel, collectives suppressed	partial	a partial term — this is what was wrong

Four defect families follow from this, spanning ParallelLinearAdapter (both fc1 and fc2 cases), GroupedExpertLinearAdapter (all three backends), and the shared-expert/MultiLoRALinear path under moe_shared_expert_overlap. Per-path detail, measured failure signatures, and a two-GPU repro: #5634.

What the fix does

Two small, autograd-visible completions, applied uniformly:

• Complete the low-rank activation after linear_in: for the fc2 case, reduce_from(copy_to(x)) over the expert-TP group — an all-reduce in both directions, the exact adjoint pair for a sum whose result every rank consumes differently. For the fc1 case, copy_to(x) — forward identity whose backward all-reduce restores the suppressed dL/dz sum. The grouped class's fc1 gather becomes all_gather_last_dim_from_tensor_parallel_region (backward = reduce-scatter-sum, the true adjoint).
• Zero-embed the output shard instead of gathering: linear_out keeps gather_output=False and the local shard is placed at this rank's offset in a full-width buffer of zeros. Each output element then has exactly one non-zero contributor, so the downstream sum (dispatcher over ETP; post_forward_comm over TP) reconstructs B @ z exactly once.

All groups are passed explicitly (group=) — the MCore default resolves the dense TP group, which is silently wrong for expert adapters.

Why this design

Constraints. A candidate fix is a choice of shard axis for the adapter's A and B, plus whatever collectives make it correct. Three constraints bind:

1. Exactness. The delta must be exactly (alpha/r)·B@A, because serving merges it into the base weight. "Close" is not a weaker form of correct — it is a delta no merged weight can represent.
2. State-matching. The adapter's output must arrive in the base's parallel state, per the table above.
3. Checkpoint compatibility. Each of A and B must keep the shard axis the stored format already records for it. That axis is checkpoint metadata (ShardedTensor.axis_fragmentations) and is what lets a TP=4 checkpoint load at TP=2; the export path pairs each tensor with the matching RowParallelMapping / ColumnParallelMapping so its per-rank pieces reassemble into the right full tensor. This is a function of the base flavor, not a single axis per matrix: A is sharded on r under a column-parallel base and on d_in under a row-parallel one, while B is always sharded on d_out. The constraint forbids changing the axis for a given flavor, not the flavor-dependence itself, which is already part of the format. Moving an axis would break existing adapter checkpoints, mis-assemble exported lora_A/lora_B, and silently change the merge.

Alternatives considered, for the row-parallel case where the base's collectives are suppressed:

Option	Exact?	Communication	Why not
Zero-embed the local shard (chosen)	Yes	one all-reduce on [tokens, rank]; output side is a local pad, no traffic	—
Keep gathering (status quo)	No	all-gather on [tokens, d_out]	this is the bug. Correct under ordinary TP where the base's output is replicated; under suppression the downstream sum counts the delta P times. Identical code, opposite correctness, which is why it went unnoticed
Gather, then scale by 1/P	Forward only	same as status quo	bakes a topology constant into the numerics: (alpha/r)·B@A no longer merges unless export applies the same factor, and the constant is silently wrong if the summing group ever changes
Reduce-scatter the low-rank activation, shard B on the rank axis	Yes	half the chosen option's boundary traffic	genuinely cheaper, but moves B's shard axis from d_out to r, violating constraint 3. Worth revisiting only alongside a checkpoint-format change
All-gather the input, then use the base's own sharding pattern	Yes	d_in / r times more traffic (roughly 100-700x)	the base's pattern needs a replicated input; re-replicating one costs far more than completing the rank-axis sum

Recommendation. Zero-embedding is the only option that is exact, state-matched, and checkpoint-compatible. It also needs no compensating factor anywhere: each output element has exactly one non-zero contributor, so the sum the dispatcher already performs reconstructs B @ z once, and the embed's adjoint is a plain slice, so gradients need no compensation either. That is what keeps stored weights and the merge path untouched. The one alternative that beats it on cost (reduce-scatter) does so by roughly one small collective per adapted expert linear and requires a checkpoint migration to adopt; at r=32 that is not a trade worth making, but it is the right thing to reconsider if the adapter checkpoint format is ever revised.

Where a correct completion does not exist without new machinery, the PR refuses loudly instead of computing silently wrong results.

What this breaks (the reason for the !)

Four configurations that previously ran while silently training wrong now raise NotImplementedError, each naming the offending config key and its workaround. This is the intended blast radius of the ! in the title.

Configs that previously ran while silently training wrong now raise NotImplementedError with messages naming the offending config:

1. DoRA on MoE expert linears, and on suppressed-comm TP-sharded bases (shared experts under moe_shared_expert_overlap).
2. experts_shared_outer_loras=True at ETP>1 (previously crashed with an unrelated shape error at first forward; export path has independent ETP>1 defects).
3. Expert-adapter dropout > 0 at ETP>1.
4. LoRALinear.weight on grouped adapter classes (previously AttributeError).

Each refusal has a practical workaround today (ETP=1, the default share_expert_adapters=True, dropout=0.0, the export-side merge path); lifting them properly is tracked in #5636.

No checkpoint format, weight layout, or export-path change: shard axes and sharded_state_dict factories are untouched, and merges still compute exactly α/r·B@A with no parallelism factor — the zero-embed design is what keeps stored weights compensation-free.

Changelog

Per-file changes

• peft/utils.py — ParallelLinearAdapter: add _reduce_expert_low_rank_activation and _embed_row_parallel_shard; constructor predicates (_expert_row_parallel, _shared_expert_row_parallel) force lin_out_gather_output=False on the two row-parallel-base cases.
• peft/utils.py — GroupedExpertLinearAdapter: replace the autograd-invisible _gather_along_last_dim with _complete_low_rank_activation / _embed_expert_row_parallel_shard across all three backends (per-expert, grouped_mm, TE-fp8); .contiguous() normalization before the activation-offloading tag; error text for the missing-group case kept byte-identical.
• peft/utils.py — construction-time NotImplementedError for SharedOuterGroupedExpertAdapter at ETP>1, and for expert-adapter dropout > 0 at ETP>1 on both adapter classes (dropout draws per-rank RNG on tokens the dispatcher replicates across ETP — not expressible by any merged weight).
• peft/dora.py — DoRA.transform refuses MoE expert linears and suppressed-comm TP-sharded bases (its dense-group adapter and per-rank magnitude norm do not compose); replicated bases still transform.
• peft/lora_layers.py — LoRALinear.weight refuses grouped adapter classes (previously an accidental AttributeError; the merge path for grouped adapters is the export-side materializer).
• peft/multi_lora_layers.py — MultiLoRALinear: same zero-embed fix for the suppressed-comm row-parallel case (its mid-GEMM reduce was already correct).
• Tests (commit 1, fail-first): test_expert_row_parallel_adapter_algebra.py and test_expert_column_parallel_adapter_backward_algebra.py (torch-only closed-form ground truth with negative controls); test_expert_lora_etp.py (CPU suite driving the real adapters through real MCore explicit_expert_comm suppression via a one-rank gloo init + mock group sizes, only the collectives faked — each fake asserts it received exactly the tensor the real collective would see); test_expert_lora_etp_distributed.py (two-rank TP=2×ETP=2 real-NCCL suite, TF32 disabled for exact tolerances) + L0_Launch_peft_expert_etp_distributed.sh; DoRA refusal tests; LoRALinear.weight refusal tests; MultiLoRA behavior-pinning tests; one existing test updated to pass a pg_collection (the completion needs a resolvable group).

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? — three layers; see verification below
• Did you add or update any necessary documentation? — docstrings state the invariants (partial-sum form, adjoints, group selection); no user-facing doc page applies to a bug fix; the new error messages name the config knob to change
• Does the PR affect components that are optional to install? — No. The fix imports only megatron.core.tensor_parallel.mappings (a hard dependency); the existing TE safe_import_from guards in peft/utils.py are untouched
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

Verification (A/B on 2×H100, real NCCL)

Verification (A/B on 2×H100, real NCCL)

Commit 1 checked out alone reproduces every defect: each demonstration is xfail(strict=True) and was verified to fail on its intended assertion (not a harness error) against the pre-fix tree; commit 2 deletes the markers, so a strict XPASS gate mechanically proves each demonstrated defect is fixed.

	Pre-fix (fail-first)	Post-fix (same tests as live assertions)
CPU suite	5 passed, 8 xfailed	13/13 passed
Distributed 2-GPU (TP=2×ETP=2)	4 passed, 7 xfailed	11/11 passed, both ranks
DoRA / weight-property / MultiLoRA regression tests	1 passed + 2 xfailed	26 passed
Closed-form algebra	6 passed	6 passed
Full tests/unit_tests/peft/ regression	—	419 passed

Measured defect signatures pre-fix (matching the closed-form predictions): fc2 combined delta = 2× a concatenation of unsummed partials; fc1 dL/dA missing all cross-rank terms under an exact forward; grouped fc2 output requires_grad == False; shared overlap and MultiLoRA deltas exactly 2× at TP=2 (MultiLoRA greatest rel. diff 1.27).

Performance and memory — no added communication at ETP=1; two output-side paths got faster

Performance and memory

Measured on the same 2x H100 setup, adapter in isolation per the A/B protocol: 50 CUDA-event-timed iterations after 10 warmup, even token counts so the pad path never fires, bf16 (the per-expert-backend row is fp32, that backend's trigger), 4096 tokens, hidden 2048, ffn 2688, rank 32, 4 local experts, TP=2 x ETP=2:

Case	fwd fixed (ms)	fwd broken (ms)	bwd fixed (ms)	bwd broken (ms)	peak fixed / broken (MiB)
Default adapter, expert fc2	0.251 ± 0.033	0.234 ± 0.058	0.361	0.259	147.1 / 147.1
Default adapter, expert fc1	0.216 ± 0.015	0.195 ± 0.014	0.369	0.370	138.8 / 138.8
Grouped adapter fc2, grouped_mm	0.370 ± 0.016	0.368 ± 0.113	0.543	n/a (no graph)	195.9 / 153.2
Grouped adapter fc1, grouped_mm	0.351 ± 0.033	0.334 ± 0.013	0.531	0.231	194.3 / 177.5
Grouped adapter fc2, per-expert backend	0.658 ± 0.035	0.559 ± 0.018	0.949	n/a (no graph)	275.9 / 205.7
Shared-expert-overlap fc2	0.210 ± 0.056	0.263 ± 0.019	0.336	0.427	224.8 / 224.8
MultiLoRA fc2 (wrapper incl. base linear)	0.527 ± 0.067	0.626 ± 0.222	0.372	0.360	249.5 / 257.5

Reading the table — with the essential caveat that the broken arm is not same-work (it computes wrong values, partial gradients, or no gradients at all, which is the bug):

• Peak memory is byte-identical wherever the two arms do comparable work (both default-adapter rows and shared-overlap), measuring the by-construction claim: the zero-embed buffer replaces the removed gather buffer one-for-one and no new collective saves tensors for backward. The grouped rows' higher fixed-arm peaks exist because the broken arm builds no (fc2) or a truncated (fc1) autograd graph — the delta is the cost of gradients existing.
• The two output-side-only cases got faster (shared-overlap −20% fwd, MultiLoRA −16% fwd): the fix removes a gather there and replaces it with a local zero-pad, while the embedded delta rides the reduce that post_forward_comm/the dispatcher already performs — zero added communication.
• Elsewhere forward overhead is +5–11%, worst +18% on the per-expert fallback backend where the completion collectives scale with local expert count; all collectives run on [tokens, rank] tensors and are latency-bound, and the whole adapter remains a sub-millisecond fraction of an MoE layer.
• Backward deltas where the broken arm has a graph reflect the restored communication plus gradient terms the broken arm never computed (e.g. grouped fc1's dL/dA).

At ETP=1 — every shipped recipe — all new branches short-circuit at group size 1, leaving one predicate check per adapter forward; dense LoRA paths are otherwise untouched. Not measured: end-to-end step-time at ETP>1 (the adapter is a small fraction of an MoE layer, but no full-model A/B was run), and a bitwise pre/post A/B at ETP=1 (asserted structurally; backed by the 419-test regression run rather than a bit comparison).

Known limits / residual risk

• The claim that the dispatcher sums partial terms across expert-TP is proven from code for the alltoall dispatcher (combine_preprocess reduce-scatter over expt_tp). For the flex/DeepEP dispatcher the combine's sum semantics are inferred from the DeepEP contract, not verified — worth a maintainer's confirmation before relying on ETP>1 with flex.
• TE-fp8 grouped path and EP>1 × ETP>1 jointly are not covered by the 2-GPU suite (the latter needs ≥4 GPUs, above the functional-test cap); the collectives sit at bf16 boundaries outside TE autograd kernels by construction.
• All shipped recipes pin expert_tensor_parallel_size=1, so these defects are latent for stock configs; they fire on custom ETP>1 topologies (where we found them) and — for the shared-expert/MultiLoRA cases — at plain TP>1 with moe_shared_expert_overlap.

Fixes #5634. Related to #5636.

[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.

[tigist-far]

Restructured the description for readability, top-down: what breaks and for whom, then the mechanism, then a new Why this design section covering the constraints and the alternatives considered, then the breaking changes. Verification, performance, and the per-file changelog moved into collapsed sections; known limits stay visible.

The per-path defect enumeration now lives only in #5634 rather than being duplicated here.

No code, test, or technical-claim changes — description only.