From 3ec5b9deac57d22fcc0a4aa47b32833ddd8b4103 Mon Sep 17 00:00:00 2001
From: Arist12 <ykzhang12@gmail.com>
Date: Fri, 21 Aug 2026 18:27:18 +0000
Subject: [PATCH] [peft] fix: support ROCm TE grouped-linear fields

---
 src/megatron/bridge/peft/utils.py   |  5 +++
 tests/unit_tests/peft/test_utils.py | 52 +++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py
index 2fa1d740b4..d3133368f7 100644
--- a/src/megatron/bridge/peft/utils.py
+++ b/src/megatron/bridge/peft/utils.py
@@ -1330,6 +1330,11 @@ def _forward_te_grouped_linear(
                 "skip_fp8_weight_update": None,
                 "save_original_input": helper.save_original_input,
                 "debug": False,
+                # ROCm TE adds these optional fields. Other TE builds do not select
+                # them because the tuple is assembled from their own field_names.
+                "m_splits_tensor": None,
+                "actual_m_splits": None,
+                "unpad_output": False,
             }
             missing = [name for name in field_names if name not in available]
             if missing:
diff --git a/tests/unit_tests/peft/test_utils.py b/tests/unit_tests/peft/test_utils.py
index b97cc2e1fa..0ad957c62d 100644
--- a/tests/unit_tests/peft/test_utils.py
+++ b/tests/unit_tests/peft/test_utils.py
@@ -958,6 +958,58 @@ def test_grouped_expert_linear_adapter_te_grouped_mlp_prefers_te_backend_over_gr
         assert mock_te_backend.call_args_list[0].kwargs["m_splits"] == [1, 2]
         assert mock_te_backend.call_args_list[1].kwargs["m_splits"] == [1, 2]
 
+    @patch(
+        "megatron.bridge.peft.utils.parallel_state.get_tensor_model_parallel_world_size",
+        return_value=1,
+    )
+    def test_te_grouped_linear_supplies_rocm_optional_fields(self, _mock_tp_world_size):
+        """ROCm-only fields should use the defaults from TE's public wrapper."""
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=MockModelParallelConfig(),
+        )
+        helper = Mock(
+            apply_bias=False,
+            fp8=False,
+            fp8_calibration=False,
+            wgrad_store=None,
+            fuse_wgrad_accumulation=False,
+            sequence_parallel=False,
+            activation_dtype=torch.float32,
+            save_original_input=False,
+        )
+        helper.prepare_forward.side_effect = lambda value, **_kwargs: value
+        helper._get_quantizers.return_value = (None, None, None, None, None, None)
+        x = torch.ones(2, 2)
+        weight = torch.ones(2, 2, 2)
+
+        def fake_forward(_ctx, value, non_tensor_args, *_tensors):
+            assert non_tensor_args == (None, None, False)
+            return value
+
+        with (
+            torch.no_grad(),
+            patch.object(adapter, "_get_te_grouped_linear_helper", return_value=helper),
+            patch(
+                "megatron.bridge.peft.utils._te_grouped_linear_contract",
+                return_value=(False, ["m_splits_tensor", "actual_m_splits", "unpad_output"]),
+            ),
+            patch(
+                "megatron.bridge.peft.utils.TEPytorchGroupedLinearAutograd.forward",
+                side_effect=fake_forward,
+            ),
+        ):
+            output = adapter._forward_te_grouped_linear(x, weight=weight, m_splits=[1, 1])
+
+        torch.testing.assert_close(output, x)
+        helper.end_forward.assert_called_once()
+
     def test_grouped_expert_linear_adapter_requires_expert_tp_group_for_gather(self):
         """Per-expert LoRA should fail clearly when expert TP is configured without initialized groups."""
         config = MockModelParallelConfig()