fix(peft): correct dense MultiLoRALinear on replicated base linears

src/megatron/bridge/peft/multi_lora_layers.py

@@ -105,8 +105,16 @@ def __init__(
         # wrapped base linear's output layout.
         self.input_is_parallel = attrs.input_is_parallel
         self.disable_sequence_parallel_comm = attrs.disable_sequence_parallel_comm
+        # False for replicated bases (TELinear parallel_mode="duplicated", e.g. MLA
+        # q/kv down-projections): their adapters are unsharded and need no TP
+        # collectives between the two grouped GEMMs.
+        self.base_linear_is_parallel = attrs.base_linear_is_parallel
         self.use_a2a = a2a_experimental
-        self._gather_output = attrs.input_is_parallel
+        # Mirrors ParallelLinearAdapter's lin_out_gather_output: row-parallel
+        # bases and replicated bases produce a full [tokens, out] tensor, so the
+        # column-sharded adapter output must be gathered before the residual
+        # add; a column-parallel base keeps the [tokens, out/tp] shard.
+        self._gather_output = attrs.input_is_parallel or not attrs.base_linear_is_parallel
 
         # ModuleList of ParallelLinearAdapters gives per-adapter optimizer state
         # isolation, clean checkpoint serialization, and bridge export compatibility.
@@ -136,6 +144,9 @@ def __init__(
         )
 
         self.tokens_per_adapter: Optional[torch.Tensor] = None
+        # Host-side sum of tokens_per_adapter (set alongside it); lets forward
+        # detect an SP-sharded input without a per-layer device sync.
+        self.tokens_per_adapter_total: Optional[int] = None
         device = next(to_wrap.parameters()).device
         dtype = next(to_wrap.parameters()).dtype
         # Non-persistent: slot lifecycle is externally managed, not checkpointed.
@@ -160,6 +171,25 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten
             x = gather_from_sequence_parallel_region(x)
 
         x_flat = x.reshape(-1, x.shape[-1])
+
+        # A replicated base (e.g. MLA q/kv down-projections) does no SP gather —
+        # under sequence parallelism it consumes the SP shard directly, so the
+        # per-slot spans must be narrowed to this rank's contiguous token window.
+        # The shard is contiguous in the same sequence-major flattening the spans
+        # address (same invariant as the MoE slot routing's SP narrow).
+        total = self.tokens_per_adapter_total
+        if total is not None and x_flat.shape[0] != total:
+            tp_size = parallel_state.get_tensor_model_parallel_world_size()
+            if x_flat.shape[0] * tp_size != total:
+                raise RuntimeError(
+                    f"{self.base_linear_name}: adapter token spans cover {total} tokens but the "
+                    f"base linear received {x_flat.shape[0]} rows, which is not the full batch "
+                    f"or its 1/{tp_size} sequence-parallel shard. Check that "
+                    f"set_tokens_per_adapter_slot() was given this micro-batch's counts."
+                )
+            start = parallel_state.get_tensor_model_parallel_rank() * x_flat.shape[0]
+            tokens_per_adapter = _narrow_token_counts_to_window(tokens_per_adapter, start, x_flat.shape[0])
+
         offsets = tokens_per_adapter.cumsum(dim=0, dtype=torch.int32)
 
         stacked_A = torch.stack([a.linear_in.weight for a in self.adapters])
@@ -168,8 +198,10 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten
         mid = torch._grouped_mm(x_flat, stacked_A.transpose(-2, -1), offsets)
 
         # TP collective between A and B: row-parallel base needs an all-reduce
-        # of the partial sums; column-parallel base needs an all-gather of the
-        # rank-sharded output to a full [tokens, dim] for the second GEMM.
+        # of the partial sums; every other base (column-parallel and replicated
+        # alike — ParallelLinearAdapter shards A on the rank axis whenever
+        # input_is_parallel is False) needs an all-gather of the rank-sharded
+        # mid to a full [tokens, dim] for the second GEMM.
         if self.input_is_parallel:
             mid = reduce_from_tensor_model_parallel_region(mid)
         else:
@@ -450,6 +482,9 @@ def __init__(
         )
 
         self.tokens_per_adapter: Optional[torch.Tensor] = None
+        # Written by set_tokens_per_adapter_slot alongside tokens_per_adapter;
+        # unused here (the slot-routing hook owns the expert-side SP narrow).
+        self.tokens_per_adapter_total: Optional[int] = None
         # Republished by the MoE-layer forward pre-hook on every experts forward
         # (including recompute replays), so it is never stale when read; the None
         # here only guards a forward that runs before install_moe_slot_routing.
@@ -560,6 +595,18 @@ def _apply_rank_mask(self, idx: int) -> None:
 _MULTI_LORA_TYPES = (MultiLoRALinear,)
 
 
+def _narrow_token_counts_to_window(counts: torch.Tensor, start: int, num_rows: int) -> torch.Tensor:
+    """Intersect contiguous per-slot token spans with the window ``[start, start + num_rows)``.
+
+    ``counts[i]`` tokens of slot ``i`` occupy the rows ``[cum[i-1], cum[i])`` of the
+    sequence-major flattened micro-batch. A base linear that consumes the
+    sequence-parallel shard sees only ``num_rows`` of those rows starting at
+    ``start``, so its spans are the per-slot overlap with that window.
+    """
+    cum = counts.cumsum(dim=0)
+    return (cum.clamp(max=start + num_rows) - (cum - counts).clamp(min=start)).clamp(min=0).to(counts.dtype)
+
+
 def _iter_multi_lora_modules(model):
     models = model if isinstance(model, list) else [model]
     for model_chunk in models:
@@ -575,8 +622,13 @@ def set_tokens_per_adapter_slot(model, tokens_per_adapter: torch.Tensor) -> None
     upcoming forward that belong to adapter slot ``i``. Must sum to the total
     token count of the micro-batch.
     """
+    # One host sync per micro-batch: layers whose base linear consumes the
+    # SP-sharded sequence (replicated bases) compare their row count against
+    # this total to narrow the spans to their shard without a per-layer sync.
+    total = int(tokens_per_adapter.sum().item())
     for module in _iter_multi_lora_modules(model):
         module.tokens_per_adapter = tokens_per_adapter
+        module.tokens_per_adapter_total = total
 
 
 def _split_sizes_to_list(splits) -> Optional[List[int]]:

tests/unit_tests/peft/test_multi_lora_layers.py

@@ -653,6 +653,42 @@ def test_load_adapter_unused_keys_raises():
         load_adapter(m, 0, over_full)
 
 
+# --------------------------------------------------------------------------- #
+# SP-shard span narrowing: a replicated base linear (e.g. MLA q/kv down-proj)
+# consumes the sequence-parallel shard directly, so the per-slot spans must be
+# intersected with this rank's contiguous token window. Getting this wrong is
+# an out-of-bounds grouped GEMM, not a shape error.
+# --------------------------------------------------------------------------- #
+def _narrow(counts, start, num_rows):
+    return multi_lora_layers_module._narrow_token_counts_to_window(
+        torch.tensor(counts, dtype=torch.int32), start, num_rows
+    ).tolist()
+
+
+def test_narrow_window_spanning_a_slot_boundary():
+    # Slots [10, 54, 0, 0] over 64 tokens; rank 0 of TP=4 sees rows [0, 16).
+    assert _narrow([10, 54, 0, 0], start=0, num_rows=16) == [10, 6, 0, 0]
+    # Rank 1 sees rows [16, 32) — entirely inside slot 1.
+    assert _narrow([10, 54, 0, 0], start=16, num_rows=16) == [0, 16, 0, 0]
+
+
+def test_narrow_full_window_is_identity():
+    assert _narrow([10, 54, 0, 0], start=0, num_rows=64) == [10, 54, 0, 0]
+
+
+def test_narrow_single_slot_batch():
+    # One active slot (the smoke-run shape): every shard lands in slot 0.
+    for rank in range(4):
+        assert _narrow([64, 0, 0, 0], start=rank * 16, num_rows=16) == [16, 0, 0, 0]
+
+
+def test_narrow_window_covering_many_small_slots():
+    # Window [3, 9) overlaps the tail of slot 0, all of slot 1, and the head of slot 2.
+    assert _narrow([4, 3, 5, 0], start=3, num_rows=6) == [1, 3, 2, 0]
+    # Counts always sum to the window size.
+    assert sum(_narrow([4, 3, 5, 0], start=3, num_rows=6)) == 6
+
+
 # --------------------------------------------------------------------------- #
 # Forward smoke / B4 reset: single-GPU integration through a real
 # ColumnParallelLinear.