From 2feb7e47aed4c621eee1b36d680a4f509a1a3ed1 Mon Sep 17 00:00:00 2001
From: Michal Futrega <mfutrega@nvidia.com>
Date: Fri, 17 Apr 2026 14:38:20 +0200
Subject: [PATCH] [tests] test: add coverage for dense_grouped_gemm in
 GPTModelProvider
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The dense_grouped_gemm feature introduces a feature-detection branch
(via inspect.signature) in transformer_engine_layer_spec that was
uncovered — codecov flagged it on the fork's PR. Adds four unit tests:
attribute default, attribute settable, and both branches of the
signature check (forwarded when supported, omitted otherwise) using
patched fakes whose signatures include or exclude the parameter.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
---
 tests/unit_tests/models/test_gpt_provider.py | 96 ++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/tests/unit_tests/models/test_gpt_provider.py b/tests/unit_tests/models/test_gpt_provider.py
index 0b63a3d77a..6fad023327 100644
--- a/tests/unit_tests/models/test_gpt_provider.py
+++ b/tests/unit_tests/models/test_gpt_provider.py
@@ -373,3 +373,99 @@ def test_default_layer_spec_default_case(self, mock_te_full_spec, mock_te_spec):
         mock_te_full_spec.assert_not_called()
         mock_te_spec.assert_called_once_with(provider)
         assert result == "te_spec"
+
+    def test_dense_grouped_gemm_defaults_to_false(self):
+        """GPTModelProvider.dense_grouped_gemm defaults to False."""
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+        )
+        assert provider.dense_grouped_gemm is False
+
+    def test_dense_grouped_gemm_can_be_enabled(self):
+        """GPTModelProvider.dense_grouped_gemm is a settable bool attribute."""
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            dense_grouped_gemm=True,
+        )
+        assert provider.dense_grouped_gemm is True
+
+    def test_transformer_engine_layer_spec_forwards_dense_grouped_gemm_when_supported(self):
+        """When the upstream spec function exposes a dense_grouped_gemm parameter,
+        transformer_engine_layer_spec must forward the provider's value to it."""
+        from megatron.bridge.models.gpt_provider import transformer_engine_layer_spec
+
+        captured: dict = {}
+
+        # Signature intentionally includes `dense_grouped_gemm` so the feature-detect
+        # branch in gpt_provider.py activates.
+        def fake_spec_supported(
+            num_experts=None,
+            moe_grouped_gemm=False,
+            qk_layernorm=False,
+            fp8=False,
+            dense_grouped_gemm=False,
+        ):
+            captured["num_experts"] = num_experts
+            captured["moe_grouped_gemm"] = moe_grouped_gemm
+            captured["qk_layernorm"] = qk_layernorm
+            captured["fp8"] = fp8
+            captured["dense_grouped_gemm"] = dense_grouped_gemm
+            return "te_spec_supported"
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            dense_grouped_gemm=True,
+        )
+
+        with patch(
+            "megatron.bridge.models.gpt_provider.get_gpt_layer_with_transformer_engine_spec",
+            new=fake_spec_supported,
+        ):
+            result = transformer_engine_layer_spec(provider)
+
+        assert result == "te_spec_supported"
+        assert captured["dense_grouped_gemm"] is True
+
+    def test_transformer_engine_layer_spec_omits_dense_grouped_gemm_when_unsupported(self):
+        """When the upstream spec function does not expose a dense_grouped_gemm
+        parameter (older Megatron-Core), transformer_engine_layer_spec must not
+        pass the kwarg — otherwise the call would raise TypeError at runtime."""
+        from megatron.bridge.models.gpt_provider import transformer_engine_layer_spec
+
+        captured: dict = {}
+
+        # Signature intentionally excludes `dense_grouped_gemm`. If the production
+        # code were to forward it, the call below would raise TypeError.
+        def fake_spec_unsupported(
+            num_experts=None,
+            moe_grouped_gemm=False,
+            qk_layernorm=False,
+            fp8=False,
+        ):
+            captured["num_experts"] = num_experts
+            captured["moe_grouped_gemm"] = moe_grouped_gemm
+            captured["qk_layernorm"] = qk_layernorm
+            captured["fp8"] = fp8
+            return "te_spec_unsupported"
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            dense_grouped_gemm=True,
+        )
+
+        with patch(
+            "megatron.bridge.models.gpt_provider.get_gpt_layer_with_transformer_engine_spec",
+            new=fake_spec_unsupported,
+        ):
+            result = transformer_engine_layer_spec(provider)
+
+        assert result == "te_spec_unsupported"
+        assert "dense_grouped_gemm" not in captured