diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f20a33f1c17..7fd101e8a56 100755 --- a/CHANGELOG.rst +++ b/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* diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 333ccbbcb3d..ae785dd6076 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -29,9 +29,6 @@ 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, @@ -39,7 +36,9 @@ 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 diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py index fbfbea57133..1b19cee5685 100644 --- a/modelopt/torch/export/plugins/mcore_common.py +++ b/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] = { diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py index 77c1ad75e1d..01156e3fbf6 100644 --- a/modelopt/torch/export/plugins/mcore_custom.py +++ b/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.""" diff --git a/modelopt/torch/export/plugins/mcore_deepseek_v4.py b/modelopt/torch/export/plugins/mcore_deepseek_v4.py new file mode 100644 index 00000000000..7cc5d2c4a93 --- /dev/null +++ b/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."), +} diff --git a/modelopt/torch/export/plugins/mcore_glm.py b/modelopt/torch/export/plugins/mcore_glm.py index 25a967860ac..11ee40d63ed 100644 --- a/modelopt/torch/export/plugins/mcore_glm.py +++ b/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."), } diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 6f3f128d1fc..8a705676e75 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -108,6 +108,13 @@ from megatron.core.transformer.torch_norm import L2Norm from megatron.core.transformer.transformer_layer import TransformerLayer + try: + from megatron.core.models.hybrid.layers.hybrid_hyper_connection import ( + HyperConnectionHybridLayer, + ) + except ImportError: # older Megatron-Core without mHC + HyperConnectionHybridLayer = None + has_mcore = True __all__ = [ @@ -183,13 +190,17 @@ def __init__( self._src_num_hidden_layers = self._hf_text_config.num_hidden_layers self._hf_text_config.num_hidden_layers = language_model.config.num_layers self._hf_text_config.hidden_size = language_model.config.hidden_size - # MLA's kv_channels is the V head dim, not HF's head_dim (e.g. glm5_next derives it from RoPE). - if not getattr(language_model.config, "multi_latent_attention", False): + is_deepseek_v4 = self._hf_config.architectures[0] == "DeepseekV4ForCausalLM" + # MLA/V4's kv_channels is not necessarily HF's native head dimension. + if not is_deepseek_v4 and not getattr( + language_model.config, "multi_latent_attention", False + ): self._hf_text_config.head_dim = language_model.config.kv_channels self._hf_text_config.num_attention_heads = language_model.config.num_attention_heads - self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups + if not is_deepseek_v4: + self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups self.is_multimodal = isinstance(model, LLaVAModel) - if not self.is_multimodal: + if not self.is_multimodal and not is_deepseek_v4: self._hf_text_config.intermediate_size = language_model.config.ffn_hidden_size self._hf_quant_config: dict = {} self._hf_extra_config = None @@ -204,17 +215,17 @@ def __init__( self.arch, LLAVA_VISION_PREFIXES if self.is_multimodal else None ) # TODO: May modify this later according to what quantization exported ckpt is, currently only support BF16. - if self.arch == "GptOssForCausalLM": + if self.arch in ("GptOssForCausalLM", "DeepseekV4ForCausalLM"): if hasattr(self._hf_config, "quantization_config"): del self._hf_config.quantization_config self.all_rules = self._populate_rule_book() self.rules = self.all_rules[self.arch] self.all_mcore_mappings = all_mcore_hf_export_mapping[self.arch] - if self.rules.get("mtp_in_decoder_layers", False) and language_model.config.mtp_num_layers: - # Only the copy from the source checkpoint is supported (the bridge builds no MTP). - raise NotImplementedError( - f"Exporting a Megatron-built MTP for {self.arch} is not supported yet." - ) + if self.rules.get("deepseek_v4", False): + self._init_deepseek_v4_export() + if self.rules.get("fold_attn_mlp_layer_pairs", False): + # Each HF decoder layer is an attention and an MLP physical layer in Megatron. + self._hf_text_config.num_hidden_layers = language_model.config.num_layers // 2 # The vision tower is copied through unquantized, so deployments must not treat it as such. self.exclude_modules = [ prefix.removesuffix(".") + "*" for prefix in self.vision_passthrough_prefixes or () @@ -583,6 +594,32 @@ def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) - if not source: warn_rank_0(f"Export self-check skipped: no tensor index found in {source_dir}.") return + if self.rules.get("deepseek_v4", False): + # Raw V4 parameters have no trailing '.weight'; ancestor module checks below + # cannot detect a missing mHC parameter, compressor APE or hash-routing table. + raw_parameters = { + key + for key in source + if not key.startswith("mtp.") + and ( + key.startswith("hc_head_") + or key.rsplit(".", 1)[-1] + in ( + "hc_attn_fn", + "hc_attn_base", + "hc_attn_scale", + "hc_ffn_fn", + "hc_ffn_base", + "hc_ffn_scale", + "attn_sink", + "ape", + "tid2eid", + ) + ) + } + missing_raw = raw_parameters - exported + if missing_raw: + raise RuntimeError(f"Export dropped V4 parameter(s): {sorted(missing_raw)[:8]}") # Narrow on purpose: compare module prefixes, not tensor names, since a quantized source # carries extras with no export counterpart, and only inside decoder layers, whose naming @@ -604,7 +641,7 @@ def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) - exported_modules.add(prefix) missing = set() for key in source: - layer = re.search(r"\.layers\.(\d+)\.", key) + layer = re.search(r"(?:^|\.)layers\.(\d+)\.", key) if layer is None: continue # see the note above: decoder layers only if int(layer.group(1)) >= hf_depth: @@ -659,12 +696,9 @@ def _get_state_dict(self): # Decoder layers for layer in model.decoder.layers: layer_id = layer.layer_number - 1 - if isinstance(layer, MambaLayer): - self._get_mamba_layer_state_dict(layer, layer_id) - elif isinstance(layer, TransformerLayer): - self._get_transformer_layer_state_dict(layer, layer_id) - else: - raise ValueError("Only TransformerLayer or MambaLayer are supported.") + if self.rules.get("fold_attn_mlp_layer_pairs", False): + layer_id //= 2 + self._get_decoder_layer_state_dict(layer, layer_id) self._layer_state_dicts[layer.layer_number] = self._state_dict if layer.layer_number != self.model.config.num_layers: @@ -677,6 +711,12 @@ def _get_state_dict(self): if hasattr(model.decoder, "final_norm") and model.decoder.final_norm: self.rules["final_norm"](model.decoder.final_norm) + if self.rules.get("deepseek_v4", False) and model.post_process: + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + value = getattr(model.decoder, name).detach() + # V4 stores every mHC control, including the head projection, in FP32. + self.rules[name](value.to(device="cpu", dtype=torch.float32)) + # Output layer if hasattr(model, "output_layer") and not model.share_embeddings_and_output_weights: self.rules["output_layer"](model.output_layer) @@ -698,6 +738,34 @@ def _get_fused_norm_weight(self, module, primary_key: str = "fused_norm"): return None, None return fused_key, weight + def _get_decoder_layer_state_dict( + self, layer, layer_id, is_mtp=False, export_hyper_connection=True + ): + """Export one decoder layer, unwrapping an mHC ``HyperConnectionHybridLayer`` first.""" + if HyperConnectionHybridLayer is not None and isinstance(layer, HyperConnectionHybridLayer): + if export_hyper_connection: + self._get_hyper_connection_state_dict(layer, layer_id) + layer = layer.inner_layer + if isinstance(layer, MambaLayer): + self._get_mamba_layer_state_dict(layer, layer_id, is_mtp=is_mtp) + elif isinstance(layer, TransformerLayer): + self._get_transformer_layer_state_dict(layer, layer_id, is_mtp=is_mtp) + else: + raise ValueError("Only TransformerLayer or MambaLayer are supported.") + + def _get_hyper_connection_state_dict(self, layer, layer_id): + """Export an mHC wrapper as HF's ``hc_{attn,ffn}_{fn,base,scale}`` of its HF layer.""" + attention = getattr(layer.inner_layer, "self_attention", None) + kind = "ffn" if attention is None or isinstance(attention, IdentityOp) else "attn" + hc = layer.hyper_connection + # Preserve V4's FP32 routing projections; other mHC exports retain their dtype policy. + dtype = torch.float32 if self.rules.get("deepseek_v4", False) else self.dtype + self.rules["hc_fn"](hc.mapping_proj.weight.detach().to(dtype), layer_id, kind) + # The HF checkpoint keeps the mHC bias and alpha scales in FP32. + self.rules["hc_base"](hc.bias.detach().float(), layer_id, kind) + alphas = torch.cat([hc.alpha_pre, hc.alpha_post, hc.alpha_res]).detach().float() + self.rules["hc_scale"](alphas, layer_id, kind) + def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.input_layernorm, IdentityOp): self.rules["input_layernorm"](layer.input_layernorm, layer_id, is_mtp=is_mtp) @@ -714,7 +782,9 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): self.rules[fused_key](norm_weight, layer_id, is_mtp=is_mtp) if not isinstance(layer.self_attention, IdentityOp): - if "MLASelfAttention" in str(type(layer.self_attention)): + if self.rules.get("deepseek_v4", False): + self._get_deepseek_v4_attention_state_dict(layer.self_attention, layer_id) + elif "MLASelfAttention" in str(type(layer.self_attention)): if hasattr(layer.self_attention, "linear_q_proj"): self.rules["linear_q_proj"]( layer.self_attention.linear_q_proj, layer_id, is_mtp=is_mtp @@ -746,6 +816,8 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): indexer = getattr(core_attention, "indexer", None) if indexer is not None: self._get_dsa_indexer_state_dict(indexer, layer_id, is_mtp) + elif "kda" in self.rules and hasattr(layer.self_attention, "in_proj"): + self._get_kda_state_dict(layer, layer_id, is_mtp=is_mtp) elif "linear_attn" in self.rules and hasattr(layer.self_attention, "in_proj"): # GatedDeltaNet (Qwen3.5 linear attention): no q/k layernorm, no core_attention. self._get_gated_delta_net_state_dict(layer, layer_id, is_mtp=is_mtp) @@ -785,9 +857,12 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.mlp, IdentityOp): if "MoE" in str(type(layer.mlp)): - self.rules["router"]( - layer.mlp.router, layer_id, dtype=self.moe_router_dtype, is_mtp=is_mtp - ) + if self.rules.get("deepseek_v4", False): + self._get_deepseek_v4_router_state_dict(layer.mlp.router, layer_id) + else: + self.rules["router"]( + layer.mlp.router, layer_id, dtype=self.moe_router_dtype, is_mtp=is_mtp + ) if hasattr(layer.mlp, "fc1_latent_proj") and layer.mlp.fc1_latent_proj is not None: self.rules["fc1_latent_proj"]( layer.mlp.fc1_latent_proj, layer_id, is_mtp=is_mtp @@ -855,6 +930,80 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id, is_mtp=is_mtp) self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id, is_mtp=is_mtp) + def _init_deepseek_v4_export(self): + """Validate the main-only TP1/PP1 V4 path and read its native optional tensor names.""" + if self.model.config.num_layers % 2: + raise ValueError("DeepSeek V4 export requires paired attention/MoE physical layers.") + if getattr(self.model.config, "mtp_num_layers", None) or getattr( + self._hf_text_config, "num_nextn_predict_layers", 0 + ): + raise NotImplementedError( + "DeepSeek V4 export currently requires a main-only HF view (MTP=0)." + ) + if torch.distributed.is_initialized() and ( + get_tensor_model_parallel_world_size() != 1 + or get_pipeline_model_parallel_world_size() != 1 + ): + raise NotImplementedError( + "DeepSeek V4 export currently requires TP1 / PP1; EP is supported." + ) + source = str(self._hf_pretrained_model_name) + self._deepseek_v4_source_keys = ( + _read_checkpoint_keys(source) + if os.path.isdir(source) + else set(get_safetensors_metadata(source).weight_map) + ) + if not self._deepseek_v4_source_keys: + raise ValueError("DeepSeek V4 export requires the source checkpoint tensor index.") + + def _get_deepseek_v4_attention_state_dict(self, attention, layer_id): + """Export V4's shared-KV attention, optional compressor and CSA indexer.""" + for rule, name in ( + ("linear_q_down_proj", "linear_q_down_proj"), + ("linear_q_layernorm", "q_layernorm"), + ("linear_q_up_proj", "linear_q_up_proj"), + ("linear_kv_proj", "linear_kv_proj"), + ("linear_kv_layernorm", "kv_layernorm"), + ("linear_proj", "linear_proj"), + ): + self.rules[rule](getattr(attention, name), layer_id) + self.rules["linear_o_group_proj"]( + attention.linear_o_group_proj.detach().to(device="cpu", dtype=self.dtype), layer_id + ) + core = attention.core_attention + self.rules["attn_sink"](core.attn_sink.detach().float().cpu(), layer_id) + compressor = getattr(core, "compressor", None) + indexer = getattr(core, "indexer", None) + for prefix, module in ( + ("compressor", compressor), + ("indexer.compressor", getattr(indexer, "compressor", None)), + ): + if module is not None: + for name in ("linear_wkv", "linear_wgate", "norm"): + self.rules[f"{prefix}.{name}"](getattr(module, name), layer_id) + self.rules[f"{prefix}.ape"](module.ape.detach().float().cpu(), layer_id) + if indexer is not None: + self.rules["indexer.linear_wq_b"](indexer.linear_wq_b, layer_id) + self.rules["indexer.linear_weights_proj"](indexer.linear_weights_proj, layer_id) + # Bridge also accepts the released checkpoint's legacy flat scorer name. + native = f"layers.{layer_id}.attn.indexer.scorer.weights_proj.weight" + legacy = native.replace(".scorer.", ".") + if ( + legacy in self._deepseek_v4_source_keys + and native not in self._deepseek_v4_source_keys + ): + self._state_dict[legacy] = self._state_dict.pop(native) + + def _get_deepseek_v4_router_state_dict(self, router, layer_id): + """Keep hash-routing tables and omit bias buffers synthesized only during import.""" + self.rules["router"](router, layer_id, dtype=self.moe_router_dtype) + prefix = f"layers.{layer_id}.ffn.gate." + if prefix + "bias" not in self._deepseek_v4_source_keys: + self._state_dict.pop(prefix + "bias", None) + tid2eid = getattr(router, "tid2eid", None) + if tid2eid is not None: + self._state_dict[prefix + "tid2eid"] = tid2eid.detach().cpu() + def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, torch.Tensor]: """Export the live MTP module, or copy it from the pretrained model if absent (and allowed).""" model = getattr(self, "model", None) @@ -866,12 +1015,20 @@ def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, to saved_state_dict = self._state_dict self._state_dict = OrderedDict() try: - for mtp_layer in mtp.layers: + for mtp_idx, mtp_layer in enumerate(mtp.layers): # Some architectures (Qwen3.5) put a single TransformerLayer here, not a container. inner = mtp_layer.mtp_model_layer inner_layers = getattr(inner, "layers", None) or [inner] - first_id = inner_layers[0].layer_number - 1 - last_id = inner_layers[-1].layer_number - 1 + if self.rules.get("mtp_in_decoder_layers", False): + # HF stores the MTP layer after the decoder layers, under the decoder's names. + first_id = last_id = self._hf_text_config.num_hidden_layers + mtp_idx + inner_ids = [first_id] * len(inner_layers) + inner_is_mtp = False + else: + first_id = inner_layers[0].layer_number - 1 + last_id = inner_layers[-1].layer_number - 1 + inner_ids = [layer.layer_number - 1 for layer in inner_layers] + inner_is_mtp = True # Outer predictor projections attach to the first inner HF index. if "mtp.enorm" in self.rules: @@ -879,19 +1036,18 @@ def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, to if "mtp.hnorm" in self.rules: self.rules["mtp.hnorm"](mtp_layer.hnorm, first_id) if "mtp.eh_proj" in self.rules: - self.rules["mtp.eh_proj"](mtp_layer.eh_proj, first_id) - - # Inner layers reuse the base decoder walker (is_mtp=True). - for inner in inner_layers: - hf_layer_id = inner.layer_number - 1 - if isinstance(inner, MambaLayer): - self._get_mamba_layer_state_dict(inner, hf_layer_id, is_mtp=True) - elif isinstance(inner, TransformerLayer): - self._get_transformer_layer_state_dict(inner, hf_layer_id, is_mtp=True) - else: - raise ValueError( - "Only TransformerLayer or MambaLayer are supported in the MTP block." - ) + if getattr(mtp_layer, "eh_proj", None) is not None: + self.rules["mtp.eh_proj"](mtp_layer.eh_proj, first_id) + else: # split projections (GLM-5.3-Flash): HF eh_proj is [e_proj | h_proj] + eh_proj = torch.cat([mtp_layer.e_proj.weight, mtp_layer.h_proj.weight], 1) + self.rules["mtp.eh_proj"](eh_proj.detach().to(self.dtype), first_id) + + # Inner layers reuse the base decoder walker. The MTP mHC weights have no HF + # counterpart (Megatron-Bridge initializes them on import), so they are dropped. + for inner, hf_layer_id in zip(inner_layers, inner_ids): + self._get_decoder_layer_state_dict( + inner, hf_layer_id, is_mtp=inner_is_mtp, export_hyper_connection=False + ) # The MTP block's own final layernorm attaches to the last inner HF index. final_layernorm = getattr(mtp_layer, "final_layernorm", None) @@ -1033,12 +1189,31 @@ def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False): self.rules["linear_attn.out_norm"](gdn.out_norm, layer_id, is_mtp=is_mtp) self.rules["linear_attn.out_proj"](gdn.out_proj, layer_id, is_mtp=is_mtp) + def _get_kda_state_dict(self, layer, layer_id, is_mtp=False): + """Export a KDA (Kimi Delta Attention) layer's ``self_attention``.""" + kda = layer.self_attention + self.rules["kda"](kda, layer_id, is_mtp=is_mtp) + for name in ("beta_proj", "f_a_proj", "f_b_proj", "g_a_proj", "g_b_proj"): + self.rules[f"kda.{name}"](getattr(kda, name), layer_id, is_mtp=is_mtp) + self.rules["kda.out_norm"](kda.out_norm, layer_id, is_mtp=is_mtp) + self.rules["kda.out_proj"](kda.out_proj, layer_id, is_mtp=is_mtp) + # The HF checkpoint keeps the kernel parameters in FP32. + self.rules["kda.A_log"](kda.A_log.detach().float(), layer_id, is_mtp=is_mtp) + self.rules["kda.dt_bias"](kda.dt_bias.detach().float(), layer_id, is_mtp=is_mtp) + def _get_dsa_indexer_state_dict(self, indexer, layer_id, is_mtp=False): - """Export the DSA indexer of a sparse MLA layer.""" + """Export the DSA kpool indexer of a sparse MLA layer.""" if "indexer.linear_wq_b" not in self.rules: raise NotImplementedError(f"No export rule for the DSA indexer of {self.arch}.") for name in ("linear_wq_b", "linear_wk", "k_norm", "linear_weights_proj"): self.rules[f"indexer.{name}"](getattr(indexer, name), layer_id, is_mtp=is_mtp) + # KPool (GLM-5.3-Flash) only; plain DSA indexers (GLM-5.2) have no compression params. + for name in ("index_kpool_compress_ape", "index_kpool_compress_gate"): + param = getattr(indexer, name, None) + if param is not None: + self.rules[f"indexer.{name}"]( + param.detach().to(self.dtype), layer_id, is_mtp=is_mtp + ) def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.norm, IdentityOp): @@ -1176,6 +1351,7 @@ def _custom_mapping_to_lambda(mapping): "self_attention_scaling": self._self_attention_scaling, "gated_mlp_slicing": self._gated_mlp_slicing, "gated_delta_net_slicing": self._gated_delta_net_slicing, + "kda_slicing": self._kda_slicing, "grouped_mlp_slicing": self._grouped_mlp_slicing, "pack_name_remapping": self._pack_name_remapping, "pack_name_remapping_gpt_oss": self._pack_name_remapping_gpt_oss, @@ -1484,22 +1660,14 @@ def _gated_mlp_slicing( up_proj_weight = weight[ffn_hidden_size:, :] if qformat in GGML_FORMATS: - self._state_dict.update( - self._get_iq_weight_state( - gate_proj_prefix + "weight", - gate_proj_weight, - qformat, - getattr(module, "weight_quantizer", None), - ) - ) - self._state_dict.update( - self._get_iq_weight_state( - up_proj_prefix + "weight", - up_proj_weight, - qformat, - getattr(module, "weight_quantizer", None), - ) + # Match the fused linear's fake-quant forward before separating gate/up. + # Packing slices independently can reuse a GPTQ pin for an unchanged + # half even when the combined, QAD-updated weight must be re-encoded. + packed = self._pack_iq_weight( + weight, qformat, getattr(module, "weight_quantizer", None) ) + self._state_dict[gate_proj_prefix + "weight"] = packed[:ffn_hidden_size].clone() + self._state_dict[up_proj_prefix + "weight"] = packed[ffn_hidden_size:].clone() elif weight_scale is None: self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight self._state_dict[up_proj_prefix + "weight"] = up_proj_weight @@ -1935,11 +2103,6 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): """ if is_mtp: prefix = self._mtp_prefix(prefix) - in_proj = module.in_proj - name_to_value, qformat, block_size = self._get_quantized_state( - in_proj, self.dtype, prefix=prefix - ) - assert tuple(module.in_proj_split_names) == ( "query", "key", @@ -1959,12 +2122,39 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): sections["alpha"], ] proj_names = ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a") - proj_prefixes = [prefix + name + "." for name in proj_names] # The recipes keep the alpha / beta gates in BF16, but Megatron fuses all six sections # behind one quantizer, so they can only be dropped here rather than by a quantizer_name. - keep_bf16 = { - p for p, n in zip(proj_prefixes, proj_names) if n in ("in_proj_a", "in_proj_b") - } + self._split_fused_projection( + module.in_proj, + prefix, + proj_names, + split_sizes, + keep_bf16_names=("in_proj_a", "in_proj_b"), + ) + + def _kda_slicing(self, module, prefix, is_mtp=False): + """Split KDA's fused q|k|v ``in_proj`` and depthwise ``conv1d`` into HF's q/k/v tensors.""" + if is_mtp: + prefix = self._mtp_prefix(prefix) + assert tuple(module.in_proj_split_names) == ("query", "key", "value"), ( + f"Unexpected KDA in_proj layout {tuple(module.in_proj_split_names)}; only the " + "two-stage-gate layout [query, key, value] is supported" + ) + split_sizes = list(module.in_proj_split_sections) + self._split_fused_projection( + module.in_proj, prefix, ("q_proj", "k_proj", "v_proj"), split_sizes + ) + conv_weights = torch.split(module.conv1d.weight.detach().to(self.dtype), split_sizes) + for name, weight in zip(("q_conv1d", "k_conv1d", "v_conv1d"), conv_weights): + self._state_dict[prefix + name + ".weight"] = weight + + def _split_fused_projection(self, fused, prefix, proj_names, split_sizes, keep_bf16_names=()): + """Export a row-fused projection as the HF projections ``prefix + proj_names[i]``.""" + name_to_value, qformat, block_size = self._get_quantized_state( + fused, self.dtype, prefix=prefix + ) + proj_prefixes = [prefix + name + "." for name in proj_names] + keep_bf16 = {p for p, n in zip(proj_prefixes, proj_names) if n in keep_bf16_names} for proj_prefix in proj_prefixes: if proj_prefix in keep_bf16: @@ -1994,7 +2184,7 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): proj_prefix + "weight", proj_weight, qformat, - getattr(in_proj, "weight_quantizer", None), + getattr(fused, "weight_quantizer", None), ) ) elif weight_scale is None: diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 83dfc2f00d5..e0d6fe91d49 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -30,6 +30,18 @@ _ROW_KEY_CHUNK = 256 +class _ReconstructedWeightSTE(torch.autograd.Function): + """Preserve the decoder's exact forward values and pass the input gradient through.""" + + @staticmethod + def forward(ctx, inputs, reconstructed): + return reconstructed + + @staticmethod + def backward(ctx, grad_output): + return grad_output, None + + class _CacheKey(NamedTuple): """What a cached payload was packed from; any change means it must be packed again.""" @@ -225,7 +237,7 @@ def fake_quantize_with_cache( dtype=inputs.dtype, block_chunk_size=decode_chunk_size, ) - return inputs + (reconstructed - inputs).detach() + return _ReconstructedWeightSTE.apply(inputs, reconstructed) @dataclass(frozen=True) diff --git a/modelopt/torch/quantization/plugins/megatron_lora.py b/modelopt/torch/quantization/plugins/megatron_lora.py index b2affb365ce..d65fb7ab502 100644 --- a/modelopt/torch/quantization/plugins/megatron_lora.py +++ b/modelopt/torch/quantization/plugins/megatron_lora.py @@ -32,10 +32,19 @@ from ..nn import QuantModuleRegistry from .megatron import _QuantMegatronMLP +# HybridStack is optional for older Megatron-Core versions. +try: + from megatron.core.models.hybrid.hybrid_block import HybridStack +except ImportError: + HybridStack = None + __all__ = [] -@QuantLoRAInputRegistry.register({TransformerBlock: "TransformerBlock"}) +@QuantLoRAInputRegistry.register( + {TransformerBlock: "TransformerBlock"} + | ({HybridStack: "HybridStack"} if HybridStack is not None else {}) +) class _MegatronQuantLoRATransformerBlock(_QuantLoRAInput): def forward(self, hidden_states, *args, **kwargs): # Reentrant checkpoints need a differentiable input even with frozen embeddings. diff --git a/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/README.md b/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/README.md new file mode 100644 index 00000000000..bd83b7fe3e0 --- /dev/null +++ b/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/README.md @@ -0,0 +1,97 @@ +# DeepSeek-V4-Flash IQ2_XXS with mergeable LoRA QAD + +Use the standard Megatron-Bridge entry points to import the official HF checkpoint, +calibrate IQ2_XXS routed experts, train adapters, and export packed IQ2 or folded BF16. +Bridge decodes the original MXFP4 expert weights and FP8 weights into BF16 during import. +Other weights stay floating point, including native FP32 mHC controls. + +## Inputs and topology + +- Prepare a main-only HF metadata view of the official checkpoint: keep the original + shards, remove `mtp.*` index entries, set `num_nextn_predict_layers=0`, and trim + per-layer metadata to the main depth. Preserve the original checkpoint separately. +- Use TP1 / PP1 / CP1 / ETP1 and SequentialMLP (`--no_moe_grouped_gemm`). The recorded + integration uses EP8 on two nodes with four GB300 GPUs each. V4 TP/PP splitting and + MTP export are not supported by this initial implementation. +- Supply a calibration JSONL and Megatron pre-tokenized QAD training data. The recorded + calibration uses a seven-source `nemotron-post-training-v3` mix, 512 windows of 4096 + tokens, and batch size 1. Pin dataset revisions and keep calibration, training and + held-out documents disjoint. + +Run the commands from the repository root on both nodes. Set `NODE_RANK` to 0 or 1 +and `MASTER_ADDR` to the first node's address. Set the input/output variables +`HF_MAIN`, `CALIB_JSONL`, `TRAIN_PREFIX`, `PTQ_MEGATRON`, `QAD_OUTPUT`, +`TRAINED_ITERATION`, `PACKED_HF`, and `FOLDED_HF` to your paths. + +```bash +LAUNCH=(torchrun --nnodes=2 --nproc_per_node=4 --node_rank="$NODE_RANK" + --master_addr="$MASTER_ADDR" --master_port="${MASTER_PORT:-29500}") +``` + +## Calibrate and save the training checkpoint + +```bash +"${LAUNCH[@]}" examples/megatron_bridge/quantize.py \ + --hf_model_name_or_path "$HF_MAIN" --trust_remote_code \ + --recipe models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq \ + --no_moe_grouped_gemm --tp_size 1 --pp_size 1 --ep_size 8 --cp_size 1 \ + --dsa_kernel_backend none \ + --calib_dataset_name "$CALIB_JSONL" --calib_num_samples 512 \ + --calib_batch_size 1 --seq_length 4096 \ + --lora_rank 8 --lora_alpha 16 \ + --lora_target_modules '*mlp.experts.local_experts.*.linear_fc1' \ + '*mlp.experts.local_experts.*.linear_fc2' \ + --skip_generate --export_megatron_path "$PTQ_MEGATRON" +``` + +Omit `--compress`: QAD needs floating-point training weights, ModelOpt quantizer state, +GPTQ payloads, and zero-initialized adapters. The recipe uses importance-weighted, +layerwise GPTQ with block size 256 and damping 0.3. It targets only routed-expert +gate/up and down projections. + +## Train with distill.py + +```bash +"${LAUNCH[@]}" examples/megatron_bridge/distill.py \ + --student_hf_path "$HF_MAIN" --teacher_hf_path "$HF_MAIN" --trust_remote_code \ + --student_megatron_path "$PTQ_MEGATRON" \ + --no_moe_grouped_gemm --tp_size 1 --pp_size 1 --ep_size 8 --cp_size 1 \ + --data_paths "$TRAIN_PREFIX" --seq_length 512 --mbs 1 --gbs 8 \ + --kd_loss_alpha 1.0 --moe_aux_loss_coeff 0 --moe_router_bias_update_rate 0 \ + --dsa_kernel_backend none --deterministic_mode --mhc_fused_backend triton \ + --recompute_granularity selective --recompute_modules mhc moe \ + --lr 0.001 --min_lr 0.001 --lr_warmup_iters 0 \ + --train_iters 4 --save_interval 2 --checkpoint_keep_last -1 \ + --eval_iters 0 --no_async_save --log_interval 1 --output_dir "$QAD_OUTPUT" +``` + +Only the routed-expert rank-8 / alpha-16 LoRA factors are trainable. The frozen teacher +runs online; leave `--logit_kl_top_k` and `--logit_kl_top_p` unset for full-vocabulary KL. +The four-step configuration is an integration smoke test, not a tuned training recipe. +Exact resume additionally requires the Bridge data-setup RNG and Core frozen-optimizer +checkpoint fixes described in the PR's dependency notes. + +## Merge and export + +Keep the unmerged training checkpoint for resume. Point `TRAINED_ITERATION` to the +saved iteration directory. The standard exporter restores and merges the adapters. + +```bash +"${LAUNCH[@]}" examples/megatron_bridge/export_quantized_megatron_to_hf.py \ + --hf_model_name_or_path "$HF_MAIN" --trust_remote_code \ + --megatron_path "$TRAINED_ITERATION" --export_unified_hf_path "$PACKED_HF" \ + --pp_size 1 --ep_size 8 --no_moe_grouped_gemm +``` + +Run a separate export with `--vllm_fakequant_export` and +`--export_unified_hf_path "$FOLDED_HF"` to store the quantized expert values as BF16. +Packed export stores IQ2_XXS expert payloads; non-experts remain floating point in both +representations. Optimized packed-IQ2 serving is a separate runtime integration. + +## Recorded validation + +The v02 stack passed full-model GPTQ, exact eight-rank PTQ reload, four QAD updates, +fresh step-two resume, adapter merge, and both exports. All 33,024 packed expert tensors +decode exactly to folded BF16; all 34,223 exported tensors were compared. Fresh folded-HF +import reproduces the merged Bridge output probe on all eight ranks. The output probe +uses one 32-token input; held-out quality and production training are separate checks. diff --git a/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq.yaml b/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq.yaml new file mode 100644 index 00000000000..12c39a198b0 --- /dev/null +++ b/modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq.yaml @@ -0,0 +1,42 @@ +# 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. + +# Bridge PTQ: IQ2_XXS only on the routed experts. Keep floating-point training +# weights (no --compress), disable grouped GEMM and prepare an MTP=0 HF view. +# Mergeable QAD attaches rank 8 / alpha 16 adapters using quantize.py's LoRA flags. +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + iq2_xxs: configs/numerics/iq2_xxs +metadata: + description: >- + DeepSeek-V4-Flash-0731 (Megatron-Bridge): routed-expert IQ2_XXS with + importance-weighted layerwise GPTQ; all other imported weights stay floating point. +quantize: + algorithm: + method: gptq + block_size: 256 + perc_damp: 0.3 + importance_weighted: true + layerwise: {enable: true} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*mlp.experts.local_experts.*.linear_fc1.weight_quantizer' + cfg: + $import: iq2_xxs + - quantizer_name: '*mlp.experts.local_experts.*.linear_fc2.weight_quantizer' + cfg: + $import: iq2_xxs + - {quantizer_name: '*mtp*', enable: false} diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml new file mode 100644 index 00000000000..a773459de40 --- /dev/null +++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml @@ -0,0 +1,61 @@ +# 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. + +# Recipe for zai-org/GLM-5.3-Flash-BF16 with Megatron-Bridge: GGML IQ2_S on all routed-expert projections and +# Q8_0 on the shared experts and the dense MLP of layers 0-2, calibrated with layerwise GPTQ. Both attention +# families (KDA and sparse MLA), the router, the vision tower, the MTP layer, embeddings and lm_head stay BF16. +# +# Run quantize.py, distill.py and export with --no_moe_grouped_gemm: GPTQ and the GGML export +# only handle per-expert linears. vLLM has no IQ2 kernels, so evaluate an export made with +# --vllm_fakequant_export, which folds the IQ2 weights into BF16. +# +# Every quantized linear's input width is a multiple of the 256-weight GGML block +# (hidden 4096, dense MLP 12288, experts 2048). + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + iq2_s: configs/numerics/iq2_s + q8_0: configs/numerics/q8_0 + +metadata: + description: >- + GLM-5.3-Flash-BF16 (Megatron-Bridge): GGML IQ2_S on the routed experts and Q8_0 on the shared experts and + dense MLP, calibrated with layerwise GPTQ. Attention, router, vision tower, MTP layer, embeddings and + lm_head stay BF16. Requires calibration data. +quantize: + algorithm: + method: gptq + # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time. + block_size: 256 + perc_damp: 0.3 + importance_weighted: true + layerwise: {enable: true} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: iq2_s + # Every token runs through the shared experts and the dense MLP, and they are ~0.5% of the MLP + # weights, so Q8_0 there costs almost nothing in size. + - quantizer_name: '*mlp.shared_experts*weight_quantizer' + cfg: + $import: q8_0 + - quantizer_name: '*mlp.linear_fc*weight_quantizer' + cfg: + $import: q8_0 + - {quantizer_name: '*visual*', enable: false} + # Megatron-Bridge builds the MTP layer (HF does not); keep it BF16 like the HF export. + - {quantizer_name: '*mtp*', enable: false} diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml index d897eb79f27..59cb92c7715 100644 --- a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml +++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml @@ -29,10 +29,10 @@ # so only 3 of the 45 layers have a plain MLP -- 9 modules in total (`mlp.gate_proj`, # `mlp.up_proj`, `mlp.down_proj` each). The other 42 layers carry `mlp.experts..*`. # -# The one rule that is not obvious is the trailing `*visual*` disable, and it is load-bearing. +# The one rule that is not obvious is the `*visual*` disable, and it is load-bearing. # The vision tower reuses the same leaf names -- `model.visual.blocks..mlp.gate_proj` and # friends, across 24 blocks -- so the dense-MLP patterns match 144 modules inside it. Entries -# apply in order, so the disable has to come LAST or the vision tower is quantized by accident. +# apply in order, so the disable has to come AFTER them or the vision tower is quantized by accident. # (`*.experts.*` matches nothing under `model.visual.*`, so only the dense-MLP patterns reach it.) # # Two model-specific notes: @@ -41,10 +41,12 @@ # layers nest under `model.language_model.layers` and layerwise_calibrate cannot find # them. # -# * The MTP (next-token-prediction) head is not quantized because it is not built. The config -# declares `num_hidden_layers: 45` (with `num_nextn_predict_layers: 1`), so the HF model class -# instantiates decoder layers 0-44 only and does not construct the MTP layer -- it is therefore -# neither quantized nor carried into the exported checkpoint. +# * The MTP (next-token-prediction) head stays BF16. The HF model class does not build it +# (decoder layers 0-44 only); Megatron-Bridge does, so the trailing `*mtp*` disable keeps it off. +# +# The same recipe also drives Megatron-Bridge `quantize.py`, where the dense MLP is +# `mlp.linear_fc1` / `mlp.linear_fc2`. Those patterns need a literal `mlp.linear_fc`, so they +# skip `mlp.shared_experts.linear_fc*` and `mlp.experts.local_experts..linear_fc*`. # # The shared `default_disabled_quantizers` unit is deliberately not imported. # Of its patterns only `*visual*` changes anything here; every other one either matches no @@ -63,8 +65,8 @@ metadata: description: >- GLM-5.3-Flash-BF16: NVFP4 (W4A4) on the routed experts and on the dense MLP of layers 0-2, plus an FP8 KV cache in cast mode using constant amax; max calibration. Shared experts, - router gate, KDA and MLA attention, vision tower, embeddings and lm_head stay BF16. The - MTP layer is not built at num_hidden_layers=45, so it is neither quantized nor exported. + router gate, KDA and MLA attention, vision tower, MTP layer, embeddings and lm_head + stay BF16. Works for both HF (hf_ptq.py) and Megatron-Bridge (quantize.py) module names. quantize: algorithm: method: max @@ -100,6 +102,21 @@ quantize: - quantizer_name: '*mlp.down_proj*input_quantizer' cfg: $import: nvfp4 + # Dense MLP, Megatron-Bridge names (fc1 = fused gate_proj + up_proj). + - quantizer_name: '*mlp.linear_fc1*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp.linear_fc1*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp.linear_fc2*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp.linear_fc2*input_quantizer' + cfg: + $import: nvfp4 - $import: kv_fp8_cast - # MUST stay last: keeps the vision tower BF16 after the dense-MLP patterns matched it. + # MUST follow the dense-MLP patterns: keeps the vision tower BF16 after they matched it. - {quantizer_name: '*visual*', enable: false} + # Megatron-Bridge builds the MTP layer (HF does not); keep it BF16 like the HF export. + - {quantizer_name: '*mtp*', enable: false} diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index bdf910981d2..f1b51a06402 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -639,6 +639,56 @@ def create_tiny_deepseek_v3_dir( ) +##### GLM-5.3-FLASH ##### +def get_tiny_glm5_next(**text_config_kwargs) -> PreTrainedModel: + """Tiny GLM-5.3-Flash VLM: one KDA + dense-MLP layer, one sparse-MLA + MoE layer, one vision block.""" + set_seed(SEED) + + # Lazy import — glm5_next is native only from transformers 5.16.1. + from transformers import Glm5NextConfig + + text_kwargs = { + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "moe_intermediate_size": 32, + "num_hidden_layers": 2, + "layer_types": ["linear_attention", "indexed_attention"], + "mlp_layer_types": ["dense", "sparse"], + "indexer_types": ["full", "full"], + "num_attention_heads": 4, + "num_key_value_heads": 4, + "q_lora_rank": 32, + "kv_lora_rank": 16, + "qk_nope_head_dim": 16, + "v_head_dim": 16, + "linear_num_heads": 4, + "linear_head_dim": 16, + "index_n_heads": 2, + "index_head_dim": 16, + "index_topk": 8, + "index_kpool": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "max_position_embeddings": 64, + "pad_token_id": 0, + } + text_kwargs.update(text_config_kwargs) + vision_kwargs = { + "depth": 1, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 2, + "image_size": 28, + "patch_size": 14, + "out_hidden_size": text_kwargs["hidden_size"], + "projection_intermediate_size": 64, + } + cfg = Glm5NextConfig(text_config=text_kwargs, vision_config=vision_kwargs) + return AutoModelForImageTextToText.from_config(cfg, dtype=torch.bfloat16) + + ##### GLM-5.3-FLASH / DEEPSEEK-V4 (config-only, for vLLM ``load_format="dummy"``) ##### def _write_config_dir(dir_path: Path, config: dict) -> Path: dir_path.mkdir(parents=True, exist_ok=True) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 6977061d282..78ff442d1bd 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -35,10 +35,37 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge")) -from distill import HAS_DIRECT_HF_SFT, _hf_source +from distill import HAS_DIRECT_HF_SFT, _hf_source, get_args from export_distilled_megatron_to_hf import _get_checkpoint_export_paths +@pytest.mark.parametrize("explicit", [False, True]) +def test_distill_qad_execution_options(monkeypatch, explicit): + command = [ + "distill.py", + "--student_hf_path", + "student", + "--teacher_hf_path", + "teacher", + "--output_dir", + "output", + "--use_mock_data", + ] + if explicit: + command += [ + "--deterministic_mode", + "--mhc_fused_backend", + "triton", + "--moe_router_bias_update_rate", + "0", + ] + monkeypatch.setattr(sys, "argv", command) + args = get_args() + assert args.deterministic_mode is explicit + assert args.mhc_fused_backend == ("triton" if explicit else None) + assert args.moe_router_bias_update_rate == (0.0 if explicit else None) + + @pytest.mark.timeout(360) # distill + inline HF export + export of every iteration def test_distill_llm_hf_export(tmp_path, num_gpus): """Distill with an inline HF export, then re-export every saved iteration standalone.""" diff --git a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py b/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py index 965c44d7c29..f8aa35ca484 100644 --- a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py +++ b/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py @@ -38,6 +38,7 @@ "Qwen3_5ForConditionalGeneration": "model.language_model.layers.{}.self_attn.", "Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.self_attn.", "NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.", + "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.self_attn.", "GlmMoeDsaForCausalLM": "model.layers.{}.self_attn.", } @@ -60,6 +61,7 @@ def test_export_mapping_emits_kv_cache_scales(arch, prefix): "Qwen3MoeForCausalLM": "model.layers.{}.mlp.experts.{}", "NemotronHForCausalLM": "backbone.layers.{}.mixer.experts.{}", "NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.experts.{}", + "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.mlp.experts.{}", "GlmMoeDsaForCausalLM": "model.layers.{}.mlp.experts.{}", } @@ -105,3 +107,33 @@ def test_new_model_nests_every_target_under_language_model(): assert vl["mtp.eh_proj"].target_name_or_prefix.startswith("language_model.mtp.") passthrough = all_mcore_hf_vision_passthrough_mapping["NemotronH_Omni_Reasoning_V3"] assert not any(p.startswith("language_model.") for p in passthrough) + + +def test_glm5_next_names_match_released_checkpoint(): + """Rule targets must format to tensor names of the released zai-org/GLM-5.3-Flash checkpoint.""" + mapping = all_mcore_hf_export_mapping["Glm5NextForConditionalGeneration"] + assert mapping["fold_attn_mlp_layer_pairs"] is True + assert mapping["mtp_in_decoder_layers"] is True + + layer = "model.language_model.layers.3." + expected = { + ("hc_fn", 3, "attn"): layer + "hc_attn_fn", + ("hc_base", 3, "ffn"): layer + "hc_ffn_base", + ("hc_scale", 3, "ffn"): layer + "hc_ffn_scale", + ("kda", 3): layer + "self_attn.", + ("kda.beta_proj", 3): layer + "self_attn.b_proj.", + ("kda.out_norm", 3): layer + "self_attn.o_norm.", + ("kda.A_log", 3): layer + "self_attn.A_log", + ("linear_kv_down_proj", 3): layer + "self_attn.kv_a_proj_with_mqa.", + ("indexer.linear_wq_b", 3): layer + "self_attn.indexer.wq_b.", + ("indexer.index_kpool_compress_ape", 3): layer + + "self_attn.indexer.index_kpool_compress_ape", + ("fused_pre_mlp_layernorm", 3): layer + "post_attention_layernorm.weight", + ("shared_experts.linear_fc2", 3): layer + "mlp.shared_experts.down_proj.", + ("mtp.eh_proj", 3): layer + "eh_proj.weight", + ("mtp.final_layernorm", 3): layer + "shared_head.norm.", + } + for (rule, *args), name in expected.items(): + assert mapping[rule].target_name_or_prefix.format(*args) == name, rule + assert mapping["kda"].func_name == "kda_slicing" + assert mapping["router"].func_kwargs["mapping"] == {"expert_bias": "e_score_correction_bias"} diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 043970ef42a..dd64dd835c8 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -300,6 +300,59 @@ def test_megatron_gated_delta_net_slicing_exports_ggml_payloads(qformat): ) +def test_megatron_kda_slicing_splits_fused_qkv_and_conv1d(): + weight = torch.randn(7, 4).bfloat16() + conv = torch.randn(7, 1, 4).bfloat16() + module = SimpleNamespace( + in_proj=object(), + in_proj_split_names=("query", "key", "value"), + in_proj_split_sections=(2, 2, 3), + conv1d=SimpleNamespace(weight=conv), + ) + exporter = _make_ggml_exporter() + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, None, None) + + exporter._kda_slicing(module, "model.layers.0.self_attn.") + + for name, rows in (("q", slice(0, 2)), ("k", slice(2, 4)), ("v", slice(4, 7))): + torch.testing.assert_close( + exporter._state_dict[f"model.layers.0.self_attn.{name}_proj.weight"], weight[rows] + ) + torch.testing.assert_close( + exporter._state_dict[f"model.layers.0.self_attn.{name}_conv1d.weight"], conv[rows] + ) + assert exporter.exclude_modules == [ + f"model.layers.0.self_attn.{name}_proj" for name in ("q", "k", "v") + ] + + +@pytest.mark.parametrize(("self_attention", "kind"), [(object(), "attn"), (None, "ffn")]) +def test_megatron_hyper_connection_exports_hf_hc_tensors(self_attention, kind): + hc = SimpleNamespace( + mapping_proj=SimpleNamespace(weight=torch.randn(24, 64)), + bias=torch.randn(24), + alpha_pre=torch.tensor([0.1]), + alpha_post=torch.tensor([0.2]), + alpha_res=torch.tensor([0.3]), + ) + layer = SimpleNamespace( + hyper_connection=hc, inner_layer=SimpleNamespace(self_attention=self_attention) + ) + exporter = _make_ggml_exporter() + exporter.rules = exporter._populate_rule_book()["Glm5NextForConditionalGeneration"] + + exporter._get_hyper_connection_state_dict(layer, 2) + + prefix = f"model.language_model.layers.2.hc_{kind}_" + assert sorted(exporter._state_dict) == [prefix + n for n in ("base", "fn", "scale")] + assert exporter._state_dict[prefix + "fn"].dtype == torch.bfloat16 + # The released checkpoint keeps the mHC bias and alpha scales in FP32. + assert exporter._state_dict[prefix + "base"].dtype == torch.float32 + torch.testing.assert_close( + exporter._state_dict[prefix + "scale"], torch.tensor([0.1, 0.2, 0.3]) + ) + + @pytest.mark.parametrize("qformat", GGML_FORMAT_NAMES) def test_megatron_packed_experts_reject_ggml_without_deployment_loader(qformat): experts = _make_ggml_experts(qformat, "linear_fc2") @@ -1143,30 +1196,6 @@ def test_mtp_state_dict_index_file(tmp_path): assert "mtp*" in exporter.exclude_modules -def _test_live_decoder_mtp_export_rejected(model_dir, rank, size): - model = get_mcore_gpt_model( - tensor_model_parallel_size=size, - pipeline_model_parallel_size=1, - initialize_megatron=True, - num_layers=2, - hidden_size=64, - num_attention_heads=4, - vocab_size=128, - max_sequence_length=32, - mtp_num_layers=1, - ).cuda() - with pytest.raises(NotImplementedError, match="Megatron-built MTP"): - GPTModelExporter(model, str(model_dir)) - - -def test_live_decoder_mtp_export_rejected(dist_workers_size_1, tmp_path): - """GLM-5 MTP is only copied from the source; a Megatron-built one must not export misnamed.""" - transformers.GlmMoeDsaConfig( - num_hidden_layers=2, architectures=["GlmMoeDsaForCausalLM"] - ).save_pretrained(tmp_path) - dist_workers_size_1.run(partial(_test_live_decoder_mtp_export_rejected, tmp_path)) - - def test_mtp_state_dict_copies_decoder_mtp_layers(tmp_path): """GLM-5 keeps MTP as an extra decoder layer; copy it dequantized when Megatron did not build it.""" model_dir = tmp_path / "fake_glm5" @@ -1188,6 +1217,8 @@ def test_mtp_state_dict_copies_decoder_mtp_layers(tmp_path): exporter._src_num_hidden_layers = 2 exporter._hf_text_config = SimpleNamespace(num_hidden_layers=1, num_nextn_predict_layers=1) + # Non-writer ranks skip the copy. + assert exporter._get_mtp_state_dict(copy_from_pretrained=False) == {} mtp_state_dict = exporter._get_mtp_state_dict() assert sorted(mtp_state_dict) == [ diff --git a/tests/gpu_megatron/torch/quantization/test_quant_lora.py b/tests/gpu_megatron/torch/quantization/test_quant_lora.py index 8b2d3a6d2ea..5b8790dcec9 100644 --- a/tests/gpu_megatron/torch/quantization/test_quant_lora.py +++ b/tests/gpu_megatron/torch/quantization/test_quant_lora.py @@ -19,7 +19,7 @@ import pytest import torch -from _test_utils.torch.megatron.models import MegatronModel +from _test_utils.torch.megatron.models import MegatronModel, get_mcore_hybrid_model from _test_utils.torch.megatron.utils import initialize_for_megatron, sharded_state_dict_test_helper from torch.nn import functional as F @@ -66,3 +66,82 @@ def _test_quant_lora(checkpoint_dir, rank, size): def test_quant_lora_distributed_checkpoint(tmp_path, request, tp_size): workers = request.getfixturevalue(f"dist_workers_size_{tp_size}") workers.run(partial(_test_quant_lora, str(tmp_path))) + + +def _test_hybrid_iq2_lora_gradients(recompute, rank, size): + initialize_for_megatron(tensor_model_parallel_size=1) + model = get_mcore_hybrid_model( + num_layers=2, + hybrid_layer_pattern="EE", + hidden_size=256, + num_attention_heads=4, + moe_ffn_hidden_size=256, + moe_shared_expert_intermediate_size=None, + num_moe_experts=2, + moe_router_topk=2, + transformer_impl="transformer_engine", + gradient_accumulation_fusion=False, + hidden_dropout=0, + attention_dropout=0, + recompute_granularity="full" if recompute else None, + recompute_method="uniform" if recompute else None, + recompute_num_layers=1 if recompute else None, + ).cuda() + targets = [ + "decoder.layers.*.mlp.experts.local_experts.*.linear_fc1", + "decoder.layers.*.mlp.experts.local_experts.*.linear_fc2", + ] + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + *[ + { + "quantizer_name": name + ".weight_quantizer", + "cfg": { + "num_bits": "iq2_xxs", + "block_sizes": {-1: 256}, + "backend": "ggml", + }, + "enable": True, + } + for name in targets + ], + ], + "algorithm": None, + }, + ) + backbone = {name: parameter.detach().clone() for name, parameter in model.named_parameters()} + mtq.enable_quant_lora(model, {"rank": 4, "target_modules": targets}) + adapters = { + name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad + } + assert len(adapters) == 16 + optimizer = torch.optim.SGD(adapters.values(), lr=0.1) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4]], device="cuda"), + "position_ids": torch.arange(4, device="cuda").unsqueeze(0), + "attention_mask": None, + } + model.train() + for step in range(2): + optimizer.zero_grad(set_to_none=True) + loss = model(**inputs).float().square().mean() + assert torch.isfinite(loss) + loss.backward() + for name, parameter in adapters.items(): + assert parameter.grad is not None and torch.isfinite(parameter.grad).all(), name + if step or name.endswith("lora_B"): + assert parameter.grad.count_nonzero(), name + optimizer.step() + for name, expected in backbone.items(): + torch.testing.assert_close(dict(model.named_parameters())[name], expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "recompute", + [False, True], +) +def test_hybrid_iq2_lora_gradients(dist_workers_size_1, recompute): + dist_workers_size_1.run(partial(_test_hybrid_iq2_lora_gradients, recompute)) diff --git a/tests/unit/recipe/test_deepseek_v4_recipe.py b/tests/unit/recipe/test_deepseek_v4_recipe.py new file mode 100644 index 00000000000..a20378e37e9 --- /dev/null +++ b/tests/unit/recipe/test_deepseek_v4_recipe.py @@ -0,0 +1,68 @@ +# 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. + +"""Keep the Flash IQ2 recipe scoped to routed SequentialMLP experts.""" + +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe + + +def _mlp(): + module = nn.Module() + module.linear_fc1 = nn.Linear(256, 512, bias=False) + module.linear_fc2 = nn.Linear(256, 256, bias=False) + return module + + +def test_flash_iq2_recipe_targets_only_routed_experts(): + model = nn.Module() + model.decoder = nn.Module() + model.decoder.layers = nn.ModuleList([nn.Module(), nn.Module()]) + for layer in model.decoder.layers: + layer.inner_layer = nn.Module() + inner = layer.inner_layer + inner.self_attention = nn.Module() + inner.self_attention.linear_proj = nn.Linear(256, 256, bias=False) + inner.mlp = nn.Module() + inner.mlp.experts = nn.Module() + inner.mlp.experts.local_experts = nn.ModuleList([_mlp(), _mlp()]) + inner.mlp.shared_experts = _mlp() + inner.mlp.router = nn.Linear(256, 2, bias=False) + model.mtp = nn.Module() + model.mtp.mlp = model.decoder.layers[0].inner_layer.mlp.__class__() + model.mtp.mlp.experts = nn.Module() + model.mtp.mlp.experts.local_experts = nn.ModuleList([_mlp()]) + config = load_recipe( + "models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq" + ).quantize.model_dump() + algorithm = config.pop("algorithm") + assert algorithm["method"] == "gptq" + assert algorithm["importance_weighted"] + assert algorithm["block_size"] == 256 + assert algorithm["perc_damp"] == 0.3 + assert algorithm["layerwise"]["enable"] + mtq.quantize(model, {**config, "algorithm": None}) + active = [] + for name, module in model.named_modules(): + if name.endswith("weight_quantizer") and module.is_enabled: + active.append(name) + assert name.startswith("decoder.") + assert ".mlp.experts.local_experts." in name + assert module.num_bits == "iq2_xxs" + elif name.endswith("input_quantizer"): + assert not module.is_enabled + assert len(active) == 8 diff --git a/tests/unit/recipe/test_glm_5_3_recipe.py b/tests/unit/recipe/test_glm_5_3_recipe.py index de0c7c73942..c225ac47eeb 100644 --- a/tests/unit/recipe/test_glm_5_3_recipe.py +++ b/tests/unit/recipe/test_glm_5_3_recipe.py @@ -22,24 +22,26 @@ *not* matched and the shared experts stay BF16. * The vision tower reuses the language MLP's leaf names (``mlp.gate_proj`` / ``up_proj`` / ``down_proj``), so the dense-MLP patterns match ``model.visual.*`` - too -- only the trailing ``*visual*`` disable (which must stay last) keeps the + too -- only the ``*visual*`` disable (which must follow them) keeps the vision tower in BF16. -* ``*mlp.gate_proj*`` must not catch the router ``mlp.gate``. This pins that behaviour so it can't silently drift. """ +import pytest import torch.nn as nn +from _test_utils.torch import transformers_models import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe _RECIPE = "models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast" +_IQ2_RECIPE = "models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq" _H = 32 class _MLP(nn.Module): - """Plain MLP leaf names, shared by the dense MLP, each routed/shared expert, and vision.""" + """Plain MLP leaf names, shared by the dense MLP, the shared experts, and vision.""" def __init__(self): super().__init__() @@ -48,78 +50,31 @@ def __init__(self): self.down_proj = nn.Linear(_H, _H, bias=False) -class _MoE(nn.Module): +class _VisionAttention(nn.Module): def __init__(self): super().__init__() - self.experts = nn.ModuleList([_MLP(), _MLP()]) - self.shared_experts = _MLP() - self.gate = nn.Linear(_H, 2, bias=False) # router - - -class _KDA(nn.Module): - """KDA linear attention (projections plus a depthwise causal conv1d).""" - - def __init__(self): - super().__init__() - self.in_proj_qkvz = nn.Linear(_H, _H, bias=False) - self.conv1d = nn.Conv1d(_H, _H, kernel_size=3, groups=_H, bias=False) - self.out_proj = nn.Linear(_H, _H, bias=False) - - -class _MLA(nn.Module): - """NoPE sparse-MLA attention (and the vision attention, which reuses the leaf names).""" - - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(_H, _H, bias=False) - self.kv_proj = nn.Linear(_H, _H, bias=False) - self.o_proj = nn.Linear(_H, _H, bias=False) - - -class _DenseLayer(nn.Module): - """Layers 0-2: KDA attention + a plain (dense) MLP.""" - - def __init__(self): - super().__init__() - self.linear_attn = _KDA() - self.mlp = _MLP() - - -class _SparseLayer(nn.Module): - """Layers 3-44: MLA attention + an MoE block.""" - - def __init__(self): - super().__init__() - self.self_attn = _MLA() - self.mlp = _MoE() + self.qkv = nn.Linear(_H, 3 * _H) + self.proj = nn.Linear(_H, _H) class _VisionBlock(nn.Module): def __init__(self): super().__init__() self.mlp = _MLP() # same gate_proj / up_proj / down_proj leaf names as the language MLP - self.attn = _MLA() - - -class _GLM53Flash(nn.Module): - """Tiny stand-in for the ``glm5_next`` VLM MoE (one dense + one sparse layer + vision).""" - - def __init__(self): - super().__init__() - self.model = nn.Module() - self.model.language_model = nn.Module() - self.model.language_model.layers = nn.ModuleList([_DenseLayer(), _SparseLayer()]) - self.model.visual = nn.Module() - self.model.visual.blocks = nn.ModuleList([_VisionBlock()]) - self.lm_head = nn.Linear(_H, _H, bias=False) + self.attn = _VisionAttention() def _nvfp4(quantizer): return quantizer.is_enabled and quantizer.num_bits == (2, 1) +def _fp8(quantizer): + return quantizer.is_enabled and quantizer.num_bits == (4, 3) + + def test_glm_5_3_recipe_quantizer_precedence(): - model = _GLM53Flash() + pytest.importorskip("transformers.models.glm5_next", reason="needs transformers>=5.16.1") + model = transformers_models.get_tiny_glm5_next() config = load_recipe(_RECIPE).quantize.model_dump() # The recipe uses plain max calibration; here we only assert quantizer placement, @@ -128,45 +83,140 @@ def test_glm_5_3_recipe_quantizer_precedence(): config["algorithm"] = None mtq.quantize(model, config) - dense = model.model.language_model.layers[0] - sparse = model.model.language_model.layers[1] + dense, sparse = model.model.language_model.layers - # Routed experts -> NVFP4 W4A4. - for expert in sparse.mlp.experts: - for proj in (expert.gate_proj, expert.up_proj, expert.down_proj): - assert _nvfp4(proj.weight_quantizer) - assert _nvfp4(proj.input_quantizer) + # Routed experts (fused into 3D params, one weight quantizer per expert) -> NVFP4 W4A4. + experts = sparse.mlp.experts + for name in ("gate_up_proj", "down_proj"): + assert _nvfp4(getattr(experts, f"{name}_input_quantizer")) + assert all(_nvfp4(q) for q in getattr(experts, f"{name}_weight_quantizers")) # Dense MLP (layers 0-2) -> NVFP4. for proj in (dense.mlp.gate_proj, dense.mlp.up_proj, dense.mlp.down_proj): assert _nvfp4(proj.weight_quantizer) assert _nvfp4(proj.input_quantizer) - # Vision tower stays BF16 -- the load-bearing case: the vision MLP reuses - # gate_proj/up_proj/down_proj, so the dense-MLP patterns match it and only the - # trailing `*visual*` disable keeps it off. - vblock = model.model.visual.blocks[0] - for proj in (vblock.mlp.gate_proj, vblock.mlp.up_proj, vblock.mlp.down_proj): - assert proj.weight_quantizer.is_enabled is False - assert proj.input_quantizer.is_enabled is False + # Sparse-MLA KV cache -> FP8. + assert _fp8(sparse.self_attn.k_bmm_quantizer) + assert _fp8(sparse.self_attn.v_bmm_quantizer) - # Shared experts and the router gate stay BF16: `*.experts.*` needs a literal - # `.experts.` (so `shared_experts` is skipped), and `*mlp.gate_proj*` doesn't match - # the router `mlp.gate`. + # The whole vision tower stays BF16 -- the load-bearing case: its blocks and merger reuse + # gate_proj/up_proj/down_proj, so the dense-MLP patterns match them and only the later + # `*visual*` disable keeps them off. + for name, module in model.model.visual.named_modules(): + if name.endswith("quantizer"): + assert module.is_enabled is False, name + + # Shared experts stay BF16: `*.experts.*` needs a literal `.experts.`, so `shared_experts` + # is skipped. (The router `mlp.gate` is not an nn.Linear, so it gets no quantizer at all.) for proj in ( sparse.mlp.shared_experts.gate_proj, sparse.mlp.shared_experts.up_proj, sparse.mlp.shared_experts.down_proj, ): assert proj.weight_quantizer.is_enabled is False - assert sparse.mlp.gate.weight_quantizer.is_enabled is False + assert proj.input_quantizer.is_enabled is False - # Both attention families stay BF16, including the KDA conv1d. - assert dense.linear_attn.conv1d.weight_quantizer.is_enabled is False - assert dense.linear_attn.in_proj_qkvz.weight_quantizer.is_enabled is False - assert dense.linear_attn.out_proj.weight_quantizer.is_enabled is False - for proj in (sparse.self_attn.q_proj, sparse.self_attn.kv_proj, sparse.self_attn.o_proj): - assert proj.weight_quantizer.is_enabled is False + # Both attention families' projections stay BF16, including the KDA conv1d and the indexer. + for attn in (dense.self_attn, sparse.self_attn): + for name, module in attn.named_modules(): + if name.endswith(("weight_quantizer", "input_quantizer")): + assert module.is_enabled is False, name - # lm_head stays BF16. + # Embeddings and lm_head stay BF16. + assert model.model.language_model.embed_tokens.weight_quantizer.is_enabled is False assert model.lm_head.weight_quantizer.is_enabled is False + + +class _McoreMLP(nn.Module): + """Megatron-Bridge MLP leaf names (dense MLP, each local expert, and the shared experts).""" + + def __init__(self): + super().__init__() + self.linear_fc1 = nn.Linear(_H, 2 * _H, bias=False) + self.linear_fc2 = nn.Linear(_H, _H, bias=False) + + +class _McoreMoE(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.Module() + self.experts.local_experts = nn.ModuleList([_McoreMLP(), _McoreMLP()]) + self.shared_experts = _McoreMLP() + + +class _McoreLayer(nn.Module): + def __init__(self, mlp): + super().__init__() + self.inner_layer = nn.Module() # mHC wraps each block as `.inner_layer` + self.inner_layer.mlp = mlp + + +class _McoreGLM53Flash(nn.Module): + """Megatron-Bridge naming: a dense and an MoE decoder layer, the MTP layer, and vision.""" + + def __init__(self): + super().__init__() + self.language_model = nn.Module() + self.language_model.decoder = nn.Module() + self.language_model.decoder.layers = nn.ModuleList( + [_McoreLayer(_McoreMLP()), _McoreLayer(_McoreMoE())] + ) + self.language_model.mtp = nn.Module() + self.language_model.mtp.layers = nn.ModuleList([_McoreLayer(_McoreMoE())]) + self.visual = nn.Module() + self.visual.blocks = nn.ModuleList([_VisionBlock()]) + + +def test_glm_5_3_recipe_megatron_names(): + model = _McoreGLM53Flash() + config = load_recipe(_RECIPE).quantize.model_dump() + config["algorithm"] = None + mtq.quantize(model, config) + + dense, sparse = (layer.inner_layer.mlp for layer in model.language_model.decoder.layers) + mtp = model.language_model.mtp.layers[0].inner_layer.mlp + + # Dense MLP and routed experts -> NVFP4 W4A4. + for mlp in (dense, *sparse.experts.local_experts): + for proj in (mlp.linear_fc1, mlp.linear_fc2): + assert _nvfp4(proj.weight_quantizer) + assert _nvfp4(proj.input_quantizer) + + # Shared experts, the whole MTP layer, and the vision tower stay BF16. + for mlp in (sparse.shared_experts, mtp.shared_experts, *mtp.experts.local_experts): + for proj in (mlp.linear_fc1, mlp.linear_fc2): + assert proj.weight_quantizer.is_enabled is False + assert proj.input_quantizer.is_enabled is False + vmlp = model.visual.blocks[0].mlp + for proj in (vmlp.gate_proj, vmlp.up_proj, vmlp.down_proj): + assert proj.weight_quantizer.is_enabled is False + + +def test_glm_5_3_iq2_recipe_megatron_names(): + model = _McoreGLM53Flash() + config = load_recipe(_IQ2_RECIPE).quantize.model_dump() + config["algorithm"] = None + mtq.quantize(model, config) + + dense, sparse = (layer.inner_layer.mlp for layer in model.language_model.decoder.layers) + mtp = model.language_model.mtp.layers[0].inner_layer.mlp + + # Routed experts -> IQ2_S; dense MLP and shared experts -> Q8_0; activations stay BF16. + for mlp, num_bits in [ + *((expert, "iq2_s") for expert in sparse.experts.local_experts), + (dense, "q8_0"), + (sparse.shared_experts, "q8_0"), + ]: + for proj in (mlp.linear_fc1, mlp.linear_fc2): + assert proj.weight_quantizer.is_enabled + assert proj.weight_quantizer.num_bits == num_bits + assert proj.input_quantizer.is_enabled is False + + # The MTP layer and the vision tower stay BF16. + for mlp in (mtp.shared_experts, *mtp.experts.local_experts): + for proj in (mlp.linear_fc1, mlp.linear_fc2): + assert proj.weight_quantizer.is_enabled is False + vmlp = model.visual.blocks[0].mlp + for proj in (vmlp.gate_proj, vmlp.up_proj, vmlp.down_proj): + assert proj.weight_quantizer.is_enabled is False diff --git a/tests/unit/torch/export/test_deepseek_v4_export.py b/tests/unit/torch/export/test_deepseek_v4_export.py new file mode 100644 index 00000000000..1aa25e03546 --- /dev/null +++ b/tests/unit/torch/export/test_deepseek_v4_export.py @@ -0,0 +1,279 @@ +# 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. + +"""Exercise native V4 naming and real IQ2 packing without requiring CUDA/MCore.""" + +import json +from collections import OrderedDict +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.plugins.vllm_fakequant_megatron import VllmFqGPTModelExporter +from modelopt.torch.export.unified_export_megatron import GPTModelExporter +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml.common import pin_packed_weight + + +def _exporter(folded=False, source_keys=()): + cls = VllmFqGPTModelExporter if folded else GPTModelExporter + exporter = cls.__new__(cls) + exporter.dtype = torch.bfloat16 + exporter.moe_router_dtype = None + exporter.arch = "DeepseekV4ForCausalLM" + exporter._state_dict = OrderedDict() + exporter.exclude_modules = [] + exporter.layer_config_dict = {} + exporter._deepseek_v4_source_keys = set(source_keys) + exporter.rules = exporter._populate_rule_book()[exporter.arch] + if folded: + exporter._quantizer_recipe_markers = [] + exporter._quantizer_validation_failure = "" + exporter._packing_experts = False + exporter._grouped_fq_modules = set() + return exporter + + +def _compressor(): + module = nn.Module() + module.linear_wkv = nn.Linear(256, 256, bias=False).bfloat16() + module.linear_wgate = nn.Linear(256, 256, bias=False).bfloat16() + module.norm = nn.LayerNorm(256, bias=False).bfloat16() + module.ape = nn.Parameter(torch.randn(4, 256)) + return module + + +def _attention(compression): + attention = nn.Module() + for name in ("linear_q_down_proj", "linear_q_up_proj", "linear_kv_proj", "linear_proj"): + setattr(attention, name, nn.Linear(256, 256, bias=False).bfloat16()) + attention.q_layernorm = nn.LayerNorm(256, bias=False).bfloat16() + attention.kv_layernorm = nn.LayerNorm(256, bias=False).bfloat16() + attention.linear_o_group_proj = nn.Parameter(torch.randn(2, 256, 256).bfloat16()) + attention.core_attention = nn.Module() + core = attention.core_attention + core.attn_sink = nn.Parameter(torch.randn(2)) + if compression: + core.compressor = _compressor() + if compression == "csa": + core.indexer = nn.Module() + core.indexer.linear_wq_b = nn.Linear(256, 256, bias=False).bfloat16() + core.indexer.linear_weights_proj = nn.Linear(256, 2, bias=False).bfloat16() + core.indexer.compressor = _compressor() + return attention + + +@pytest.mark.parametrize("compression", [None, "csa", "hca"]) +@pytest.mark.parametrize("legacy", [False, True]) +def test_v4_attention_exports_compressors_and_optional_indexer(compression, legacy): + prefix = "layers.3.attn." + scorer = prefix + ( + "indexer.weights_proj.weight" if legacy else "indexer.scorer.weights_proj.weight" + ) + exporter = _exporter(source_keys=[scorer]) + attention = _attention(compression) + exporter._get_deepseek_v4_attention_state_dict(attention, 3) + state = exporter._state_dict + assert len(state) == {None: 8, "hca": 12, "csa": 18}[compression] + assert torch.equal(state[prefix + "wo_a.weight"], attention.linear_o_group_proj) + assert state[prefix + "attn_sink"].dtype == torch.float32 + if compression: + assert state[prefix + "compressor.ape"].dtype == torch.float32 + if compression == "csa": + assert scorer in state + assert state[scorer].dtype == torch.bfloat16 + + +@pytest.mark.parametrize("source_bias", [False, True]) +def test_v4_router_preserves_integer_hash_table_and_source_bias_policy(source_bias): + prefix = "layers.0.ffn.gate." + exporter = _exporter(source_keys=[prefix + "bias"] if source_bias else []) + router = nn.Linear(256, 4, bias=False).bfloat16() + router.register_buffer("expert_bias", torch.arange(4, dtype=torch.float32)) + router.register_buffer("tid2eid", torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)) + exporter._get_deepseek_v4_router_state_dict(router, 0) + assert (prefix + "bias" in exporter._state_dict) == source_bias + assert torch.equal(exporter._state_dict[prefix + "tid2eid"], router.tid2eid) + assert exporter._state_dict[prefix + "tid2eid"].dtype == torch.int32 + + +@pytest.mark.parametrize("folded", [False, True]) +def test_v4_expert_export_splits_gate_up_and_reuses_packed_codes(folded): + torch.manual_seed(6105) + module = nn.Linear(256, 512, bias=False).bfloat16() + module.config = SimpleNamespace(ffn_hidden_size=256) + mtq.quantize( + module, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": "iq2_xxs", "backend": "ggml", "block_sizes": {-1: 256}}, + }, + ], + "algorithm": None, + }, + ) + expected = module.weight_quantizer(module.weight).detach() + fmt = GGML_FORMAT_REGISTRY["iq2_xxs"] + packed = fmt.pack(module.weight, module.weight_quantizer).reshape(512, 1, 66) + exporter = _exporter(folded=folded) + exporter.rules["local_experts.linear_fc1"](module, 2, 17) + state = exporter._state_dict + if folded: + recipe = exporter._extract_quantizer_recipe_markers({0: state}) + assert len(recipe) == 6 # input/weight/output quantizers for each split projection + assert not exporter._quantizer_validation_failure + assert set(state) == {"layers.2.ffn.experts.17.w1.weight", "layers.2.ffn.experts.17.w3.weight"} + for key, start in (("w1", 0), ("w3", 256)): + weight = state[f"layers.2.ffn.experts.17.{key}.weight"] + if folded: + assert weight.dtype == torch.bfloat16 + assert torch.equal(weight, expected[start : start + 256]) + else: + assert weight.dtype == torch.uint8 + assert torch.equal(weight.reshape(256, 1, 66), packed[start : start + 256]) + + +@pytest.mark.parametrize("folded", [False, True]) +def test_v4_expert_export_matches_forward_when_only_up_weights_change(folded): + """An unchanged gate must follow the fused layer's updated quantization decision.""" + torch.manual_seed(41) + fmt = GGML_FORMAT_REGISTRY["iq2_xxs"] + packed, shape = fmt.quantize(torch.randn(4, 256, dtype=torch.bfloat16)) + module = nn.Linear(256, 4, bias=False).bfloat16() + module.config = SimpleNamespace(ffn_hidden_size=2) + mtq.quantize( + module, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": "iq2_xxs", "backend": "ggml", "block_sizes": {-1: 256}}, + }, + ], + "algorithm": None, + }, + ) + with torch.no_grad(): + module.weight.copy_(fmt.dequantize(packed, shape, dtype=torch.bfloat16)) + module.weight[2:].add_(0.125) + pin_packed_weight(module.weight_quantizer, "iq2_xxs", packed) + expected = module.weight_quantizer(module.weight).detach() + exporter = _exporter(folded=folded) + exporter.rules["local_experts.linear_fc1"](module, 0, 0) + weights = [ + exporter._state_dict[f"layers.0.ffn.experts.0.{part}.weight"] for part in ("w1", "w3") + ] + if not folded: + weights = [ + fmt.dequantize(value, torch.tensor([2, 256]), dtype=torch.bfloat16) for value in weights + ] + assert torch.equal(torch.cat(weights), expected) + + +@pytest.mark.parametrize("deepseek_v4", [False, True]) +def test_mhc_export_preserves_v4_fp32_projection_and_existing_dtype_policy(deepseek_v4): + exporter = _exporter() + exporter.rules["deepseek_v4"] = deepseek_v4 + projection = torch.tensor([[1.000001, 0.1234567]], dtype=torch.float32) + layer = SimpleNamespace( + inner_layer=SimpleNamespace(), + hyper_connection=SimpleNamespace( + mapping_proj=SimpleNamespace(weight=projection), + bias=torch.tensor([0.1234567]), + alpha_pre=torch.tensor([1.0]), + alpha_post=torch.tensor([2.0]), + alpha_res=torch.tensor([3.0]), + ), + ) + exporter._get_hyper_connection_state_dict(layer, 0) + state = exporter._state_dict + expected = projection if deepseek_v4 else projection.bfloat16() + assert state["layers.0.hc_ffn_fn"].dtype == expected.dtype + assert torch.equal(state["layers.0.hc_ffn_fn"], expected) + assert state["layers.0.hc_ffn_base"].dtype == torch.float32 + assert state["layers.0.hc_ffn_scale"].dtype == torch.float32 + + +def test_v4_export_preserves_fp32_head_controls(): + exporter = _exporter() + controls = { + name: torch.tensor([[1.000001, 0.1234567]], dtype=torch.float32) + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale") + } + exporter.model = SimpleNamespace( + pre_process=False, + post_process=True, + decoder=SimpleNamespace(layers=[], **controls), + ) + exporter._get_state_dict() + for name, expected in controls.items(): + assert exporter._state_dict[name].dtype == torch.float32 + assert torch.equal(exporter._state_dict[name], expected) + + +def test_v4_source_coverage_checks_root_layer_names(tmp_path): + source = tmp_path / "source" + exported = tmp_path / "exported" + source.mkdir() + exported.mkdir() + missing = "layers.0.attn.indexer.compressor.wkv.weight" + (source / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {missing: "x.safetensors"}}) + ) + (exported / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"layers.0.attn.wkv.weight": "x.safetensors"}}) + ) + exporter = _exporter() + exporter._hf_text_config = SimpleNamespace(num_hidden_layers=1) + exporter._src_num_hidden_layers = 1 + with pytest.raises(RuntimeError, match="Export dropped 1 tensor"): + exporter._verify_exported_keys(exported, source) + + +@pytest.mark.parametrize( + "missing", ["hc_head_scale", "layers.0.hc_attn_base", "layers.0.ffn.gate.tid2eid"] +) +def test_v4_source_coverage_checks_raw_parameters(tmp_path, missing): + source = tmp_path / "source" + exported = tmp_path / "exported" + source.mkdir() + exported.mkdir() + (source / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {missing: "x.safetensors"}}) + ) + (exported / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"layers.0.attn.wkv.weight": "x.safetensors"}}) + ) + exporter = _exporter() + with pytest.raises(RuntimeError, match="Export dropped V4 parameter"): + exporter._verify_exported_keys(exported, source) + + +@pytest.mark.parametrize(("num_layers", "mtp"), [(3, 0), (4, 1)]) +def test_v4_export_rejects_unpaired_layers_and_mtp(num_layers, mtp): + exporter = _exporter() + exporter.model = SimpleNamespace( + config=SimpleNamespace(num_layers=num_layers, mtp_num_layers=mtp) + ) + exporter._hf_text_config = SimpleNamespace(num_nextn_predict_layers=mtp) + with pytest.raises((ValueError, NotImplementedError), match=r"paired|main-only"): + exporter._init_deepseek_v4_export() diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 136ce92f6ed..c3506f1c7c9 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -283,3 +283,22 @@ def test_format_fake_quant_rejects_another_formats_quantizer(num_bits): GGML_FORMAT_REGISTRY[num_bits].fake_quant( torch.ones(1, block_size), SimpleNamespace(num_bits=other) ) + + +@pytest.mark.parametrize("num_bits", FORMAT_NAMES) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fake_quant_matches_packed_decode_exactly_and_preserves_gradient(num_bits, dtype): + """Subtract/add STE rounding must not change the values exported as packed weights.""" + torch.manual_seed(6105) + weight = torch.randn(4, 256, dtype=dtype, requires_grad=True) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + ) + fmt = GGML_FORMAT_REGISTRY[num_bits] + actual = quantizer(weight) + packed = fmt.pack(weight, quantizer) + decoded = fmt.dequantize(packed, torch.tensor(weight.shape), dtype=dtype) + assert torch.equal(actual, decoded) + gradient = torch.randn_like(weight) + actual.backward(gradient) + assert torch.equal(weight.grad, gradient)