fix: Remove experimental_attention_variant arg from get_gpt_layer call in mlm_compat

[chtruong814]

What does this PR do ?

This removes the kwarg experimental_attention_variant from the get_gpt_layer calls in mlm_compat. That kwarg was removed in this PR on the MCore main branch recently.
NVIDIA/Megatron-LM#1989

We initially encountered a problem with this with Export-Deploy after attempting to use the latest version of MBridge and MCore.

  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/model_load_save.py", line 328, in build_and_load_model
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     return _load_checkpoint()
  (ServeReplica:llama:MegatronRayDeployable pid=47125)            ^^^^^^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/model_load_save.py", line 308, in _load_checkpoint
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     model = _call_model_provider(model_cfg)
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/model_load_save.py", line 292, in _call_model_provider
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     return _get_model(megatron_args, provider, model_cfg)
  (ServeReplica:llama:MegatronRayDeployable pid=47125)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/mlm_compat/model.py", line 220, in _get_model
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     model = build_model()
  (ServeReplica:llama:MegatronRayDeployable pid=47125)             ^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/mlm_compat/model.py", line 212, in build_model
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     model = model_provider_func(args, model_cfg, pre_process=pre_process, post_process=post_process)
  (ServeReplica:llama:MegatronRayDeployable pid=47125)             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/mlm_compat/model.py", line 113, in _gpt_provider
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     transformer_layer_spec = _get_transformer_layer_spec(args, use_te, config.use_kitchen)
  (ServeReplica:llama:MegatronRayDeployable pid=47125)                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125)   File "/opt/venv/lib/python3.12/site-packages/megatron/bridge/training/mlm_compat/model.py", line 57, in _get_transformer_layer_spec
  (ServeReplica:llama:MegatronRayDeployable pid=47125)     return get_gpt_layer_with_transformer_engine_spec(
  (ServeReplica:llama:MegatronRayDeployable pid=47125)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  (ServeReplica:llama:MegatronRayDeployable pid=47125) TypeError: get_gpt_layer_with_transformer_engine_spec() got an unexpected keyword argument 'experimental_attention_variant'

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

• Refactor
  • Internal optimization to training components with no impact on user-facing functionality or public interfaces.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Removed the experimental_attention_variant parameter from two function calls within the MLM compatibility model's transformer layer specification method. The parameter is no longer passed to either the Transformer Engine or local layer specification paths, simplifying the function signatures while maintaining existing control flow.

Changes

Cohort / File(s)	Change Summary
MLM Attention Variant Parameter Removal src/megatron/bridge/training/mlm_compat/model.py	Removed experimental_attention_variant argument from calls to get_gpt_layer_with_transformer_engine_spec() and get_gpt_layer_local_spec() in _get_transformer_layer_spec()

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: removing the experimental_attention_variant argument from get_gpt_layer calls in mlm_compat, which directly matches the changeset.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Test Results For Major Changes	✅ Passed	PR contains a minor bug fix removing outdated parameter arguments from two function calls to restore MCore API compatibility.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}