diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml index 3cf948399f..c1fd5daf0e 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml @@ -124,6 +124,7 @@ validation_dataloader: optimizer: _target_: torch.optim.AdamW + fused: true betas: - 0.9 - 0.95 diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml index 3e99c07f81..d4b7ef5f86 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml @@ -59,7 +59,7 @@ model: load_base_model: true backend: _target_: nemo_automodel.components.models.common.BackendConfig - attn: sdpa + attn: tilelang linear: torch rms_norm: torch_fp32 rope_fusion: false @@ -122,6 +122,7 @@ validation_dataloader: optimizer: _target_: torch.optim.AdamW + fused: true betas: [0.9, 0.95] eps: 1e-8 lr: 1e-5 diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml new file mode 100644 index 0000000000..a5252e9b90 --- /dev/null +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml @@ -0,0 +1,168 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +# LoRA fine-tuning of deepseek-ai/DeepSeek-V4-Flash on HellaSwag with the frozen +# routed experts kept resident in mxfp4 (fp4-e2m1 + e8m0 block scales) and +# dequantized on the fly in the grouped-GEMM forward/backward. +# +# Only the experts are quantized; every other weight (MLA attention, dense MLP, +# embeddings, MTP, lm_head) stays bf16. The routed experts are ~90%+ of the +# parameters, so packing them to ~4 bits cuts steady-state base-weight memory by +# roughly 4x relative to the bf16 LoRA recipe. +# +# v1 constraint: mxfp4 experts require the torch_mm GroupedExperts backend, so +# this recipe uses dispatcher=torch (NOT deepep). DeepEP + mxfp4 is a follow-up. + +recipe: TrainFinetuneRecipeForNextTokenPrediction + +seed: 1234 + +step_scheduler: + global_batch_size: 128 + local_batch_size: 1 + ckpt_every_steps: 500 + val_every_steps: 500 + gc_every_steps: 10 + num_epochs: 1 + max_steps: 100 + +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 1 + pp_size: 1 + # Single node, 8x H200: ep_size must divide dp_size*cp_size (= world_size here). + ep_size: 8 + + sequence_parallel: false + # Recommended on: the 43-layer activation footprint, not the (now-packed) expert + # weights, is the steady-state pressure once experts are mxfp4. + activation_checkpointing: true + + moe: + reshard_after_forward: false + wrap_outer_model: false + +dist_env: + backend: nccl + timeout_minutes: 30 + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config + config: + _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + name_or_path: deepseek-ai/DeepSeek-V4-Flash + num_nextn_predict_layers: 0 + # Required by FusedLinearCrossEntropy: the trainer calls the model with + # logits_to_keep=1 and reads out.hidden_states to apply the fused lm_head + CE + # without materializing full [seq, vocab] logits. The forward only populates + # hidden_states when config.output_hidden_states is set. + output_hidden_states: true + trust_remote_code: false + load_base_model: true + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: tilelang + linear: torch + rms_norm: torch_fp32 + rope_fusion: false + # mxfp4 experts run on the torch_mm grouped-GEMM path (DeepEP not yet supported). + dispatcher: torch + experts: torch_mm + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: + # attention + - "*wq_a" + - "*wq_b" + - "*wkv" + - "*wo_b" + # routed experts: GroupedExperts -> GroupedExpertsLoRAMXFP4 + # (frozen base stays mxfp4-packed; only the LoRA adapters train) + - "*mlp.experts" + # shared experts: dense MLP nn.Linear (base stays bf16, LoRA adapters bf16) + - "*shared_experts.gate_proj" + - "*shared_experts.up_proj" + - "*shared_experts.down_proj" + dim: 8 + alpha: 32 + use_triton: True + # fp32 LoRA adapters: fp32 master weights + fp32 AdamW state for stability (small-update + # swamping), while FSDP2's param_dtype=bf16 still runs the adapter matmuls in bf16. Adapters + # are tiny so the fp32 optimizer-state cost is negligible. Matches HF PEFT's default. + lora_dtype: float32 + # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly. + expert_weight_format: mxfp4 + +checkpoint: + enabled: false + # The DSV4-Flash checkpoint is quantized: fp4 routed experts + fp8 non-expert + # projections. This enables quant-aware loading -- the fp8 non-expert weights are + # dequantized to bf16, while the fp4 experts are passed through packed (mxfp4, + # never materialized in bf16). Required to load this checkpoint. + dequantize_base_checkpoint: true + +loss_fn: + # Fused linear cross-entropy (Apple cut_cross_entropy): applies the lm_head + CE without + # materializing the [seq, vocab=129280] logits, removing the ~16 GiB fp32 logits spike and + # raising the single-node context ceiling. Uses the DSV4 forward's logits_to_keep / + # hidden_states path; loss stays fp32-quality (cut_cross_entropy accumulates the logsumexp in + # fp32 internally). Swap back to masked_ce.MaskedCrossEntropy if cut_cross_entropy is unavailable. + _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: train + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + +packed_sequence: + packed_sequence_size: 0 + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + pad_seq_len_divisible: 64 + shuffle: true + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: validation + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + pad_seq_len_divisible: 64 + shuffle: false + drop_last: true + +optimizer: + _target_: torch.optim.AdamW + fused: true + betas: [0.9, 0.95] + eps: 1e-8 + lr: 1e-5 + weight_decay: 0.1 diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml new file mode 100644 index 0000000000..1134fdbc57 --- /dev/null +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml @@ -0,0 +1,196 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +# LoRA fine-tuning of deepseek-ai/DeepSeek-V4-Flash on HellaSwag with the frozen +# routed experts kept resident in mxfp4 (fp4-e2m1 + e8m0 block scales) and +# dequantized on the fly in the grouped-GEMM forward/backward, dispatched with HybridEP. +# +# This is the HybridEP counterpart of deepseek_v4_flash_hellaswag_lora_mxfp4.yaml. +# mxfp4 quantizes only the expert *weights* (local to each EP rank); HybridEP only +# governs token *dispatch/combine* (fused all-to-all of bf16 activations + fp32 probs). +# They intersect only at the two post-dispatch grouped GEMMs, which read the packed +# weights via MXFP4GroupedMM. Unlike the torch dispatcher (which all-gathers every +# token to every rank), HybridEP routes each rank only the tokens for its local experts, +# so this is the scalable EP path for DSV4 at EP=8/64. +# On an 8-GPU NVLink node, set NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8. +# +# Only the experts are quantized; every other weight (MLA attention, dense MLP, +# embeddings, MTP, lm_head) stays bf16. The routed experts are ~90%+ of the +# parameters, so packing them to ~4 bits cuts steady-state base-weight memory by +# roughly 4x relative to the bf16 LoRA recipe. +# +# Constraint: mxfp4 experts require the torch_mm experts backend (the grouped_gemm +# 'gmm' path has no packed variant), so this recipe pairs dispatcher=hybridep with +# experts=torch_mm. + +recipe: TrainFinetuneRecipeForNextTokenPrediction + +seed: 1234 + +step_scheduler: + global_batch_size: 128 + local_batch_size: 1 + ckpt_every_steps: 500 + val_every_steps: 500 + gc_every_steps: 10 + num_epochs: 1 + max_steps: 100 + +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 1 + pp_size: 1 + # Single node, 8x H200: ep_size must divide dp_size*cp_size (= world_size here). + ep_size: 8 + + sequence_parallel: false + # Recommended on: the 43-layer activation footprint, not the (now-packed) expert + # weights, is the steady-state pressure once experts are mxfp4. + activation_checkpointing: true + + moe: + reshard_after_forward: false + wrap_outer_model: false + +dist_env: + backend: nccl + timeout_minutes: 30 + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config + config: + _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + name_or_path: deepseek-ai/DeepSeek-V4-Flash + num_nextn_predict_layers: 0 + # Required by FusedLinearCrossEntropy: the trainer calls the model with + # logits_to_keep=1 and reads out.hidden_states to apply the fused lm_head + CE + # without materializing full [seq, vocab] logits. The forward only populates + # hidden_states when config.output_hidden_states is set. + output_hidden_states: true + trust_remote_code: false + load_base_model: true + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: tilelang + linear: torch + rms_norm: torch_fp32 + rope_fusion: false + # mxfp4 experts run on the torch_mm grouped-GEMM path, dispatched with HybridEP. + dispatcher: hybridep + experts: torch_mm + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: + # Direct attention projections only. The more permissive "*wkv"/"*wq_b" + # patterns also match the FP32 HCA compressor/indexer projections. + - "*.self_attn.wq_a" + - "*.self_attn.wq_b" + - "*.self_attn.wkv" + - "*.self_attn.wo_b" + # routed experts: GroupedExpertsDeepEP -> GroupedExpertsDeepEPLoRAMXFP4 + # (frozen base stays mxfp4-packed; only the LoRA adapters train) + - "*mlp.experts" + # shared experts: dense MLP nn.Linear (base stays bf16, LoRA adapters bf16) + - "*shared_experts.gate_proj" + - "*shared_experts.up_proj" + - "*shared_experts.down_proj" + dim: 8 + alpha: 32 + use_triton: True + # BF16 adapters match the base compute dtype and keep every patched linear's + # FSDP2 storage dtype uniform. FP32 adapters would require separate parameter-owning + # modules for the adapter weights under the current FSDP2 dtype isolation rules. + lora_dtype: bfloat16 + # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly. + expert_weight_format: mxfp4 + +checkpoint: + enabled: false + # The DSV4-Flash checkpoint is quantized: fp4 routed experts + fp8 non-expert + # projections. This enables quant-aware loading -- the fp8 non-expert weights are + # dequantized to bf16, while the fp4 experts are passed through packed (mxfp4, + # never materialized in bf16). Required to load this checkpoint. + dequantize_base_checkpoint: true + +loss_fn: + # Fused linear cross-entropy (Apple cut_cross_entropy): applies the lm_head + CE without + # materializing the [seq, vocab=129280] logits, removing the ~16 GiB fp32 logits spike and + # raising the single-node context ceiling (~30k -> ~36-38k tokens on one 8xH200). Uses the + # DSV4 forward's logits_to_keep / hidden_states path; loss stays fp32-quality (cut_cross_entropy + # accumulates the logsumexp in fp32 internally). Swap back to masked_ce.MaskedCrossEntropy if + # cut_cross_entropy is unavailable. + _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: train + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + +packed_sequence: + packed_sequence_size: 0 + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + pad_seq_len_divisible: 64 + shuffle: true + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag + path_or_dataset: rowan/hellaswag + split: validation + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash + +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + pad_seq_len_divisible: 64 + shuffle: false + drop_last: true + +optimizer: + # Keep the BF16 adapter parameters/checkpoint while TE FusedAdam reconstructs + # FP32 master weights from the BF16 values plus their stored 16-bit remainders. + _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam + betas: [0.9, 0.95] + eps: 1e-8 + lr: 1e-5 + weight_decay: 0.1 + adam_w_mode: true + bias_correction: true + master_weights: true + master_weight_dtype: torch.float32 + store_param_remainders: true + exp_avg_dtype: torch.float32 + exp_avg_sq_dtype: torch.float32 + +wandb: + enable: false + project: automodel-dsv4 + name: deepseek-v4-flash-mxfp4-lora-hybridep-te-fusedadam + group: pr-2548 + tags: [deepseek-v4-flash, mxfp4, lora, hybridep, te-fusedadam] + mode: online diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 920555963b..d2be37be86 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -95,6 +95,15 @@ def _ensure_tied_lm_heads(model) -> None: ensure_tied_lm_head(model_part) +def _call_model_hook(model, hook_name: str, *args) -> None: + """Call an explicitly implemented hook on each model part, if present.""" + model_parts = model.parts if hasattr(model, "parts") else [model] + for model_part in model_parts: + hook = getattr(type(model_part), hook_name, None) + if hook is not None: + hook(model_part, *args) + + def _safe_moe_tp_parts(model) -> list[torch.nn.Module]: """Return model parts using the conservative custom-MoE TP plan.""" model_parts = model.parts if hasattr(model, "parts") else [model] @@ -158,6 +167,7 @@ def _apply_peft_and_lower_precision( peft_config.use_triton = False # Skip freeze here - will do global freeze after checkpoint loading apply_lora_to_linear_modules(model, peft_config, quantization_config=quantization_config, skip_freeze=True) + _call_model_hook(model, "prepare_peft_checkpoint_load", peft_config) # FP8 if fp8_config is not None: @@ -790,6 +800,9 @@ def apply_model_infrastructure( "check freeze_config and the PEFT configuration." ) + if peft_config is not None: + _call_model_hook(model, "finalize_peft_checkpoint_load", peft_config) + if autopipeline is None: print_trainable_parameters(model) # Once model's been sharded # Ensure model is on the correct device. diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index b07c1bd5cc..fe516794ed 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -22,7 +22,14 @@ import torch.nn.functional as F from torch.distributed.tensor import DTensor -from nemo_automodel.components._peft.lora_experts import GroupedExpertsDeepEPLoRA, GroupedExpertsLoRA +from nemo_automodel.components._peft.lora_experts import ( + GroupedExpertsDeepEPLoRA, + GroupedExpertsLoRA, +) +from nemo_automodel.components._peft.lora_experts_mxfp4 import ( + GroupedExpertsDeepEPLoRAMXFP4, + GroupedExpertsLoRAMXFP4, +) from nemo_automodel.components._peft.lora_kernel import ( lora_da_dx_update_wrapper, lora_db_update_wrapper, @@ -31,6 +38,7 @@ from nemo_automodel.components._peft.module_matcher import ModuleMatcher from nemo_automodel.components.moe.layers import GroupedExperts, GroupedExpertsDeepEP, GroupedExpertsTE from nemo_automodel.components.moe.mok_experts import GroupedExpertsMoK +from nemo_automodel.components.moe.quantized_experts import GroupedExpertsMXFP4, MXFP4ExpertStorageMixin from nemo_automodel.shared.import_utils import safe_import, safe_import_te from nemo_automodel.shared.tp_linear import tp_linear_forward from nemo_automodel.shared.utils import dtype_from_str @@ -57,6 +65,10 @@ class PeftConfig: use_memory_efficient_lora: bool = True use_triton: bool = False moe_rank_scaling: bool = False + # "mxfp4" keeps frozen MoE expert base weights packed as fp4-e2m1 + e8m0 block + # scales, dequantized on the fly in forward/backward. Experts only; requires + # the torch_mm experts backend. + expert_weight_format: Literal["bf16", "mxfp4"] = "bf16" def to_dict(self): return self.__dict__.copy() @@ -77,6 +89,7 @@ def from_dict(cls, d: dict[str, Any]): use_memory_efficient_lora=d.get("use_memory_efficient_lora", True), use_triton=d.get("use_triton", False), moe_rank_scaling=d.get("moe_rank_scaling", False), + expert_weight_format=d.get("expert_weight_format", "bf16"), ) @@ -526,6 +539,8 @@ def patch_moe_module( alpha=32, lora_A_init_method="xavier", lora_dtype=None, + expert_weight_format="bf16", + passthrough=False, ): """ Patches a custom MoE module (GroupedExperts or GroupedExpertsDeepEP) with LoRA. @@ -536,30 +551,34 @@ def patch_moe_module( alpha (int, optional): LoRA scaling factor. Defaults to 32. lora_A_init_method (str, optional): Initialization method for LoRA A matrix. Defaults to "xavier". lora_dtype (torch.dtype or str, optional): Data type for LoRA weights. Defaults to None. + expert_weight_format (str, optional): "bf16" keeps frozen base expert weights in + floating point; "mxfp4" keeps them packed as fp4-e2m1 + e8m0 block scales with + on-the-fly dequantization. Defaults to "bf16". + passthrough (bool, optional): Only used with expert_weight_format="mxfp4". When True, + the frozen base is registered as packed placeholders at init so a packed fp4 + checkpoint loads straight in (no bf16 expert materialization). Defaults to False. Returns: - nn.Module: The LoRA-wrapped MoE module (GroupedExpertsLoRA or GroupedExpertsDeepEPLoRA). + nn.Module: The LoRA-wrapped MoE module. """ + if expert_weight_format not in ("bf16", "mxfp4"): + raise ValueError(f"Unsupported expert_weight_format: {expert_weight_format}") + common = dict(lora_dim=dim, alpha=alpha, lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype) + mxfp4 = expert_weight_format == "mxfp4" if isinstance(orig_module, GroupedExpertsMoK): raise NotImplementedError("LoRA is not supported for Mixture-of-Kittens expert modules.") if isinstance(orig_module, GroupedExpertsTE): raise NotImplementedError("LoRA is not supported for Transformer Engine (TE) expert modules.") elif isinstance(orig_module, GroupedExpertsDeepEP): - new_module = GroupedExpertsDeepEPLoRA( - orig_module, - lora_dim=dim, - alpha=alpha, - lora_A_init_method=lora_A_init_method, - lora_dtype=lora_dtype, - ) + if mxfp4: + new_module = GroupedExpertsDeepEPLoRAMXFP4(orig_module, passthrough=passthrough, **common) + else: + new_module = GroupedExpertsDeepEPLoRA(orig_module, **common) elif isinstance(orig_module, GroupedExperts): - new_module = GroupedExpertsLoRA( - orig_module, - lora_dim=dim, - alpha=alpha, - lora_A_init_method=lora_A_init_method, - lora_dtype=lora_dtype, - ) + if mxfp4: + new_module = GroupedExpertsLoRAMXFP4(orig_module, passthrough=passthrough, **common) + else: + new_module = GroupedExpertsLoRA(orig_module, **common) else: raise NotImplementedError(f"Unsupported MoE module type: {type(orig_module)}") @@ -646,13 +665,18 @@ def apply_lora_to_linear_modules( moe_dim, ) - # Replace the module in the model + # Replace the module in the model. For mxfp4, build LoRA experts in + # passthrough mode (packed base placeholders) so the packed fp4 checkpoint + # loads straight in — matching the model-wide packed adapter mode set in + # _apply_peft_and_lower_precision and avoiding bf16 expert materialization. new_module = patch_moe_module( module, dim=moe_dim, alpha=peft_config.alpha, lora_A_init_method=peft_config.lora_A_init, lora_dtype=lora_dtype, + expert_weight_format=peft_config.expert_weight_format, + passthrough=(peft_config.expert_weight_format == "mxfp4"), ) # Find parent and replace @@ -699,6 +723,84 @@ def apply_lora_to_linear_modules( return num_modules_matched +def convert_frozen_experts_to_mxfp4(model: nn.Module, passthrough: bool = False) -> int: + """Swap frozen ``GroupedExperts`` modules to mxfp4-resident ``GroupedExpertsMXFP4``. + + Applies to routed experts that are NOT LoRA-targeted (those that received a + LoRA adapter are already ``GroupedExpertsLoRAMXFP4``). This is the path that + delivers the storage win for the common case of LoRA on attention with frozen + experts. Must be called with the base weights frozen. + + Args: + model: Model to convert in place. + passthrough: When True, build the new modules in packed-storage mode at + init (no bf16 weights) so a packed fp4 checkpoint loads straight in, + capping the load-time peak. The state-dict adapter must also be put in + ``expert_storage_format='mxfp4'`` so it emits packed keys. When False + (default), weights load as bf16 and pack after load (higher load peak, + but works with any checkpoint and is the validated path). + + Returns: + Number of expert modules converted. + """ + # Import here to avoid a hard dependency at module import time. + from nemo_automodel.components.moe.experts import GroupedExpertsDeepEP, GroupedExpertsTE + from nemo_automodel.components.moe.quantized_experts import GroupedExpertsDeepEPMXFP4 + + # Exact-type → mxfp4-resident replacement. Exact `type(...) is` (not isinstance) so the + # LoRA-on-experts subclasses (already MXFP4ExpertStorageMixin, skipped above) and any + # bf16 LoRA experts are left untouched — only genuinely frozen modules are converted. + frozen_conversions = { + GroupedExperts: GroupedExpertsMXFP4, + GroupedExpertsDeepEP: GroupedExpertsDeepEPMXFP4, + } + + num_converted = 0 + unsupported = 0 + for name, module in list(model.named_modules()): + # Already mxfp4-resident (frozen or LoRA-targeted) — skip. + if isinstance(module, MXFP4ExpertStorageMixin): + continue + if isinstance(module, GroupedExpertsTE): + unsupported += 1 + continue + new_cls = frozen_conversions.get(type(module)) + if new_cls is not None: + new_module = new_cls(module, passthrough=passthrough) + parent_name, _, child_name = name.rpartition(".") + parent = model.get_submodule(parent_name) if parent_name else model + setattr(parent, child_name, new_module) + num_converted += 1 + + if unsupported: + logger.warning( + "expert_weight_format='mxfp4' skipped %d Transformer Engine expert module(s); TE experts have no " + "packed variant. Use backend.experts='torch_mm' (with backend.dispatcher='torch' or 'deepep').", + unsupported, + ) + return num_converted + + +def pack_mxfp4_expert_base_weights(model: nn.Module) -> int: + """Pack any deferred mxfp4-resident expert modules after base weights are loaded. + + Both ``GroupedExpertsLoRAMXFP4`` and frozen ``GroupedExpertsMXFP4`` modules + created on the meta device defer packing until their base weights are + materialized from the checkpoint. Call this after checkpoint load to convert + them; modules pack one at a time so the bf16 weights of at most one expert + module coexist with their packed copy. + + Returns: + Number of modules packed. + """ + num_packed = 0 + for module in model.modules(): + if isinstance(module, MXFP4ExpertStorageMixin) and not module._mxfp4_resident: + module.pack_base_weights() + num_packed += 1 + return num_packed + + class LoRATritonFunction(torch.autograd.Function): """ Autograd function that avoids saving the LoRA A activation. diff --git a/nemo_automodel/components/_peft/lora_experts_mxfp4.py b/nemo_automodel/components/_peft/lora_experts_mxfp4.py new file mode 100644 index 0000000000..68dc8d2c55 --- /dev/null +++ b/nemo_automodel/components/_peft/lora_experts_mxfp4.py @@ -0,0 +1,302 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""MXFP4-resident expert LoRA implementations.""" + +import torch +from torch.distributed.tensor import DTensor, Partial, Shard + +from nemo_automodel.components._peft.lora_experts import ( + GroupedExpertsDeepEPLoRA, + GroupedExpertsLoRA, + _to_local, +) +from nemo_automodel.components.moe.experts import ( + GroupedExperts, + GroupedExpertsDeepEP, + _apply_bias, + _permute_tokens_for_grouped_mm, +) +from nemo_automodel.components.moe.quantized_experts import MXFP4ExpertStorageMixin + + +class GroupedExpertsLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsLoRA): + """GroupedExperts + LoRA with the frozen base weights resident in packed mxfp4. + + The base gate/up and down projections are stored as packed fp4-e2m1 int8 plus + ``float8_e8m0fnu`` block scales (checkpoint orientation ``[n_experts, out_dim, + in_dim]``) and dequantized on the fly inside ``MXFP4GroupedMM`` during forward + and backward (see ``MXFP4ExpertStorageMixin``). Only the LoRA adapters (and + optional expert biases) remain in floating point. + + Two load paths: + - ``passthrough=False`` (default): packing is deferred. The base loads as bf16 and + is packed after the checkpoint load (``pack_base_weights()``). Works with any + checkpoint, but materializes bf16 experts at load (high peak). + - ``passthrough=True``: register packed base placeholders at init (no bf16 storage) + so a packed fp4 checkpoint loads straight into them via the adapter's packed path, + never materializing bf16 experts. The scale-out path for the full DeepSeek-V4-Flash. + """ + + def __init__( + self, + orig_module: GroupedExperts, + lora_dim=8, + alpha=32, + lora_A_init_method="xavier", + lora_dtype=None, + passthrough=False, + ): + super().__init__( + orig_module, + lora_dim=lora_dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) + if passthrough: + # Swap the frozen bf16 base placeholders (from super().__init__) for packed + # mxfp4 placeholders; the LoRA adapters just built are left untouched. A packed + # checkpoint then loads straight in with no bf16 expert materialization. + self._init_packed_placeholders() + else: + self._init_mxfp4_storage() + + def forward(self, x: torch.Tensor, token_mask: torch.Tensor, weights: torch.Tensor, indices: torch.Tensor): + """Forward pass with mxfp4 base weights and LoRA injection. + + Mirrors GroupedExpertsLoRA.forward, replacing the base grouped GEMMs with + MXFP4GroupedMM over the packed weights. Falls back to the parent (bf16) + path while packing is still deferred. + """ + if not self._mxfp4_resident: + return super().forward(x, token_mask, weights, indices) + + assert not isinstance(x, DTensor) + input_dtype = x.dtype + + if isinstance(self.gate_and_up_projs_packed, DTensor): + ep_mesh = self.gate_and_up_projs_packed.device_mesh + assert ep_mesh is not None + assert ep_mesh.ndim == 1 + ep_size = ep_mesh.size() + ep_rank = ep_mesh.get_local_rank() + else: + ep_mesh = None + ep_size = 1 + ep_rank = 0 + + assert self.n_routed_experts % ep_size == 0 + + if ep_size > 1: + x = DTensor.from_local(x, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor( + grad_placements=[Partial()] + ) + weights = DTensor.from_local(weights.float(), device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor( + grad_placements=[Partial()] + ) + indices = DTensor.from_local(indices, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor() + token_mask = DTensor.from_local(token_mask, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor() + + n_local_experts = self.n_routed_experts // ep_size + experts_start_idx = ep_rank * n_local_experts + + y = self._forward_grouped_mm_mxfp4(x, token_mask, weights, indices, n_local_experts, experts_start_idx) + + if ep_size > 1: + y = DTensor.from_local(y, device_mesh=ep_mesh, placements=[Partial()]) + y = y.redistribute(placements=[Shard(0)]).to_local() + + return y.to(input_dtype) + + def _forward_grouped_mm_mxfp4(self, x, token_mask, weights, indices, n_local_experts, experts_start_idx): + """Grouped GEMM forward path over packed mxfp4 base weights with LoRA injection.""" + sorted_token_ids, sorted_weights, tokens_per_expert, offs = _permute_tokens_for_grouped_mm( + indices, + weights, + token_mask, + n_local_experts, + experts_start_idx, + ) + + # Match the activation dtype for the LoRA grouped GEMMs. The frozen base is + # dequantized to x.dtype inside MXFP4GroupedMM, but the adapters may be a + # different dtype (GroupedExperts allocates its base — hence the adapter dtype — + # as fp32 when no backend dtype is set), which would mismatch torch._grouped_mm. + lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A).to(x.dtype) + lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B).to(x.dtype) + lora_down_A = _to_local(self.lora_down_A).to(x.dtype) + lora_down_B = _to_local(self.lora_down_B).to(x.dtype) + + y = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + + if tokens_per_expert.sum() > 0: + permuted_x = x[sorted_token_ids] + permuted_probs = sorted_weights.unsqueeze(-1) + + if self.expert_bias: + gate_up_proj_bias = _to_local(self.gate_up_proj_bias) + down_proj_bias = _to_local(self.down_proj_bias) + + # Gate+Up projection + LoRA + output1 = self._mxfp4_base_mm(permuted_x, "gate_and_up_projs", offs) + lora_out1_A = torch._grouped_mm(permuted_x, lora_gate_and_up_A, offs=offs) + lora_out1 = torch._grouped_mm(lora_out1_A, lora_gate_and_up_B, offs=offs) + output1 = output1 + lora_out1 * self.scale + + if self.expert_bias: + output1 = _apply_bias(output1, gate_up_proj_bias, tokens_per_expert) + + output1 = self.expert_activation_grouped(output1, permuted_probs) + + # Down projection + LoRA + output2 = self._mxfp4_base_mm(output1, "down_projs", offs) + lora_out2_A = torch._grouped_mm(output1, lora_down_A, offs=offs) + lora_out2 = torch._grouped_mm(lora_out2_A, lora_down_B, offs=offs) + output2 = output2 + lora_out2 * self.scale + + if self.expert_bias: + output2 = _apply_bias(output2, down_proj_bias, tokens_per_expert, permuted_probs) + + scatter_ids = sorted_token_ids.unsqueeze(1).expand_as(output2) + y.scatter_add_(0, scatter_ids, output2.float()) + else: + # Dummy computation for gradient flow; dequantize only expert 0. + gate_up_w0 = self._mxfp4_dequant_expert0("gate_and_up_projs", x.dtype) + down_w0 = self._mxfp4_dequant_expert0("down_projs", x.dtype) + output1 = torch.matmul(x[0] * 0, gate_up_w0) + output1 = ( + output1 + + torch.matmul(torch.matmul(x[0] * 0, lora_gate_and_up_A[0]), lora_gate_and_up_B[0]) * self.scale + ) + output1_ = self.expert_activation_grouped(output1, weights[0, 0, None].unsqueeze(0)) + output2 = torch.matmul(output1_, down_w0) + output2 = output2 + torch.matmul(torch.matmul(output1_ * 0, lora_down_A[0]), lora_down_B[0]) * self.scale + y[0] += output2[0] + + return y + + +class GroupedExpertsDeepEPLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsDeepEPLoRA): + """GroupedExpertsDeepEP + LoRA with the frozen base weights resident in packed mxfp4. + + The DeepEP fused all-to-all token dispatch is reused unchanged from + ``GroupedExpertsDeepEPLoRA``; only the two frozen base grouped GEMMs read the packed + fp4-e2m1 + e8m0 base weights via ``MXFP4GroupedMM`` instead of bf16. The LoRA A/B + adapters (and optional expert biases) stay in floating point and their grouped GEMMs + are unchanged. + + Requires the torch_mm experts backend; the grouped_gemm (``gmm``) path has no packed + variant. When constructed from a module whose base weights are still on the meta device, + packing is deferred until ``pack_base_weights()`` runs after the checkpoint is loaded. + """ + + def __init__( + self, + orig_module: GroupedExpertsDeepEP, + lora_dim=8, + alpha=32, + lora_A_init_method="xavier", + lora_dtype=None, + passthrough=False, + ): + super().__init__( + orig_module, + lora_dim=lora_dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) + if passthrough: + # Packed base placeholders (no bf16) so a packed fp4 checkpoint loads straight + # in; the LoRA adapters from super().__init__ are untouched. See + # GroupedExpertsLoRAMXFP4 for the passthrough vs deferred distinction. + self._init_packed_placeholders() + else: + self._init_mxfp4_storage() + + def forward( + self, + x: torch.Tensor, + token_mask: torch.Tensor, + weights: torch.Tensor, + indices: torch.Tensor, + ): + """Forward with mxfp4 base weights, DeepEP dispatch, and LoRA injection. + + Mirrors ``GroupedExpertsDeepEPLoRA.forward`` (torch_mm branch), replacing the base + grouped GEMMs with ``MXFP4GroupedMM`` over the packed weights. Falls back to the + bf16 parent while packing is still deferred. + """ + if not self._mxfp4_resident: + return super().forward(x, token_mask, weights, indices) + + assert not isinstance(x, DTensor) + assert self.use_torch_mm, "mxfp4-resident DeepEP experts require the torch_mm experts backend." + assert self.n_routed_experts % self.ep_size == 0 + + indices = indices.masked_fill(~token_mask.unsqueeze(-1), -1) + (permuted_local_hidden_states, tokens_per_expert, permuted_probs) = self.token_dispatcher.token_permutation2( + hidden_states=x, + num_local_tokens=x.size(0), + token_probs=weights, + token_indices=indices, + ) + permuted_probs = permuted_probs.unsqueeze(-1) + + # Match the activation dtype for the LoRA grouped GEMMs (the base dequantizes to + # x.dtype inside MXFP4GroupedMM; adapters may be fp32 — see GroupedExpertsLoRAMXFP4). + lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A).to(x.dtype) + lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B).to(x.dtype) + lora_down_A = _to_local(self.lora_down_A).to(x.dtype) + lora_down_B = _to_local(self.lora_down_B).to(x.dtype) + + if torch.count_nonzero(tokens_per_expert) > 0: + tokens_per_expert_gpu = tokens_per_expert.to(device=permuted_local_hidden_states.device, non_blocking=True) + offs = tokens_per_expert_gpu.cumsum(dim=0).to(torch.int32) + + # Gate+Up projection (mxfp4 base) + LoRA + output1 = self._mxfp4_base_mm(permuted_local_hidden_states, "gate_and_up_projs", offs) + lora_out1_A = torch._grouped_mm(permuted_local_hidden_states, lora_gate_and_up_A, offs=offs) + lora_out1 = torch._grouped_mm(lora_out1_A, lora_gate_and_up_B, offs=offs) + output1 = output1 + lora_out1 * self.scale + + if self.expert_bias: + output1 = _apply_bias(output1, _to_local(self.gate_up_proj_bias), tokens_per_expert) + + output1 = self.expert_activation(output1, permuted_probs) + + # Down projection (mxfp4 base) + LoRA + output2 = self._mxfp4_base_mm(output1, "down_projs", offs) + lora_out2_A = torch._grouped_mm(output1, lora_down_A, offs=offs) + lora_out2 = torch._grouped_mm(lora_out2_A, lora_down_B, offs=offs) + output2 = output2 + lora_out2 * self.scale + + if self.expert_bias: + output2 = _apply_bias(output2, _to_local(self.down_proj_bias), tokens_per_expert, permuted_probs) + else: + # Dummy computation for gradient flow; dequantize only expert 0. + gate_up_w0 = self._mxfp4_dequant_expert0("gate_and_up_projs", x.dtype) + down_w0 = self._mxfp4_dequant_expert0("down_projs", x.dtype) + output1 = torch.matmul(x[0] * 0, gate_up_w0) + output1 = ( + output1 + + torch.matmul(torch.matmul(x[0] * 0, lora_gate_and_up_A[0]), lora_gate_and_up_B[0]) * self.scale + ) + output1_ = self.expert_activation(output1, permuted_probs) + output2 = torch.matmul(output1_, down_w0) + output2 = output2 + torch.matmul(torch.matmul(output1_ * 0, lora_down_A[0]), lora_down_B[0]) * self.scale + + y = self.token_dispatcher.token_unpermutation(output2) + return y diff --git a/nemo_automodel/components/loss/linear_ce.py b/nemo_automodel/components/loss/linear_ce.py index 4c6552d3d1..f4e42301e2 100644 --- a/nemo_automodel/components/loss/linear_ce.py +++ b/nemo_automodel/components/loss/linear_ce.py @@ -247,6 +247,13 @@ def forward( if self.logit_softcapping == 0: self.logit_softcapping = None + # cut_cross_entropy fuses e @ c.T in the input precision (with fp32 logsumexp + # accumulation internally), so e (hidden_states) and c (lm_weight) must share a + # dtype. Cast the classifier weight to the activation dtype to support models whose + # lm_head is kept in fp32 while activations are bf16 (e.g. DeepSeek V4). + if lm_weight.dtype != hidden_states.dtype: + lm_weight = lm_weight.to(hidden_states.dtype) + # Compute loss with shift=False to match PyTorch behavior # Set filter_eps=None to avoid any token filtering loss = linear_cross_entropy( diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index 7b09407a86..79ed0c2767 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -429,7 +429,9 @@ def forward( scores = F.linear(x.float(), self.weight.float()) if self.score_func == "sqrtsoftplus": - scores = F.softplus(scores).sqrt() + # clamp_min: softplus underflows to 0.0 for very negative logits and sqrt'(0)=inf + # makes the backward NaN; bound it here with a negligible forward change. + scores = F.softplus(scores).clamp_min(1e-12).sqrt() elif self.score_func == "sigmoid": scores = scores.sigmoid() else: @@ -1204,6 +1206,18 @@ def get_output_embeddings(self): def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings + def prepare_peft_checkpoint_load(self, peft_config) -> None: + """Configure model-owned PEFT storage before loading a checkpoint.""" + from nemo_automodel.components.models.deepseek_v4.mxfp4_peft import prepare_mxfp4_peft_checkpoint_load + + prepare_mxfp4_peft_checkpoint_load(self, peft_config) + + def finalize_peft_checkpoint_load(self, peft_config) -> None: + """Finalize model-owned PEFT storage after loading a checkpoint.""" + from nemo_automodel.components.models.deepseek_v4.mxfp4_peft import finalize_mxfp4_peft_checkpoint_load + + finalize_mxfp4_peft_checkpoint_load(self, peft_config) + def customize_pipeline_stage_modules( self, module_names_per_stage: list[list[str]], diff --git a/nemo_automodel/components/models/deepseek_v4/mxfp4_peft.py b/nemo_automodel/components/models/deepseek_v4/mxfp4_peft.py new file mode 100644 index 0000000000..b6221003c4 --- /dev/null +++ b/nemo_automodel/components/models/deepseek_v4/mxfp4_peft.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""DeepSeek V4 integration for MXFP4-resident experts under PEFT.""" + +import logging + +import torch.nn as nn + +from nemo_automodel.components.distributed.init_utils import get_world_size_safe + +logger = logging.getLogger(__name__) + + +def prepare_mxfp4_peft_checkpoint_load(model: nn.Module, peft_config) -> None: + """Prepare a DeepSeek V4 model to load routed experts directly in MXFP4.""" + if getattr(peft_config, "expert_weight_format", "bf16") != "mxfp4": + return + + from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4 + + # MXFP4-resident experts rely on the MoE parallelizer to shard both packed + # values and scales. Without it, the scale tensors are not applied correctly. + if get_world_size_safe() == 1: + raise ValueError( + "peft.expert_weight_format='mxfp4' requires expert parallelism " + "(multi-GPU with distributed.ep_size>1); it is not supported on a single GPU " + "(the packed expert scales are not applied without the MoE parallelizer). " + "Use ep_size>1, or set expert_weight_format='bf16'." + ) + + adapter = getattr(model, "state_dict_adapter", None) + if adapter is not None and hasattr(adapter, "expert_storage_format"): + adapter.expert_storage_format = "mxfp4" + + num_converted = convert_frozen_experts_to_mxfp4(model, passthrough=True) + logger.info("Converted %d frozen expert module(s) to mxfp4-resident storage (passthrough)", num_converted) + + +def finalize_mxfp4_peft_checkpoint_load(model: nn.Module, peft_config) -> None: + """Pack any deferred DeepSeek V4 expert weights after checkpoint loading.""" + if getattr(peft_config, "expert_weight_format", "bf16") != "mxfp4": + return + + from nemo_automodel.components._peft.lora import pack_mxfp4_expert_base_weights + + num_packed = pack_mxfp4_expert_base_weights(model) + if num_packed: + logger.info("Packed %d MoE expert modules to mxfp4-resident storage", num_packed) diff --git a/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py b/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py index e7bfaef4ad..0991d75c75 100644 --- a/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py +++ b/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py @@ -229,11 +229,17 @@ def __init__( moe_config: MoEConfig, backend: BackendConfig, dtype: torch.dtype = torch.float32, + expert_storage_format: str = "bf16", ): self.config = config self.moe_config = moe_config self.backend = backend self.dtype = dtype + # "bf16": dequantize routed experts to bf16 on load (default). + # "mxfp4": passthrough — keep the checkpoint's packed fp4 (int8) + e8m0 + # scales and aggregate them into ``*_packed`` / ``*_scales`` params so + # experts are never materialized in bf16 (caps the load-time peak). + self.expert_storage_format = expert_storage_format self._checkpoint_expert_quant_layout_cache: _ExpertQuantLayout | None = None # ------------------------------------------------------------------ @@ -345,6 +351,11 @@ def _dequantize(self, state_dict: dict[str, Any]) -> dict[str, Any]: if scale_key is not None: scale = state_dict[scale_key] if self._is_expert_weight_key(key): + # mxfp4 passthrough: leave packed expert weight + scale in place; + # _aggregate_experts_packed consumes both. Only fp4-layout + # checkpoints are eligible — fall back to dequant otherwise. + if self.expert_storage_format == "mxfp4" and self._expert_key_is_fp4(weight, scale): + continue state_dict[key] = self._dequantize_expert_weight(key, weight, scale) else: state_dict[key] = dequantize_from_fp8(weight, scale, dtype=self.dtype, name=key) @@ -360,6 +371,8 @@ def _aggregate_experts( device_mesh: DeviceMesh | None, ) -> dict[str, Any]: """Aggregate per-expert weights (w1/w2/w3) into stacked gate_and_up/down tensors.""" + if self.expert_storage_format == "mxfp4": + return self._aggregate_experts_packed(state_dict, device_mesh) n_experts = self.moe_config.n_routed_experts if device_mesh is not None: @@ -442,6 +455,86 @@ def _aggregate_experts( return out + def _expert_key_is_fp4(self, weight: torch.Tensor, scale: torch.Tensor) -> bool: + return self._expert_quant_layout_from_tensors(weight, scale) is _ExpertQuantLayout.FP4 + + def _aggregate_experts_packed( + self, + state_dict: dict[str, Any], + device_mesh: DeviceMesh | None, + ) -> dict[str, Any]: + """Aggregate per-expert packed fp4 weights + e8m0 scales WITHOUT dequantizing. + + Concatenates gate (w1) and up (w3) along the output dim and stacks experts + on dim 0 — packing is along the contraction (input) dim, so both operations + are layout-preserving and no unpacking is needed. Emits ``*_packed`` / + ``*_scales`` keys matching ``MXFP4ExpertStorageMixin``'s packed params. + """ + n_experts = self.moe_config.n_routed_experts + if device_mesh is not None: + rank = ( + get_submesh(device_mesh, ("ep",)).get_rank() + if "ep" in device_mesh.mesh_dim_names + else device_mesh.get_rank() + ) + start_expert, end_expert = get_expert_range_for_rank_from_mesh(device_mesh, n_experts) + expected_per_rank = end_expert - start_expert + else: + rank = None + expected_per_rank = n_experts + + # layer -> {"gate_and_up": {eid: {"w1": (packed, scale), "w3": (...)}}, "down": {eid: (packed, scale)}} + by_layer: dict[str, dict] = {} + out: dict[str, Any] = {} + + for key in list(state_dict.keys()): + m = _EXPERT_PATTERN.match(key) + if m is None: + # Drop orphaned expert scale keys (consumed via their weight key); pass the rest through. + if ".ffn.experts." in key and key.endswith(".scale"): + continue + out[key] = state_dict[key] + continue + + layer_num, expert_num, which = m.group(1), int(m.group(2)), m.group(3) + if not should_load_expert_for_rank(expert_num, device_mesh, n_experts): + continue + + weight = state_dict[key] + scale = state_dict.get(key[: -len(".weight")] + ".scale") + assert scale is not None, f"missing scale for packed expert weight {key}" + packed = weight.to_local() if is_dtensor(weight) else weight + scale_local = scale.to_local() if is_dtensor(scale) else scale + + layer = by_layer.setdefault(layer_num, {"gate_and_up": {}, "down": {}}) + if which in ("w1", "w3"): + layer["gate_and_up"].setdefault(expert_num, {})[which] = (packed, scale_local) + else: # w2 = down + layer["down"][expert_num] = (packed, scale_local) + + gu = layer.get("gate_and_up") + if gu is not None and len(gu) == expected_per_rank and all("w1" in d and "w3" in d for d in gu.values()): + eids = sorted(gu.keys()) + # cat(gate, up) along the output dim (dim 0 of [out, in//2] / [out, in//block]). + packed_stack = torch.stack([torch.cat([gu[e]["w1"][0], gu[e]["w3"][0]], dim=0) for e in eids], dim=0) + scale_stack = torch.stack([torch.cat([gu[e]["w1"][1], gu[e]["w3"][1]], dim=0) for e in eids], dim=0) + base = f"model.layers.{layer_num}.mlp.experts.gate_and_up_projs" + out[base + "_packed"] = create_dtensor_from_local(packed_stack, device_mesh, rank) + out[base + "_scales"] = create_dtensor_from_local(scale_stack, device_mesh, rank) + del layer["gate_and_up"] + + down = layer.get("down") + if down is not None and len(down) == expected_per_rank: + eids = sorted(down.keys()) + packed_stack = torch.stack([down[e][0] for e in eids], dim=0) + scale_stack = torch.stack([down[e][1] for e in eids], dim=0) + base = f"model.layers.{layer_num}.mlp.experts.down_projs" + out[base + "_packed"] = create_dtensor_from_local(packed_stack, device_mesh, rank) + out[base + "_scales"] = create_dtensor_from_local(scale_stack, device_mesh, rank) + del layer["down"] + + return out + def _rename_all(self, state_dict: dict[str, Any]) -> dict[str, Any]: """Apply the HF->internal rename table to every key.""" return {_rename_hf_key(k): v for k, v in state_dict.items()} @@ -672,6 +765,12 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t if quantization: quantized = [] for key, value in result: + # mxfp4 passthrough: experts are already split into packed int8 + # weight + e8m0 scale by _split_merged_expert; do not rebuild + # placeholders (the bf16->fp4 path assumes a bf16 input). + if self.expert_storage_format == "mxfp4" and self._is_expert_weight_key(key): + quantized.append((key, value)) + continue if key.endswith(".weight") and not self._is_non_quantized(key): base = key[: -len(".weight")] if self._is_expert_weight_key(key): @@ -949,6 +1048,36 @@ def _split_merged_expert(self, fqn: str, tensor: Any) -> list[tuple[str, Any]]: Handles DTensor inputs (EP-sharded) by working on the local shard only, emitting keys only for the experts owned by the current rank. """ + # mxfp4 passthrough: packed params are already in checkpoint orientation + # [E, out, in//{2,block}]; split per-expert (dim 0) then gate||up along the + # output dim (dim 0 of the per-expert tensor). No transpose — unlike the + # bf16 path which stores compute layout and must transpose. + gate_up_packed_pat = re.compile(r"^(model\.layers\.(\d+)\.mlp\.experts)\.gate_and_up_projs_(packed|scales)$") + down_packed_pat = re.compile(r"^(model\.layers\.(\d+)\.mlp\.experts)\.down_projs_(packed|scales)$") + + m = gate_up_packed_pat.match(fqn) + if m: + layer_num, suffix = m.group(2), m.group(3) + disk_suffix = "weight" if suffix == "packed" else "scale" + expert_tensors, expert_ids = split_experts_weights_dtensor_aware(tensor, self.moe_config.n_routed_experts) + result = [] + for t, eid in zip(expert_tensors, expert_ids): + out_dim = t.shape[0] // 2 + gate_t, up_t = t.split(out_dim, dim=0) + result.append((f"layers.{layer_num}.ffn.experts.{eid}.w1.{disk_suffix}", gate_t)) + result.append((f"layers.{layer_num}.ffn.experts.{eid}.w3.{disk_suffix}", up_t)) + return result + + m = down_packed_pat.match(fqn) + if m: + layer_num, suffix = m.group(2), m.group(3) + disk_suffix = "weight" if suffix == "packed" else "scale" + expert_tensors, expert_ids = split_experts_weights_dtensor_aware(tensor, self.moe_config.n_routed_experts) + return [ + (f"layers.{layer_num}.ffn.experts.{eid}.w2.{disk_suffix}", t) + for t, eid in zip(expert_tensors, expert_ids) + ] + gate_up_pat = re.compile(r"^(model\.layers\.(\d+)\.mlp\.experts)\.gate_and_up_projs$") down_pat = re.compile(r"^(model\.layers\.(\d+)\.mlp\.experts)\.down_projs$") diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 2ebea56d40..9c96e90680 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -1695,6 +1695,11 @@ def to_local(tensor): else: return tensor + # mxfp4-resident experts hold packed base weights (no gate_and_up_projs / + # down_projs); those are filled from the checkpoint, nothing to init here. + if getattr(module, "_mxfp4_resident", False): + return + with torch.device(buffer_device): if isinstance(module, (GroupedExperts, GroupedExpertsDeepEP)): to_local(module.gate_and_up_projs).normal_(mean=0.0, std=init_std) diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index 709733a4ff..8358d82275 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -303,8 +303,13 @@ def _partition_fn(self, name, module, device_mesh): assert device_mesh.ndim == 1 for name, param in module.named_parameters(recurse=False): - dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)])) - dist_param.requires_grad = param.requires_grad + # Pass requires_grad at construction: nn.Parameter defaults to + # requires_grad=True, which raises for non-floating dtypes (e.g. the + # int8 / e8m0 packed tensors of mxfp4-resident experts) before a + # later assignment could fix it. + dist_param = nn.Parameter( + distribute_tensor(param, device_mesh, [Shard(0)]), requires_grad=param.requires_grad + ) module.register_parameter(name, dist_param) if isinstance(module, (GroupedExpertsDeepEP, GroupedExpertsMoK)): diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py new file mode 100644 index 0000000000..e5445387ad --- /dev/null +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -0,0 +1,397 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""MXFP4-resident expert storage for frozen MoE experts. + +``GroupedExpertsMXFP4`` keeps the frozen routed-expert base weights packed as +fp4-e2m1 + e8m0 block scales (the DeepSeek V4 Flash checkpoint format) and +dequantizes on the fly inside the grouped GEMM, instead of holding them in +bf16. This is the storage win for LoRA / frozen-base training of large MoE +models, where the routed experts dominate parameter memory. + +The format-specific pack/unpack/GEMM logic lives in ``MXFP4ExpertStorageMixin`` +so a future integer-int4 (e.g. GLM) variant can reuse the same module wiring by +swapping the mixin's primitives. +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor + +from nemo_automodel.components.moe.experts import ( + GroupedExperts, + GroupedExpertsDeepEP, + _apply_bias, + _permute_tokens_for_grouped_mm, +) +from nemo_automodel.components.quantization.mxfp4 import ( + MXFP4_BLOCK_SIZE, + MXFP4GroupedMM, + dequantize_mxfp4, + quantize_mxfp4, +) + + +def _to_local(t): + """Return the local shard of a DTensor, or the tensor unchanged.""" + return t.to_local() if isinstance(t, DTensor) else t + + +class MXFP4ExpertStorageMixin: + """Packed-mxfp4 base-weight storage and grouped GEMM for routed experts. + + Mixed into a ``GroupedExperts`` (or ``GroupedExpertsLoRA``) subclass. The base + projections ``gate_and_up_projs`` / ``down_projs`` are stored as packed fp4 + (int8, two e2m1 nibbles per byte) plus ``float8_e8m0fnu`` block scales, in + checkpoint orientation ``[n_experts, out_dim, in_dim]`` so the block scales run + along the contraction dim. The bf16 parameters are dropped once packed. + + Packing is deferred when the base weights are still on the meta device: the + module behaves like its bf16 parent until ``pack_base_weights()`` runs (after + the checkpoint is loaded). + """ + + _MXFP4_BASE_NAMES: tuple[str, ...] = ("gate_and_up_projs", "down_projs") + # Storage-parameter suffixes, in pack/unpack order. Kept as a tuple so the + # registration helper is format-driven rather than hardcoding two names. + _PACKED_SUFFIXES: tuple[str, ...] = ("_packed", "_scales") + + def _init_mxfp4_storage(self) -> None: + """Validate the backend and pack immediately if base weights are materialized.""" + if not self.use_torch_mm: + raise NotImplementedError( + "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm'). " + "The grouped_gemm path (backend.experts='gmm') has no packed variant; with DeepEP dispatch use " + "backend.dispatcher='deepep' together with backend.experts='torch_mm'." + ) + self._mxfp4_resident = False + if not _to_local(getattr(self, self._MXFP4_BASE_NAMES[0])).is_meta: + self.pack_base_weights() + + @torch.no_grad() + def _init_packed_placeholders(self) -> None: + """Register meta packed storage params from config shapes (no bf16 weights). + + Used by the passthrough path so a packed fp4 checkpoint loads straight into + these params without ever materializing bf16 experts. Config-driven, so it is + shared by the torch (``GroupedExpertsMXFP4``) and DeepEP + (``GroupedExpertsDeepEPMXFP4``) frozen variants. + """ + cfg = self.config + block = MXFP4_BLOCK_SIZE + up_proj_dim = cfg.moe_inter_dim * 2 if self.is_gated else cfg.moe_inter_dim + expert_dim = cfg.expert_dim + moe_inter = cfg.moe_inter_dim + e = cfg.n_routed_experts + assert expert_dim % block == 0 and moe_inter % block == 0, ( + f"expert dims must be divisible by {block} for mxfp4 (expert_dim={expert_dim}, moe_inter={moe_inter})" + ) + # Checkpoint orientation [E, out, in], packed along the contraction (in) dim. + shapes = { + "gate_and_up_projs": ((e, up_proj_dim, expert_dim // 2), (e, up_proj_dim, expert_dim // block)), + "down_projs": ((e, expert_dim, moe_inter // 2), (e, expert_dim, moe_inter // block)), + } + for name, (packed_shape, scale_shape) in shapes.items(): + packed = torch.empty(packed_shape, dtype=torch.int8, device="meta") + scales = torch.empty(scale_shape, dtype=torch.float8_e8m0fnu, device="meta") + self.register_packed_base_weight(name, (packed, scales)) + self._mxfp4_resident = True + + @torch.no_grad() + def register_packed_base_weight(self, name: str, tensors: tuple[torch.Tensor, ...], reference=None) -> None: + """Register packed storage params for base projection ``name``. + + Decoupled from quantization so it can run either as a post-load conversion + (``pack_base_weights`` passes freshly quantized tensors) or at module init + (a chunk-loader passes meta placeholders, then loads the quantized + checkpoint straight into them — the path that avoids ever materializing + bf16 experts at GLM-744B scale). Replaces the bf16 parameter ``name`` if + present. + + Args: + name: Base projection name (e.g. ``"gate_and_up_projs"``). + tensors: Storage tensors in ``_PACKED_SUFFIXES`` order. + reference: Optional DTensor whose mesh/placements the storage tensors + inherit (use the pre-pack bf16 param, or a meta DTensor at init). + """ + assert len(tensors) == len(self._PACKED_SUFFIXES), ( + f"expected {len(self._PACKED_SUFFIXES)} tensors {self._PACKED_SUFFIXES}, got {len(tensors)}" + ) + if isinstance(reference, DTensor): + tensors = tuple(DTensor.from_local(t, reference.device_mesh, reference.placements) for t in tensors) + if name in self._parameters: + del self._parameters[name] + for suffix, tensor in zip(self._PACKED_SUFFIXES, tensors): + self.register_parameter(name + suffix, nn.Parameter(tensor, requires_grad=False)) + + @torch.no_grad() + def pack_base_weights(self) -> None: + """Pack the frozen base projections to mxfp4 and free the bf16 tensors. + + No-op when already packed. Requires the base weights to be materialized. + """ + if self._mxfp4_resident: + return + for name in self._MXFP4_BASE_NAMES: + param = getattr(self, name) + local = _to_local(param) + assert not local.is_meta, f"pack_base_weights requires materialized '{name}'" + # [E, in, out] (compute layout) -> [E, out, in] (checkpoint layout) so the + # mx block scales run along the contraction dim. + tensors = quantize_mxfp4(local.transpose(-2, -1).contiguous()) + self.register_packed_base_weight(name, tensors, reference=param) + self._mxfp4_resident = True + + def _mxfp4_base_mm(self, x: torch.Tensor, name: str, offs: torch.Tensor) -> torch.Tensor: + """Grouped GEMM ``x @ W`` over the packed base weight ``name`` (dequant on the fly).""" + packed = _to_local(getattr(self, name + "_packed")) + scales = _to_local(getattr(self, name + "_scales")) + return MXFP4GroupedMM.apply(x, packed, scales, offs) + + def _mxfp4_dequant_expert0(self, name: str, dtype: torch.dtype) -> torch.Tensor: + """Dequantize expert 0 of base weight ``name`` to compute layout ``[in, out]``.""" + packed = _to_local(getattr(self, name + "_packed"))[0] + scales = _to_local(getattr(self, name + "_scales"))[0] + return dequantize_mxfp4(packed, scales, dtype).transpose(-2, -1) + + +class GroupedExpertsMXFP4(MXFP4ExpertStorageMixin, GroupedExperts): + """Frozen routed experts with mxfp4-resident base weights and no adapter. + + Drop-in replacement for ``GroupedExperts`` when the experts are frozen (e.g. + LoRA training that targets only attention). Forward mirrors + ``GroupedExperts._forward_grouped_mm`` but reads the packed base weights. + """ + + def __init__(self, orig_module: GroupedExperts, passthrough: bool = False): + """ + Args: + orig_module: The bf16 GroupedExperts to replace. + passthrough: When True, register packed storage placeholders at init + (no bf16 weights) so a quantized checkpoint loads straight into + them — experts are never materialized in bf16. Requires the base + weights to be meta (i.e. loaded later from a packed checkpoint). + """ + super().__init__(orig_module.config, backend=None) + if not self.use_torch_mm and not orig_module.use_torch_mm: + raise NotImplementedError( + "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm')." + ) + self.use_torch_mm = orig_module.use_torch_mm + + if passthrough: + # The bf16 base params from super().__init__ are placeholders only + # (meta under init_empty_weights); _init_packed_placeholders deletes + # them and registers meta packed storage, so no bf16 experts are ever + # materialized — the packed checkpoint loads straight into them. + if self.expert_bias: + self.gate_up_proj_bias.requires_grad_(False) + self.down_proj_bias.requires_grad_(False) + self._init_packed_placeholders() + return + + if not getattr(orig_module, "gate_and_up_projs", None).is_meta: + self.gate_and_up_projs.data = _to_local(orig_module.gate_and_up_projs).clone() + self.down_projs.data = _to_local(orig_module.down_projs).clone() + if self.expert_bias: + self.gate_up_proj_bias.data = _to_local(orig_module.gate_up_proj_bias).clone() + self.down_proj_bias.data = _to_local(orig_module.down_proj_bias).clone() + self.gate_and_up_projs.requires_grad_(False) + self.down_projs.requires_grad_(False) + self._init_mxfp4_storage() + + def forward( + self, + x: torch.Tensor, + token_mask: torch.Tensor, + weights: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + """Forward over mxfp4 base weights. Falls back to bf16 until packing is done.""" + if not self._mxfp4_resident: + return super().forward(x, token_mask, weights, indices) + + assert not isinstance(x, DTensor) + input_dtype = x.dtype + + if isinstance(self.gate_and_up_projs_packed, DTensor): + ep_mesh = self.gate_and_up_projs_packed.device_mesh + assert ep_mesh is not None and ep_mesh.ndim == 1, "We only support 1D mesh for MoE" + ep_size = ep_mesh.size() + ep_rank = ep_mesh.get_local_rank() + else: + ep_mesh = None + ep_size = 1 + ep_rank = 0 + + assert self.n_routed_experts % ep_size == 0 + + if ep_size > 1: + from torch.distributed.tensor import Partial, Shard + + x = DTensor.from_local(x, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor( + grad_placements=[Partial()] + ) + weights = DTensor.from_local(weights.float(), device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor( + grad_placements=[Partial()] + ) + indices = DTensor.from_local(indices, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor() + token_mask = DTensor.from_local(token_mask, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor() + + n_local_experts = self.n_routed_experts // ep_size + experts_start_idx = ep_rank * n_local_experts + + y = self._forward_grouped_mm_mxfp4(x, token_mask, weights, indices, n_local_experts, experts_start_idx) + + if ep_size > 1: + from torch.distributed.tensor import Partial, Shard + + y = DTensor.from_local(y, device_mesh=ep_mesh, placements=[Partial()]) + y = y.redistribute(placements=[Shard(0)]).to_local() + + return y.to(input_dtype) + + def _forward_grouped_mm_mxfp4(self, x, token_mask, weights, indices, n_local_experts, experts_start_idx): + sorted_token_ids, sorted_weights, tokens_per_expert, offs = _permute_tokens_for_grouped_mm( + indices, weights, token_mask, n_local_experts, experts_start_idx + ) + y = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + + if tokens_per_expert.sum() > 0: + permuted_x = x[sorted_token_ids] + permuted_probs = sorted_weights.unsqueeze(-1) + + output1 = self._mxfp4_base_mm(permuted_x, "gate_and_up_projs", offs) + if self.expert_bias: + output1 = _apply_bias(output1, _to_local(self.gate_up_proj_bias), tokens_per_expert) + output1 = self.expert_activation_grouped(output1, permuted_probs) + + output2 = self._mxfp4_base_mm(output1, "down_projs", offs) + if self.expert_bias: + output2 = _apply_bias(output2, _to_local(self.down_proj_bias), tokens_per_expert, permuted_probs) + + scatter_ids = sorted_token_ids.unsqueeze(1).expand_as(output2) + y.scatter_add_(0, scatter_ids, output2.float()) + else: + # Dummy computation for gradient flow when no tokens routed locally. + gate_up_w0 = self._mxfp4_dequant_expert0("gate_and_up_projs", x.dtype) + down_w0 = self._mxfp4_dequant_expert0("down_projs", x.dtype) + output1 = torch.matmul(x[0] * 0, gate_up_w0) + output1_ = self.expert_activation_grouped(output1, weights[0, 0, None].unsqueeze(0)) + output2 = torch.matmul(output1_, down_w0) + y[0] += output2[0] + + return y + + +class GroupedExpertsDeepEPMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsDeepEP): + """Frozen routed experts with mxfp4-resident base weights under DeepEP dispatch. + + Drop-in replacement for ``GroupedExpertsDeepEP`` when the experts are frozen + (e.g. LoRA on attention only). The DeepEP fused all-to-all token dispatch is reused + unchanged — mxfp4 only changes the two post-dispatch grouped GEMMs, which read the + packed base weights via ``MXFP4GroupedMM`` instead of bf16 ``torch._grouped_mm``. + + Requires the torch_mm experts backend (``backend.experts='torch_mm'``); the + grouped_gemm (``gmm``) path has no packed variant. + """ + + def __init__(self, orig_module: GroupedExpertsDeepEP, passthrough: bool = False): + """ + Args: + orig_module: The bf16 GroupedExpertsDeepEP to replace. + passthrough: When True, register packed storage placeholders at init (no + bf16 weights) so a quantized checkpoint loads straight into them. + """ + super().__init__( + orig_module.config, + backend=None, + dispatcher_backend=orig_module.dispatcher_backend, + dispatcher_num_sms=orig_module.dispatcher_num_sms, + dispatcher_share_token_dispatcher=orig_module.dispatcher_share_token_dispatcher, + dispatcher_async_dispatch=orig_module.dispatcher_async_dispatch, + ) + # backend=None leaves use_torch_mm False; inherit the original's choice so the + # mxfp4 storage guard enforces torch_mm (set before _init_mxfp4_storage runs). + self.use_torch_mm = orig_module.use_torch_mm + + if passthrough: + if self.expert_bias: + self.gate_up_proj_bias.requires_grad_(False) + self.down_proj_bias.requires_grad_(False) + self._init_packed_placeholders() + return + + if not _to_local(orig_module.gate_and_up_projs).is_meta: + self.gate_and_up_projs.data = _to_local(orig_module.gate_and_up_projs).clone() + self.down_projs.data = _to_local(orig_module.down_projs).clone() + if self.expert_bias: + self.gate_up_proj_bias.data = _to_local(orig_module.gate_up_proj_bias).clone() + self.down_proj_bias.data = _to_local(orig_module.down_proj_bias).clone() + self.gate_and_up_projs.requires_grad_(False) + self.down_projs.requires_grad_(False) + self._init_mxfp4_storage() + + def forward( + self, + x: torch.Tensor, + token_mask: torch.Tensor, + weights: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + """Forward over mxfp4 base weights with DeepEP dispatch. + + Mirrors ``GroupedExpertsDeepEP.forward``, replacing the two base + ``torch._grouped_mm`` calls with ``MXFP4GroupedMM`` over the packed weights. + Falls back to the bf16 parent while packing is still deferred. + """ + if not self._mxfp4_resident: + return super().forward(x, token_mask, weights, indices) + + assert not isinstance(x, DTensor) + assert self.use_torch_mm, "mxfp4-resident DeepEP experts require the torch_mm experts backend." + assert self.n_routed_experts % self.ep_size == 0, ( + f"Number of experts must be divisible by ep_size (ep_size={self.ep_size})" + ) + + indices = indices.masked_fill(~token_mask.unsqueeze(-1), -1) + (permuted_local_hidden_states, tokens_per_expert, permuted_probs) = self.token_dispatcher.token_permutation2( + hidden_states=x, + num_local_tokens=x.size(0), + token_probs=weights, + token_indices=indices, + ) + permuted_probs = permuted_probs.unsqueeze(-1) + + if torch.count_nonzero(tokens_per_expert) > 0: + tokens_per_expert_gpu = tokens_per_expert.to(device=permuted_local_hidden_states.device, non_blocking=True) + offs = tokens_per_expert_gpu.cumsum(dim=0).to(torch.int32) + + output1 = self._mxfp4_base_mm(permuted_local_hidden_states, "gate_and_up_projs", offs) + if self.expert_bias: + output1 = _apply_bias(output1, _to_local(self.gate_up_proj_bias), tokens_per_expert) + output1 = self.expert_activation(output1, permuted_probs) + output2 = self._mxfp4_base_mm(output1, "down_projs", offs) + if self.expert_bias: + output2 = _apply_bias(output2, _to_local(self.down_proj_bias), tokens_per_expert, permuted_probs) + else: + # Dummy computation for gradient flow when no tokens routed locally. + gate_up_w0 = self._mxfp4_dequant_expert0("gate_and_up_projs", x.dtype) + down_w0 = self._mxfp4_dequant_expert0("down_projs", x.dtype) + output1 = torch.matmul(x[0] * 0, gate_up_w0) + output1_ = self.expert_activation(output1, permuted_probs) + output2 = torch.matmul(output1_, down_w0) + + y = self.token_dispatcher.token_unpermutation(output2) + return y diff --git a/nemo_automodel/components/quantization/__init__.py b/nemo_automodel/components/quantization/__init__.py index 044fe286ad..d25e28f796 100644 --- a/nemo_automodel/components/quantization/__init__.py +++ b/nemo_automodel/components/quantization/__init__.py @@ -6,6 +6,12 @@ create_fp8_config_from_dict, verify_fp8_conversion, ) +from .mxfp4 import ( + MXFP4_BLOCK_SIZE, + MXFP4GroupedMM, + dequantize_mxfp4, + quantize_mxfp4, +) from .qlora import ( HAS_BNB, create_bnb_config, @@ -27,6 +33,10 @@ "HAS_BNB", "create_bnb_config", "verify_qlora_quantization", + "MXFP4_BLOCK_SIZE", + "MXFP4GroupedMM", + "dequantize_mxfp4", + "quantize_mxfp4", ] if HAVE_TORCHAO: diff --git a/nemo_automodel/components/quantization/mxfp4.py b/nemo_automodel/components/quantization/mxfp4.py new file mode 100644 index 0000000000..de6c25d9ca --- /dev/null +++ b/nemo_automodel/components/quantization/mxfp4.py @@ -0,0 +1,151 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""MXFP4 (fp4 e2m1 + e8m0 block scales) pack/unpack utilities for MoE expert weights. + +The packed layout matches the DeepSeek V4 Flash routed-expert checkpoint format: +two e2m1 values per int8 byte (low nibble at even column index, high nibble at the +following odd column) with one ``float8_e8m0fnu`` scale per 32 contiguous columns. +``MXFP4GroupedMM`` provides a grouped GEMM over packed weights that re-dequantizes +in backward instead of saving the dequantized tensor, so frozen expert weights stay +packed at steady state during LoRA training. +""" + +import torch + +MXFP4_BLOCK_SIZE = 32 + +# FP4 e2m1 value table: low 3 bits -> magnitude, MSB -> sign. +# Layout: [positive values for codes 0-7, negative values for codes 8-15]. +_FP4_E2M1_TABLE = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) + +# Midpoints between consecutive positive e2m1 magnitudes, used to round to nearest. +_FP4_E2M1_MIDPOINTS = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=torch.float32) + +# Per-byte expansion of the two packed e2m1 nibbles: row ``b`` holds +# ``(value(low nibble), value(high nibble))``. A single gather over this 256-row +# table decodes both fp4 values per byte with one int64 index tensor (half the +# gathers of indexing the 16-entry table twice). +_FP4_BYTE_TABLE = torch.stack( + [ + _FP4_E2M1_TABLE[torch.arange(256) & 0x0F], + _FP4_E2M1_TABLE[(torch.arange(256) >> 4) & 0x0F], + ], + dim=-1, +) # [256, 2] + + +def dequantize_mxfp4(packed: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Unpack fp4 e2m1 packed-int8 values and apply the per-32-column e8m0 scale. + + Args: + packed: int8 tensor of shape ``[..., K // 2]`` holding two e2m1 values per byte. + scales: ``float8_e8m0fnu`` tensor of shape ``[..., K // 32]``. + dtype: Output dtype. + + Returns: + Dequantized tensor of shape ``[..., K]`` in ``dtype``. + """ + packed_u8 = packed.contiguous().view(torch.uint8) + # Single gather over the 256-row byte table yields both e2m1 values per byte, + # interleaved (low, high) so column indices match the original layout. + pairs = _FP4_BYTE_TABLE.to(packed_u8.device)[packed_u8.long()] # [..., K // 2, 2] + fp4_vals = pairs.flatten(-2) # [..., K] + + # float8_e8m0fnu casts straight to its 2^(e - 127) value (OCP MX spec: byte 0x00 + # decodes to 2^-127, not zero). No all-zero-block special case is needed -- such a + # block's fp4 codes are already 0, so the product is 0 regardless of the scale. + scale_f32 = scales.to(torch.float32) # [..., K // 32] + + # Stay blocked and broadcast the per-32-column scale instead of materializing a + # full [..., K] scale tensor (cheaper, and fuses better under torch.compile). + blocked = fp4_vals.view(*fp4_vals.shape[:-1], scale_f32.shape[-1], MXFP4_BLOCK_SIZE) + return (blocked * scale_f32.unsqueeze(-1)).flatten(-2).to(dtype) + + +def quantize_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize along the last dim to the packed mxfp4 layout used by ``dequantize_mxfp4``. + + Block scales are computed as ``2^(floor(log2(amax)) - 2)`` so that values that are + already exactly representable (e.g. a dequantized fp4 checkpoint) round-trip + value-exactly. + + Args: + weight: Floating-point tensor of shape ``[..., K]`` with ``K`` divisible by 32. + + Returns: + Tuple of (int8 packed tensor ``[..., K // 2]``, ``float8_e8m0fnu`` scales ``[..., K // 32]``). + """ + k = weight.shape[-1] + assert k % MXFP4_BLOCK_SIZE == 0, f"last dim {k} must be divisible by {MXFP4_BLOCK_SIZE}" + + w = weight.float() + blocks = w.view(*w.shape[:-1], k // MXFP4_BLOCK_SIZE, MXFP4_BLOCK_SIZE) + amax = blocks.abs().amax(dim=-1) + + # e2m1 max magnitude is 6 = 1.5 * 2^2, so the shared block exponent is + # floor(log2(amax)) - 2. frexp gives floor(log2(amax)) == exponent - 1 exactly, + # avoiding log2 round-off near powers of two (the common case for a dequantized-fp4 + # checkpoint); frexp(0).exponent == 0 also removes the need for a zero guard. + exp_field = torch.frexp(amax).exponent - 3 # (exponent - 1) - 2 + scale_bytes = (exp_field + 127).clamp(1, 254).to(torch.uint8) + scales = scale_bytes.view(torch.float8_e8m0fnu) + scale = scales.to(torch.float32) # exact power of two; same decode as dequantize_mxfp4 + + # Round each scaled magnitude to the nearest e2m1 magnitude code. bucketize rounds + # half away from zero rather than to-nearest-even (PTX cvt / OCP recommendation); + # irrelevant for the exact-round-trip path, where scaled values land exactly on + # codes and never on a midpoint tie. + scaled = blocks.abs() / scale.unsqueeze(-1) + midpoints = _FP4_E2M1_MIDPOINTS.to(w.device) + codes = torch.bucketize(scaled, midpoints).to(torch.uint8) + codes = codes | torch.where(blocks < 0, torch.full_like(codes, 0x08), torch.zeros_like(codes)) + codes = codes.reshape(*w.shape[:-1], k) + + packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).view(torch.int8) + return packed.contiguous(), scales.contiguous() + + +class MXFP4GroupedMM(torch.autograd.Function): + """Grouped GEMM over mxfp4-packed frozen weights with dequantization on the fly. + + Saves only the packed weights for backward and re-dequantizes there, so the + bf16 weight tensor is a transient in both passes instead of being kept alive + by autograd. Weights are stored as ``[E, N, K]`` packed along ``K``, which is + the natural dequantization output and the operand the backward GEMM needs + directly (``grad_x = grad_out @ W``). The forward needs ``[E, K, N]``, which + ``torch._grouped_mm`` consumes as a transposed view (cuBLAS transB) — no + contiguous copy required. + + No weight gradient is produced — the base weights are frozen under LoRA. + """ + + @staticmethod + def forward(ctx, x: torch.Tensor, packed: torch.Tensor, scales: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + w_t = dequantize_mxfp4(packed, scales, x.dtype) # [E, N, K] + # Pass the transposed view directly; torch._grouped_mm handles the + # strided mat2 (transB), avoiding a full bf16 weight copy per forward. + out = torch._grouped_mm(x, w_t.transpose(-2, -1), offs=offs) + ctx.save_for_backward(packed, scales, offs) + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + packed, scales, offs = ctx.saved_tensors + w_t = dequantize_mxfp4(packed, scales, grad_out.dtype) # [E, N, K] == W^T + grad_x = torch._grouped_mm(grad_out.contiguous(), w_t, offs=offs) + return grad_x, None, None, None diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 6e6ca2fa8a..2e3dda68d3 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -1367,7 +1367,32 @@ def _run_validation_epoch(self, val_dataloader): total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device) total_num_label_tokens = 0 - for batch in val_dataloader: + # Keep every rank in lockstep over validation batches. The MoE expert + # forward issues collectives over the EP group; with uneven per-rank + # validation shard sizes, ranks would issue those collectives a + # different number of times and deadlock (observed at end-of-training + # validation on EP MoE models). Drive the loop by a global-MIN + # "does every rank still have a batch?" all-reduce so all ranks run + # exactly the same number of forwards (= global-min batch count). + # ``len(dataloader)`` is not relied upon (StatefulDataLoader may not + # define it). + dist_active = torch.distributed.is_available() and torch.distributed.is_initialized() + val_iter = iter(val_dataloader) + while True: + try: + batch = next(val_iter) + has_batch = 1 + except StopIteration: + batch = None + has_batch = 0 + if dist_active: + flag = torch.tensor(has_batch, dtype=torch.long, device=self.dist_env.device) + torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN) + if int(flag.item()) == 0: + break + elif has_batch == 0: + break + loss_buffer = [] num_label_tokens = (batch["labels"] != -100).sum().item() self._forward_backward_step( @@ -1610,8 +1635,28 @@ def main(config_path=None): config_path = pathlib.Path(__file__).parent.resolve() / "llama_3_2_1b_hellaswag.yaml" cfg = parse_args_and_load_config(config_path) trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) - trainer.setup() - trainer.run_train_validation_loop() + try: + trainer.setup() + trainer.run_train_validation_loop() + finally: + # Tear down the distributed process group so the process exits cleanly. + # Without this, NCCL/the elastic agent waits on the still-initialized + # group and the run hangs at shutdown after the final step. + _destroy_process_group_if_initialized() + + +def _destroy_process_group_if_initialized() -> None: + """Best-effort barrier + ``destroy_process_group`` so ranks exit together.""" + if not (torch.distributed.is_available() and torch.distributed.is_initialized()): + return + try: + torch.distributed.barrier() + except Exception as exc: # a rank may have already exited; destroy anyway + logger.warning("Barrier before process-group teardown failed: %s", exc) + try: + torch.distributed.destroy_process_group() + except Exception as exc: + logger.warning("destroy_process_group failed during teardown: %s", exc) if __name__ == "__main__": diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py new file mode 100644 index 0000000000..f59ba6521e --- /dev/null +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -0,0 +1,354 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +import pytest +import torch + +from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4, patch_moe_module +from nemo_automodel.components._peft.lora_experts import GroupedExpertsLoRA +from nemo_automodel.components._peft.lora_experts_mxfp4 import GroupedExpertsLoRAMXFP4 +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.layers import GroupedExperts +from nemo_automodel.components.moe.quantized_experts import GroupedExpertsMXFP4 +from nemo_automodel.components.quantization.mxfp4 import dequantize_mxfp4, quantize_mxfp4 + + +@pytest.fixture +def device(): + if torch.cuda.is_available(): + return torch.device(f"cuda:{torch.cuda.current_device()}") + return torch.device("cpu") + + +@pytest.fixture +def moe_config(): + # Dims divisible by 32 so both contraction dims are mxfp4-blockable. + return MoEConfig( + n_routed_experts=4, + n_shared_experts=0, + n_activated_experts=2, + n_expert_groups=1, + n_limited_groups=1, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="softmax", + route_scale=1.0, + dim=64, + inter_dim=128, + moe_inter_dim=64, + norm_topk_prob=False, + expert_activation="swiglu", + dtype=torch.bfloat16, + ) + + +def test_quantize_dequantize_idempotent(): + """A second quantize/dequantize round-trip must reproduce the first exactly.""" + torch.manual_seed(0) + w = torch.randn(3, 8, 64, dtype=torch.bfloat16) + packed, scales = quantize_mxfp4(w) + assert packed.dtype == torch.int8 + assert packed.shape == (3, 8, 32) + assert scales.dtype == torch.float8_e8m0fnu + assert scales.shape == (3, 8, 2) + + dq1 = dequantize_mxfp4(packed, scales, torch.bfloat16) + packed2, scales2 = quantize_mxfp4(dq1) + dq2 = dequantize_mxfp4(packed2, scales2, torch.bfloat16) + assert torch.equal(dq1, dq2) + + +def test_quantize_zero_block(): + """All-zero blocks must encode to scale byte 0 and decode back to exact zeros.""" + w = torch.zeros(2, 64, dtype=torch.bfloat16) + w[1, 32:] = 1.5 # one nonzero block to confirm mixed handling + packed, scales = quantize_mxfp4(w) + dq = dequantize_mxfp4(packed, scales, torch.bfloat16) + assert torch.equal(dq[0], torch.zeros(64, dtype=torch.bfloat16)) + assert torch.equal(dq[1, :32], torch.zeros(32, dtype=torch.bfloat16)) + assert torch.equal(dq[1, 32:], torch.full((32,), 1.5, dtype=torch.bfloat16)) + + +def test_dequantize_matches_dsv4_adapter(): + """fp4_utils dequant must agree exactly with the DeepSeek V4 state dict adapter.""" + from nemo_automodel.components.models.deepseek_v4.state_dict_adapter import DeepSeekV4StateDictAdapter + + torch.manual_seed(1) + w = torch.randn(4, 16, 96, dtype=torch.bfloat16) + packed, scales = quantize_mxfp4(w) + dq = dequantize_mxfp4(packed, scales, torch.bfloat16) + ref = DeepSeekV4StateDictAdapter._dequantize_expert_fp4( + packed.view(-1, packed.shape[-1]), scales.view(-1, scales.shape[-1]), torch.bfloat16 + ) + assert torch.equal(ref.view_as(dq), dq) + + +def _make_representable_(experts: GroupedExperts) -> None: + """Replace expert weights in-place with their mxfp4 round-trip so storage is exact.""" + with torch.no_grad(): + for name in ("gate_and_up_projs", "down_projs"): + param = getattr(experts, name) + w_t = param.data.transpose(-2, -1).contiguous() + packed, scales = quantize_mxfp4(w_t) + param.data.copy_(dequantize_mxfp4(packed, scales, param.dtype).transpose(-2, -1)) + + +def _routing_inputs(moe_config, num_tokens, device, dtype): + torch.manual_seed(7) + x = torch.randn(num_tokens, moe_config.dim, dtype=dtype, device=device, requires_grad=True) + indices = torch.stack( + [ + torch.randperm(moe_config.n_routed_experts, device=device)[: moe_config.n_activated_experts] + for _ in range(num_tokens) + ] + ) + weights = torch.rand(num_tokens, moe_config.n_activated_experts, dtype=dtype, device=device) + weights = weights / weights.sum(dim=-1, keepdim=True) + token_mask = torch.ones(num_tokens, dtype=torch.bool, device=device) + return x, token_mask, weights, indices + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mxfp4_module_packs_base_weights(moe_config, device): + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + + lora = GroupedExpertsLoRAMXFP4(orig, lora_dim=8, alpha=16) + + assert lora._mxfp4_resident + assert not hasattr(lora, "gate_and_up_projs") + assert not hasattr(lora, "down_projs") + assert lora.gate_and_up_projs_packed.dtype == torch.int8 + assert lora.gate_and_up_projs_scales.dtype == torch.float8_e8m0fnu + assert not lora.gate_and_up_projs_packed.requires_grad + # Checkpoint orientation: [E, out, in/2] for gate+up ([E, 2*inter, dim/2]). + assert lora.gate_and_up_projs_packed.shape == (4, 2 * moe_config.moe_inter_dim, moe_config.dim // 2) + assert lora.down_projs_packed.shape == (4, moe_config.dim, moe_config.moe_inter_dim // 2) + # Only LoRA params are trainable. + trainable = {n for n, p in lora.named_parameters() if p.requires_grad} + assert trainable == {"lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mxfp4_forward_backward_matches_bf16(moe_config, device): + """With fp4-representable base weights, the packed module must match the bf16 module.""" + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + _make_representable_(orig) + + ref = GroupedExpertsLoRA(orig, lora_dim=8, alpha=16).to(device) + mx = GroupedExpertsLoRAMXFP4(orig, lora_dim=8, alpha=16).to(device) + with torch.no_grad(): + for name in ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"): + getattr(mx, name).data.copy_(getattr(ref, name).data) + + x_ref, token_mask, weights, indices = _routing_inputs(moe_config, 32, device, torch.bfloat16) + x_mx = x_ref.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref, token_mask, weights, indices) + y_mx = mx(x_mx, token_mask, weights, indices) + torch.testing.assert_close(y_mx, y_ref, atol=1e-6, rtol=1e-6) + + y_ref.float().pow(2).sum().backward() + y_mx.float().pow(2).sum().backward() + torch.testing.assert_close(x_mx.grad, x_ref.grad, atol=1e-5, rtol=1e-5) + for name in ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"): + torch.testing.assert_close(getattr(mx, name).grad, getattr(ref, name).grad, atol=1e-5, rtol=1e-5) + # Base weights are frozen; no grads anywhere else. + assert mx.gate_and_up_projs_packed.grad is None + assert mx.down_projs_packed.grad is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mxfp4_no_routed_tokens(moe_config, device): + """The all-masked dummy path must run and keep LoRA gradients flowing.""" + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + + mx = GroupedExpertsLoRAMXFP4(orig, lora_dim=8, alpha=16).to(device) + x, _, weights, indices = _routing_inputs(moe_config, 8, device, torch.bfloat16) + token_mask = torch.zeros(8, dtype=torch.bool, device=device) + + y = mx(x, token_mask, weights, indices) + assert y.shape == x.shape + y.float().sum().backward() + assert x.grad is not None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mxfp4_requires_torch_mm_backend(moe_config, device): + orig = GroupedExperts(moe_config).to(device) # default per-expert loop backend + with torch.no_grad(): + orig.init_weights(buffer_device=device) + with pytest.raises(NotImplementedError, match="torch_mm"): + GroupedExpertsLoRAMXFP4(orig, lora_dim=8, alpha=16) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_patch_moe_module_mxfp4(moe_config, device): + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + + patched = patch_moe_module(orig, dim=4, alpha=8, expert_weight_format="mxfp4") + assert isinstance(patched, GroupedExpertsLoRAMXFP4) + + patched_bf16 = patch_moe_module(orig, dim=4, alpha=8) + assert isinstance(patched_bf16, GroupedExpertsLoRA) + assert not isinstance(patched_bf16, GroupedExpertsLoRAMXFP4) + + +# --- frozen experts (no adapter) --------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_frozen_mxfp4_packs_and_freezes(moe_config, device): + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + + mx = GroupedExpertsMXFP4(orig) + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert mx.down_projs_scales.dtype == torch.float8_e8m0fnu + # Nothing trainable: experts are fully frozen. + assert [n for n, p in mx.named_parameters() if p.requires_grad] == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_frozen_mxfp4_forward_matches_bf16(moe_config, device): + orig = GroupedExperts(moe_config).to(device) + orig.use_torch_mm = True + with torch.no_grad(): + orig.init_weights(buffer_device=device) + _make_representable_(orig) + + mx = GroupedExpertsMXFP4(orig) + x, token_mask, weights, indices = _routing_inputs(moe_config, 32, device, torch.bfloat16) + + # Reference: bf16 grouped_mm path on the same (fp4-representable) weights. + y_ref = orig(x.detach(), token_mask, weights, indices) + y_mx = mx(x.detach(), token_mask, weights, indices) + torch.testing.assert_close(y_mx, y_ref, atol=1e-6, rtol=1e-6) + + +def test_register_packed_base_weight_is_init_capable(moe_config): + """Lock the API the GLM chunk-loader depends on: packed params can be registered + decoupled from quantization (here, with meta placeholders), replacing the bf16 + param without ever materializing it.""" + orig = GroupedExperts(moe_config) + orig.use_torch_mm = True # avoid the post-init pack; weights are still on meta + + mx = GroupedExpertsMXFP4.__new__(GroupedExpertsMXFP4) + GroupedExperts.__init__(mx, moe_config, backend=None) + mx.use_torch_mm = True + mx._mxfp4_resident = False + + # Register meta placeholders of the packed shapes, as a chunk-loader would at init. + e = moe_config.n_routed_experts + gate_up_out = 2 * moe_config.moe_inter_dim + placeholders = { + "gate_and_up_projs": ( + torch.empty(e, gate_up_out, moe_config.dim // 2, dtype=torch.int8, device="meta"), + torch.empty(e, gate_up_out, moe_config.dim // 32, dtype=torch.float8_e8m0fnu, device="meta"), + ), + "down_projs": ( + torch.empty(e, moe_config.dim, moe_config.moe_inter_dim // 2, dtype=torch.int8, device="meta"), + torch.empty(e, moe_config.dim, moe_config.moe_inter_dim // 32, dtype=torch.float8_e8m0fnu, device="meta"), + ), + } + for name, tensors in placeholders.items(): + mx.register_packed_base_weight(name, tensors) + + assert not hasattr(mx, "gate_and_up_projs") + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert mx.gate_and_up_projs_packed.is_meta + assert mx.down_projs_scales.dtype == torch.float8_e8m0fnu + assert not mx.gate_and_up_projs_packed.requires_grad + + +def test_passthrough_init_registers_packed_params_on_meta(moe_config): + """Phase-B packed-at-init: build in packed mode from meta weights, no bf16.""" + with torch.device("meta"): + orig = GroupedExperts(moe_config) + orig.use_torch_mm = True + + mx = GroupedExpertsMXFP4(orig, passthrough=True) + + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") # never created bf16 storage + up_proj_dim = 2 * moe_config.moe_inter_dim # gated + assert tuple(mx.gate_and_up_projs_packed.shape) == (moe_config.n_routed_experts, up_proj_dim, moe_config.dim // 2) + assert tuple(mx.gate_and_up_projs_scales.shape) == (moe_config.n_routed_experts, up_proj_dim, moe_config.dim // 32) + assert tuple(mx.down_projs_packed.shape) == ( + moe_config.n_routed_experts, + moe_config.dim, + moe_config.moe_inter_dim // 2, + ) + assert mx.gate_and_up_projs_packed.is_meta + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert mx.down_projs_scales.dtype == torch.float8_e8m0fnu + assert [n for n, p in mx.named_parameters() if p.requires_grad] == [] + + +def test_lora_passthrough_init_packed_base_with_trainable_adapters(moe_config): + """LoRA-on-experts passthrough: packed-at-init base (no bf16) + trainable LoRA + adapters. This is the path a packed fp4 checkpoint loads into when experts are + LoRA-targeted (previously size-mismatched against the packed checkpoint keys).""" + # Mirror the real flow (infra wraps PEFT swap in init_empty_weights/meta) so + # the GroupedExpertsLoRA base-weight copy is a meta->meta no-op. + with torch.device("meta"): + orig = GroupedExperts(moe_config) + orig.use_torch_mm = True + mx = GroupedExpertsLoRAMXFP4(orig, lora_dim=8, alpha=16, passthrough=True) + + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") # base is packed, never bf16 + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert mx.gate_and_up_projs_packed.is_meta + # Only the LoRA adapters are trainable; packed base is frozen. + trainable = {n for n, p in mx.named_parameters() if p.requires_grad} + assert trainable == {"lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_convert_frozen_experts_to_mxfp4(moe_config, device): + import torch.nn as nn + + class TinyModel(nn.Module): + def __init__(self): + super().__init__() + e = GroupedExperts(moe_config) + e.use_torch_mm = True + self.experts = e + + model = TinyModel().to(device) + with torch.no_grad(): + model.experts.init_weights(buffer_device=device) + + n = convert_frozen_experts_to_mxfp4(model) + assert n == 1 + assert isinstance(model.experts, GroupedExpertsMXFP4) + # Idempotent: an already-converted module is not re-wrapped. + assert convert_frozen_experts_to_mxfp4(model) == 0 diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py new file mode 100644 index 0000000000..0ea9171276 --- /dev/null +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py @@ -0,0 +1,295 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""mxfp4-resident experts under DeepEP dispatch. + +The DeepEP token all-to-all is mocked (``MockDeepEPDispatcher``) so the forward path is +exercised without a real DeepEP backend or process group — only the post-dispatch +grouped GEMM differs between bf16 and mxfp4, and that is what these tests pin down. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4, patch_moe_module +from nemo_automodel.components._peft.lora_experts import GroupedExpertsDeepEPLoRA +from nemo_automodel.components._peft.lora_experts_mxfp4 import GroupedExpertsDeepEPLoRAMXFP4 +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.experts import GroupedExpertsDeepEP +from nemo_automodel.components.moe.quantized_experts import GroupedExpertsDeepEPMXFP4, GroupedExpertsMXFP4 +from nemo_automodel.components.quantization.mxfp4 import dequantize_mxfp4, quantize_mxfp4 + + +@pytest.fixture +def device(): + if torch.cuda.is_available(): + return torch.device(f"cuda:{torch.cuda.current_device()}") + return torch.device("cpu") + + +@pytest.fixture +def moe_config(): + # Dims divisible by 32 so both contraction dims are mxfp4-blockable. + return MoEConfig( + n_routed_experts=4, + n_shared_experts=0, + n_activated_experts=2, + n_expert_groups=1, + n_limited_groups=1, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="softmax", + route_scale=1.0, + dim=64, + inter_dim=128, + moe_inter_dim=64, + norm_topk_prob=False, + expert_activation="swiglu", + dtype=torch.bfloat16, + ) + + +class MockDeepEPDispatcher: + """Mock dispatcher that returns pre-set permuted tensors (no comms).""" + + def __init__(self, permuted_x, tokens_per_expert, permuted_probs): + self.token_permutation2 = MagicMock(return_value=(permuted_x, tokens_per_expert, permuted_probs)) + + def token_unpermutation(self, hidden_states): + return hidden_states + + +def _make_representable_(experts) -> None: + """Replace expert weights in-place with their mxfp4 round-trip so storage is exact.""" + with torch.no_grad(): + for name in ("gate_and_up_projs", "down_projs"): + param = getattr(experts, name) + w_t = param.data.transpose(-2, -1).contiguous() + packed, scales = quantize_mxfp4(w_t) + param.data.copy_(dequantize_mxfp4(packed, scales, param.dtype).transpose(-2, -1)) + + +def _make_deepep(moe_config, device, *, use_torch_mm=True, representable=False): + """Build a materialized GroupedExpertsDeepEP with single-rank dispatcher state injected.""" + orig = GroupedExpertsDeepEP(moe_config).to(device).to(torch.bfloat16) + with torch.no_grad(): + orig.init_weights(device) + if representable: + _make_representable_(orig) + orig.n_routed_experts = moe_config.n_routed_experts + orig.ep_size = 1 + orig.ep_rank = 0 + orig.use_torch_mm = use_torch_mm + return orig + + +def _inject_dispatcher(module, num_tokens, dim, device): + """Attach single-rank dispatcher state + a mock that routes all tokens to expert 0.""" + module.n_routed_experts = 4 + module.ep_size = 1 + module.ep_rank = 0 + tokens_per_expert = torch.tensor([num_tokens, 0, 0, 0], dtype=torch.long, device="cpu") + permuted_x = torch.randn(num_tokens, dim, device=device, dtype=torch.bfloat16) + permuted_probs = torch.ones(num_tokens, device=device, dtype=torch.bfloat16) + module.token_dispatcher = MockDeepEPDispatcher(permuted_x, tokens_per_expert, permuted_probs) + return permuted_x, tokens_per_expert, permuted_probs + + +# --- wiring / construction ---------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_frozen_deepep_mxfp4_packs_and_freezes(moe_config, device): + orig = _make_deepep(moe_config, device) + mx = GroupedExpertsDeepEPMXFP4(orig) + + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") + assert not hasattr(mx, "down_projs") + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert mx.down_projs_scales.dtype == torch.float8_e8m0fnu + # Checkpoint orientation [E, out, in/2]. + assert mx.gate_and_up_projs_packed.shape == (4, 2 * moe_config.moe_inter_dim, moe_config.dim // 2) + # DeepEP dispatcher knobs are carried over. + assert mx.dispatcher_backend == orig.dispatcher_backend + assert mx.use_torch_mm is True + # Fully frozen. + assert [n for n, p in mx.named_parameters() if p.requires_grad] == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_lora_deepep_mxfp4_only_lora_trainable(moe_config, device): + orig = _make_deepep(moe_config, device) + mx = GroupedExpertsDeepEPLoRAMXFP4(orig, lora_dim=8, alpha=16) + + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + trainable = {n for n, p in mx.named_parameters() if p.requires_grad} + assert trainable == {"lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_patch_moe_module_deepep_mxfp4(moe_config, device): + orig = _make_deepep(moe_config, device) + patched = patch_moe_module(orig, dim=8, alpha=16, expert_weight_format="mxfp4") + assert isinstance(patched, GroupedExpertsDeepEPLoRAMXFP4) + + orig2 = _make_deepep(moe_config, device) + patched_bf16 = patch_moe_module(orig2, dim=8, alpha=16) + assert isinstance(patched_bf16, GroupedExpertsDeepEPLoRA) + assert not isinstance(patched_bf16, GroupedExpertsDeepEPLoRAMXFP4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_deepep_mxfp4_requires_torch_mm_backend(moe_config, device): + orig = _make_deepep(moe_config, device, use_torch_mm=False) # gmm path -> unsupported + with pytest.raises(NotImplementedError, match="torch_mm"): + GroupedExpertsDeepEPMXFP4(orig) + with pytest.raises(NotImplementedError, match="torch_mm"): + GroupedExpertsDeepEPLoRAMXFP4(orig, lora_dim=8, alpha=16) + + +def test_passthrough_deepep_registers_packed_params_on_meta(moe_config): + with torch.device("meta"): + orig = GroupedExpertsDeepEP(moe_config) + orig.use_torch_mm = True + + mx = GroupedExpertsDeepEPMXFP4(orig, passthrough=True) + assert mx._mxfp4_resident + assert not hasattr(mx, "gate_and_up_projs") # never created bf16 storage + up_proj_dim = 2 * moe_config.moe_inter_dim # gated + assert tuple(mx.gate_and_up_projs_packed.shape) == (moe_config.n_routed_experts, up_proj_dim, moe_config.dim // 2) + assert mx.gate_and_up_projs_packed.is_meta + assert mx.gate_and_up_projs_packed.dtype == torch.int8 + assert [n for n, p in mx.named_parameters() if p.requires_grad] == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_convert_frozen_experts_to_mxfp4_handles_deepep(moe_config, device): + import torch.nn as nn + + class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.experts = _make_deepep(moe_config, device) + + model = TinyModel() + n = convert_frozen_experts_to_mxfp4(model) + assert n == 1 + assert isinstance(model.experts, GroupedExpertsDeepEPMXFP4) + # Not the torch-path class. + assert not isinstance(model.experts, GroupedExpertsMXFP4) + # Idempotent. + assert convert_frozen_experts_to_mxfp4(model) == 0 + + +# --- forward numerics (mock dispatcher) --------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_frozen_deepep_mxfp4_forward_matches_bf16(moe_config, device): + """Frozen mxfp4 DeepEP forward must match the bf16 DeepEP forward on representable weights.""" + orig = _make_deepep(moe_config, device, representable=True) + mx = GroupedExpertsDeepEPMXFP4(orig) + + num_tokens = 8 + permuted_x, tokens_per_expert, permuted_probs = _inject_dispatcher(mx, num_tokens, moe_config.dim, device) + # Share the exact same mock return on the bf16 reference. + orig.token_dispatcher = MockDeepEPDispatcher(permuted_x, tokens_per_expert, permuted_probs) + + x = torch.randn(num_tokens, moe_config.dim, device=device, dtype=torch.bfloat16) + weights = torch.ones(num_tokens, 1, device=device, dtype=torch.bfloat16) + indices = torch.zeros(num_tokens, 1, dtype=torch.long, device=device) + token_mask = torch.ones(num_tokens, dtype=torch.bool, device=device) + + y_mx = mx(x, token_mask, weights, indices) + # GroupedExpertsDeepEP.forward calls .to_local() on plain Parameters; patch for the test. + with torch.no_grad(), patch.object(torch.Tensor, "to_local", new=lambda self: self, create=True): + y_ref = orig(x, token_mask, weights, indices) + + assert y_mx.shape == (num_tokens, moe_config.dim) + torch.testing.assert_close(y_mx, y_ref, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_lora_deepep_mxfp4_forward_matches_bf16(moe_config, device): + """LoRA mxfp4 DeepEP forward must match the bf16 DeepEP LoRA forward with the same adapters.""" + orig = _make_deepep(moe_config, device, representable=True) + + ref = GroupedExpertsDeepEPLoRA(orig, lora_dim=8, alpha=16).to(device).to(torch.bfloat16) + ref.use_torch_mm = True + # Pack on CPU-allocated base params, then move packed storage + adapters to device. + mx = GroupedExpertsDeepEPLoRAMXFP4(orig, lora_dim=8, alpha=16).to(device) + # Give LoRA non-trivial values (B is zero-init by default) and sync both modules. The + # DeepEP base allocates fp32 (it relies on a later .to(bf16)); assign the bf16 ref tensor + # outright so the mxfp4 adapters are bf16 — .to(bf16) on the module would corrupt the + # int8 packed base. + with torch.no_grad(): + for name in ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"): + getattr(ref, name).data.normal_(0, 0.02) + getattr(mx, name).data = getattr(ref, name).data.clone().to(device) + + num_tokens = 8 + permuted_x, tokens_per_expert, permuted_probs = _inject_dispatcher(mx, num_tokens, moe_config.dim, device) + ref.token_dispatcher = MockDeepEPDispatcher(permuted_x, tokens_per_expert, permuted_probs) + ref.n_routed_experts = 4 + ref.ep_size = 1 + + x = torch.randn(num_tokens, moe_config.dim, device=device, dtype=torch.bfloat16) + weights = torch.ones(num_tokens, 1, device=device, dtype=torch.bfloat16) + indices = torch.zeros(num_tokens, 1, dtype=torch.long, device=device) + token_mask = torch.ones(num_tokens, dtype=torch.bool, device=device) + + y_mx = mx(x, token_mask, weights, indices) + with torch.no_grad(), patch.object(torch.Tensor, "to_local", new=lambda self: self, create=True): + y_ref = ref(x, token_mask, weights, indices) + + assert y_mx.shape == (num_tokens, moe_config.dim) + torch.testing.assert_close(y_mx, y_ref, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_lora_deepep_mxfp4_no_routed_tokens(moe_config, device): + """The all-zero-tokens dummy path must run and keep LoRA gradients flowing.""" + orig = _make_deepep(moe_config, device) + mx = GroupedExpertsDeepEPLoRAMXFP4(orig, lora_dim=8, alpha=16).to(device) + # DeepEP base allocates fp32; make adapters bf16 to match activations (see note above). + with torch.no_grad(): + for name in ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B"): + p = getattr(mx, name) + p.data = p.data.to(torch.bfloat16) + + num_tokens = 8 + # tokens_per_expert all zeros -> dummy path. + permuted_x = torch.randn(num_tokens, moe_config.dim, device=device, dtype=torch.bfloat16) + permuted_probs = torch.ones(num_tokens, device=device, dtype=torch.bfloat16) + tokens_per_expert = torch.zeros(4, dtype=torch.long, device="cpu") + mx.n_routed_experts = 4 + mx.ep_size = 1 + mx.token_dispatcher = MockDeepEPDispatcher(permuted_x, tokens_per_expert, permuted_probs) + + x = torch.randn(num_tokens, moe_config.dim, device=device, dtype=torch.bfloat16, requires_grad=True) + weights = torch.ones(num_tokens, 1, device=device, dtype=torch.bfloat16) + indices = torch.zeros(num_tokens, 1, dtype=torch.long, device=device) + token_mask = torch.ones(num_tokens, dtype=torch.bool, device=device) + + y = mx(x, token_mask, weights, indices) + assert y.shape == x.shape + assert torch.isfinite(y).all() + y.float().sum().backward() + assert mx.lora_gate_and_up_A.grad is not None diff --git a/tests/unit_tests/_transformers/test_auto_model.py b/tests/unit_tests/_transformers/test_auto_model.py index 4c2abbf6ce..42068366e1 100644 --- a/tests/unit_tests/_transformers/test_auto_model.py +++ b/tests/unit_tests/_transformers/test_auto_model.py @@ -33,7 +33,11 @@ _patch_remote_code_compat, _resolve_distributed_setup, ) -from nemo_automodel._transformers.infrastructure import _apply_peft_and_lower_precision, instantiate_infrastructure +from nemo_automodel._transformers.infrastructure import ( + _apply_peft_and_lower_precision, + _call_model_hook, + instantiate_infrastructure, +) from nemo_automodel._transformers.model_init import ( _filter_kwargs_for_init, _filter_meta_device_from_init_context, @@ -799,6 +803,51 @@ def test_apply_peft_disables_triton_with_autopipeline(self, caplog): assert mock_peft_config.use_triton is False assert "Disabling Triton with Pipeline Parallelism" in caplog.text + def test_apply_peft_calls_explicit_model_hook(self): + """Model-owned PEFT setup runs after generic LoRA patching.""" + + class HookedModel: + def __init__(self): + self.prepared_with = None + + def prepare_peft_checkpoint_load(self, peft_config): + self.prepared_with = peft_config + + model = HookedModel() + peft_config = MagicMock() + + with patch("nemo_automodel._transformers.infrastructure.apply_lora_to_linear_modules") as apply_lora: + _apply_peft_and_lower_precision( + model, + tp_size=1, + autopipeline=None, + peft_config=peft_config, + quantization_config=None, + fp8_config=None, + qat_quantizer=None, + ) + + apply_lora.assert_called_once() + assert model.prepared_with is peft_config + + def test_model_hook_calls_each_explicit_pipeline_part(self): + """Model-owned post-load setup is preserved for pipeline model parts.""" + + class HookedPart: + def __init__(self): + self.finalized_with = None + + def finalize_peft_checkpoint_load(self, peft_config): + self.finalized_with = peft_config + + parts = [HookedPart(), HookedPart()] + model = types.SimpleNamespace(parts=parts) + peft_config = object() + + _call_model_hook(model, "finalize_peft_checkpoint_load", peft_config) + + assert all(part.finalized_with is peft_config for part in parts) + def test_apply_fp8_when_configured(self): """When fp8_config provided, calls apply_fp8_to_model.""" mock_model = MagicMock() diff --git a/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py new file mode 100644 index 0000000000..837b525816 --- /dev/null +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py @@ -0,0 +1,211 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""mxfp4 passthrough load path for DeepSeek-V4: experts load packed, never bf16.""" + +import json +import os +from unittest.mock import Mock + +import pytest +import torch + +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.deepseek_v4.config import DeepseekV4Config +from nemo_automodel.components.models.deepseek_v4.state_dict_adapter import DeepSeekV4StateDictAdapter +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.quantization.mxfp4 import dequantize_mxfp4, quantize_mxfp4 + +# Real DeepSeek-V4-Flash checkpoint (fp4 experts). Tests using it self-skip when absent. +_REAL_CKPT = "/raid0/data/models/DeepSeek-V4-Flash" + +HIDDEN = 64 +MOE_INTER = 32 +N_EXPERTS = 4 + + +def _make_adapter(expert_storage_format="bf16"): + config = DeepseekV4Config( + vocab_size=256, + hidden_size=HIDDEN, + num_hidden_layers=2, + num_attention_heads=4, + head_dim=16, + qk_rope_head_dim=8, + q_lora_rank=32, + o_lora_rank=32, + o_groups=2, + n_routed_experts=N_EXPERTS, + num_experts_per_tok=2, + moe_intermediate_size=MOE_INTER, + num_nextn_predict_layers=0, + ) + moe_config = Mock(spec=MoEConfig) + moe_config.n_routed_experts = N_EXPERTS + moe_config.moe_inter_dim = MOE_INTER + return DeepSeekV4StateDictAdapter( + config, moe_config, BackendConfig(), dtype=torch.bfloat16, expert_storage_format=expert_storage_format + ) + + +def _synthetic_fp4_checkpoint(seed=0): + """Per-expert packed fp4 (int8) weights + e8m0 scales for one MoE layer. + + w1/w3 (gate/up): [moe_inter, hidden] checkpoint orientation; w2 (down): + [hidden, moe_inter]. Returns the HF-format dict plus the bf16 weights each + packed tensor decodes to (for reference). + """ + torch.manual_seed(seed) + sd = {} + ref = {} # (expert, which) -> dequantized bf16 weight [out, in] + for e in range(N_EXPERTS): + for which, (out_dim, in_dim) in { + "w1": (MOE_INTER, HIDDEN), + "w3": (MOE_INTER, HIDDEN), + "w2": (HIDDEN, MOE_INTER), + }.items(): + w = torch.randn(out_dim, in_dim, dtype=torch.bfloat16) + packed, scales = quantize_mxfp4(w) + base = f"layers.0.ffn.experts.{e}.{which}" + sd[base + ".weight"] = packed + sd[base + ".scale"] = scales + ref[(e, which)] = dequantize_mxfp4(packed, scales, torch.bfloat16) + return sd, ref + + +def test_passthrough_emits_packed_keys_no_bf16(): + sd, _ = _synthetic_fp4_checkpoint() + out = _make_adapter("mxfp4").from_hf(dict(sd)) + + gu_packed = out["model.layers.0.mlp.experts.gate_and_up_projs_packed"] + gu_scales = out["model.layers.0.mlp.experts.gate_and_up_projs_scales"] + dn_packed = out["model.layers.0.mlp.experts.down_projs_packed"] + dn_scales = out["model.layers.0.mlp.experts.down_projs_scales"] + + # Packed storage, never materialized to bf16. + assert gu_packed.dtype == torch.int8 and gu_scales.dtype == torch.float8_e8m0fnu + assert dn_packed.dtype == torch.int8 and dn_scales.dtype == torch.float8_e8m0fnu + # gate||up concatenated along the output dim: [E, 2*moe_inter, hidden//2]. + assert tuple(gu_packed.shape) == (N_EXPERTS, 2 * MOE_INTER, HIDDEN // 2) + assert tuple(gu_scales.shape) == (N_EXPERTS, 2 * MOE_INTER, HIDDEN // 32) + assert tuple(dn_packed.shape) == (N_EXPERTS, HIDDEN, MOE_INTER // 2) + # No bf16 expert weight keys, and no orphaned scale keys leaked through. + assert not any(k.endswith("gate_and_up_projs") or k.endswith("down_projs") for k in out) + assert not any(".ffn.experts." in k for k in out) + + +def test_passthrough_decodes_to_same_weights_as_bf16_path(): + """Unpacking the passthrough output must equal the bf16 dequant+aggregate path.""" + sd, _ = _synthetic_fp4_checkpoint() + out_mx = _make_adapter("mxfp4").from_hf(dict(sd)) + out_bf16 = _make_adapter("bf16").from_hf(dict(sd)) + + # bf16 path: gate_and_up_projs is compute layout [E, in=hidden, 2*moe_inter]. + gu_bf16 = out_bf16["model.layers.0.mlp.experts.gate_and_up_projs"] + dn_bf16 = out_bf16["model.layers.0.mlp.experts.down_projs"] + + # passthrough: unpack [E, 2*moe_inter, hidden] then transpose to compute layout. + gu_unpacked = dequantize_mxfp4( + out_mx["model.layers.0.mlp.experts.gate_and_up_projs_packed"], + out_mx["model.layers.0.mlp.experts.gate_and_up_projs_scales"], + torch.bfloat16, + ).transpose(-2, -1) + dn_unpacked = dequantize_mxfp4( + out_mx["model.layers.0.mlp.experts.down_projs_packed"], + out_mx["model.layers.0.mlp.experts.down_projs_scales"], + torch.bfloat16, + ).transpose(-2, -1) + + assert torch.equal(gu_unpacked, gu_bf16) + assert torch.equal(dn_unpacked, dn_bf16) + + +def test_passthrough_roundtrip_to_hf_recovers_checkpoint_keys(): + """from_hf (aggregate) -> to_hf (split) must reproduce the per-expert packed + checkpoint keys with matching dtypes and bit-exact values — this is the path + the DCP loader uses to enumerate destination tensors.""" + sd, _ = _synthetic_fp4_checkpoint() + adapter = _make_adapter("mxfp4") + model_sd = adapter.from_hf(dict(sd)) + hf = adapter.to_hf(model_sd, quantization=True) + + for e in range(N_EXPERTS): + for w in ("w1", "w2", "w3"): + wk = f"layers.0.ffn.experts.{e}.{w}.weight" + sk = f"layers.0.ffn.experts.{e}.{w}.scale" + assert wk in hf and sk in hf, f"missing {wk}/{sk}" + assert hf[wk].dtype == torch.int8 + assert hf[sk].dtype == torch.float8_e8m0fnu + # cat-then-split is identity: recovered packed bytes equal the original. + assert torch.equal(hf[wk], sd[wk]), f"{wk} value mismatch" + assert torch.equal(hf[sk].view(torch.uint8), sd[sk].view(torch.uint8)), f"{sk} value mismatch" + + +@pytest.mark.skipif( + not os.path.isdir(_REAL_CKPT) or not os.path.isfile(f"{_REAL_CKPT}/model.safetensors.index.json"), + reason=f"real DeepSeek-V4-Flash checkpoint not present at {_REAL_CKPT}", +) +def test_passthrough_against_real_checkpoint(): + """Validate passthrough on real fp4 expert tensors: correct dtypes/shapes and + bit-exact decode vs the bf16 dequant path on actual checkpoint bytes.""" + from safetensors import safe_open + + layer, n_exp = 3, 8 + weight_map = json.load(open(f"{_REAL_CKPT}/model.safetensors.index.json"))["weight_map"] + needed = [ + f"layers.{layer}.ffn.experts.{e}.{w}.{suffix}" + for e in range(n_exp) + for w in ("w1", "w2", "w3") + for suffix in ("weight", "scale") + ] + by_file: dict[str, list[str]] = {} + for k in needed: + by_file.setdefault(weight_map[k], []).append(k) + sd = {} + for fname, keys in by_file.items(): + with safe_open(f"{_REAL_CKPT}/{fname}", framework="pt") as h: + for k in keys: + sd[k] = h.get_tensor(k) + + # Real layout: int8 packed weights + float8_e8m0fnu scales. + w1 = sd[f"layers.{layer}.ffn.experts.0.w1.weight"] + assert w1.dtype == torch.int8 + assert sd[f"layers.{layer}.ffn.experts.0.w1.scale"].dtype == torch.float8_e8m0fnu + + cfg = DeepseekV4Config.from_pretrained(_REAL_CKPT) + moe_config = Mock(spec=MoEConfig) + moe_config.n_routed_experts = n_exp + moe_config.moe_inter_dim = cfg.moe_intermediate_size + + def adapter(fmt): + return DeepSeekV4StateDictAdapter( + cfg, moe_config, BackendConfig(), dtype=torch.bfloat16, expert_storage_format=fmt + ) + + out_mx = adapter("mxfp4").from_hf(dict(sd)) + out_bf16 = adapter("bf16").from_hf({k: v.clone() for k, v in sd.items()}) + + base = f"model.layers.{layer}.mlp.experts" + assert out_mx[f"{base}.gate_and_up_projs_packed"].dtype == torch.int8 + assert out_mx[f"{base}.gate_and_up_projs_packed"].shape == ( + n_exp, + 2 * cfg.moe_intermediate_size, + cfg.hidden_size // 2, + ) + + for proj in ("gate_and_up_projs", "down_projs"): + unpacked = dequantize_mxfp4( + out_mx[f"{base}.{proj}_packed"], out_mx[f"{base}.{proj}_scales"], torch.bfloat16 + ).transpose(-2, -1) + assert torch.equal(unpacked, out_bf16[f"{base}.{proj}"]), f"{proj} decode mismatch vs bf16 path" diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index 77dbbe7fe4..9be0749b53 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -100,8 +100,9 @@ def _install_torch_and_layers_stubs(monkeypatch): nn_stub = types.ModuleType("torch.nn") class Parameter: - def __init__(self, data=None): + def __init__(self, data=None, requires_grad=True): self.data = data + self.requires_grad = requires_grad class Module: pass