[HollowMan6]

Pending on:

• Add KDA (Kimi Delta Attention) hybrid layers with GLM-style two-stage gates NVIDIA/Megatron-LM#7054
• Add mHC precision controls, mean output contract, and full activation recompute with mHC NVIDIA/Megatron-LM#7521
• Add KPool DSA indexer, NoPE MLA support, and FP8 precision controls for the indexer and MLA attention NVIDIA/Megatron-LM#7522
• Add DSA indexer weights projection precision controls NVIDIA/Megatron-LM#6842

What does this PR do ?

Add HF ↔ Megatron support for GLM-5.3-Flash (model_type glm5_next, Glm5NextForConditionalGeneration): a unified bridge, provider, and VLM wrapper for its hybrid KDA/DSA + MoE + mHC language model with a replicated HF vision tower, plus the shared conversion, PEFT, and packed-sequence pieces it needs.

Changelog

GLM-5.3-Flash (src/megatron/bridge/models/glm5next/)

• Add GLM53FlashBridge, GLM53FlashModelProvider, and GLM53FlashModel, registered for Glm5NextForConditionalGeneration with model_type="glm5_next" and exported from megatron.bridge.models.
• Build hybrid_layer_pattern from the HF layer_types / mlp_layer_types: each HF decoder layer becomes two physical layers (KDA K or DSA D attention at 2N, dense - or MoE E FFN at 2N+1), with moe_layer_freq derived per physical module and num_hidden_layers recovered from the pattern on export.
• Configure MLA NoPE geometry (qk_pos_emb_head_dim=0, kv_channels=v_head_dim), the kpool DSA indexer (FP8 kpool; index_kpool / index_kpool_always_select_tail read from the HF config), KDA two-stage low-rank gates with gate_lower_bound, mHC multi-stream hyper-connections with fp32 mappings, sigmoid routing with expert bias and an always-on shared expert (moe_shared_expert_gate=False), grouped GEMM with the HybridEP flex dispatcher, and a one-layer DE MTP using the repeated-layer spec.
• Weight mappings: fused KDA in_proj / conv1d built from separate HF Q/K/V tensors with section-wise TP sharding (_ColumnParallelConcatMapping), MLA projections and indexer weights, dense and MoE FFN (grouped-GEMM and sequential expert layouts), mHC hc_{attn,ffn}_{fn,base,scale}, and the MTP eh_proj split into e_proj / h_proj (_EhProjSplitMapping, import only).
• Dequantize blockwise FP8 checkpoint weights on import; keep Megatron init for the MTP inner-layer hyper-connection params, which have no HF source.
• VLM wrapper: replicated HF Glm5NextVisionModel on pre-process ranks, image/video features scattered into the embeddings before THD packing and the sequence-parallel scatter, text-only fallback when no pixel inputs are given, and freeze_{language_model,vision_model,vision_projection} options.

Hybrid providers

• Add HybridMLAModelProvider (HybridModelProvider + MLATransformerConfig). With pipeline_model_parallel_size > 1 and DSA layers in the pattern, it inserts balanced | separators only before full-indexer DSA layers so IndexShare groups never cross a stage, validates user-provided separators, and rejects num_layers_in_first/last_pipeline_stage.
• Add scatter_embedding_sequence_parallel to HybridModelProvider (default True, matching HybridModel) and pass it through; the VLM provider sets it to False so vision features are injected before the SP scatter.

Shared conversion infrastructure

• Add HCAlphaMapping, mapping an HF [3] hc_*_scale tensor to Megatron's alpha_pre / alpha_post / alpha_res (index 0 exports all three), and migrate DeepSeekV4Bridge from its private _HCAlphaMapping / _HCAlphaSecondaryMapping to it with equivalent import/export behavior.
• Add InitOnlyMapping for Megatron params with no HF counterpart (import keeps init, export is a no-op).
• Add MegatronModelBridge._prefix_mapping_registry to re-key a text-model registry under a VLM prefix while keeping nested _tp_mapping names in sync.

PEFT

• Split fused KDA in_proj LoRA linear_out weights by Q/K/V section sizes when kda_two_stage_gates is set, instead of the GQA-style split_qkv_weights path, so KDA LoRA adapters export correctly.

Packed sequences

• Move preprocess_packed_seqs from models/qwen_vl/modelling_qwen3_vl/utils.py to training/utils/packed_seq_utils.py (shared by Qwen3-VL and GLM-5.3-Flash) and add use_fp8_padding, which aligns each sequence to lcm(16, align) and the packed total to align * 128.

Tests

• tests/unit_tests/models/glm5next/test_glm53flash_mapping.py: provider precision contract and multimodal config copy, language_model. prefixing of text mappings, HCAlphaMapping import/export and wildcard resolution, section-wise KDA TP sharding round-trips (TP 1 to 16).
• tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py: packed-sequence alignment and gradient flow with and without FP8 padding.
• tests/unit_tests/models/test_model_bridge_lora.py: KDA Q/K/V LoRA linear_out split.

GitHub Actions CI

See the CI section in 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? (no model page yet; see follow-ups below)
• 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?

Glm5NextForConditionalGeneration is imported behind a try/except that falls back to the class name, Glm5NextModel / Glm5NextVisionModel are imported lazily inside GLM53FlashModel.__init__, and the new unit tests use pytest.importorskip on the GLM-5.3-Flash config module.

Additional Information

• No linked issue.
• Dependency: transformers. The glm5_next model family ships in transformers v5.17.0; pyproject.toml currently pins transformers<=5.15.0, so on the pinned version the bridge only resolves by class name, the VLM wrapper cannot be instantiated, and the new unit tests are skipped. The pin bump should go in a separate build: PR per the contributing guide.
• preprocess_packed_seqs moved without a compatibility re-export from its old Qwen3-VL module.
• Follow-ups not in this PR: docs/models/glm/ page and README / docs model tables, recipes and examples, verification card, and a functional HF ↔ Megatron round-trip test.

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

[kevalmorabia97]

Bug: the MTP final_layernorm is imported from the main norm instead of the MTP layer's shared_head.norm

In glm53flash_bridge.py#L727-L730 the mapping is:

AutoMapping(
    f"{mtp_mg}.final_layernorm.weight",
    "model.language_model.norm.weight",
),

The GLM-5.3-Flash checkpoint has its own model.language_model.layers.45.shared_head.norm.weight, which this mapping never reads:

• On import, the MTP head's norm is silently replaced with the decoder's final norm. That hurts MTP / speculative-decoding acceptance.
• On export, the main norm is written back out, so the checkpoint doesn't round-trip.

The GLM-4.5 bridge maps the same tensor from shared_head.norm (glm45_bridge.py#L295-L297).

Suggested fix:

                     AutoMapping(
                         f"{mtp_mg}.final_layernorm.weight",
-                        "model.language_model.norm.weight",
+                        f"{mtp_hf}.shared_head.norm.weight",
                     ),

How we found it: we round-tripped a tiny GLM-5.3-Flash through this bridge. It had 4 HF layers built from the real config, with random weights and every norm randomized so no two were equal. We imported it, then exported it with ModelOpt's Megatron→HF exporter (NVIDIA/Model-Optimizer#2539).

• With the current mapping: the exported layers.4.shared_head.norm.weight equalled the source model.language_model.norm.weight, not the source shared_head.norm. It was the only one of 197 unquantized tensors that didn't round-trip bit-exactly.
• With the fix above: all 197 round-trip bit-exactly.

This was on the 26.10.rc2 container with Megatron-LM #7054 / #7521 / #7522 and transformers==5.17.0.

[HollowMan6]

Thank you @kevalmorabia97, this is now fixed in e6f7cf2

[kevalmorabia97]

Two training-loop failures with GLM-5.3-Flash (e6f7cf28), found while testing QAD; both reproduce with plain distillation (no quantization)

Setup:

• nvcr.io/nvidian/nemo:26.10.rc2 plus this PR at e6f7cf28, Megatron-LM #7054 / #7521 / #7522, and transformers==5.17.0.
• A tiny random GLM-5.3-Flash built from the real config: 4 HF layers (3 KDA + 1 DSA), hidden 512, 8 routed experts, hc_mult=4, plus the MTP layer.
• One GPU, TP=PP=EP=1, driven by the Megatron-Bridge distillation API (megatron.bridge.training.distill.distill) through ModelOpt's examples/megatron_bridge/distill.py.
• Settings: --use_mock_data --seq_length 64 --mbs 1 --gbs 4 --train_iters 3, with --teacher_hf_path and --student_hf_path both pointing at the tiny model.

Iteration 1 trains fine. Loss is finite and grad norm 1.92. Then:

1. Saving a checkpoint with the distributed optimizer asserts

use_distributed_optimizer=True, ckpt_format="torch_dist", overlap_grad_reduce / overlap_param_gather. Fails at the first save:

megatron/core/optimizer/distrib_optimizer.py:1979, in sharded_param_state_dp_reshardable
    assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), (
AssertionError: (torch.Size([49152]), 0, 79298560)

• The only 49,152-element parameters in the model are the mHC hyper_connection.mapping_proj.weight tensors (24 × 2048), in the decoder layers and in the MTP layer. The optimizer state for one of them is being matched against the whole 79,298,560-element grad-buffer range.
• The provider keeps mHC in FP32, so this may be the mixed-dtype buffer handling around #7521.
• dist_ckpt_optim_fully_reshardable=True fails at the same point: RuntimeError: The size of tensor a (79298560) must match the size of tensor b (49152).

2. With the non-distributed optimizer, iteration 2 asserts in the grad buffer

With use_distributed_optimizer=False (in both the DDP and optimizer configs), the checkpoint saves, but the next step fails:

megatron/core/distributed/param_and_grad_buffer.py:338, in reset   (from DistributedDataParallel.zero_grad_buffer)
    assert len(self.per_param_grad_ready_counts) == len(self.params)

That looks like some parameters never receive a gradient in iteration 1. I haven't isolated which ones. The DSA indexer, whose training goes through its own loss, is a candidate.

Both happen without any quantization. Running the same job with a plain BF16 student, no ModelOpt quantizers involved, gives the identical assertion (torch.Size([49152]), 0, 79298560) at the first save.

Workaround we used: one iteration, --no_async_save, and the non-distributed optimizer. That was enough to validate a train → save → export round-trip, but not a real multi-step training run.

[HollowMan6]

@kevalmorabia97 For 1, you should be able to fix it by patching this dev commit to main NVIDIA/Megatron-LM@ceaa6c5, for 2, that's a good catch and now it should be fixed with NVIDIA/Megatron-LM@bfb1773

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[kevalmorabia97]

LGTM from the Model Optimizer side. Tested 22e09a8b with NVIDIA/Megatron-LM#7521 + NVIDIA/Megatron-LM#7054 in a NeMo 26.10 container on GB300: ModelOpt GPTQ PTQ, distillation with save + resume, and quantized HF export on a 4-layer GLM-5.3-Flash, plus full-size GLM-5.3-Flash QAD on 32 nodes (EP32 x CP4, packed SFT) with distributed-optimizer save + resume, matching the loss and grad norms of our earlier full PTQ + QAD + vLLM eval study on older heads of this stack.

GLM-5.3-Flash support in Model Optimizer (NVIDIA/Model-Optimizer#2539) is waiting on this PR, so we'd appreciate a review and merge once the Megatron-LM PRs land. Thanks @HollowMan6!

[kevalmorabia97]

@cuichenx can we bump Megatron-LM submodule to include NVIDIA/Megatron-LM#7054 to unblock this PR?

[cuichenx]

@cuichenx can we bump Megatron-LM submodule to include NVIDIA/Megatron-LM#7054 to unblock this PR?

yep, monitoring #6717

[HollowMan6]

Rebased! Would you mind triggering the CI? Thanks!

cc: @kevalmorabia97 @cuichenx

[kevalmorabia97]

Thanks for the rebase! Two things before CI:

1. The branch is based on 5bc5076, which predates #6717. The PR's 3rdparty/Megatron-LM is still a296523 (without NVIDIA/Megatron-LM#7054), so CI on the head commit won't have KDA. Could you rebase onto current main? It's conflict-free.
2. tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py:49 still imports preprocess_packed_seqs from modelling_qwen3_vl.utils, which this PR removes. The module then fails to import, including the new test_packed_sequence_alignment_and_gradients. Importing it from megatron.bridge.training.utils.packed_seq_utils fixes it. With that change, the glm5next, qwen3_vl, LoRA, deepseek, hybrid and packed-seq unit tests pass on main + this PR with Megatron-LM 5de526f2.

Code-wise the commits are unchanged from my last review; happy to approve once these are in.

[HollowMan6]

Thank you @kevalmorabia97! Sorry I forgot to update my local main branch, now it's rebased and the test case is fixed.

[kevalmorabia97]

Thanks! Both are in: the branch is on main with Megatron-LM 5de526f2, and the 10 earlier commits are unchanged. On e342d93, the glm5next, qwen3_vl, LoRA bridge, deepseek, hybrid, packed-seq and peft utils unit tests pass (618 passed). The only failures are Transformer Engine kernel errors from running on Ada GPUs locally. LGTM.

[kevalmorabia97]

/ok to test e342d93

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.