diff --git a/src/megatron/bridge/peft/multi_lora_layers.py b/src/megatron/bridge/peft/multi_lora_layers.py index 693028ddbe..4dfb3c3bd1 100644 --- a/src/megatron/bridge/peft/multi_lora_layers.py +++ b/src/megatron/bridge/peft/multi_lora_layers.py @@ -19,8 +19,9 @@ via per-layer ``tokens_per_adapter`` set by :func:`set_tokens_per_adapter_slot`. Forward stacks the raw weights of all adapters and uses ``torch._grouped_mm`` -for a single fused kernel; TP/SP collectives are issued once around the two -GEMMs to match the layout of the wrapped base linear. +as an eligible fast path, with per-slot linear operations as the fallback; +TP/SP collectives are issued once around the two projections to match the +layout of the wrapped base linear. :class:`MultiLoRAGroupedExpertLinear` is the MoE counterpart, wrapping a grouped expert linear (``mlp.experts.linear_fc{1,2}`` of a ``TEGroupedMLP``) with one @@ -36,6 +37,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from megatron.core import parallel_state from megatron.core.tensor_parallel.mappings import ( all_to_all, @@ -55,12 +57,70 @@ ) +_GROUPED_MM_ALIGNMENT_BYTES = 16 +_GROUPED_MM_AUTOGRAD_DTYPES = (torch.bfloat16, torch.float16) +_GROUPED_MM_MIN_CUDA_CAPABILITY = (8, 0) + + +def _has_grouped_mm_layout(tensor: torch.Tensor) -> bool: + """Return whether a tensor satisfies PyTorch's CUDA grouped-MM layout contract.""" + if tensor.ndim not in (2, 3) or tensor.numel() == 0: + return False + + alignment = _GROUPED_MM_ALIGNMENT_BYTES // tensor.element_size() + if tensor.data_ptr() % _GROUPED_MM_ALIGNMENT_BYTES != 0: + return False + if tensor.ndim == 3 and tensor.stride(0) % alignment != 0: + return False + + rows, columns = tensor.shape[-2:] + row_stride, column_stride = tensor.stride()[-2:] + row_major = column_stride == 1 and row_stride >= max(1, columns) and row_stride % alignment == 0 + column_major = row_stride == 1 and column_stride >= max(1, rows) and column_stride % alignment == 0 + return row_major or column_major + + +def _can_use_grouped_mm(x: torch.Tensor, grouped_weights: torch.Tensor) -> bool: + """Return whether grouped MM is safe for forward and backward.""" + if not hasattr(torch, "_grouped_mm") or not x.is_cuda or not grouped_weights.is_cuda: + return False + if x.dtype not in _GROUPED_MM_AUTOGRAD_DTYPES or grouped_weights.dtype != x.dtype: + return False + if torch.cuda.get_device_capability(x.device) < _GROUPED_MM_MIN_CUDA_CAPABILITY: + return False + + # Backward applies the same 16-byte stride validation to its contiguous + # output gradient, whose row stride is the projection's output width. + output_width = grouped_weights.shape[-1] + if output_width * x.element_size() % _GROUPED_MM_ALIGNMENT_BYTES != 0: + return False + return _has_grouped_mm_layout(x) and _has_grouped_mm_layout(grouped_weights) + + +def _dense_multi_lora_mm( + x: torch.Tensor, + stacked_weights: torch.Tensor, + *, + token_splits: Tuple[int, ...], + offsets: torch.Tensor, +) -> torch.Tensor: + """Apply one dense Multi-LoRA projection with an eligible grouped-MM fast path.""" + grouped_weights = stacked_weights.transpose(-2, -1) + if all(split > 0 for split in token_splits) and _can_use_grouped_mm(x, grouped_weights): + return torch._grouped_mm(x, grouped_weights, offsets) + + inputs = x.split(token_splits, dim=0) + return torch.cat( + [F.linear(adapter_input, adapter_weight) for adapter_input, adapter_weight in zip(inputs, stacked_weights)] + ) + + class MultiLoRALinear(AdapterWrapper): """Megatron parallel linear wrapped with *N* concurrent LoRA adapters. Each adapter slot is a :class:`ParallelLinearAdapter` stored in an - ``nn.ModuleList``. Forward uses grouped GEMM with a single set of - TP/SP comms for efficiency. + ``nn.ModuleList``. Forward uses grouped GEMM when eligible and otherwise + falls back to per-slot linear operations, with one set of TP/SP comms. For bridge export compatibility, use :func:`expose_adapter_slot` to temporarily expose one slot as ``.adapter``. @@ -80,11 +140,11 @@ def __init__( a2a_experimental: bool = False, ) -> None: nn.Module.__init__(self) - # The grouped-GEMM forward below never runs each adapter's own + # The fused Multi-LoRA forward below never runs each adapter's own # ParallelLinearAdapter.forward, so adapter dropout would be silently # dropped. Reject dropout>0 loudly instead of pretending to apply it. assert dropout == 0.0, ( - f"MultiLoRALinear grouped-GEMM path does not apply adapter dropout " + f"MultiLoRALinear fused projection path does not apply adapter dropout " f"(got dropout={dropout}); set dropout/--lora-dropout to 0." ) self.to_wrap = to_wrap @@ -100,14 +160,14 @@ def __init__( # input_is_parallel distinguishes column-parallel base (False, e.g. linear_qkv, # linear_fc1) from row-parallel base (True, e.g. linear_proj, linear_fc2). - # It controls which TP collective runs between the two grouped GEMMs and + # It controls which TP collective runs between the two adapter projections and # whether the second GEMM's output needs to be all-gathered to match the # 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. + # collectives between the two adapter projections. self.base_linear_is_parallel = attrs.base_linear_is_parallel self.use_a2a = a2a_experimental # Mirrors ParallelLinearAdapter's lin_out_gather_output: row-parallel @@ -144,8 +204,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. + # Immutable host metadata is cached alongside tokens_per_adapter so + # forward never synchronizes each layer to recover split sizes. + self.tokens_per_adapter_splits: Optional[Tuple[int, ...]] = None self.tokens_per_adapter_total: Optional[int] = None device = next(to_wrap.parameters()).device dtype = next(to_wrap.parameters()).dtype @@ -162,6 +223,8 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten return linear_output, bias tokens_per_adapter = self.tokens_per_adapter + token_splits = self.tokens_per_adapter_splits + assert tokens_per_adapter is not None and token_splits is not None x = layernorm_output.contiguous() # SP gather (once) — for column-parallel base layers without an LN-fused @@ -188,14 +251,15 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten 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]) + token_splits = _narrow_token_counts_to_window(token_splits, start, x_flat.shape[0]) + tokens_per_adapter = tokens_per_adapter.new_tensor(token_splits) offsets = tokens_per_adapter.cumsum(dim=0, dtype=torch.int32) stacked_A = torch.stack([a.linear_in.weight for a in self.adapters]) stacked_B = torch.stack([a.linear_out.weight for a in self.adapters]) - mid = torch._grouped_mm(x_flat, stacked_A.transpose(-2, -1), offsets) + mid = _dense_multi_lora_mm(x_flat, stacked_A, token_splits=token_splits, offsets=offsets) # TP collective between A and B: row-parallel base needs an all-reduce # of the partial sums; every other base (column-parallel and replicated @@ -207,7 +271,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten else: mid = gather_from_tensor_model_parallel_region(mid) - out = torch._grouped_mm(mid, stacked_B.transpose(-2, -1), offsets) + out = _dense_multi_lora_mm(mid, stacked_B, token_splits=token_splits, offsets=offsets) # Per-token scaling is applied *before* the output-side TP/SP comms. # ``per_token_scaling`` is indexed by the full token count @@ -595,7 +659,7 @@ 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: +def _narrow_token_counts_to_window(counts: Sequence[int], start: int, num_rows: int) -> Tuple[int, ...]: """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 @@ -603,8 +667,14 @@ def _narrow_token_counts_to_window(counts: torch.Tensor, start: int, num_rows: i 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) + end = start + num_rows + slot_start = 0 + narrowed = [] + for count in counts: + slot_end = slot_start + count + narrowed.append(max(0, min(slot_end, end) - max(slot_start, start))) + slot_start = slot_end + return tuple(narrowed) def _iter_multi_lora_modules(model): @@ -622,12 +692,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()) + # One host sync per micro-batch: cache immutable split sizes for every + # layer's fallback and any sequence-parallel narrowing. + token_splits = tuple(int(count) for count in tokens_per_adapter.tolist()) + total = sum(token_splits) for module in _iter_multi_lora_modules(model): module.tokens_per_adapter = tokens_per_adapter + module.tokens_per_adapter_splits = token_splits module.tokens_per_adapter_total = total diff --git a/tests/unit_tests/peft/test_multi_lora_layers.py b/tests/unit_tests/peft/test_multi_lora_layers.py index 2e02999d66..43f1c40e07 100644 --- a/tests/unit_tests/peft/test_multi_lora_layers.py +++ b/tests/unit_tests/peft/test_multi_lora_layers.py @@ -45,6 +45,7 @@ import pytest import torch import torch.nn as nn +import torch.nn.functional as F from megatron.bridge.models.conversion.peft_bridge import MegatronPeftBridge from megatron.bridge.peft import multi_lora as multi_lora_mod @@ -151,6 +152,125 @@ def adapter_deps_patch() -> ExitStack: return stack +def test_grouped_mm_layout_requires_16_byte_alignment() -> None: + aligned = torch.empty(4, 8, dtype=torch.bfloat16) + misaligned = torch.empty(33, dtype=torch.bfloat16)[1:].view(4, 8) + unaligned_stride = torch.empty(4, 9, dtype=torch.bfloat16) + + assert multi_lora_layers_module._has_grouped_mm_layout(aligned) + assert not multi_lora_layers_module._has_grouped_mm_layout(misaligned) + assert not multi_lora_layers_module._has_grouped_mm_layout(unaligned_stride) + + +def test_grouped_mm_requires_aligned_backward_output_width() -> None: + x = MagicMock(is_cuda=True, dtype=torch.bfloat16) + x.device = torch.device("cuda", 0) + grouped_weights = MagicMock(is_cuda=True, dtype=torch.bfloat16) + x.element_size.return_value = 2 + + with ( + patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), + patch.object(multi_lora_layers_module, "_has_grouped_mm_layout", return_value=True), + ): + grouped_weights.shape = (2, 8, 2) + assert not multi_lora_layers_module._can_use_grouped_mm(x, grouped_weights) + + grouped_weights.shape = (2, 8, 8) + assert multi_lora_layers_module._can_use_grouped_mm(x, grouped_weights) + + +def test_grouped_mm_requires_sm80_or_newer() -> None: + x = MagicMock(is_cuda=True, dtype=torch.bfloat16) + x.device = torch.device("cuda", 0) + grouped_weights = MagicMock(is_cuda=True, dtype=torch.bfloat16) + + with ( + patch.object(torch.cuda, "get_device_capability", return_value=(7, 5)), + patch.object(multi_lora_layers_module, "_has_grouped_mm_layout") as has_layout, + ): + assert not multi_lora_layers_module._can_use_grouped_mm(x, grouped_weights) + + has_layout.assert_not_called() + + +def test_dense_multi_lora_mm_preserves_eligible_grouped_mm_path() -> None: + x = torch.randn(4, 8) + weights = torch.randn(2, 8, 8) + token_counts = torch.tensor([2, 2], dtype=torch.int32) + offsets = token_counts.cumsum(dim=0) + expected = torch.cat([F.linear(x[:2], weights[0]), F.linear(x[2:], weights[1])]) + + with ( + patch.object(multi_lora_layers_module, "_can_use_grouped_mm", return_value=True), + patch.object(torch, "_grouped_mm", return_value=expected) as grouped_mm, + ): + actual = multi_lora_layers_module._dense_multi_lora_mm(x, weights, token_splits=(2, 2), offsets=offsets) + + grouped_mm.assert_called_once() + grouped_args = grouped_mm.call_args.args + assert grouped_args[0] is x + torch.testing.assert_close(grouped_args[1], weights.transpose(-2, -1)) + assert grouped_args[2] is offsets + torch.testing.assert_close(actual, expected) + + +def test_dense_multi_lora_mm_fallback_handles_empty_slots() -> None: + x = torch.empty(0, 8, requires_grad=True) + weights = torch.randn(2, 6, 8, requires_grad=True) + token_counts = torch.tensor([0, 0], dtype=torch.int32) + offsets = token_counts.cumsum(dim=0) + + with patch.object(torch, "_grouped_mm", side_effect=AssertionError("grouped MM must not run")) as grouped_mm: + actual = multi_lora_layers_module._dense_multi_lora_mm(x, weights, token_splits=(0, 0), offsets=offsets) + + grouped_mm.assert_not_called() + assert actual.shape == (0, 6) + actual.sum().backward() + assert x.grad is not None + assert weights.grad is not None + + +def test_dense_multi_lora_mm_fallback_preserves_mixed_empty_slot_gradients() -> None: + x = torch.arange(1, 33, dtype=torch.float64).reshape(4, 8).requires_grad_() + weights = torch.arange(1, 145, dtype=torch.float64).reshape(3, 6, 8).requires_grad_() + token_counts = torch.tensor([2, 0, 2], dtype=torch.int32) + offsets = token_counts.cumsum(dim=0) + expected = torch.cat([F.linear(x[:2], weights[0]), F.linear(x[2:], weights[2])]) + + with patch.object(torch, "_grouped_mm", side_effect=AssertionError("grouped MM must not run")) as grouped_mm: + actual = multi_lora_layers_module._dense_multi_lora_mm(x, weights, token_splits=(2, 0, 2), offsets=offsets) + + grouped_mm.assert_not_called() + torch.testing.assert_close(actual, expected) + actual.sum().backward() + assert x.grad is not None + assert weights.grad is not None + assert torch.count_nonzero(weights.grad[0]) > 0 + assert torch.count_nonzero(weights.grad[1]) == 0 + assert torch.count_nonzero(weights.grad[2]) > 0 + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), + reason="requires an SM80+ CUDA device", +) +def test_dense_multi_lora_mm_cuda_mixed_empty_slot_backward() -> None: + x = torch.arange(1, 33, device="cuda", dtype=torch.bfloat16).reshape(4, 8).requires_grad_() + weights = torch.arange(1, 193, device="cuda", dtype=torch.bfloat16).reshape(3, 8, 8).requires_grad_() + token_counts = torch.tensor([2, 0, 2], device="cuda", dtype=torch.int32) + + actual = multi_lora_layers_module._dense_multi_lora_mm( + x, + weights, + token_splits=(2, 0, 2), + offsets=token_counts.cumsum(dim=0), + ) + actual.float().sum().backward() + + assert weights.grad is not None + assert torch.count_nonzero(weights.grad[1]) == 0 + + # ====================================================================== # MultiLoRALinear: per-slot rank/alpha bookkeeping + rank masking # ====================================================================== @@ -170,6 +290,7 @@ def test_slot_defaults_after_construction(self) -> None: assert layer.n_adapters == 3 assert layer.max_rank == 8 assert layer.tokens_per_adapter is None + assert layer.tokens_per_adapter_splits is None assert torch.equal(layer.alpha_values, torch.ones(3)) assert torch.equal(layer.rank_values, torch.full((3,), 8.0)) @@ -196,6 +317,67 @@ def test_slot_metadata_registered_as_buffers(self) -> None: assert layer.alpha_values.dtype == torch.float64 assert layer.rank_values.dtype == torch.float64 + def test_forward_falls_back_when_grouped_mm_layout_is_ineligible(self) -> None: + """A TP-sharded BF16 local rank of two must avoid grouped MM.""" + layer = _build_multi_lora_linear(in_features=8, out_features=6, n_adapters=2, dim=2) + layer.to(dtype=torch.bfloat16) + layer.init_adapter_slot(0, rank=2, alpha=2) + layer.init_adapter_slot(1, rank=2, alpha=2) + set_tokens_per_adapter_slot([layer], torch.tensor([2, 2], dtype=torch.int32)) + + with torch.no_grad(): + layer.to_wrap.weight.zero_() + layer.to_wrap.bias.zero_() + for adapter in layer.adapters: + adapter.linear_in.weight.normal_() + adapter.linear_out.weight.normal_() + + x = torch.randn(4, 8, dtype=torch.bfloat16, requires_grad=True) + expected = torch.cat( + [ + layer.adapters[0].linear_out(layer.adapters[0].linear_in(x[:2])), + layer.adapters[1].linear_out(layer.adapters[1].linear_in(x[2:])), + ] + ) + + with ( + patch.object( + torch, + "_grouped_mm", + side_effect=RuntimeError("strides should be multiple of 16 bytes"), + ) as grouped_mm, + patch.object( + multi_lora_layers_module, "gather_from_tensor_model_parallel_region", side_effect=lambda t: t + ), + ): + actual, _ = layer(x) + + grouped_mm.assert_not_called() + torch.testing.assert_close(actual, expected) + actual.sum().backward() + assert x.grad is not None + + def test_forward_reuses_host_splits_and_grouped_offsets(self) -> None: + layer = _build_multi_lora_linear(in_features=8, out_features=8, n_adapters=2, dim=8) + set_tokens_per_adapter_slot([layer], torch.tensor([2, 2], dtype=torch.int32)) + calls = [] + + def fake_projection(x, stacked_weights, *, token_splits, offsets): + calls.append((token_splits, offsets)) + return x.new_zeros((x.shape[0], stacked_weights.shape[1])) + + with ( + patch.object(multi_lora_layers_module, "_dense_multi_lora_mm", side_effect=fake_projection), + patch.object( + multi_lora_layers_module, "gather_from_tensor_model_parallel_region", side_effect=lambda t: t + ), + ): + layer(torch.randn(4, 8)) + + assert len(calls) == 2 + assert calls[0][0] is calls[1][0] is layer.tokens_per_adapter_splits + assert calls[0][1] is calls[1][1] + def test_init_adapter_slot_sets_rank_alpha_and_masks(self) -> None: layer = _build_multi_lora_linear(dim=8) with torch.no_grad(): @@ -325,6 +507,8 @@ def test_set_tokens_per_adapter_slot(self) -> None: for module in container.mods: assert module.tokens_per_adapter is tokens + assert module.tokens_per_adapter_splits == (3, 5) + assert module.tokens_per_adapter_total == 8 def test_init_and_clear_adapter_slot_across_model(self) -> None: container = _MultiLoRAContainer(n_layers=2) @@ -660,9 +844,7 @@ def test_load_adapter_unused_keys_raises(): # 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() + return list(multi_lora_layers_module._narrow_token_counts_to_window(tuple(counts), start, num_rows)) def test_narrow_window_spanning_a_slot_boundary():