Address CodeRabbit review feedback on GDN FLOPs calculator

- Make GDN config strict: raise ValueError when linear_attention_freq is
  None, raise TypeError on invalid types, assert list length matches
  num_layers (mirrors existing MoE validation pattern)
- Replace getattr fallbacks for GDN dims with direct attribute access so
  missing fields fail explicitly
- Add @pytest.mark.unit to TestGDNLayerFlops class
- Strengthen test_gdn_layer_freq_list to verify 6/2 split by comparing
  against equivalent int freq=3 and against pure-standard baseline

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/training/utils/flop_utils.py

@@ -388,45 +388,55 @@ def transformer_flops():
         # weighted sum of GDN and standard-attention per-layer costs.
         experimental_attention_variant = getattr(cfg.model, "experimental_attention_variant", None)
         if experimental_attention_variant == "gated_delta_net":
-            linear_attention_freq = getattr(cfg.model, "linear_attention_freq", None)
-            if linear_attention_freq is not None:
-                if isinstance(linear_attention_freq, int):
-                    linear_attention_pattern = [
-                        0 if ((i + 1) % linear_attention_freq == 0) else 1 for i in range(num_layers)
-                    ]
-                elif isinstance(linear_attention_freq, list):
-                    linear_attention_pattern = list(linear_attention_freq)
-                else:
-                    linear_attention_pattern = [1] * num_layers
-
-                num_gdn_layers = sum(linear_attention_pattern)
-                num_standard_attn_layers = num_layers - num_gdn_layers
-
-                standard_self_attn_per_layer = self_attn_term / num_layers if num_layers > 0 else 0
-
-                qk_head_dim = getattr(cfg.model, "linear_key_head_dim", 128)
-                v_head_dim = getattr(cfg.model, "linear_value_head_dim", 128)
-                num_qk_heads = getattr(cfg.model, "linear_num_key_heads", 16)
-                num_v_heads = getattr(cfg.model, "linear_num_value_heads", 48)
-                conv_kernel_dim = getattr(cfg.model, "linear_conv_kernel_dim", 4)
-
-                qk_dim = qk_head_dim * num_qk_heads
-                v_dim = v_head_dim * num_v_heads
-
-                gdn_self_attn_per_layer = (
-                    3
-                    * 2
-                    * (
-                        cfg.model.hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads)
-                        + conv_kernel_dim * (2 * qk_dim + v_dim)
-                        + num_v_heads * (v_head_dim**2) * 4
-                        + cfg.model.hidden_size * v_dim
-                    )
+            linear_attention_freq = cfg.model.linear_attention_freq
+            if linear_attention_freq is None:
+                raise ValueError(
+                    "linear_attention_freq must be set when experimental_attention_variant='gated_delta_net'"
+                )
+            if isinstance(linear_attention_freq, int):
+                linear_attention_pattern = [
+                    0 if ((i + 1) % linear_attention_freq == 0) else 1 for i in range(num_layers)
+                ]
+            elif isinstance(linear_attention_freq, list):
+                linear_attention_pattern = list(linear_attention_freq)
+                assert len(linear_attention_pattern) == num_layers, (
+                    f"Invalid length of linear_attention_pattern: {len(linear_attention_pattern)}, "
+                    f"expected {num_layers}, "
+                    f"current linear_attention_freq: {linear_attention_freq}"
+                )
+            else:
+                raise TypeError(
+                    f"linear_attention_freq must be int or list, got {type(linear_attention_freq).__name__}"
                 )
 
-                self_attn_term = (
-                    gdn_self_attn_per_layer * num_gdn_layers + standard_self_attn_per_layer * num_standard_attn_layers
+            num_gdn_layers = sum(linear_attention_pattern)
+            num_standard_attn_layers = num_layers - num_gdn_layers
+
+            standard_self_attn_per_layer = self_attn_term / num_layers if num_layers > 0 else 0
+
+            qk_head_dim = cfg.model.linear_key_head_dim
+            v_head_dim = cfg.model.linear_value_head_dim
+            num_qk_heads = cfg.model.linear_num_key_heads
+            num_v_heads = cfg.model.linear_num_value_heads
+            conv_kernel_dim = cfg.model.linear_conv_kernel_dim
+
+            qk_dim = qk_head_dim * num_qk_heads
+            v_dim = v_head_dim * num_v_heads
+
+            gdn_self_attn_per_layer = (
+                3
+                * 2
+                * (
+                    cfg.model.hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads)
+                    + conv_kernel_dim * (2 * qk_dim + v_dim)
+                    + num_v_heads * (v_head_dim**2) * 4
+                    + cfg.model.hidden_size * v_dim
                 )
+            )
+
+            self_attn_term = (
+                gdn_self_attn_per_layer * num_gdn_layers + standard_self_attn_per_layer * num_standard_attn_layers
+            )
 
         padded_vocab_size = calculate_padded_vocab_size(
             cfg.model.vocab_size,

tests/unit_tests/training/utils/test_flop_utils.py

@@ -388,6 +388,7 @@ def test_swiglu_scaling_factor(self):
         assert flops_swiglu == expected_swiglu, f"SwiGLU: expected {expected_swiglu:.2e} but got {flops_swiglu:.2e}"
 
 
+@pytest.mark.unit
 class TestGDNLayerFlops:
     """Tests for Gated DeltaNet (GDN) FLOPs calculation in transformer_flops path."""
 
@@ -487,10 +488,10 @@ def test_gdn_only_layers(self):
         )
 
     def test_gdn_layer_freq_list(self):
-        """Test GDN with linear_attention_freq as a list pattern."""
+        """Test GDN with linear_attention_freq as a list pattern (6 GDN, 2 standard)."""
         batch_size = 1
-        # [1, 1, 0, 1, 1, 0, 1, 1] means 6 GDN, 2 standard
-        freq_list = [1, 1, 0, 1, 1, 0, 1, 1]
+        freq_list = [1, 1, 0, 1, 1, 0, 1, 1]  # 6 GDN, 2 standard
+        assert sum(freq_list) == 6
         model_cfg = self._qwen35_27b_config(
             num_layers=8,
             linear_attention_freq=freq_list,
@@ -499,6 +500,20 @@ def test_gdn_layer_freq_list(self):
         flops = num_floating_point_operations(cfg, batch_size=batch_size)
         assert flops > 0
 
+        # Verify the mask is actually applied: must differ from pure standard attention.
+        baseline_cfg = MockConfigContainer(
+            model=self._qwen35_27b_config(num_layers=8, experimental_attention_variant=None)
+        )
+        baseline_flops = num_floating_point_operations(baseline_cfg, batch_size=batch_size)
+        assert flops != baseline_flops, "List-based GDN mask should differ from pure standard attention"
+
+        # freq_list [1,1,0,1,1,0,1,1] is identical to the pattern generated by int freq=3.
+        int_freq_cfg = MockConfigContainer(model=self._qwen35_27b_config(num_layers=8, linear_attention_freq=3))
+        int_freq_flops = num_floating_point_operations(int_freq_cfg, batch_size=batch_size)
+        assert flops == int_freq_flops, (
+            "List [1,1,0,1,1,0,1,1] should produce the same FLOPs as int freq=3 (equivalent 6/2 split)"
+        )
+
     def test_gdn_exact_self_attn_term(self):
         """Verify the GDN self_attn_term matches the expected formula from Megatron-LM."""
         batch_size = 1