Skip to content

Commit b8b23fd

Browse files
tigist-fartaufeeque9claude
committed
fix(peft)!: make expert LoRA exact under suppressed tensor-parallel communication
Complete the adapter arithmetic that explicit_expert_comm suppression leaves undone, using autograd-visible collectives on the groups each adapter already carries: reduce_from(copy_to(x)) over expert-TP for the row-parallel-input case (all-reduce on both passes), copy_to(x) for the column-parallel case (backward-only all-reduce restoring the suppressed dL/dz sum), and a differentiable last-dim all-gather (backward reduce-scatter-sum) replacing the grouped adapter's autograd-invisible bare all_gather in all three backends. On the output side, keep linear_out unsharded-gather-free and zero-embed the local shard at its rank offset, so the downstream sum (the dispatcher over expert-TP; SharedExpertMLP.post_forward_comm over dense TP) counts the delta exactly once with no scaling factor; MultiLoRALinear gets the same treatment on suppressed-comm row-parallel bases. Where no correct completion exists without new machinery, refuse loudly instead of training wrong: DoRA on expert or suppressed-comm TP-sharded linears, SharedOuterGroupedExpertAdapter at ETP>1, expert-adapter dropout at ETP>1 (per-rank RNG on dispatcher-replicated tokens), and LoRALinear.weight on grouped adapter classes. BREAKING CHANGE: the four configurations above now raise NotImplementedError where they previously ran with silently wrong numerics (or an accidental AttributeError). Checkpoint format, weight layout, and export paths are unchanged. The strict-xfail markers from the previous commit are deleted here; every demonstration now asserts for real. Verified on 2x H100 (--runxfail, TF32 off): CPU suite 13/13, distributed 11/11 on both ranks, refusal and wiring tests 26 passed, full peft regression 419 passed. At ETP=1 every new branch short-circuits at group size 1; peak memory is unchanged by construction (the embed buffer replaces the removed gather buffer). The fix design originates from Mohammad Taufeeque's earlier implementation of the same defect class. Co-authored-by: Taufeeque <9taufeeque9@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Tigist Diriba <tigist@far.ai>
1 parent e0b1a27 commit b8b23fd

10 files changed

Lines changed: 410 additions & 122 deletions

‎src/megatron/bridge/peft/dora.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from megatron.bridge.peft.base import PEFT
2222
from megatron.bridge.peft.dora_layers import DoRALinear, ParallelLinearDoRAAdapter
2323
from megatron.bridge.peft.module_matcher import ModuleMatcher
24-
from megatron.bridge.peft.utils import get_adapter_attributes_from_linear
24+
from megatron.bridge.peft.utils import get_adapter_attributes_from_linear, is_expert_linear
2525

2626

2727
logger = logging.getLogger(__name__)
@@ -90,7 +90,22 @@ def transform(self, m: nn.Module, name: Optional[str] = None, prefix: Optional[s
9090

9191
if (ans := self.match(m, name, prefix)) is not None:
9292
_, full_name = ans
93+
if is_expert_linear(full_name):
94+
raise NotImplementedError(
95+
f"DoRA does not support MoE expert linears (matched {full_name}): the DoRA "
96+
"adapter is built without is_expert, so its sharding and collectives do not "
97+
"compose with expert parallelism (TP, ETP, or EP > 1). Exclude expert "
98+
"modules from target_modules to proceed."
99+
)
93100
attrs = get_adapter_attributes_from_linear(m)
101+
if attrs.disable_tensor_parallel_comm and attrs.base_linear_is_parallel:
102+
raise NotImplementedError(
103+
f"DoRA does not support TP-sharded linears whose own tensor-parallel "
104+
f"communication is suppressed (matched {full_name}; e.g. shared experts "
105+
"under moe_shared_expert_overlap): DoRA's per-rank magnitude norm does not "
106+
"compose with the downstream TP reduction. Exclude the module from "
107+
"target_modules. A replicated base (e.g. a duplicated TELinear) is fine."
108+
)
94109
logger.info(f"Adding DoRA to: {full_name}")
95110
adapter = ParallelLinearDoRAAdapter(
96111
attrs.in_features,

‎src/megatron/bridge/peft/lora_layers.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,19 @@ class to provide a specific implementation of the forward method.
3535
@property
3636
def weight(self) -> torch.Tensor:
3737
"""Return the effective base weight, including the LoRA delta when enabled."""
38+
from megatron.bridge.peft.utils import GroupedExpertLinearAdapter, SharedOuterGroupedExpertAdapter
39+
40+
if isinstance(self.adapter, (GroupedExpertLinearAdapter, SharedOuterGroupedExpertAdapter)):
41+
# Grouped bases expose weight0..N (no single 2D to_wrap.weight), and these
42+
# adapters carry 3D per-expert weights with no `tp_group` attribute — the
43+
# group-less merge below would silently merge only the local shard. The
44+
# merge path for grouped adapters is the export-side materializer
45+
# (models/conversion/peft_bridge.py), not this property.
46+
raise NotImplementedError(
47+
f"LoRALinear.weight is not supported for grouped expert adapters "
48+
f"({type(self.adapter).__name__} on {self.adapter.base_linear_name}); "
49+
"use the checkpoint/export merge path instead."
50+
)
3851
base_weight = self.to_wrap.weight
3952
if not self._adapter_enabled:
4053
return base_weight

‎src/megatron/bridge/peft/multi_lora_layers.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,18 @@ def __init__(
115115
# column-sharded adapter output must be gathered before the residual
116116
# add; a column-parallel base keeps the [tokens, out/tp] shard.
117117
self._gather_output = attrs.input_is_parallel or not attrs.base_linear_is_parallel
118+
# Shared-expert fc2 under moe_shared_expert_overlap: the base's own TP
119+
# collectives are suppressed (attrs.disable_tensor_parallel_comm) and
120+
# SharedExpertMLP.post_forward_comm sums the full-width partials across
121+
# the dense TP group. A gathered delta would be counted once per rank;
122+
# zero-embedding this rank's shard instead gives each output element
123+
# exactly one non-zero contributor, so the downstream sum counts it once.
124+
# Mirrors ParallelLinearAdapter._shared_expert_row_parallel.
125+
self._suppressed_comm_row_parallel = bool(
126+
attrs.input_is_parallel and attrs.disable_tensor_parallel_comm and attrs.base_linear_is_parallel
127+
)
128+
if self._suppressed_comm_row_parallel:
129+
self._gather_output = False
118130

119131
# ModuleList of ParallelLinearAdapters gives per-adapter optimizer state
120132
# isolation, clean checkpoint serialization, and bridge export compatibility.
@@ -155,6 +167,21 @@ def __init__(
155167
"rank_values", torch.full((n_adapters,), dim, dtype=dtype, device=device), persistent=False
156168
)
157169

170+
def _embed_row_parallel_shard(self, out: torch.Tensor) -> torch.Tensor:
171+
"""Zero-embed this rank's output shard at its dense-TP offset.
172+
173+
Padding is differentiable and its adjoint is the matching slice, so the
174+
downstream ``post_forward_comm`` sum reproduces the delta exactly once with
175+
no ``1/tp`` factor and no gradient compensation.
176+
"""
177+
tp_size = parallel_state.get_tensor_model_parallel_world_size()
178+
if tp_size <= 1:
179+
return out
180+
rank = parallel_state.get_tensor_model_parallel_rank()
181+
shard_width = out.shape[-1]
182+
left = rank * shard_width
183+
return nn.functional.pad(out, (left, (tp_size - 1) * shard_width - left))
184+
158185
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
159186
linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs)
160187

@@ -226,8 +253,13 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten
226253
# Match the wrapped base linear's output layout: row-parallel base
227254
# produces a fully-summed [tokens, h_out] tensor (which we then SP
228255
# scatter); column-parallel base keeps the [tokens, h_out/tp] shard.
256+
# A suppressed-comm row-parallel base (shared experts under overlap)
257+
# emits a full-width partial that post_forward_comm sums across TP, so
258+
# the delta must be contributed as a zero-embedded local shard instead.
229259
if self._gather_output:
230260
out = gather_from_tensor_model_parallel_region(out)
261+
elif self._suppressed_comm_row_parallel:
262+
out = self._embed_row_parallel_shard(out)
231263

232264
if not self.disable_sequence_parallel_comm and self.input_is_parallel:
233265
if self.use_a2a:

0 commit comments

Comments
 (0)