diff --git a/src/megatron/bridge/peft/multi_lora_layers.py b/src/megatron/bridge/peft/multi_lora_layers.py index f5d4e484cb..693028ddbe 100644 --- a/src/megatron/bridge/peft/multi_lora_layers.py +++ b/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]]: diff --git a/tests/unit_tests/peft/test_multi_lora_layers.py b/tests/unit_tests/peft/test_multi_lora_layers.py index 4486803c5a..2e02999d66 100644 --- a/tests/unit_tests/peft/test_multi_lora_layers.py +++ b/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.