fix: Remove experimental_attention_variant arg from get_gpt_layer call in mlm_compat (#2073)

Signed-off-by: Charlie Truong <[email protected]>
Co-authored-by: oliver könig <[email protected]>

Author: chtruong814

src/megatron/bridge/training/mlm_compat/model.py

@@ -59,7 +59,6 @@ def _get_transformer_layer_spec(args: argparse.Namespace, use_te: bool, use_kitc
             moe_grouped_gemm=args.moe_grouped_gemm,
             qk_layernorm=args.qk_layernorm,
             multi_latent_attention=args.multi_latent_attention,
-            experimental_attention_variant=getattr(args, "experimental_attention_variant", None),
             moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             qk_l2_norm=args.qk_l2_norm,
             use_kitchen=use_kitchen,
@@ -70,7 +69,6 @@ def _get_transformer_layer_spec(args: argparse.Namespace, use_te: bool, use_kitc
             moe_grouped_gemm=args.moe_grouped_gemm,
             qk_layernorm=args.qk_layernorm,
             multi_latent_attention=args.multi_latent_attention,
-            experimental_attention_variant=getattr(args, "experimental_attention_variant", None),
             moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             normalization=args.normalization,
             use_kitchen=use_kitchen,

tests/unit_tests/Launch_Unit_Tests.sh

@@ -26,7 +26,7 @@ fi
 echo ""
 
 CUDA_VISIBLE_DEVICES="0,1" uv run coverage run -a --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ -m pytest \
-    --timeout=0.75 \
+    --timeout=1 \
     -o log_cli=true \
     -o log_cli_level=INFO \
     --disable-warnings \

tests/unit_tests/training/mlm_compat/test_model.py

@@ -105,7 +105,6 @@ def test_te_spec(self, mock_te_spec_func, mock_args):
             moe_grouped_gemm=False,
             qk_layernorm=False,
             multi_latent_attention=False,
-            experimental_attention_variant=None,
             moe_use_legacy_grouped_gemm=False,
             qk_l2_norm=False,
             use_kitchen=False,
@@ -121,7 +120,6 @@ def test_local_spec(self, mock_local_spec_func, mock_args):
             moe_grouped_gemm=False,
             qk_layernorm=False,
             multi_latent_attention=False,
-            experimental_attention_variant=None,
             moe_use_legacy_grouped_gemm=False,
             normalization="LayerNorm",
             use_kitchen=True,