peft: fix multi-LoRA activation-recompute replay and grouped-expert adapter DDP flagging

src/megatron/bridge/peft/recompute.py

@@ -75,9 +75,14 @@ def maybe_enable_recompute_inputs_grad(model, peft_recompute_patched: Set[int] |
                 continue
 
             params = list(unwrapped_model.named_parameters())
-            trainable_adapter = any(p.requires_grad and ".adapter." in n.lower() for n, p in params)
+            # Multi-LoRA slots live under ".adapters.<slot>."; single-LoRA under ".adapter.".
+            trainable_adapter = any(
+                p.requires_grad and (".adapter." in n.lower() or ".adapters." in n.lower()) for n, p in params
+            )
             trainable_base = any(
-                p.requires_grad and (".to_wrap." not in n.lower() and ".adapter." not in n.lower()) for n, p in params
+                p.requires_grad
+                and (".to_wrap." not in n.lower() and ".adapter." not in n.lower() and ".adapters." not in n.lower())
+                for n, p in params
             )
 
             if not (trainable_adapter and not trainable_base):

src/megatron/bridge/peft/utils.py

@@ -27,6 +27,7 @@
 from megatron.core import ModelParallelConfig, parallel_state
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict, ShardedTensor, ShardedTensorFactory
 from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
+from megatron.core.tensor_parallel.layers import set_tensor_model_parallel_attributes
 from megatron.core.tensor_parallel.mappings import (
     gather_from_sequence_parallel_region,
     scatter_to_sequence_parallel_region,
@@ -1119,23 +1120,31 @@ def __init__(
         ParallelLinearAdapter._get_init_fn(self, column_init_method)(linear_in_weight)
         ParallelLinearAdapter._get_init_fn(self, row_init_method)(linear_out_weight)
 
-        expert_parallel = (
-            parallel_state.get_expert_model_parallel_world_size() or model_parallel_config.expert_model_parallel_size
-        ) > 1
         self._linear_in_tp_axis = linear_in_tp_axis
         self._linear_out_tp_axis = linear_out_tp_axis
         self.linear_in = nn.Module()
         self.linear_in.weight = nn.Parameter(linear_in_weight)
         self.linear_out = nn.Module()
         self.linear_out.weight = nn.Parameter(linear_out_weight)
+        # Ported from upstream (NVIDIA-NeMo/Megatron-Bridge main): adapter
+        # gradients must reduce over the expert process groups whenever the
+        # experts are EP-sharded or the expert TP size differs from the dense
+        # TP size — with ETP < TP the weights are TP-duplicated and only the
+        # expert DP group folds those replicas' gradients together. The dense
+        # bucket is safe only when both groups coincide (EP == 1, ETP == TP);
+        # deriving the flag from EP size alone left EP=1, TP>1 runs unsynced.
+        use_expert_process_groups = (
+            parallel_state.get_expert_model_parallel_world_size() or model_parallel_config.expert_model_parallel_size
+        ) > 1 or expert_tp_size != (
+            parallel_state.get_tensor_model_parallel_world_size() or model_parallel_config.tensor_model_parallel_size
+        )
         for weight, tp_axis in (
             (self.linear_in.weight, linear_in_tp_axis),
             (self.linear_out.weight, linear_out_tp_axis),
         ):
-            setattr(weight, "allreduce", not expert_parallel)
+            setattr(weight, "allreduce", not use_expert_process_groups)
             if tp_axis is not None:
-                setattr(weight, "partition_dim", tp_axis)
-                setattr(weight, "partition_stride", 1)
+                set_tensor_model_parallel_attributes(weight, True, tp_axis, 1)
 
         if dropout > 0.0:
             self.dropout = nn.Dropout(dropout)

tests/unit_tests/peft/test_recompute.py

@@ -59,6 +59,25 @@ def modules(self):
             yield module
 
 
+class DummyMultiLoRAModel(torch.nn.Module):
+    """Model whose only trainable params are multi-LoRA slots (".adapters.<i>.")."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.config = SimpleNamespace(recompute_method="uniform")
+        self.block = DummyTransformerBlock()
+
+        # Frozen base parameter (not trainable)
+        self.base = torch.nn.Linear(1, 1, bias=False)
+        self.base.weight.requires_grad = False
+
+        # Multi-LoRA wrappers hold one adapter per slot in an ``adapters``
+        # ModuleList, so parameter names contain ".adapters.<i>." rather than
+        # ".adapter.". Nest under a ModuleDict so the full name carries the
+        # leading dot (e.g. "linear_fc1.adapters.0.weight").
+        self.linear_fc1 = torch.nn.ModuleDict({"adapters": torch.nn.ModuleList([DummyAdapter(), DummyAdapter()])})
+
+
 def _patch_transformer_block(monkeypatch):
     import megatron.core.transformer.transformer_block as transformer_block
 
@@ -90,3 +109,46 @@ def test_maybe_enable_recompute_inputs_grad_patches_block(monkeypatch):
     # Second invocation should be a no-op (no duplicate patch)
     maybe_enable_recompute_inputs_grad(model, patched_registry)
     assert model.block.forward is patched_forward
+
+
+def test_maybe_enable_recompute_inputs_grad_patches_block_multi_lora(monkeypatch):
+    """Multi-LoRA slot params (".adapters.<i>.") must be recognized as adapters.
+
+    Regression test: they used to be classified as trainable base weights, the
+    patch was skipped, and under full activation recompute the checkpointed
+    region never replayed in backward — every adapter grad stayed zero.
+    """
+    _patch_transformer_block(monkeypatch)
+    recompute_mod.PEFT_RECOMPUTE_PATCHED.clear()
+
+    model = DummyMultiLoRAModel()
+    param_names = [n for n, _ in model.named_parameters()]
+    assert any(".adapters." in n for n in param_names)
+    assert not any(".adapter." in n for n in param_names)
+
+    patched_registry = maybe_enable_recompute_inputs_grad(model, set())
+
+    assert id(model) in patched_registry
+
+    input_tensor = torch.zeros(2, 2)
+    assert input_tensor.requires_grad is False
+
+    model.block(input_tensor)
+    assert model.block.last_input_requires_grad is True
+
+
+def test_maybe_enable_recompute_inputs_grad_skips_trainable_base(monkeypatch):
+    """A genuinely trainable base weight must still disable the patch."""
+    _patch_transformer_block(monkeypatch)
+    recompute_mod.PEFT_RECOMPUTE_PATCHED.clear()
+
+    model = DummyMultiLoRAModel()
+    model.base.weight.requires_grad = True
+
+    patched_registry = maybe_enable_recompute_inputs_grad(model, set())
+
+    assert id(model) not in patched_registry
+
+    input_tensor = torch.zeros(2, 2)
+    model.block(input_tensor)
+    assert model.block.last_input_requires_grad is False

tests/unit_tests/peft/test_utils.py

@@ -670,6 +670,45 @@ def test_parallel_linear_adapter_sharded_state_dict_fc1_special_case(
 class TestGroupedExpertLinearAdapter:
     """Tests for grouped-expert per-expert LoRA adapters."""
 
+    @pytest.mark.parametrize(
+        ("tp_size", "expected_allreduce"),
+        [(2, False), (1, True)],
+    )
+    def test_grouped_expert_linear_adapter_weight_flags_follow_expert_process_groups(
+        self, tp_size, expected_allreduce
+    ):
+        """allreduce mirrors upstream's use_expert_process_groups condition.
+
+        With ETP != TP (here ETP=1, TP=2) the TP-duplicated adapter weights
+        must route to DDP's expert bucket (allreduce=False) whose reduce group
+        folds in the TP peers; the old EP-only condition left them in the
+        dense bucket, so cross-TP gradients never synchronized and replicas
+        silently diverged. When the groups coincide (EP=1, ETP == TP) the
+        dense bucket is equivalent and allreduce stays True.
+        """
+        config = MockModelParallelConfig()
+        config.tensor_model_parallel_size = tp_size
+        adapter = GroupedExpertLinearAdapter(
+            in_features=4,
+            out_features=4,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=True,
+            model_parallel_config=config,
+        )
+
+        for weight, expected_axis in (
+            (adapter.linear_in.weight, 2),
+            (adapter.linear_out.weight, 1),
+        ):
+            assert weight.allreduce is expected_allreduce
+            assert weight.tensor_model_parallel is True
+            assert weight.partition_dim == expected_axis
+            assert weight.partition_stride == 1
+        assert adapter.is_expert is True
+
     @pytest.mark.parametrize("split_kwarg", ["m_splits", "tokens_per_expert"])
     def test_grouped_expert_linear_adapter_accepts_tensor_split_kwargs(self, split_kwarg):
         """Tensor-valued split kwargs should not trigger ambiguous truth-value errors."""