[model] feat: add GLM-5 (MoE + MLA + DSA) bridge and provider

[yaoyu-33]

What does this PR do?

Takes over #2469 and rebases onto latest main.

Adds support for the GLM-5 model (MoE + MLA + DSA architecture).

Changes

• GLM5Bridge: bidirectional HF ↔ Megatron-Core checkpoint conversion for GlmMoeDsaForCausalLM
• GLM5ModelProvider: config/arch mapping for the glm_moe_dsa architecture
• Overrides maybe_modify_loaded_hf_weight / maybe_modify_converted_hf_weight to handle fused expert tensors (3D) introduced in transformers>=5.2.0
• Exposes GLM5Bridge and GLM5ModelProvider in the glm_moe_dsa package
• Functional tests for GLM-5 conversion across parallelism configs

Note: Requires transformers>=5.2.0 locally to use this feature, due to the new fused expert weight layout.

Related

Closes #2343
Takes over #2469

GitHub Actions CI

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?
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for GLM-5 models with Mixture of Experts (MoE) and Distributed Sparse Attention (DSA) capabilities, including HuggingFace-to-Megatron conversion.

• Updates

  • Reorganized model provider exports for improved API structure, including expanded support for Llama, DeepSeek, NeMoTron, and GLM variants.
  • Updated upstream Megatron-LM dependency.

• Tests

  • Added functional test suite for GLM-5 model conversion validation across multiple parallelism configurations.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 6e1bf75

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request adds GLM-5 MoE+DSA model support to the Megatron bridge framework by introducing a new bridge implementation and model provider with comprehensive conversion logic. Additionally, the main models package exports are refactored to transition from bridge-based to provider-based classes, and the Megatron-LM submodule pointer is updated.

Changes

Cohort / File(s)	Summary
Submodule Update 3rdparty/Megatron-LM	Updated Megatron-LM submodule commit reference from 9ed8b0c to 0e74689.
Models Package Export Refactoring src/megatron/bridge/models/__init__.py	Transitioned exports from bridge-based classes to provider-based classes across multiple model families (DeepSeek, GLM, GPT-OSS, Llama, NeMoTron, etc.). Removed obsolete imports for Qwen Omni, Sarvam, and Qwen-VL-35 variants while adding granular provider exports for model size variants.
GLM-5 MoE+DSA Bridge Implementation src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py	Implements GLM5Bridge for HuggingFace-to-Megatron conversion, featuring HF state enumeration, extensive layer-wise conversion mapping for MLA/DSA attention variants, expert-based MoE handling with fused/per-expert tensor strategies, and conditional MTP-layer support. Includes expert tensor merging logic for roundtrip export.
GLM-5 Model Provider Configuration src/megatron/bridge/models/glm_moe_dsa/glm5_provider.py	Defines GLM5ModelProvider dataclass with MLATransformerConfig integration, conditional transformer layer-spec selection (experimental variant with fallback), and comprehensive default hyperparameters covering RoPE, MLA ranks/dimensions, DSA indexer settings, MoE configuration, and MTP auxiliary parameters.
GLM-5 Package Initialization src/megatron/bridge/models/glm_moe_dsa/__init__.py	Package initializer that exports GLM5Bridge and GLM5ModelProvider from their respective modules.
GLM-5 Functional Tests tests/functional_tests/models/glm_moe_dsa/*.py	Functional test suite for GLM-5 conversion, including toy model creation fixture, roundtrip conversion validation across multiple parallelism strategies (TP/PP/EP), and optional disk-local model testing. Verifies model structure, configuration preservation, and conversion correctness under GPU-only execution.

Sequence Diagram(s)

sequenceDiagram
    actor User
    participant HFModel as HuggingFace<br/>GlmMoeDsaForCausalLM
    participant GLM5Bridge as GLM5Bridge<br/>Converter
    participant MappingRegistry as Conversion<br/>Mapping Registry
    participant MegatronModel as Megatron-Core<br/>GPTModel

    User->>GLM5Bridge: Initiate conversion<br/>(hf_pretrained, megatron_model)
    GLM5Bridge->>HFModel: Enumerate state keys<br/>& load config
    HFModel-->>GLM5Bridge: HF state dict & config
    
    GLM5Bridge->>GLM5Bridge: Build GLM5ModelProvider<br/>with MLA/DSA/MoE params
    GLM5Bridge->>MappingRegistry: Construct mapping registry<br/>(global + layer-wise)
    
    MappingRegistry->>MappingRegistry: Detect expert layout<br/>(fused vs per-expert)
    MappingRegistry->>MappingRegistry: Define attention/MLA/DSA<br/>& MoE expert mappings
    MappingRegistry-->>GLM5Bridge: Configured registry
    
    GLM5Bridge->>GLM5Bridge: Execute conversion tasks<br/>Apply mappings to weights
    GLM5Bridge->>GLM5Bridge: Merge per-expert tensors<br/>for HF export format
    
    GLM5Bridge->>MegatronModel: Load converted weights
    MegatronModel-->>User: Converted model ready

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 53.85% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Test Results For Major Changes	⚠️ Warning	PR adds major GLM-5 feature but lacks test execution results and evidence of passing tests; unresolved review issues identified.	Add actual test execution results, confirmation of model validation against source, and address all review comments before merge.

✅ 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 '[model] feat: add GLM-5 (MoE + MLA + DSA) bridge and provider' directly and concisely summarizes the main change: adding GLM-5 bridge and provider with specific technical features.
Linked Issues check	✅ Passed	The PR successfully implements GLM-5 bridge and provider support as requested in issue #2343, including HF↔Megatron conversion, architecture mapping, expert tensor handling, and functional tests [#2343].
Out of Scope Changes check	✅ Passed	All code changes directly support GLM-5 bridge/provider implementation. The submodule update to Megatron-LM is a dependency refresh likely needed for GLM-5 support, and export changes align with integrating new providers into the public API.

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

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch yuya-add-glm5

📝 Coding Plan

• Generate coding plan for human review comments

---

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.}

[coderabbitai]

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@3rdparty/Megatron-LM`:
- Line 1: The PR points the 3rdparty/Megatron-LM submodule at an unreachable
commit (0e7468950283aaa89d36c3ef105892db8dd00f56); update the submodule to a
published, fetchable commit: check out a known public commit in the
3rdparty/Megatron-LM repo, update the gitlink in the superproject (stage and
commit the change), and ensure .gitmodules (if modified) is consistent so git
submodule update works in CI; reference the submodule path 3rdparty/Megatron-LM
and replace the bad commit hash with the valid published commit hash in the
superproject commit.

In `@src/megatron/bridge/models/__init__.py`:
- Around line 135-156: Two different modules export NemotronNano12Bv2Provider
and the second import overwrites the first; fix by aliasing one or both imports
(e.g., import NemotronNano12Bv2Provider from megatron.bridge.models.nemotron_vl
as NemotronNano12Bv2ProviderVL and from
megatron.bridge.models.nemotronh.nemotron_h_provider as
NemotronNano12Bv2ProviderH), then update the top-level exports that currently
reference NemotronNano12Bv2Provider to point to the correct alias (choose which
provider should be exported under the public name) and update any internal
references to use the new alias names (check usages of
NemotronNano12Bv2Provider, NemotronVLBridge, and entries in the module's __all__
if present).
- Around line 104-111: The import currently brings in LlamaNemotronBridge but it
is not exported in __all__, causing a flake8 F401; either remove
LlamaNemotronBridge from the import tuple in the
megatron.bridge.models.llama_nemotron import block or add "LlamaNemotronBridge"
to the module's __all__ list to explicitly re-export it; locate the import
statement that lists Llama31Nemotron70BProvider, Llama31NemotronNano8BProvider,
Llama31NemotronUltra253BProvider, Llama33NemotronSuper49BProvider,
LlamaNemotronBridge, LlamaNemotronHeterogeneousProvider and update it
accordingly so the import/export surface is consistent.
- Around line 67-70: The module re-exports GLM-5 symbols unguarded causing
ImportError on transformers<5.3.0; wrap the import of GLM5Bridge and
GLM5ModelProvider from megatron.bridge.models.glm_moe_dsa in a try/except
ImportError (catching the missing GlmMoeDsaForCausalLM) and on failure set those
names to None or omit them from __all__ so importing megatron.bridge.models
won't raise; alternatively update the package requirement to transformers>=5.3.0
in pyproject.toml if you prefer to enforce availability instead of guarding the
import.

In `@src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py`:
- Around line 127-135: Replace the warning in the
ImportError/ModuleNotFoundError handler with a fail-fast exception: instead of
logger.warning("DSA spec not available; falling back to standard GPT decoder
block spec."), raise a NotImplementedError (or RuntimeError) that clearly states
the DSA experimental_attention variant is required because subsequent conversion
registers DSA-specific target mappings (e.g., self_attn.indexer.*) and will fail
or skip weights if the spec is absent; reference provider.transformer_layer_spec
and get_transformer_block_with_experimental_attention_variant_spec in the
message so maintainers know which component is missing.
- Around line 400-449: This method implements a bridge-local fused-expert merge
that duplicates grouping/transpose behavior already handled by
GLMExpertGateUpProjMapping/GLMExpertDownProjMapping and must be removed; replace
the body of maybe_modify_converted_hf_weight so it no longer performs any
per-expert caching/stacking or special-case mapping checks (remove all uses of
self.hf_weights_cache and the if isinstance(...
GLMExpertGateUpProjMapping/GLMExpertDownProjMapping) branch and the merging
logic), and instead simply return converted_weights_dict unchanged (or the
incoming dict) so grouped export/transpose behavior is driven solely by the
mapping/param_mapping code (keep the method signature unchanged).

In `@tests/functional_tests/models/glm_moe_dsa/test_glm5_conversion.py`:
- Around line 138-145: Add an explicit 2-GPU requirement check and
documentation: in the module containing _build_roundtrip_cmd, import pytest and
torch, add a module-level docstring or test docstrings noting "requires 2 CUDA
devices", and either decorate affected tests with
`@pytest.mark.skipif`(torch.cuda.device_count() < 2, reason="requires 2 CUDA
devices") or call pytest.skip at the start of those test functions when
torch.cuda.device_count() < 2; ensure the _build_roundtrip_cmd usage remains
unchanged but tests are skipped early with a clear message about the 2-GPU
requirement.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 511bd143-5667-4f5c-8f39-5458ec7f6045

📥 Commits

Reviewing files that changed from the base of the PR and between f565adb and 6e1bf75.

📒 Files selected for processing (7)

• 3rdparty/Megatron-LM
• src/megatron/bridge/models/__init__.py
• src/megatron/bridge/models/glm_moe_dsa/__init__.py
• src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py
• src/megatron/bridge/models/glm_moe_dsa/glm5_provider.py
• tests/functional_tests/models/glm_moe_dsa/__init__.py
• tests/functional_tests/models/glm_moe_dsa/test_glm5_conversion.py

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Point the submodule at a published Megatron-LM commit.

Both reported pipelines fail during git submodule update with fatal: ... not our ref 0e7468950283aaa89d36c3ef105892db8dd00f56. As-is, the repo cannot fetch this revision, so the PR is not buildable or testable.

🧰 Tools

🪛 GitHub Actions: Build docs

[error] 1-1: Submodule update failed during 'git submodule update --init --force --depth=1 --recursive'. Git error: 'fatal: remote error: upload-pack: not our ref 0e74689'.

🪛 GitHub Actions: CICD NeMo

[error] 1-1: Submodule update failed: "fatal: remote error: upload-pack: not our ref 0e74689" during 'git -c protocol.version=2 submodule update --init --force --depth=1 --recursive'.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@3rdparty/Megatron-LM` at line 1, The PR points the 3rdparty/Megatron-LM
submodule at an unreachable commit (0e7468950283aaa89d36c3ef105892db8dd00f56);
update the submodule to a published, fetchable commit: check out a known public
commit in the 3rdparty/Megatron-LM repo, update the gitlink in the superproject
(stage and commit the change), and ensure .gitmodules (if modified) is
consistent so git submodule update works in CI; reference the submodule path
3rdparty/Megatron-LM and replace the bad commit hash with the valid published
commit hash in the superproject commit.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Show the unguarded import chain for the new GLM-5 exports.
rg -n -C2 'from megatron\.bridge\.models\.glm_moe_dsa import|from transformers import GlmMoeDsaForCausalLM' \
  src/megatron/bridge/models/__init__.py \
  src/megatron/bridge/models/glm_moe_dsa/__init__.py \
  src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 919

---

🏁 Script executed:

# Check how other optional models are handled in __init__.py
# Look for try-except patterns and imports
head -100 src/megatron/bridge/models/__init__.py | tail -50

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 1374

---

🏁 Script executed:

# Search for try-except patterns around model imports in __init__.py
rg -n 'try:|except|ImportError|ModuleNotFoundError' src/megatron/bridge/models/__init__.py -A 2 -B 2

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 259

---

🏁 Script executed:

# Check if glm_moe_dsa has an __init__.py that does the import
cat src/megatron/bridge/models/glm_moe_dsa/__init__.py 2>/dev/null || echo "No __init__.py found"

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 889

---

🏁 Script executed:

# Check how other transformers imports are handled in model bridges
rg -n 'from transformers import' src/megatron/bridge/models/*/\*.py | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 144

---

🏁 Script executed:

# Check full structure of __init__.py for try-except guards around imports
rg -n 'try:|except|ImportError' src/megatron/bridge/models/__init__.py | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 139

---

🏁 Script executed:

# Count total lines in __init__.py to understand scope
wc -l src/megatron/bridge/models/__init__.py

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 113

---

🏁 Script executed:

# Search for transformers imports in model bridge files
find src/megatron/bridge/models -name "*.py" -type f -exec grep -l "from transformers import" {} \; | head -10

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 1280

---

🏁 Script executed:

# Look at the full __init__.py to check for try-except patterns
cat src/megatron/bridge/models/__init__.py | grep -n "try\|except" | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 264

---

🏁 Script executed:

# Check pyproject.toml or requirements files for transformers version constraints
find . -maxdepth 2 -name "pyproject.toml" -o -name "requirements*.txt" | xargs cat 2>/dev/null | grep -i transformers | head -10

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 105

---

🏁 Script executed:

# Check if GlmMoeDsaForCausalLM is mentioned in any config or doc
rg -i "GlmMoeDsaForCausalLM" --type py --type md --type toml --type txt | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 1184

---

🏁 Script executed:

# Check Megatron transformers requirement to see if GlmMoeDsaForCausalLM is available
rg "transformers" pyproject.toml setup.py setup.cfg 2>/dev/null | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 120

---

🌐 Web query:

GlmMoeDsaForCausalLM transformers version history available since when

💡 Result:

GlmMoeDsaForCausalLM was added to Hugging Face Transformers via PR #43858 ("Add GlmMoeDsa"), which merged approximately 36 days before 2026-03-20, placing it around mid-February 2026. The model documentation source file indicates it was added to Transformers on 2026-02-08. It supports GLM-5 (zai-org/GLM-5), a 744B parameter MoE model released around 2026-02-09. The class is available in Transformers v5.3.0 and later, as evidenced by code references in the docs.

Citations:

• 1: https://huggingface.co/docs/transformers/main/model_doc/glm_moe_dsa
• 2: https://github.com/huggingface/transformers/blob/main/docs/source/en/model_doc/glm_moe_dsa.md
• 3: https://app.semanticdiff.com/gh/huggingface/transformers/commit/33a8e68d4955c13cf6d6c629e76f12fc81974fc5
• 4: https://github.com/zai-org/GLM-5

---

🏁 Script executed:

# Check if there's any conditional import logic in test files that might hint at version requirements
rg -n "GlmMoeDsaForCausalLM" tests/functional_tests/models/glm_moe_dsa/test_glm5_conversion.py -B5 -A5 | head -30

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 1393

---

🏁 Script executed:

# Check if any other model bridges use try-except for transformers imports
rg -n "try:|except" src/megatron/bridge/models/*/\*_bridge.py 2>/dev/null | head -20

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 53

---

Guard the GLM-5 re-export behind version checking or raise the minimum transformers version to 5.3.0.

GlmMoeDsaForCausalLM was added to Transformers on 2026-02-08 and is only available in v5.3.0+. However, pyproject.toml allows transformers>=5.0.0, meaning users on 5.0.0–5.2.x will encounter ImportError: cannot import name 'GlmMoeDsaForCausalLM' when importing megatron.bridge.models, even if they never use GLM-5. Wrap the import in a try-except block or bump the minimum transformers version constraint to >=5.3.0.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/models/__init__.py` around lines 67 - 70, The module
re-exports GLM-5 symbols unguarded causing ImportError on transformers<5.3.0;
wrap the import of GLM5Bridge and GLM5ModelProvider from
megatron.bridge.models.glm_moe_dsa in a try/except ImportError (catching the
missing GlmMoeDsaForCausalLM) and on failure set those names to None or omit
them from __all__ so importing megatron.bridge.models won't raise; alternatively
update the package requirement to transformers>=5.3.0 in pyproject.toml if you
prefer to enforce availability instead of guarding the import.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Remove or re-export LlamaNemotronBridge.

It's still imported in this block, but it's no longer listed in __all__. That matches the flake8 F401 failure already reported for this file and leaves the intended top-level export surface ambiguous.

🧰 Tools

🪛 Flake8 (7.3.0)

[error] 104-104: 'megatron.bridge.models.llama_nemotron.LlamaNemotronBridge' imported but unused

(F401)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/models/__init__.py` around lines 104 - 111, The import
currently brings in LlamaNemotronBridge but it is not exported in __all__,
causing a flake8 F401; either remove LlamaNemotronBridge from the import tuple
in the megatron.bridge.models.llama_nemotron import block or add
"LlamaNemotronBridge" to the module's __all__ list to explicitly re-export it;
locate the import statement that lists Llama31Nemotron70BProvider,
Llama31NemotronNano8BProvider, Llama31NemotronUltra253BProvider,
Llama33NemotronSuper49BProvider, LlamaNemotronBridge,
LlamaNemotronHeterogeneousProvider and update it accordingly so the
import/export surface is consistent.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Resolve the NemotronNano12Bv2Provider name collision.

This module imports NemotronNano12Bv2Provider from both megatron.bridge.models.nemotron_vl and megatron.bridge.models.nemotronh.nemotron_h_provider. The second import overwrites the first, so the top-level export at Line 345 can only point at one provider and the other becomes unreachable.

Also applies to: 345-346

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/models/__init__.py` around lines 135 - 156, Two different
modules export NemotronNano12Bv2Provider and the second import overwrites the
first; fix by aliasing one or both imports (e.g., import
NemotronNano12Bv2Provider from megatron.bridge.models.nemotron_vl as
NemotronNano12Bv2ProviderVL and from
megatron.bridge.models.nemotronh.nemotron_h_provider as
NemotronNano12Bv2ProviderH), then update the top-level exports that currently
reference NemotronNano12Bv2Provider to point to the correct alias (choose which
provider should be exported under the public name) and update any internal
references to use the new alias names (check usages of
NemotronNano12Bv2Provider, NemotronVLBridge, and entries in the module's __all__
if present).

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Drop the bridge-local fused-expert merge override.

GLMExpertGateUpProjMapping / GLMExpertDownProjMapping already opt into grouped export in src/megatron/bridge/models/conversion/param_mapping.py, and GLMExpertDownProjMapping already carries its transpose-on-export behavior. Re-merging here reintroduces per-bridge caching/order logic and risks double-stacking or misordering expert tensors on export.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py` around lines 400 -
449, This method implements a bridge-local fused-expert merge that duplicates
grouping/transpose behavior already handled by
GLMExpertGateUpProjMapping/GLMExpertDownProjMapping and must be removed; replace
the body of maybe_modify_converted_hf_weight so it no longer performs any
per-expert caching/stacking or special-case mapping checks (remove all uses of
self.hf_weights_cache and the if isinstance(...
GLMExpertGateUpProjMapping/GLMExpertDownProjMapping) branch and the merging
logic), and instead simply return converted_weights_dict unchanged (or the
incoming dict) so grouped export/transpose behavior is driven solely by the
mapping/param_mapping code (keep the method signature unchanged).

[yaoyu-33]

/ok to test 951dc36

[yaoyu-33]

revert this change, not convince it's required. If it's toy model issue, should fix in toy model

[yaoyu-33]

check other model? this part should use super to do

[yaoyu-33]

just combine param mapping and layerspec mapping

[yaoyu-33]

/ok to test 7d055e4