[OMNIML-6104] Enable DeepSeek-V4-Flash IQ2_XXS PTQ and mergeable LoRA QAD

CHANGELOG.rst

@@ -18,9 +18,11 @@ Changelog
 
 *Megatron Framework (M-LM / M-Bridge)*
 
+- Add DeepSeek-V4-Flash routed-expert IQ2_XXS GPTQ and mergeable LoRA QAD through the standard Megatron-Bridge entry points, with packed IQ2 and folded BF16 HF export. The initial export path supports main-only models with TP1/PP1 and expert parallelism.
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
 - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
+- Add GLM-5.3-Flash (``glm5_next``) support to ``examples/megatron_bridge`` PTQ and unified HF export with the ``models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast`` recipe, matching ``nvidia/GLM-5.3-Flash-NVFP4``, and a GGML IQ2_S recipe for its routed experts (``models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq``). Requires a Megatron-Bridge and Megatron-Core with GLM-5.3-Flash support.
 
 *Misc*
 

examples/megatron_bridge/distill.py

@@ -29,17 +29,16 @@
 from export_distilled_megatron_to_hf import export_llm_to_hf, save_vlm_to_hf
 from megatron.bridge import AutoBridge
 from megatron.bridge.models.distillation_provider import convert_to_distillation_provider
-from megatron.bridge.recipes.utils.optimizer_utils import (
-    distributed_fused_adam_with_cosine_annealing,
-)
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
     FinetuningDatasetConfig,
     GPTDatasetConfig,
     LoggerConfig,
     MockGPTDatasetConfig,
+    OptimizerConfig,
     RNGConfig,
+    SchedulerConfig,
     TokenizerConfig,
     TrainingConfig,
     ValidationConfig,
@@ -280,6 +279,23 @@ def get_args():
         action="store_true",
         help="Keep the student's MTP head at its input weights: skip its loss and freeze it.",
     )
+    parser.add_argument(
+        "--deterministic_mode",
+        action="store_true",
+        help="Use Bridge/Core deterministic execution for exact training and resume checks.",
+    )
+    parser.add_argument(
+        "--mhc_fused_backend",
+        choices=["auto", "native", "triton", "cutile"],
+        default=None,
+        help="Select the existing Core mHC backend; use triton for matched exact-resume controls.",
+    )
+    parser.add_argument(
+        "--moe_router_bias_update_rate",
+        type=float,
+        default=None,
+        help="Override router expert-bias updates; set 0 to preserve the frozen router during adapter QAD.",
+    )
     parser.add_argument(
         "--moe_aux_loss_coeff",
         type=float,
@@ -497,6 +513,11 @@ def _tokenizer_prepends_bos(args) -> bool:
 
 
 def main(args: argparse.Namespace, owns_the_run: bool = True):
+    if args.deterministic_mode:
+        # Set library requirements before any checkpoint import can run a CUDA kernel.
+        from megatron.training.determinism import apply_determinism_env
+
+        apply_determinism_env(os.environ)
     student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
         args.student_megatron_path
     )
@@ -533,11 +554,20 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
+        if args.mhc_fused_backend is not None:
+            provider.use_fused_mhc = True
+            provider.mhc_fused_backend = args.mhc_fused_backend
+        if args.deterministic_mode:
+            provider.deterministic_mode = True
+            provider.cross_entropy_loss_fusion = False
+            provider.moe_router_aux_loss_fusion = False
         if args.dsa_kernel_backend is not None:
             provider.dsa_kernel_backend = args.dsa_kernel_backend
         set_moe_expert_layout(provider, moe_grouped_gemm)
         if args.moe_aux_loss_coeff is not None:
             provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
+        if args.moe_router_bias_update_rate is not None:
+            provider.moe_router_bias_update_rate = args.moe_router_bias_update_rate
         if per_token_loss:
             provider.calculate_per_token_loss = True
         if args.recompute_granularity is not None:
@@ -631,11 +661,31 @@ def _restore_student_hook(model_chunks):
         distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True)
 
     # Build optimizer and scheduler
-    optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing(
-        lr_warmup_iters=args.lr_warmup_iters,
-        max_lr=args.lr,
-        min_lr=args.min_lr,
+    # Importing the recipe helper eagerly loads unrelated optional model dependencies.
+    optimizer_config = OptimizerConfig(
+        optimizer="adam",
+        lr=args.lr,
+        min_lr=args.min_lr if args.min_lr is not None else 0.1 * args.lr,
+        weight_decay=0.1,
+        bf16=True,
+        fp16=False,
+        adam_beta1=0.9,
         adam_beta2=0.95,
+        adam_eps=1e-8,
+        use_distributed_optimizer=True,
+        clip_grad=1.0,
+    )
+    scheduler_config = SchedulerConfig(
+        start_weight_decay=0.033,
+        end_weight_decay=0.033,
+        weight_decay_incr_style="constant",
+        lr_decay_style="cosine",
+        lr_wsd_decay_style="minus_sqrt",
+        lr_wsd_decay_iters=None,
+        lr_warmup_iters=args.lr_warmup_iters,
+        lr_warmup_init=0.0,
+        lr_decay_iters=None,
+        override_opt_param_scheduler=True,
     )
 
     # Build dataset config

modelopt/torch/export/plugins/mcore_common.py

@@ -18,7 +18,12 @@
 from typing import Any
 
 from .mcore_deepseek import deepseek_causal_lm_export, deepseek_causal_lm_import
-from .mcore_glm import glm_moe_dsa_causal_lm_export
+from .mcore_deepseek_v4 import deepseek_v4_causal_lm_export
+from .mcore_glm import (
+    GLM5_NEXT_VISION_PREFIXES,
+    glm5_next_causal_lm_export,
+    glm_moe_dsa_causal_lm_export,
+)
 from .mcore_gptoss import gptoss_causal_lm_export, gptoss_causal_lm_import
 from .mcore_llama import (
     eagle3_deep_llama_causal_lm_export,
@@ -51,6 +56,7 @@
 all_mcore_hf_export_mapping: dict[str, Any] = {
     "DeepseekV2ForCausalLM": deepseek_causal_lm_export,
     "DeepseekV3ForCausalLM": deepseek_causal_lm_export,
+    "DeepseekV4ForCausalLM": deepseek_v4_causal_lm_export,
     "LlamaForCausalLM": llama_causal_lm_export,
     "Llama4ForConditionalGeneration": llama4_causal_lm_export,
     "NemotronForCausalLM": nemotron_causal_lm_export,
@@ -66,6 +72,7 @@
     "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export,
     "Qwen3_5ForConditionalGeneration": qwen3_5_vl_causal_lm_export,
     "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export,
+    "Glm5NextForConditionalGeneration": glm5_next_causal_lm_export,
     "GlmMoeDsaForCausalLM": glm_moe_dsa_causal_lm_export,
 }
 
@@ -77,6 +84,7 @@
     "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "NemotronH_Omni_Reasoning_V3": NEMOTRON_H_OMNI_VISION_PREFIXES,
+    "Glm5NextForConditionalGeneration": GLM5_NEXT_VISION_PREFIXES,
 }
 
 all_mcore_hf_import_mapping: dict[str, Any] = {

modelopt/torch/export/plugins/mcore_custom.py

@@ -200,6 +200,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
         )
 
 
+class KimiDeltaAttentionSlicing(CustomModuleMapping):
+    """A custom module mapping that splits KDA's fused q|k|v ``in_proj`` and ``conv1d``."""
+
+    def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
+        """Create a custom module mapping that splits the fused KDA projections."""
+        super().__init__(
+            func_name="kda_slicing",
+            target_name_or_prefix=target_name_or_prefix,
+            func_kwargs=func_kwargs,
+        )
+
+
 class PackNameRemapping(CustomModuleMapping):
     """A custom module mapping that packs module after name remapping."""
 

modelopt/torch/export/plugins/mcore_deepseek_v4.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Native DeepSeek V4 HF names for the paired attention/MoE HybridModel.
+
+Bridge supplies import and MXFP4/FP8 decoding. These rules export its floating-point
+training weights through the existing packed and folded ModelOpt exporters.
+"""
+
+from .mcore_custom import GatedMLPSlicing, NameRemapping
+
+__all__ = ["deepseek_v4_causal_lm_export"]
+
+deepseek_v4_causal_lm_export = {
+    "deepseek_v4": True,
+    "fold_attn_mlp_layer_pairs": True,
+    "word_embeddings": NameRemapping("embed."),
+    "final_norm": NameRemapping("norm."),
+    "output_layer": NameRemapping("head."),
+    "hc_head_fn": NameRemapping("hc_head_fn"),
+    "hc_head_base": NameRemapping("hc_head_base"),
+    "hc_head_scale": NameRemapping("hc_head_scale"),
+    "hc_fn": NameRemapping("layers.{}.hc_{}_fn"),
+    "hc_base": NameRemapping("layers.{}.hc_{}_base"),
+    "hc_scale": NameRemapping("layers.{}.hc_{}_scale"),
+    "input_layernorm": NameRemapping("layers.{}.attn_norm."),
+    "pre_mlp_layernorm": NameRemapping("layers.{}.ffn_norm."),
+    "linear_q_down_proj": NameRemapping("layers.{}.attn.wq_a."),
+    "linear_q_layernorm": NameRemapping("layers.{}.attn.q_norm."),
+    "linear_q_up_proj": NameRemapping("layers.{}.attn.wq_b."),
+    "linear_kv_proj": NameRemapping("layers.{}.attn.wkv."),
+    "linear_kv_layernorm": NameRemapping("layers.{}.attn.kv_norm."),
+    "linear_o_group_proj": NameRemapping("layers.{}.attn.wo_a.weight"),
+    "linear_proj": NameRemapping("layers.{}.attn.wo_b."),
+    "attn_sink": NameRemapping("layers.{}.attn.attn_sink"),
+    "compressor.linear_wkv": NameRemapping("layers.{}.attn.compressor.wkv."),
+    "compressor.linear_wgate": NameRemapping("layers.{}.attn.compressor.wgate."),
+    "compressor.ape": NameRemapping("layers.{}.attn.compressor.ape"),
+    "compressor.norm": NameRemapping("layers.{}.attn.compressor.norm."),
+    "indexer.linear_wq_b": NameRemapping("layers.{}.attn.indexer.wq_b."),
+    "indexer.linear_weights_proj": NameRemapping("layers.{}.attn.indexer.scorer.weights_proj."),
+    "indexer.compressor.linear_wkv": NameRemapping("layers.{}.attn.indexer.compressor.wkv."),
+    "indexer.compressor.linear_wgate": NameRemapping("layers.{}.attn.indexer.compressor.wgate."),
+    "indexer.compressor.ape": NameRemapping("layers.{}.attn.indexer.compressor.ape"),
+    "indexer.compressor.norm": NameRemapping("layers.{}.attn.indexer.compressor.norm."),
+    "router": NameRemapping("layers.{}.ffn.gate.", {"mapping": {"expert_bias": "bias"}}),
+    "shared_experts.linear_fc1": GatedMLPSlicing(
+        "layers.{}.ffn.shared_experts.", {"gate_proj_name": "w1", "up_proj_name": "w3"}
+    ),
+    "shared_experts.linear_fc2": NameRemapping("layers.{}.ffn.shared_experts.w2."),
+    "local_experts.linear_fc1": GatedMLPSlicing(
+        "layers.{}.ffn.experts.{}.", {"gate_proj_name": "w1", "up_proj_name": "w3"}
+    ),
+    "local_experts.linear_fc2": NameRemapping("layers.{}.ffn.experts.{}.w2."),
+}

modelopt/torch/export/plugins/mcore_glm.py

@@ -15,19 +15,95 @@
 
 """Custom mappings from Megatron Core models to GLM-5.x Hugging Face models.
 
-GLM-5 / GLM-5.2 (``glm_moe_dsa``) is DeepSeek-V3-style MLA with a DSA indexer. It keeps the MTP
-layer at HF index ``num_hidden_layers`` under the decoder's names; Megatron-Bridge does not build
-it, so the exporter copies it from the source checkpoint.
+GLM-5 / GLM-5.2 (``glm_moe_dsa``) is DeepSeek-V3-style MLA with a DSA indexer.
+
+For GLM-5.3-Flash (``glm5_next``), Megatron-Bridge builds each HF decoder layer as two physical
+layers wrapped in mHC hyper-connections: attention (KDA or NoPE-MLA DSA) at ``2N`` and the dense MLP /
+MoE at ``2N+1``. The vision tower is copied from HF.
+
+Both keep the MTP layer at HF index ``num_hidden_layers`` under the decoder's names.
 """
 
 from .mcore_custom import (
+    GatedMLPSlicing,
     GroupedGatedMLPSlicing,
     GroupedMLPSlicing,
+    KimiDeltaAttentionSlicing,
     NameRemapping,
     SelfAttentionScaling,
+    with_language_model_prefix,
 )
 from .mcore_deepseek import deepseek_causal_lm_export
 
+# Vision-tower weights copied straight from the HF checkpoint (never quantized).
+GLM5_NEXT_VISION_PREFIXES = ("model.visual.",)
+
+_glm5_next_causal_lm_export: dict = {
+    # Layer-level flags read by the exporter (see the module docstring).
+    "fold_attn_mlp_layer_pairs": True,
+    "mtp_in_decoder_layers": True,
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "final_norm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+    # mHC hyper-connections: formatted with (hf_layer_id, "attn" | "ffn").
+    "hc_fn": NameRemapping("model.layers.{}.hc_{}_fn"),
+    "hc_base": NameRemapping("model.layers.{}.hc_{}_base"),
+    "hc_scale": NameRemapping("model.layers.{}.hc_{}_scale"),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    # KDA linear attention (fused q|k|v in_proj and conv1d are split by ``kda``).
+    "kda": KimiDeltaAttentionSlicing("model.layers.{}.self_attn."),
+    "kda.beta_proj": NameRemapping("model.layers.{}.self_attn.b_proj."),
+    "kda.f_a_proj": NameRemapping("model.layers.{}.self_attn.f_a_proj."),
+    "kda.f_b_proj": NameRemapping("model.layers.{}.self_attn.f_b_proj."),
+    "kda.g_a_proj": NameRemapping("model.layers.{}.self_attn.g_a_proj."),
+    "kda.g_b_proj": NameRemapping("model.layers.{}.self_attn.g_b_proj."),
+    "kda.A_log": NameRemapping("model.layers.{}.self_attn.A_log"),
+    "kda.dt_bias": NameRemapping("model.layers.{}.self_attn.dt_bias"),
+    "kda.out_norm": NameRemapping("model.layers.{}.self_attn.o_norm."),
+    "kda.out_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    # NoPE MLA with the DSA kpool indexer
+    "linear_q_down_proj": NameRemapping("model.layers.{}.self_attn.q_a_proj."),
+    "linear_q_layernorm": NameRemapping("model.layers.{}.self_attn.q_a_layernorm."),
+    "linear_q_up_proj": NameRemapping("model.layers.{}.self_attn.q_b_proj."),
+    "linear_kv_down_proj": NameRemapping("model.layers.{}.self_attn.kv_a_proj_with_mqa."),
+    "linear_kv_layernorm": NameRemapping("model.layers.{}.self_attn.kv_a_layernorm."),
+    "linear_kv_up_proj": NameRemapping("model.layers.{}.self_attn.kv_b_proj."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "core_attention": SelfAttentionScaling("model.layers.{}.self_attn."),
+    "indexer.linear_wq_b": NameRemapping("model.layers.{}.self_attn.indexer.wq_b."),
+    "indexer.linear_wk": NameRemapping("model.layers.{}.self_attn.indexer.wk."),
+    "indexer.k_norm": NameRemapping("model.layers.{}.self_attn.indexer.k_norm."),
+    "indexer.linear_weights_proj": NameRemapping("model.layers.{}.self_attn.indexer.weights_proj."),
+    "indexer.index_kpool_compress_ape": NameRemapping(
+        "model.layers.{}.self_attn.indexer.index_kpool_compress_ape"
+    ),
+    "indexer.index_kpool_compress_gate": NameRemapping(
+        "model.layers.{}.self_attn.indexer.index_kpool_compress_gate"
+    ),
+    # Dense MLP (the pre-MLP norm is fused into linear_fc1)
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "fused_pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.weight"),
+    "linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+    # MoE
+    "router": NameRemapping(
+        "model.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
+    ),
+    "shared_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.shared_experts."),
+    "shared_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.shared_experts.down_proj."),
+    "local_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.experts.{}."),
+    "local_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.experts.{}.down_proj."),
+    "experts.linear_fc1": GroupedGatedMLPSlicing("model.layers.{}.mlp.experts.{{}}"),
+    "experts.linear_fc2": GroupedMLPSlicing("model.layers.{}.mlp.experts.{{}}.down_proj"),
+    # MTP (split e_proj / h_proj are concatenated back into eh_proj)
+    "mtp.enorm": NameRemapping("model.layers.{}.enorm."),
+    "mtp.hnorm": NameRemapping("model.layers.{}.hnorm."),
+    "mtp.eh_proj": NameRemapping("model.layers.{}.eh_proj.weight"),
+    "mtp.final_layernorm": NameRemapping("model.layers.{}.shared_head.norm."),
+}
+
+glm5_next_causal_lm_export = with_language_model_prefix(_glm5_next_causal_lm_export)
+
 glm_moe_dsa_causal_lm_export: dict = {
     **deepseek_causal_lm_export,
     "mtp_in_decoder_layers": True,
@@ -38,4 +114,8 @@
     "indexer.linear_weights_proj": NameRemapping("model.layers.{}.self_attn.indexer.weights_proj."),
     "experts.linear_fc1": GroupedGatedMLPSlicing("model.layers.{}.mlp.experts.{{}}"),
     "experts.linear_fc2": GroupedMLPSlicing("model.layers.{}.mlp.experts.{{}}.down_proj"),
+    "mtp.enorm": NameRemapping("model.layers.{}.enorm."),
+    "mtp.hnorm": NameRemapping("model.layers.{}.hnorm."),
+    "mtp.eh_proj": NameRemapping("model.layers.{}.eh_proj."),
+    "mtp.final_layernorm": NameRemapping("model.layers.{}.shared_head.norm."),
 }