From f04ff6903f23b73d01def5b3e43717fd010ba0f2 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 4 Jun 2026 23:30:46 +0000 Subject: [PATCH 01/29] chore(dsv4): fused AdamW + lora attn=tilelang Signed-off-by: Daniel --- .../llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml | 1 + .../deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) 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 4018225f99..556915acf4 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 839872ea99..073f0c559e 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 From f1543b0ce76752698bccf8c98d1aaf882333e018 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 4 Jun 2026 23:30:46 +0000 Subject: [PATCH 02/29] fix(moe): lazy DeepEP buffer alloc (single-node OOM workaround) Signed-off-by: Daniel --- nemo_automodel/components/moe/experts.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 559435180c..1218f3d3c0 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -693,8 +693,15 @@ def init_token_dispatcher(self, ep_mesh: DeviceMesh): config=config, ep_group=ep_group, ) - if self.dispatcher_backend == "deepep": - self._init_deepep_buffer(ep_group) + # NOTE: previously called `self._init_deepep_buffer(ep_group)` here to + # eagerly allocate the DeepEP NVSHMEM buffer at model construction + # (introduced in #2076, e42584e3). On single-node EP=8 ep_shard=1 + # DSv4-Flash, the eager allocation collides with the ~135 GB load-time + # peak and OOMs at `_aggregate_experts` torch.stack. Revert to the + # original lazy allocation in FusedDispatch.forward (fused_a2a.py:136 + # via the global `_buffer` cache). Both code paths produce the same + # buffer; only the *timing* differs. + # _init_deepep_buffer remains defined below for explicit callers. def _init_deepep_buffer(self, ep_group: dist.ProcessGroup) -> None: """Initialize DeepEP communication buffers before activation checkpointing.""" From 09c6288b03ed94b30645bdbbdba41e49f6966865 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 00:47:45 +0000 Subject: [PATCH 03/29] feat(peft): mxfp4-resident MoE expert base weights for LoRA Keep frozen routed-expert base weights packed as fp4-e2m1 + e8m0 block scales (the DeepSeek V4 Flash checkpoint format) at steady state during LoRA training, dequantizing on the fly in forward and backward via a custom grouped-GEMM autograd function that saves only the packed tensors. Opt in with peft.expert_weight_format: mxfp4 (torch_mm experts backend only). Packing happens at PEFT-swap time when weights are materialized, or after checkpoint load for meta-initialized models. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- .../_transformers/infrastructure.py | 9 + nemo_automodel/components/_peft/lora.py | 44 +++- .../components/_peft/lora_experts.py | 202 ++++++++++++++++ nemo_automodel/components/moe/fp4_utils.py | 139 +++++++++++ .../_peft/test_lora_experts_mxfp4.py | 215 ++++++++++++++++++ 5 files changed, 606 insertions(+), 3 deletions(-) create mode 100644 nemo_automodel/components/moe/fp4_utils.py create mode 100644 tests/unit_tests/_peft/test_lora_experts_mxfp4.py diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 2accc46d5f..574e4afd4a 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -576,6 +576,15 @@ def apply_model_infrastructure( if "lora_" not in name and param.requires_grad: param.requires_grad_(False) + # Pack deferred mxfp4-resident expert base weights now that the checkpoint is loaded. + if peft_config is not None and getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": + from nemo_automodel.components._peft.lora import pack_mxfp4_expert_base_weights + + for mp in model.parts if hasattr(model, "parts") else [model]: + num_packed = pack_mxfp4_expert_base_weights(mp) + if num_packed: + logger.info("Packed %d MoE expert modules to mxfp4-resident storage", num_packed) + 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 663f8a1643..5f63162c65 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -23,7 +23,11 @@ from torch.distributed.tensor import DTensor from torch.distributed.tensor.placement_types import Shard as _Shard -from nemo_automodel.components._peft.lora_experts import GroupedExpertsDeepEPLoRA, GroupedExpertsLoRA +from nemo_automodel.components._peft.lora_experts import ( + GroupedExpertsDeepEPLoRA, + GroupedExpertsLoRA, + GroupedExpertsLoRAMXFP4, +) from nemo_automodel.components._peft.lora_kernel import ( lora_da_dx_update_wrapper, lora_db_update_wrapper, @@ -55,6 +59,10 @@ class PeftConfig: lora_dtype: Optional[torch.dtype] = None 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() @@ -74,6 +82,7 @@ def from_dict(cls, d: dict[str, Any]): lora_dtype=d.get("lora_dtype", None), 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"), ) @@ -459,6 +468,7 @@ def patch_moe_module( alpha=32, lora_A_init_method="xavier", lora_dtype=None, + expert_weight_format="bf16", ): """ Patches a custom MoE module (GroupedExperts or GroupedExpertsDeepEP) with LoRA. @@ -469,13 +479,20 @@ 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". 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}") if isinstance(orig_module, GroupedExpertsTE): raise NotImplementedError("LoRA is not supported for Transformer Engine (TE) expert modules.") elif isinstance(orig_module, GroupedExpertsDeepEP): + if expert_weight_format == "mxfp4": + raise NotImplementedError("expert_weight_format='mxfp4' is not supported for DeepEP expert modules yet.") new_module = GroupedExpertsDeepEPLoRA( orig_module, lora_dim=dim, @@ -484,7 +501,8 @@ def patch_moe_module( lora_dtype=lora_dtype, ) elif isinstance(orig_module, GroupedExperts): - new_module = GroupedExpertsLoRA( + lora_cls = GroupedExpertsLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsLoRA + new_module = lora_cls( orig_module, lora_dim=dim, alpha=alpha, @@ -577,6 +595,7 @@ def apply_lora_to_linear_modules( 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, ) # Find parent and replace @@ -612,6 +631,25 @@ def apply_lora_to_linear_modules( return num_modules_matched +def pack_mxfp4_expert_base_weights(model: nn.Module) -> int: + """Pack any deferred mxfp4-resident expert modules after base weights are loaded. + + GroupedExpertsLoRAMXFP4 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 layer coexist with their packed copy. + + Returns: + Number of modules packed. + """ + num_packed = 0 + for module in model.modules(): + if isinstance(module, GroupedExpertsLoRAMXFP4) and not module._mxfp4_resident: + module.pack_base_weights() + num_packed += 1 + return num_packed + + class LoRATritonFunction(torch.autograd.Function): """ Autograd function that calls the triton kernel wrappers for the LoRA forward and backward passes. diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index 18b783ea9c..d2ba47b77a 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -24,6 +24,7 @@ _apply_bias, _permute_tokens_for_grouped_mm, ) +from nemo_automodel.components.moe.fp4_utils import MXFP4GroupedMM, dequantize_mxfp4, quantize_mxfp4 from nemo_automodel.shared.utils import dtype_from_str try: @@ -352,6 +353,207 @@ def _forward_grouped_mm( return y +class GroupedExpertsLoRAMXFP4(GroupedExpertsLoRA): + """GroupedExperts + LoRA with the frozen base weights resident in packed mxfp4. + + The base gate/up and down projections are stored transposed relative to the + compute layout (``[n_experts, out_dim, in_dim]``, matching the DeepSeek V4 Flash + checkpoint orientation) as packed fp4-e2m1 int8 plus ``float8_e8m0fnu`` block + scales, and are dequantized on the fly inside ``MXFP4GroupedMM`` during forward + and backward. Only the LoRA adapters (and optional expert biases) remain in + floating point. + + When constructed from a module whose weights are still on the meta device, + packing is deferred: the module behaves exactly like ``GroupedExpertsLoRA`` + until ``pack_base_weights()`` is called (after the base checkpoint is loaded). + """ + + def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_init_method="xavier", lora_dtype=None): + super().__init__( + orig_module, + lora_dim=lora_dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) + if not self.use_torch_mm: + raise NotImplementedError( + "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm')." + ) + self._mxfp4_resident = False + if not _to_local(self.gate_and_up_projs).is_meta: + self.pack_base_weights() + + @torch.no_grad() + def pack_base_weights(self) -> None: + """Pack the frozen base weights to mxfp4 and free the floating-point tensors. + + No-op when already packed. Must be called after the base weights are + materialized (i.e. after checkpoint load in the deferred-packing flow). + """ + if self._mxfp4_resident: + return + for name in ("gate_and_up_projs", "down_projs"): + param = getattr(self, name) + local = _to_local(param) + assert not local.is_meta, "pack_base_weights requires materialized base weights" + # [E, in, out] -> [E, out, in] so the mx block scales run along the + # contraction dim, matching the checkpoint orientation. + packed, scales = quantize_mxfp4(local.transpose(-2, -1).contiguous()) + if isinstance(param, DTensor): + packed = DTensor.from_local(packed, param.device_mesh, param.placements) + scales = DTensor.from_local(scales, param.device_mesh, param.placements) + del self._parameters[name] + self.register_parameter(name + "_packed", nn.Parameter(packed, requires_grad=False)) + self.register_parameter(name + "_scales", nn.Parameter(scales, requires_grad=False)) + self._mxfp4_resident = True + + 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 + + gate_and_up_packed = _to_local(self.gate_and_up_projs_packed) + gate_and_up_scales = _to_local(self.gate_and_up_projs_scales) + down_packed = _to_local(self.down_projs_packed) + down_scales = _to_local(self.down_projs_scales) + lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) + lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) + lora_down_A = _to_local(self.lora_down_A) + lora_down_B = _to_local(self.lora_down_B) + + 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, + gate_and_up_packed, + gate_and_up_scales, + down_packed, + down_scales, + lora_gate_and_up_A, + lora_gate_and_up_B, + lora_down_A, + lora_down_B, + 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, + gate_and_up_packed, + gate_and_up_scales, + down_packed, + down_scales, + lora_gate_and_up_A, + lora_gate_and_up_B, + lora_down_A, + lora_down_B, + 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, + ) + + 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 = MXFP4GroupedMM.apply(permuted_x, gate_and_up_packed, gate_and_up_scales, 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 = MXFP4GroupedMM.apply(output1, down_packed, down_scales, 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 = dequantize_mxfp4(gate_and_up_packed[0], gate_and_up_scales[0], x.dtype).transpose(-2, -1) + down_w0 = dequantize_mxfp4(down_packed[0], down_scales[0], x.dtype).transpose(-2, -1) + 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 GroupedExpertsDeepEPLoRA(GroupedExpertsDeepEP): """ GroupedExpertsDeepEP + LoRA. diff --git a/nemo_automodel/components/moe/fp4_utils.py b/nemo_automodel/components/moe/fp4_utils.py new file mode 100644 index 0000000000..e7527becf9 --- /dev/null +++ b/nemo_automodel/components/moe/fp4_utils.py @@ -0,0 +1,139 @@ +# 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) + + +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) + low = (packed_u8 & 0x0F).long() + high = ((packed_u8 >> 4) & 0x0F).long() + table = _FP4_E2M1_TABLE.to(packed_u8.device) + # Interleave (low, high) per byte so column indices match the original layout. + fp4_vals = torch.stack([table[low], table[high]], dim=-1).flatten(-2) + + # Decode e8m0 to fp32: 2^(e - 127), with byte 0 mapping to 0 (all-zero block). + scale_u8 = scales.contiguous().view(torch.uint8).int() + scale_f32 = torch.where( + scale_u8 == 0, + torch.zeros_like(scale_u8, dtype=torch.float32), + torch.pow(2.0, (scale_u8 - 127).float()), + ) + + scale_expanded = scale_f32.repeat_interleave(MXFP4_BLOCK_SIZE, dim=-1) + scale_expanded = scale_expanded[..., : fp4_vals.shape[-1]] + return (fp4_vals * scale_expanded).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. amax == 0 maps to scale byte 0 (decoded as 0). + nonzero = amax > 0 + exp = torch.zeros_like(amax) + exp[nonzero] = torch.floor(torch.log2(amax[nonzero])) - 2.0 + scale_bytes = torch.where( + nonzero, + (exp + 127.0).clamp(1.0, 254.0), + torch.zeros_like(exp), + ).to(torch.uint8) + # Zero-amax blocks divide by 1 instead of 0; their codes are all zero anyway. + scale = torch.where(nonzero, torch.pow(2.0, (scale_bytes.int() - 127).float()), torch.ones_like(exp)) + + # Round each scaled magnitude to the nearest e2m1 magnitude code. + 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(), scale_bytes.view(torch.float8_e8m0fnu).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 transposed relative to the GEMM operand + (``[E, N, K]`` packed along ``K``), which makes the backward GEMM operand + (``W^T``) the natural dequantization output; forward pays one transpose-copy. + + 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] + out = torch._grouped_mm(x, w_t.transpose(-2, -1).contiguous(), 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/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..aac1ebafc8 --- /dev/null +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -0,0 +1,215 @@ +# 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 patch_moe_module +from nemo_automodel.components._peft.lora_experts import GroupedExpertsLoRA, GroupedExpertsLoRAMXFP4 +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.fp4_utils import dequantize_mxfp4, quantize_mxfp4 +from nemo_automodel.components.moe.layers import GroupedExperts + + +@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) From ae75364360f61850d2fd92f1e9df1cf43db4d107 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 17:44:29 +0000 Subject: [PATCH 04/29] feat(peft): mxfp4-resident frozen MoE experts for DeepSeek-V4-Flash LoRA Extend mxfp4-resident expert storage to FROZEN (non-LoRA-targeted) routed experts, which is where the memory win actually lands for LoRA-on-attention recipes. Adds GroupedExpertsMXFP4 (frozen, packed base, dequant-on-the-fly) and convert_frozen_experts_to_mxfp4(), invoked from the PEFT-application step when peft.expert_weight_format=mxfp4. Packing of both frozen and LoRA-targeted experts is deferred until after checkpoint load. Format-specific pack/unpack/GEMM logic is factored into MXFP4ExpertStorageMixin (shared by the frozen and LoRA variants) as the seam for a future int4 codec. Only experts are quantized; all other weights stay bf16. v1 supports the torch_mm GroupedExperts backend (DeepEP experts are skipped with a warning). Adds an example DSV4-Flash LoRA+mxfp4 config. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- ...eepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 144 ++++++++++++ .../_transformers/infrastructure.py | 9 + nemo_automodel/components/_peft/lora.py | 51 ++++- .../components/_peft/lora_experts.py | 104 ++------- .../components/moe/quantized_experts.py | 209 ++++++++++++++++++ .../_peft/test_lora_experts_mxfp4.py | 61 ++++- 6 files changed, 487 insertions(+), 91 deletions(-) create mode 100644 examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml create mode 100644 nemo_automodel/components/moe/quantized_experts.py 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..b469e569d7 --- /dev/null +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml @@ -0,0 +1,144 @@ +# 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 + 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: + - "*wq_a" + - "*wq_b" + - "*wkv" + - "*wo_b" + dim: 8 + alpha: 32 + use_triton: True + # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly. + expert_weight_format: mxfp4 + +checkpoint: + enabled: false + # The DSV4-Flash checkpoint stores experts as fp4; they are dequantized to bf16 + # on load and then repacked to mxfp4 (Phase A). Required for the current flow. + dequantize_base_checkpoint: true + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +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/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 574e4afd4a..1caf56a67a 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -89,6 +89,15 @@ def _apply_peft_and_lower_precision( # 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) + # Convert frozen (non-LoRA-targeted) routed experts to mxfp4-resident storage. + # LoRA-targeted experts are already GroupedExpertsLoRAMXFP4 from the call above. + # Packing of both is deferred until after the checkpoint is loaded. + if getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": + from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4 + + num_converted = convert_frozen_experts_to_mxfp4(model) + logger.info("Converted %d frozen expert module(s) to mxfp4-resident storage", num_converted) + # FP8 if fp8_config is not None: model = apply_fp8_to_model(model, config=fp8_config) diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 5f63162c65..0391de5780 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -35,6 +35,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.quantized_experts import GroupedExpertsMXFP4, MXFP4ExpertStorageMixin from nemo_automodel.shared.import_utils import safe_import, safe_import_te from nemo_automodel.shared.utils import dtype_from_str @@ -631,20 +632,60 @@ def apply_lora_to_linear_modules( return num_modules_matched +def convert_frozen_experts_to_mxfp4(model: nn.Module) -> 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. + + 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 + + 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, (GroupedExpertsDeepEP, GroupedExpertsTE)): + unsupported += 1 + continue + if type(module) is GroupedExperts: + new_module = GroupedExpertsMXFP4(module) + 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 DeepEP/TE expert module(s); only the torch_mm " + "GroupedExperts backend is supported. Set backend.dispatcher='torch' and backend.experts='torch_mm'.", + 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. - GroupedExpertsLoRAMXFP4 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 layer coexist with their packed copy. + 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, GroupedExpertsLoRAMXFP4) and not module._mxfp4_resident: + if isinstance(module, MXFP4ExpertStorageMixin) and not module._mxfp4_resident: module.pack_base_weights() num_packed += 1 return num_packed diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index d2ba47b77a..438129c79d 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -24,7 +24,7 @@ _apply_bias, _permute_tokens_for_grouped_mm, ) -from nemo_automodel.components.moe.fp4_utils import MXFP4GroupedMM, dequantize_mxfp4, quantize_mxfp4 +from nemo_automodel.components.moe.quantized_experts import MXFP4ExpertStorageMixin from nemo_automodel.shared.utils import dtype_from_str try: @@ -353,15 +353,14 @@ def _forward_grouped_mm( return y -class GroupedExpertsLoRAMXFP4(GroupedExpertsLoRA): +class GroupedExpertsLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsLoRA): """GroupedExperts + LoRA with the frozen base weights resident in packed mxfp4. - The base gate/up and down projections are stored transposed relative to the - compute layout (``[n_experts, out_dim, in_dim]``, matching the DeepSeek V4 Flash - checkpoint orientation) as packed fp4-e2m1 int8 plus ``float8_e8m0fnu`` block - scales, and are dequantized on the fly inside ``MXFP4GroupedMM`` during forward - and backward. Only the LoRA adapters (and optional expert biases) remain in - floating point. + 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. When constructed from a module whose weights are still on the meta device, packing is deferred: the module behaves exactly like ``GroupedExpertsLoRA`` @@ -376,37 +375,7 @@ def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_ini lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype, ) - if not self.use_torch_mm: - raise NotImplementedError( - "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm')." - ) - self._mxfp4_resident = False - if not _to_local(self.gate_and_up_projs).is_meta: - self.pack_base_weights() - - @torch.no_grad() - def pack_base_weights(self) -> None: - """Pack the frozen base weights to mxfp4 and free the floating-point tensors. - - No-op when already packed. Must be called after the base weights are - materialized (i.e. after checkpoint load in the deferred-packing flow). - """ - if self._mxfp4_resident: - return - for name in ("gate_and_up_projs", "down_projs"): - param = getattr(self, name) - local = _to_local(param) - assert not local.is_meta, "pack_base_weights requires materialized base weights" - # [E, in, out] -> [E, out, in] so the mx block scales run along the - # contraction dim, matching the checkpoint orientation. - packed, scales = quantize_mxfp4(local.transpose(-2, -1).contiguous()) - if isinstance(param, DTensor): - packed = DTensor.from_local(packed, param.device_mesh, param.placements) - scales = DTensor.from_local(scales, param.device_mesh, param.placements) - del self._parameters[name] - self.register_parameter(name + "_packed", nn.Parameter(packed, requires_grad=False)) - self.register_parameter(name + "_scales", nn.Parameter(scales, requires_grad=False)) - self._mxfp4_resident = True + 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. @@ -434,15 +403,6 @@ def forward(self, x: torch.Tensor, token_mask: torch.Tensor, weights: torch.Tens assert self.n_routed_experts % ep_size == 0 - gate_and_up_packed = _to_local(self.gate_and_up_projs_packed) - gate_and_up_scales = _to_local(self.gate_and_up_projs_scales) - down_packed = _to_local(self.down_projs_packed) - down_scales = _to_local(self.down_projs_scales) - lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) - lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) - lora_down_A = _to_local(self.lora_down_A) - lora_down_B = _to_local(self.lora_down_B) - if ep_size > 1: x = DTensor.from_local(x, device_mesh=ep_mesh, placements=[Shard(0)]).full_tensor( grad_placements=[Partial()] @@ -456,22 +416,7 @@ def forward(self, x: torch.Tensor, token_mask: torch.Tensor, weights: torch.Tens 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, - gate_and_up_packed, - gate_and_up_scales, - down_packed, - down_scales, - lora_gate_and_up_A, - lora_gate_and_up_B, - lora_down_A, - lora_down_B, - n_local_experts, - experts_start_idx, - ) + 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()]) @@ -479,23 +424,7 @@ def forward(self, x: torch.Tensor, token_mask: torch.Tensor, weights: torch.Tens return y.to(input_dtype) - def _forward_grouped_mm_mxfp4( - self, - x, - token_mask, - weights, - indices, - gate_and_up_packed, - gate_and_up_scales, - down_packed, - down_scales, - lora_gate_and_up_A, - lora_gate_and_up_B, - lora_down_A, - lora_down_B, - n_local_experts, - experts_start_idx, - ): + 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, @@ -505,6 +434,11 @@ def _forward_grouped_mm_mxfp4( experts_start_idx, ) + lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) + lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) + lora_down_A = _to_local(self.lora_down_A) + lora_down_B = _to_local(self.lora_down_B) + y = torch.zeros(x.shape, dtype=torch.float32, device=x.device) if tokens_per_expert.sum() > 0: @@ -516,7 +450,7 @@ def _forward_grouped_mm_mxfp4( down_proj_bias = _to_local(self.down_proj_bias) # Gate+Up projection + LoRA - output1 = MXFP4GroupedMM.apply(permuted_x, gate_and_up_packed, gate_and_up_scales, offs) + 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 @@ -527,7 +461,7 @@ def _forward_grouped_mm_mxfp4( output1 = self.expert_activation_grouped(output1, permuted_probs) # Down projection + LoRA - output2 = MXFP4GroupedMM.apply(output1, down_packed, down_scales, offs) + 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 @@ -539,8 +473,8 @@ def _forward_grouped_mm_mxfp4( y.scatter_add_(0, scatter_ids, output2.float()) else: # Dummy computation for gradient flow; dequantize only expert 0. - gate_up_w0 = dequantize_mxfp4(gate_and_up_packed[0], gate_and_up_scales[0], x.dtype).transpose(-2, -1) - down_w0 = dequantize_mxfp4(down_packed[0], down_scales[0], x.dtype).transpose(-2, -1) + 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 diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py new file mode 100644 index 0000000000..c6ce08830c --- /dev/null +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -0,0 +1,209 @@ +# 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, + _apply_bias, + _permute_tokens_for_grouped_mm, +) +from nemo_automodel.components.moe.fp4_utils import 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") + + 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')." + ) + 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 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. + packed, scales = quantize_mxfp4(local.transpose(-2, -1).contiguous()) + if isinstance(param, DTensor): + packed = DTensor.from_local(packed, param.device_mesh, param.placements) + scales = DTensor.from_local(scales, param.device_mesh, param.placements) + del self._parameters[name] + self.register_parameter(name + "_packed", nn.Parameter(packed, requires_grad=False)) + self.register_parameter(name + "_scales", nn.Parameter(scales, requires_grad=False)) + 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): + super().__init__(orig_module.config, backend=None) + self.use_torch_mm = orig_module.use_torch_mm + 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 diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index aac1ebafc8..770ba08d60 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -15,11 +15,12 @@ import pytest import torch -from nemo_automodel.components._peft.lora import patch_moe_module +from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4, patch_moe_module from nemo_automodel.components._peft.lora_experts import GroupedExpertsLoRA, GroupedExpertsLoRAMXFP4 from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.fp4_utils import dequantize_mxfp4, quantize_mxfp4 from nemo_automodel.components.moe.layers import GroupedExperts +from nemo_automodel.components.moe.quantized_experts import GroupedExpertsMXFP4 @pytest.fixture @@ -213,3 +214,61 @@ def test_patch_moe_module_mxfp4(moe_config, device): 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) + + +@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 From 1f4b5c0ff592f653e999a93737d77c1e903272c2 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 18:06:45 +0000 Subject: [PATCH 05/29] feat(moe): init-capable packed-weight registration for mxfp4 experts Make MXFP4ExpertStorageMixin register its packed storage parameters at module init (not only post-load), driven by a _PACKED_SUFFIXES list so the helper is format-driven rather than hardcoding the two tensor names. Enables building the model packed-at-init for the passthrough load path. Signed-off-by: Daniel --- .../components/moe/quantized_experts.py | 39 +++++++++++++++---- .../_peft/test_lora_experts_mxfp4.py | 35 +++++++++++++++++ 2 files changed, 67 insertions(+), 7 deletions(-) diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py index c6ce08830c..d64d420989 100644 --- a/nemo_automodel/components/moe/quantized_experts.py +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -57,6 +57,9 @@ class MXFP4ExpertStorageMixin: """ _MXFP4_BASE_NAMES: tuple[str, ...] = ("gate_and_up_projs", "down_projs") + # Storage-parameter suffixes, in pack/unpack order; kept here 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.""" @@ -68,6 +71,33 @@ def _init_mxfp4_storage(self) -> None: if not _to_local(getattr(self, self._MXFP4_BASE_NAMES[0])).is_meta: self.pack_base_weights() + @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. @@ -82,13 +112,8 @@ def pack_base_weights(self) -> None: 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. - packed, scales = quantize_mxfp4(local.transpose(-2, -1).contiguous()) - if isinstance(param, DTensor): - packed = DTensor.from_local(packed, param.device_mesh, param.placements) - scales = DTensor.from_local(scales, param.device_mesh, param.placements) - del self._parameters[name] - self.register_parameter(name + "_packed", nn.Parameter(packed, requires_grad=False)) - self.register_parameter(name + "_scales", nn.Parameter(scales, requires_grad=False)) + 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: diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index 770ba08d60..14b04714c2 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -252,6 +252,41 @@ def test_frozen_mxfp4_forward_matches_bf16(moe_config, device): 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 + + @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 From 28874c09c970d1176dbcdc5d3d2710153431129e Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 18:51:12 +0000 Subject: [PATCH 06/29] feat(moe): mxfp4 passthrough load path for DeepSeek-V4 (no bf16 experts) Phase B: load experts directly as packed fp4 so they are never materialized in bf16, capping the load-time peak (the actual blocker for one-node DSV4-Flash). - DeepSeekV4StateDictAdapter gains expert_storage_format='mxfp4': from_hf skips expert dequant and aggregates per-expert packed int8 + e8m0 scales into *_packed/*_scales keys. Concatenating gate||up along the output dim and stacking experts on dim 0 is layout-preserving (packing is along the contraction dim), so no unpacking is needed. - GroupedExpertsMXFP4(passthrough=True) registers meta packed placeholders from config at init (via the init-capable register_packed_base_weight), so the packed checkpoint loads straight in with no bf16 storage. - convert_frozen_experts_to_mxfp4(passthrough=...) plumbs the mode; kept opt-in so the validated bf16-then-pack path stays the default until the real checkpoint confirms Phase B end to end. Tests: synthetic fp4 checkpoint proves passthrough emits packed keys (no bf16, no orphaned scales) and decodes bit-identically to the bf16 dequant+aggregate path; packed-at-init registers correct meta shapes/dtypes. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/components/_peft/lora.py | 13 +- .../models/deepseek_v4/state_dict_adapter.py | 93 +++++++++++++ .../components/moe/quantized_experts.py | 51 ++++++- .../_peft/test_lora_experts_mxfp4.py | 24 ++++ .../test_dsv4_mxfp4_passthrough.py | 125 ++++++++++++++++++ 5 files changed, 302 insertions(+), 4 deletions(-) create mode 100644 tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 0391de5780..25210247d0 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -632,7 +632,7 @@ def apply_lora_to_linear_modules( return num_modules_matched -def convert_frozen_experts_to_mxfp4(model: nn.Module) -> int: +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 @@ -640,6 +640,15 @@ def convert_frozen_experts_to_mxfp4(model: nn.Module) -> int: 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. """ @@ -656,7 +665,7 @@ def convert_frozen_experts_to_mxfp4(model: nn.Module) -> int: unsupported += 1 continue if type(module) is GroupedExperts: - new_module = GroupedExpertsMXFP4(module) + new_module = GroupedExpertsMXFP4(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) 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 d00f826b68..a3498c5511 100644 --- a/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py +++ b/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py @@ -225,11 +225,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 # ------------------------------------------------------------------ @@ -341,6 +347,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) @@ -356,6 +367,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: @@ -438,6 +451,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()} diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py index d64d420989..c7a07770e1 100644 --- a/nemo_automodel/components/moe/quantized_experts.py +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -34,7 +34,7 @@ _apply_bias, _permute_tokens_for_grouped_mm, ) -from nemo_automodel.components.moe.fp4_utils import MXFP4GroupedMM, dequantize_mxfp4, quantize_mxfp4 +from nemo_automodel.components.moe.fp4_utils import MXFP4_BLOCK_SIZE, MXFP4GroupedMM, dequantize_mxfp4, quantize_mxfp4 def _to_local(t): @@ -137,9 +137,33 @@ class GroupedExpertsMXFP4(MXFP4ExpertStorageMixin, GroupedExperts): ``GroupedExperts._forward_grouped_mm`` but reads the packed base weights. """ - def __init__(self, orig_module: GroupedExperts): + 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() @@ -150,6 +174,29 @@ def __init__(self, orig_module: GroupedExperts): self.down_projs.requires_grad_(False) self._init_mxfp4_storage() + @torch.no_grad() + def _init_packed_placeholders(self) -> None: + """Register meta packed storage params from config shapes (no bf16 weights).""" + 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 = self.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 + def forward( self, x: torch.Tensor, diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index 14b04714c2..7152dd9f12 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -287,6 +287,30 @@ def test_register_packed_base_weight_is_init_capable(moe_config): 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] == [] + + @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 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..6919b557d3 --- /dev/null +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py @@ -0,0 +1,125 @@ +# 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.""" + +from unittest.mock import Mock + +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.moe.fp4_utils import dequantize_mxfp4, quantize_mxfp4 + +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) From ac16ac6ef2907bde1b25928c38405447cd22faf0 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 19:01:05 +0000 Subject: [PATCH 07/29] refactor(quantization): move mxfp4 primitives into components/quantization Address organization: the quantization primitives belong alongside fp8/qat/qlora, not under moe/. Mirrors how fp8.py holds the format/config while the layers that use it live elsewhere. - Move fp4_utils.py -> quantization/mxfp4.py. - Export mxfp4 from quantization/__init__. - GroupedExpertsMXFP4 / MXFP4ExpertStorageMixin stay in moe/ (GroupedExperts subclasses = model layers); GroupedExpertsLoRAMXFP4 stays in _peft (LoRA module, delegates all precision logic to the mixin). - Add a real-checkpoint-gated test validating passthrough against /raid0/data/models/DeepSeek-V4-Flash: confirms int8+e8m0 layout and bit-exact decode vs the bf16 path on actual checkpoint bytes (self-skips when absent). Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- .../components/moe/quantized_experts.py | 7 +- .../components/quantization/__init__.py | 10 +++ .../fp4_utils.py => quantization/mxfp4.py} | 0 .../_peft/test_lora_experts_mxfp4.py | 2 +- .../test_dsv4_mxfp4_passthrough.py | 67 ++++++++++++++++++- 5 files changed, 83 insertions(+), 3 deletions(-) rename nemo_automodel/components/{moe/fp4_utils.py => quantization/mxfp4.py} (100%) diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py index c7a07770e1..11b74bc78f 100644 --- a/nemo_automodel/components/moe/quantized_experts.py +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -34,7 +34,12 @@ _apply_bias, _permute_tokens_for_grouped_mm, ) -from nemo_automodel.components.moe.fp4_utils import MXFP4_BLOCK_SIZE, MXFP4GroupedMM, dequantize_mxfp4, quantize_mxfp4 +from nemo_automodel.components.quantization.mxfp4 import ( + MXFP4_BLOCK_SIZE, + MXFP4GroupedMM, + dequantize_mxfp4, + quantize_mxfp4, +) def _to_local(t): 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/moe/fp4_utils.py b/nemo_automodel/components/quantization/mxfp4.py similarity index 100% rename from nemo_automodel/components/moe/fp4_utils.py rename to nemo_automodel/components/quantization/mxfp4.py diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index 7152dd9f12..96d18600ec 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -18,9 +18,9 @@ from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4, patch_moe_module from nemo_automodel.components._peft.lora_experts import GroupedExpertsLoRA, GroupedExpertsLoRAMXFP4 from nemo_automodel.components.moe.config import MoEConfig -from nemo_automodel.components.moe.fp4_utils import dequantize_mxfp4, quantize_mxfp4 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 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 index 6919b557d3..e2320ff202 100644 --- a/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py @@ -14,15 +14,21 @@ """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.moe.fp4_utils import dequantize_mxfp4, quantize_mxfp4 +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 @@ -123,3 +129,62 @@ def test_passthrough_decodes_to_same_weights_as_bf16_path(): assert torch.equal(gu_unpacked, gu_bf16) assert torch.equal(dn_unpacked, dn_bf16) + + +@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" From a00a2956f6372a61effcd00eefa8a54ef28326a2 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 11 Jun 2026 21:44:24 +0000 Subject: [PATCH 08/29] feat(dsv4): wire mxfp4 passthrough end-to-end + to_hf packed split - DeepSeekV4StateDictAdapter.to_hf now splits *_packed/*_scales params back into per-expert checkpoint keys (mirror of the from_hf aggregation), so the DCP loader can enumerate destination tensors; quantization placeholder step is bypassed for already-packed experts. - infrastructure: when peft.expert_weight_format=mxfp4, set the adapter to passthrough mode and convert experts packed-at-init before load, so the fp4 checkpoint loads straight into packed params (no bf16 experts). - Tests: from_hf->to_hf round-trip recovers per-expert packed keys bit-exactly. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- .../_transformers/infrastructure.py | 16 ++++++-- .../models/deepseek_v4/state_dict_adapter.py | 40 +++++++++++++++++++ .../test_dsv4_mxfp4_passthrough.py | 21 ++++++++++ 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 1caf56a67a..945fdf2eab 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -91,12 +91,22 @@ def _apply_peft_and_lower_precision( # Convert frozen (non-LoRA-targeted) routed experts to mxfp4-resident storage. # LoRA-targeted experts are already GroupedExpertsLoRAMXFP4 from the call above. - # Packing of both is deferred until after the checkpoint is loaded. + # Passthrough mode (packed-at-init + adapter emits packed keys) loads the + # fp4 checkpoint straight into packed params, never materializing bf16 + # experts — capping the load-time peak. if getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4 - num_converted = convert_frozen_experts_to_mxfp4(model) - logger.info("Converted %d frozen expert module(s) to mxfp4-resident storage", num_converted) + # Put the state-dict adapter(s) in passthrough mode BEFORE the checkpoint + # load so both to_hf (destination keys) and from_hf (aggregation) keep + # experts packed. + for part in getattr(model, "parts", [model]): + adapter = getattr(part, "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) # FP8 if fp8_config is not None: 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 a3498c5511..65709ad5ce 100644 --- a/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py +++ b/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py @@ -753,6 +753,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): @@ -1026,6 +1032,40 @@ 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/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py index e2320ff202..837b525816 100644 --- a/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_mxfp4_passthrough.py @@ -131,6 +131,27 @@ def test_passthrough_decodes_to_same_weights_as_bf16_path(): 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}", From 73c69b6757570d4eebc40ee64dbe122171fc26c6 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 01:10:25 +0000 Subject: [PATCH 09/29] fix(moe): skip random init for mxfp4-resident experts _init_weights touched module.gate_and_up_projs, which passthrough GroupedExpertsMXFP4 modules don't have (only *_packed/*_scales, filled from the checkpoint). Skip init for mxfp4-resident experts. Found by an end-to-end single-GPU run on real DeepSeek-V4-Flash weights (2 train steps, finite loss). Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/components/moe/experts.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 1218f3d3c0..1db5088ea8 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -1289,6 +1289,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) From 9c2facfc5244be41d5d9ffa9403a91e20a04704a Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 03:03:41 +0000 Subject: [PATCH 10/29] fix(moe): set requires_grad at Parameter construction in ExpertParallel nn.Parameter() defaults requires_grad=True, which raises 'only Tensors of floating point dtype can require gradients' for the int8/e8m0 packed tensors of mxfp4-resident experts before the subsequent requires_grad assignment runs. Pass requires_grad at construction. Identical behavior for float params; unblocks EP sharding of packed experts. Validated by a 2-GPU ep_size=2 run on real DeepSeek-V4-Flash (both steps, finite loss, 55 GiB/rank vs 124 single-GPU). Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/components/moe/parallelizer.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index 5a8d4f4a48..457cdb2917 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -91,8 +91,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): From 5424533319f9a503dc65c48fd7da7ac2ad2cd0ce Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 04:05:12 +0000 Subject: [PATCH 11/29] fix(dsv4): require expert parallelism for mxfp4; guard single-GPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root-caused the single-GPU correctness bug: at world_size=1 the MoE parallelizer is skipped, the packed e8m0 expert scales are never applied, and experts decode to unscaled fp4 (~100x too large) — silently wrong (loss 19.60 vs bf16 17.28). Diagnosed via dequant-norm debug: ep-sharded experts std=0.025 (correct) vs world=1 std=2.48 (raw fp4 grid, scale==1). mxfp4 is only correct when experts are EP-sharded (validated: ep_size=2 matches the bf16 baseline within 0.04%). Guard world_size=1 with a clear error instead of training on garbage. Multi-GPU ep_size>1 is the supported (and only sensible, given model size) path. Removed debug instrumentation. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/_transformers/infrastructure.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 945fdf2eab..b55cfae528 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -97,6 +97,20 @@ def _apply_peft_and_lower_precision( if getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4 + # mxfp4-resident experts require expert parallelism: the packed scales + # are only loaded/applied correctly when the MoE parallelizer shards + # the experts (world_size>1, ep_size>1). At world_size=1 parallelization + # is skipped and the packed scales are not applied — the experts decode + # to unscaled fp4 (~100x too large), silently corrupting results. Fail + # loudly rather than train on garbage. + 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'." + ) + # Put the state-dict adapter(s) in passthrough mode BEFORE the checkpoint # load so both to_hf (destination keys) and from_hf (aggregation) keep # experts packed. From 746671fba53df611ec13b6f13bdd8c96433701bf Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 18:54:30 +0000 Subject: [PATCH 12/29] fix(recipe): clean process exit + deadlock-safe distributed validation Two issues made training appear to hang at the end: 1. No process-group teardown. run_train_validation_loop returned but the distributed process group was never destroyed, so NCCL / the elastic agent waited and the process never exited. Add a best-effort barrier + destroy_process_group in main()'s finally. 2. End-of-training validation could deadlock under expert parallelism. The MoE expert forward issues EP collectives; if DP ranks see uneven validation shard sizes they call those collectives a different number of times and hang. Drive the validation loop by a global-MIN "does every rank still have a batch?" all-reduce so all ranks run the same number of forwards. Validated on 2xGPU DSV4-Flash mxfp4 LoRA: validation completes and the process exits cleanly (exit 0). Note: validation still runs over the full val set at the last step (is_ckpt_step -> is_last_step); bounding it (eval_iters cap) is a possible follow-up for snappier end-of-training. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/recipes/llm/train_ft.py | 51 ++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 59eeebc884..aa2b98b915 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -1649,7 +1649,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( @@ -1875,8 +1900,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__": From fdbec62b57aa3caa18382f23592a886c26564285 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 22:12:30 +0000 Subject: [PATCH 13/29] perf(mxfp4): drop redundant contiguous copy in expert grouped GEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MXFP4GroupedMM.forward dequantized to [E,N,K] then did .transpose(-2,-1).contiguous() to feed torch._grouped_mm — a full bf16 weight copy per forward. torch._grouped_mm accepts the transposed view directly (cuBLAS transB, verified bit-identical), so pass the view and skip the copy. Microbench (E=32,N=K=4096): the op goes 4.1ms -> 0.4ms (copy was pure waste). End-to-end DSV4-Flash 8-GPU ep8, 4096-token packed seq: ~4630 -> ~4970 tps (+7%); mxfp4 slowdown vs bf16 improves from ~1.67x to ~1.56x. Loss unchanged. Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/components/quantization/mxfp4.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/nemo_automodel/components/quantization/mxfp4.py b/nemo_automodel/components/quantization/mxfp4.py index e7527becf9..c4d9098dfe 100644 --- a/nemo_automodel/components/quantization/mxfp4.py +++ b/nemo_automodel/components/quantization/mxfp4.py @@ -117,9 +117,11 @@ class MXFP4GroupedMM(torch.autograd.Function): 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 transposed relative to the GEMM operand - (``[E, N, K]`` packed along ``K``), which makes the backward GEMM operand - (``W^T``) the natural dequantization output; forward pays one transpose-copy. + 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. """ @@ -127,7 +129,9 @@ class MXFP4GroupedMM(torch.autograd.Function): @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] - out = torch._grouped_mm(x, w_t.transpose(-2, -1).contiguous(), offs=offs) + # 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 From 3e2de26f9094f5090be3883746eace364b4d01c5 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 23:12:35 +0000 Subject: [PATCH 14/29] fix(peft): support mxfp4 passthrough for LoRA-targeted experts When experts are LoRA-targeted under expert_weight_format=mxfp4, patch_moe_module built GroupedExpertsLoRAMXFP4 with deferred bf16 base params, but the adapter's passthrough mode emits packed fp4 checkpoint keys -> 'Size mismatch [2048,2048] vs [2048,4096]' at load. Add packed-at-init (passthrough) to GroupedExpertsLoRAMXFP4 (shared _init_packed_placeholders moved to MXFP4ExpertStorageMixin) and thread the flag through patch_moe_module / apply_lora_to_linear_modules / infrastructure. Validated on 8-GPU DSV4-Flash, experts LoRA-targeted: loads packed (1.28B trainable LoRA params, 0.82%), 3 steps train, step-0 loss 1.9474 == the frozen-expert run (LoRA B=0 no-op confirms the mxfp4 base decodes identically under LoRA). Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- ...eepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 4 ++ .../_transformers/infrastructure.py | 15 +++++- nemo_automodel/components/_peft/lora.py | 33 +++++++++--- .../components/_peft/lora_experts.py | 23 ++++++++- .../components/moe/quantized_experts.py | 51 ++++++++++--------- .../_peft/test_lora_experts_mxfp4.py | 20 ++++++++ 6 files changed, 111 insertions(+), 35 deletions(-) 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 index b469e569d7..0aaec76338 100644 --- 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 @@ -86,6 +86,10 @@ peft: - "*wq_b" - "*wkv" - "*wo_b" + - "*mlp.experts" # routed experts: mxfp4-resident base + LoRA (GroupedExpertsLoRAMXFP4) + - "*shared_experts.gate_proj" # shared experts: standard bf16 LoRA + - "*shared_experts.up_proj" + - "*shared_experts.down_proj" dim: 8 alpha: 32 use_triton: True diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index b55cfae528..a2f0575dff 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -86,8 +86,19 @@ def _apply_peft_and_lower_precision( logger.info("Enabling PEFT with Pipeline Parallelism") logger.info("Disabling Triton with Pipeline Parallelism Enabled.") 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) + # Skip freeze here - will do global freeze after checkpoint loading. + # For mxfp4, LoRA-targeted experts are built packed-at-init (passthrough) + # so the packed fp4 checkpoint loads straight in (matches the adapter's + # passthrough mode set below); otherwise they'd carry bf16 base params and + # size-mismatch against the packed checkpoint keys. + expert_passthrough = getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4" + apply_lora_to_linear_modules( + model, + peft_config, + quantization_config=quantization_config, + skip_freeze=True, + expert_passthrough=expert_passthrough, + ) # Convert frozen (non-LoRA-targeted) routed experts to mxfp4-resident storage. # LoRA-targeted experts are already GroupedExpertsLoRAMXFP4 from the call above. diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 25210247d0..6de1a78b1e 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -470,6 +470,7 @@ def patch_moe_module( lora_A_init_method="xavier", lora_dtype=None, expert_weight_format="bf16", + passthrough=False, ): """ Patches a custom MoE module (GroupedExperts or GroupedExpertsDeepEP) with LoRA. @@ -483,6 +484,9 @@ def patch_moe_module( 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): For mxfp4, build the packed base at init (no bf16 + materialization) so a packed fp4 checkpoint loads straight in. Set by the + checkpoint-load flow. Defaults to False (pack from materialized weights). Returns: nn.Module: The LoRA-wrapped MoE module. @@ -502,14 +506,23 @@ def patch_moe_module( lora_dtype=lora_dtype, ) elif isinstance(orig_module, GroupedExperts): - lora_cls = GroupedExpertsLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsLoRA - new_module = lora_cls( - orig_module, - lora_dim=dim, - alpha=alpha, - lora_A_init_method=lora_A_init_method, - lora_dtype=lora_dtype, - ) + if expert_weight_format == "mxfp4": + new_module = GroupedExpertsLoRAMXFP4( + orig_module, + lora_dim=dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + passthrough=passthrough, + ) + else: + new_module = GroupedExpertsLoRA( + orig_module, + lora_dim=dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) else: raise NotImplementedError(f"Unsupported MoE module type: {type(orig_module)}") @@ -522,6 +535,7 @@ def apply_lora_to_linear_modules( peft_config: PeftConfig, quantization_config=None, skip_freeze: bool = False, + expert_passthrough: bool = False, ) -> int: """ Replace selected nn.Linear layers with LinearLoRA layers (in-place). @@ -531,6 +545,8 @@ def apply_lora_to_linear_modules( peft_config: PEFT configuration for LoRA parameters. quantization_config: Optional separate QLoRA quantization configuration. skip_freeze: If True, skip the global parameter freeze (caller will handle it later). + expert_passthrough: For mxfp4 expert LoRA, build packed-at-init so a packed fp4 + checkpoint loads straight in (set by the checkpoint-load flow). Returns: Number of modules that were modified with LoRA. @@ -597,6 +613,7 @@ def apply_lora_to_linear_modules( lora_A_init_method=peft_config.lora_A_init, lora_dtype=lora_dtype, expert_weight_format=peft_config.expert_weight_format, + passthrough=expert_passthrough, ) # Find parent and replace diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index 438129c79d..a520507953 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -367,7 +367,15 @@ class GroupedExpertsLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsLoRA): until ``pack_base_weights()`` is called (after the base checkpoint is loaded). """ - def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_init_method="xavier", lora_dtype=None): + 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, @@ -375,7 +383,18 @@ def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_ini lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype, ) - self._init_mxfp4_storage() + if passthrough: + # Build packed-at-init (no bf16 base) so a packed fp4 checkpoint loads + # straight into the frozen base; the LoRA adapters stay floating point. + if not self.use_torch_mm: + raise NotImplementedError( + "mxfp4-resident expert weights require the torch_mm experts backend " + "(backend.experts='torch_mm')." + ) + self._mxfp4_resident = False + 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. diff --git a/nemo_automodel/components/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py index 11b74bc78f..931e2d80cf 100644 --- a/nemo_automodel/components/moe/quantized_experts.py +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -133,6 +133,34 @@ def _mxfp4_dequant_expert0(self, name: str, dtype: torch.dtype) -> torch.Tensor: scales = _to_local(getattr(self, name + "_scales"))[0] return dequantize_mxfp4(packed, scales, dtype).transpose(-2, -1) + @torch.no_grad() + def _init_packed_placeholders(self) -> None: + """Register meta packed storage params from config shapes (no bf16 weights). + + Used by the passthrough load flow (both frozen ``GroupedExpertsMXFP4`` and + LoRA-targeted ``GroupedExpertsLoRAMXFP4``): the module is built packed-at-init + so a packed fp4 checkpoint loads straight in with no bf16 materialization. + """ + 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 = self.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 + class GroupedExpertsMXFP4(MXFP4ExpertStorageMixin, GroupedExperts): """Frozen routed experts with mxfp4-resident base weights and no adapter. @@ -179,29 +207,6 @@ def __init__(self, orig_module: GroupedExperts, passthrough: bool = False): self.down_projs.requires_grad_(False) self._init_mxfp4_storage() - @torch.no_grad() - def _init_packed_placeholders(self) -> None: - """Register meta packed storage params from config shapes (no bf16 weights).""" - 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 = self.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 - def forward( self, x: torch.Tensor, diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index 96d18600ec..d3f17189a1 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -311,6 +311,26 @@ def test_passthrough_init_registers_packed_params_on_meta(moe_config): 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 From 0982a1e4f9c602102825a0978541218c68826155 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 13 Jun 2026 00:15:05 +0000 Subject: [PATCH 15/29] style(mxfp4): apply ruff format to expert LoRA + DSV4 adapter ruff 0.9.x format-only reflow of lines added in the mxfp4 expert work; no behavior change. (ruff was not installed in the dev venv when these landed.) Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- nemo_automodel/components/_peft/lora_experts.py | 3 +-- .../components/models/deepseek_v4/state_dict_adapter.py | 8 ++------ 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index a520507953..aa638b8aa7 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -388,8 +388,7 @@ def __init__( # straight into the frozen base; the LoRA adapters stay floating point. if not self.use_torch_mm: raise NotImplementedError( - "mxfp4-resident expert weights require the torch_mm experts backend " - "(backend.experts='torch_mm')." + "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm')." ) self._mxfp4_resident = False self._init_packed_placeholders() 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 65709ad5ce..6ddea266b4 100644 --- a/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py +++ b/nemo_automodel/components/models/deepseek_v4/state_dict_adapter.py @@ -1043,9 +1043,7 @@ def _split_merged_expert(self, fqn: str, tensor: Any) -> list[tuple[str, Any]]: 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 - ) + 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 @@ -1058,9 +1056,7 @@ def _split_merged_expert(self, fqn: str, tensor: Any) -> list[tuple[str, Any]]: 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 - ) + 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) From bd3f9d000fb353551ab743014d4ad9a23cb5d002 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 13 Jun 2026 00:15:16 +0000 Subject: [PATCH 16/29] perf(mxfp4): faster, spec-correct expert dequant (~2.4x kernel, ~2.5x e2e) Rework dequantize_mxfp4 / quantize_mxfp4 on the training hot path: - e8m0 scale via scales.to(float32) direct cast (OCP MX spec: byte 0x00 is 2^-127, not zero). Drops the old zero special-case, which silently zeroed any block legitimately scaled by 2^-127; all-zero blocks still decode to 0 since their fp4 codes are 0. - Decode both packed nibbles with a single gather over a 256x2 byte table instead of two int64 gathers over the 16-entry table. - Broadcast the per-32-col scale over a blocked view instead of repeat_interleave materializing a full [..., K] scale tensor. - Compute the block exponent with torch.frexp (exponent-1 == floor(log2) exactly) instead of floor(log2(amax)), avoiding round-off at powers of two -- the common case for a dequantized-fp4 checkpoint -- and removing the amax==0 guard (frexp(0).exponent == 0). Bit-exact round-trip preserved (max_diff 0.0); 18/18 mxfp4 unit tests pass. Overall speedup: - GPU microbench at DSV4-Flash per-rank expert shapes [32,4096,4096] / [32,4096,2048]: dequant 2.39x faster; dequant is ~95-98% of the expert-op cost at 128 tok/expert, so the grouped-GEMM op itself speeds up ~2.31x. - 8-GPU DSV4-Flash LoRA end-to-end (ep_size=8, attn=torch): ~2.5x tokens/sec in the short-sequence, dequant-bound regime (per-step ratios 2.3-2.7x at matched label-token counts), with loss bit-identical to the old kernel at every step. - The gain is largest where dequant dominates (short seq) and shrinks as the grouped-GEMM compute grows. Measured at the 4096-token packed-seq regime (8xH200, attn=tilelang, dispatcher=torch): new mxfp4 6,567 tps vs old 5,040 tps = 1.30x, slightly beating the earlier ~6,300/~1.27x projection. The old kernel here (5,040 tps) reproduces the historical 4,970-tps baseline within 1.4%, confirming the regime; the bf16 gap narrows from 1.56x to ~1.18x (bf16 ~7,740 tps), and the new kernel also uses ~6 GiB/rank less (48.9 vs 54.9 GiB). Under attn=tilelang the forward is non-deterministic, so step-0 loss is not bit-identical across kernels (8.9743 vs 8.9569, ~0.2%); the dequant itself stays bit-exact (round-trip max_diff 0.0). bucketize still rounds half-away-from-zero rather than to-nearest-even; a no-op for the exact-round-trip path, documented as the known divergence if a from-scratch bf16 quantizer is ever added. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel --- .../components/quantization/mxfp4.py | 66 +++++++++++-------- 1 file changed, 37 insertions(+), 29 deletions(-) diff --git a/nemo_automodel/components/quantization/mxfp4.py b/nemo_automodel/components/quantization/mxfp4.py index c4d9098dfe..de6c25d9ca 100644 --- a/nemo_automodel/components/quantization/mxfp4.py +++ b/nemo_automodel/components/quantization/mxfp4.py @@ -36,6 +36,18 @@ # 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. @@ -49,23 +61,20 @@ def dequantize_mxfp4(packed: torch.Tensor, scales: torch.Tensor, dtype: torch.dt Dequantized tensor of shape ``[..., K]`` in ``dtype``. """ packed_u8 = packed.contiguous().view(torch.uint8) - low = (packed_u8 & 0x0F).long() - high = ((packed_u8 >> 4) & 0x0F).long() - table = _FP4_E2M1_TABLE.to(packed_u8.device) - # Interleave (low, high) per byte so column indices match the original layout. - fp4_vals = torch.stack([table[low], table[high]], dim=-1).flatten(-2) + # 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] - # Decode e8m0 to fp32: 2^(e - 127), with byte 0 mapping to 0 (all-zero block). - scale_u8 = scales.contiguous().view(torch.uint8).int() - scale_f32 = torch.where( - scale_u8 == 0, - torch.zeros_like(scale_u8, dtype=torch.float32), - torch.pow(2.0, (scale_u8 - 127).float()), - ) + # 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] - scale_expanded = scale_f32.repeat_interleave(MXFP4_BLOCK_SIZE, dim=-1) - scale_expanded = scale_expanded[..., : fp4_vals.shape[-1]] - return (fp4_vals * scale_expanded).to(dtype) + # 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]: @@ -89,19 +98,18 @@ def quantize_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 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. amax == 0 maps to scale byte 0 (decoded as 0). - nonzero = amax > 0 - exp = torch.zeros_like(amax) - exp[nonzero] = torch.floor(torch.log2(amax[nonzero])) - 2.0 - scale_bytes = torch.where( - nonzero, - (exp + 127.0).clamp(1.0, 254.0), - torch.zeros_like(exp), - ).to(torch.uint8) - # Zero-amax blocks divide by 1 instead of 0; their codes are all zero anyway. - scale = torch.where(nonzero, torch.pow(2.0, (scale_bytes.int() - 127).float()), torch.ones_like(exp)) - - # Round each scaled magnitude to the nearest e2m1 magnitude code. + # 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) @@ -109,7 +117,7 @@ def quantize_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: codes = codes.reshape(*w.shape[:-1], k) packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).view(torch.int8) - return packed.contiguous(), scale_bytes.view(torch.float8_e8m0fnu).contiguous() + return packed.contiguous(), scales.contiguous() class MXFP4GroupedMM(torch.autograd.Function): From 064f8e7040f1b404af2da0ccbfc88c468e454939 Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Sun, 14 Jun 2026 06:38:39 +0000 Subject: [PATCH 17/29] feat(moe): support mxfp4-resident experts under DeepEP dispatch Add GroupedExpertsDeepEPMXFP4 (frozen) and GroupedExpertsDeepEPLoRAMXFP4 (LoRA-on-experts) so DeepSeek V4 routed experts can stay packed as fp4-e2m1 + e8m0 block scales while using the DeepEP fused all-to-all token dispatch. mxfp4 only changes the two post-dispatch grouped GEMMs (dequant on the fly via MXFP4GroupedMM); dispatch/combine are unchanged. Requires backend.experts=torch_mm. - Lift the mxfp4+DeepEP guards in patch_moe_module and convert_frozen_experts_to_mxfp4 (TE experts still unsupported). - Hoist _init_packed_placeholders into MXFP4ExpertStorageMixin so the torch and DeepEP frozen variants share the passthrough placeholder path. - Add example recipe deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml. - Add unit tests: guard lifting + GEMM-substitution numerics via a mock dispatcher. Validated on 8xH200 EP=8: 12-step train + validation + checkpoint, finite loss/grad_norm, ~40 GiB/rank. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- ..._v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 153 +++++++++ nemo_automodel/components/_peft/lora.py | 59 ++-- .../components/_peft/lora_experts.py | 122 +++++-- .../components/moe/quantized_experts.py | 166 ++++++++-- .../_peft/test_lora_experts_mxfp4_deepep.py | 297 ++++++++++++++++++ 5 files changed, 714 insertions(+), 83 deletions(-) create mode 100644 examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml create mode 100644 tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py 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..fcdeeca38c --- /dev/null +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml @@ -0,0 +1,153 @@ +# 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 DeepEP. +# +# This is the DeepEP counterpart of deepseek_v4_flash_hellaswag_lora_mxfp4.yaml. +# mxfp4 quantizes only the expert *weights* (local to each EP rank); DeepEP 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), DeepEP routes each rank only the tokens for its local experts, +# so this is the scalable EP path for DSV4 at EP=8/64. +# +# 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=deepep 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 + 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 DeepEP. + dispatcher: deepep + experts: torch_mm + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: + - "*wq_a" + - "*wq_b" + - "*wkv" + - "*wo_b" + dim: 8 + alpha: 32 + use_triton: True + # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly. + expert_weight_format: mxfp4 + +checkpoint: + enabled: false + # The DSV4-Flash checkpoint stores experts as fp4; they are dequantized to bf16 + # on load and then repacked to mxfp4 (Phase A). Required for the current flow. + dequantize_base_checkpoint: true + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +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/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 812e528c94..2f7f797c65 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -25,6 +25,7 @@ from nemo_automodel.components._peft.lora_experts import ( GroupedExpertsDeepEPLoRA, + GroupedExpertsDeepEPLoRAMXFP4, GroupedExpertsLoRA, GroupedExpertsLoRAMXFP4, ) @@ -517,7 +518,6 @@ def patch_moe_module( lora_A_init_method="xavier", lora_dtype=None, expert_weight_format="bf16", - passthrough=False, ): """ Patches a custom MoE module (GroupedExperts or GroupedExpertsDeepEP) with LoRA. @@ -531,9 +531,6 @@ def patch_moe_module( 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): For mxfp4, build the packed base at init (no bf16 - materialization) so a packed fp4 checkpoint loads straight in. Set by the - checkpoint-load flow. Defaults to False (pack from materialized weights). Returns: nn.Module: The LoRA-wrapped MoE module. @@ -543,9 +540,8 @@ def patch_moe_module( if isinstance(orig_module, GroupedExpertsTE): raise NotImplementedError("LoRA is not supported for Transformer Engine (TE) expert modules.") elif isinstance(orig_module, GroupedExpertsDeepEP): - if expert_weight_format == "mxfp4": - raise NotImplementedError("expert_weight_format='mxfp4' is not supported for DeepEP expert modules yet.") - new_module = GroupedExpertsDeepEPLoRA( + deepep_lora_cls = GroupedExpertsDeepEPLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsDeepEPLoRA + new_module = deepep_lora_cls( orig_module, lora_dim=dim, alpha=alpha, @@ -553,23 +549,14 @@ def patch_moe_module( lora_dtype=lora_dtype, ) elif isinstance(orig_module, GroupedExperts): - if expert_weight_format == "mxfp4": - new_module = GroupedExpertsLoRAMXFP4( - orig_module, - lora_dim=dim, - alpha=alpha, - lora_A_init_method=lora_A_init_method, - lora_dtype=lora_dtype, - passthrough=passthrough, - ) - else: - new_module = GroupedExpertsLoRA( - orig_module, - lora_dim=dim, - alpha=alpha, - lora_A_init_method=lora_A_init_method, - lora_dtype=lora_dtype, - ) + lora_cls = GroupedExpertsLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsLoRA + new_module = lora_cls( + orig_module, + lora_dim=dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) else: raise NotImplementedError(f"Unsupported MoE module type: {type(orig_module)}") @@ -582,7 +569,6 @@ def apply_lora_to_linear_modules( peft_config: PeftConfig, quantization_config=None, skip_freeze: bool = False, - expert_passthrough: bool = False, ) -> int: """ Replace selected nn.Linear layers with LinearLoRA layers (in-place). @@ -592,8 +578,6 @@ def apply_lora_to_linear_modules( peft_config: PEFT configuration for LoRA parameters. quantization_config: Optional separate QLoRA quantization configuration. skip_freeze: If True, skip the global parameter freeze (caller will handle it later). - expert_passthrough: For mxfp4 expert LoRA, build packed-at-init so a packed fp4 - checkpoint loads straight in (set by the checkpoint-load flow). Returns: Number of modules that were modified with LoRA. @@ -666,7 +650,6 @@ def apply_lora_to_linear_modules( lora_A_init_method=peft_config.lora_A_init, lora_dtype=lora_dtype, expert_weight_format=peft_config.expert_weight_format, - passthrough=expert_passthrough, ) # Find parent and replace @@ -734,6 +717,15 @@ def convert_frozen_experts_to_mxfp4(model: nn.Module, passthrough: bool = False) """ # 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 @@ -741,11 +733,12 @@ def convert_frozen_experts_to_mxfp4(model: nn.Module, passthrough: bool = False) # Already mxfp4-resident (frozen or LoRA-targeted) — skip. if isinstance(module, MXFP4ExpertStorageMixin): continue - if isinstance(module, (GroupedExpertsDeepEP, GroupedExpertsTE)): + if isinstance(module, GroupedExpertsTE): unsupported += 1 continue - if type(module) is GroupedExperts: - new_module = GroupedExpertsMXFP4(module, passthrough=passthrough) + 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) @@ -753,8 +746,8 @@ def convert_frozen_experts_to_mxfp4(model: nn.Module, passthrough: bool = False) if unsupported: logger.warning( - "expert_weight_format='mxfp4' skipped %d DeepEP/TE expert module(s); only the torch_mm " - "GroupedExperts backend is supported. Set backend.dispatcher='torch' and backend.experts='torch_mm'.", + "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 diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index aa638b8aa7..42d6f4a9c2 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -367,15 +367,7 @@ class GroupedExpertsLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsLoRA): until ``pack_base_weights()`` is called (after the base checkpoint is loaded). """ - def __init__( - self, - orig_module: GroupedExperts, - lora_dim=8, - alpha=32, - lora_A_init_method="xavier", - lora_dtype=None, - passthrough=False, - ): + def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_init_method="xavier", lora_dtype=None): super().__init__( orig_module, lora_dim=lora_dim, @@ -383,17 +375,7 @@ def __init__( lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype, ) - if passthrough: - # Build packed-at-init (no bf16 base) so a packed fp4 checkpoint loads - # straight into the frozen base; the LoRA adapters stay floating point. - if not self.use_torch_mm: - raise NotImplementedError( - "mxfp4-resident expert weights require the torch_mm experts backend (backend.experts='torch_mm')." - ) - self._mxfp4_resident = False - self._init_packed_placeholders() - else: - self._init_mxfp4_storage() + 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. @@ -711,3 +693,103 @@ def forward( y = self.token_dispatcher.token_unpermutation(output2) 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 + ): + super().__init__( + orig_module, + lora_dim=lora_dim, + alpha=alpha, + lora_A_init_method=lora_A_init_method, + lora_dtype=lora_dtype, + ) + 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) + + lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) + lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) + lora_down_A = _to_local(self.lora_down_A) + lora_down_B = _to_local(self.lora_down_B) + + 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/moe/quantized_experts.py b/nemo_automodel/components/moe/quantized_experts.py index 931e2d80cf..e5445387ad 100644 --- a/nemo_automodel/components/moe/quantized_experts.py +++ b/nemo_automodel/components/moe/quantized_experts.py @@ -31,6 +31,7 @@ from nemo_automodel.components.moe.experts import ( GroupedExperts, + GroupedExpertsDeepEP, _apply_bias, _permute_tokens_for_grouped_mm, ) @@ -62,7 +63,7 @@ class MXFP4ExpertStorageMixin: """ _MXFP4_BASE_NAMES: tuple[str, ...] = ("gate_and_up_projs", "down_projs") - # Storage-parameter suffixes, in pack/unpack order; kept here so the + # 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") @@ -70,12 +71,43 @@ 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')." + "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``. @@ -133,34 +165,6 @@ def _mxfp4_dequant_expert0(self, name: str, dtype: torch.dtype) -> torch.Tensor: scales = _to_local(getattr(self, name + "_scales"))[0] return dequantize_mxfp4(packed, scales, dtype).transpose(-2, -1) - @torch.no_grad() - def _init_packed_placeholders(self) -> None: - """Register meta packed storage params from config shapes (no bf16 weights). - - Used by the passthrough load flow (both frozen ``GroupedExpertsMXFP4`` and - LoRA-targeted ``GroupedExpertsLoRAMXFP4``): the module is built packed-at-init - so a packed fp4 checkpoint loads straight in with no bf16 materialization. - """ - 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 = self.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 - class GroupedExpertsMXFP4(MXFP4ExpertStorageMixin, GroupedExperts): """Frozen routed experts with mxfp4-resident base weights and no adapter. @@ -289,3 +293,105 @@ def _forward_grouped_mm_mxfp4(self, x, token_mask, weights, indices, n_local_exp 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/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..ffa3632235 --- /dev/null +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py @@ -0,0 +1,297 @@ +# 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, + 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 From 51af07c4334d107d085ece9b7a6a99354162855f Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Sun, 14 Jun 2026 07:36:26 +0000 Subject: [PATCH 18/29] feat(peft): enable LoRA on routed+shared experts under mxfp4 Makes routed-expert (and shared-expert) LoRA actually work with mxfp4-resident base weights on the full DeepSeek-V4-Flash, fixing three issues that blocked it: - Load wiring: GroupedExperts*LoRAMXFP4 now support a passthrough mode that registers packed base placeholders at init (instead of deferred bf16-load-then-pack), so the packed fp4 checkpoint loads straight in via _aggregate_experts_packed. This avoids the bf16 _aggregate_experts re-stack that OOM'd at load (~137 GiB/rank) and keeps the steady-state footprint at the packed ~49 GiB/rank. Threaded through patch_moe_module / apply_lora_to_linear_modules; infra keeps the adapter in packed mode for both frozen and LoRA experts. - Adapter dtype: the LoRA grouped GEMMs now cast adapters to the activation dtype (GroupedExpertsDeepEP allocates its base, hence adapter sizing, as fp32 when no backend dtype is set), fixing "mat1 and mat2 have the same dtype, BFloat16 != float". - Gate NaN guard: clamp_min(1e-12) under the sqrtsoftplus gate sqrt (both the generic Gate and the DSV4 hash gate). softplus underflows to 0.0 for very negative logits and sqrt'(0)=inf makes the backward NaN; the clamp bounds the gradient with a negligible forward change. - Enable expert LoRA in both deepseek_v4 mxfp4 recipes (target *mlp.experts and *shared_experts.*proj). Validated on 8xH200 EP=8 (packed 4096): routed+shared expert LoRA trains 15 steps, loss 8.96->5.72 monotonic, grad_norm finite, no NaN, 48.9 GiB/rank, val loss 2.32. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- ...eepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 8 ++- ..._v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 8 +++ .../_transformers/infrastructure.py | 21 ++---- nemo_automodel/components/_peft/lora.py | 36 +++++----- .../components/_peft/lora_experts.py | 66 ++++++++++++++----- .../components/models/deepseek_v4/model.py | 5 +- nemo_automodel/components/moe/layers.py | 6 +- 7 files changed, 100 insertions(+), 50 deletions(-) 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 index 0aaec76338..162a19c455 100644 --- 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 @@ -82,12 +82,16 @@ model: peft: _target_: nemo_automodel.components._peft.lora.PeftConfig target_modules: + # attention - "*wq_a" - "*wq_b" - "*wkv" - "*wo_b" - - "*mlp.experts" # routed experts: mxfp4-resident base + LoRA (GroupedExpertsLoRAMXFP4) - - "*shared_experts.gate_proj" # shared experts: standard bf16 LoRA + # 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 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 index fcdeeca38c..b6f46476e4 100644 --- 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 @@ -91,10 +91,18 @@ model: peft: _target_: nemo_automodel.components._peft.lora.PeftConfig target_modules: + # attention - "*wq_a" - "*wq_b" - "*wkv" - "*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 diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index ab0a203678..ce838f3448 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -88,19 +88,8 @@ def _apply_peft_and_lower_precision( logger.info("Enabling PEFT with Pipeline Parallelism") logger.info("Disabling Triton with Pipeline Parallelism Enabled.") peft_config.use_triton = False - # Skip freeze here - will do global freeze after checkpoint loading. - # For mxfp4, LoRA-targeted experts are built packed-at-init (passthrough) - # so the packed fp4 checkpoint loads straight in (matches the adapter's - # passthrough mode set below); otherwise they'd carry bf16 base params and - # size-mismatch against the packed checkpoint keys. - expert_passthrough = getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4" - apply_lora_to_linear_modules( - model, - peft_config, - quantization_config=quantization_config, - skip_freeze=True, - expert_passthrough=expert_passthrough, - ) + # 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) # Convert frozen (non-LoRA-targeted) routed experts to mxfp4-resident storage. # LoRA-targeted experts are already GroupedExpertsLoRAMXFP4 from the call above. @@ -126,7 +115,11 @@ def _apply_peft_and_lower_precision( # Put the state-dict adapter(s) in passthrough mode BEFORE the checkpoint # load so both to_hf (destination keys) and from_hf (aggregation) keep - # experts packed. + # experts packed. Both frozen (convert_frozen_experts_to_mxfp4 passthrough) + # and LoRA-targeted experts (GroupedExperts*LoRAMXFP4 built passthrough in + # apply_lora_to_linear_modules) load the packed fp4 keys straight into packed + # params — no bf16 expert materialization, so the load-time _aggregate_experts + # bf16 re-stack OOM is avoided. for part in getattr(model, "parts", [model]): adapter = getattr(part, "state_dict_adapter", None) if adapter is not None and hasattr(adapter, "expert_storage_format"): diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 2f7f797c65..76add3b1d0 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -518,6 +518,7 @@ def patch_moe_module( lora_A_init_method="xavier", lora_dtype=None, expert_weight_format="bf16", + passthrough=False, ): """ Patches a custom MoE module (GroupedExperts or GroupedExpertsDeepEP) with LoRA. @@ -531,32 +532,29 @@ def patch_moe_module( 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. """ 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, GroupedExpertsTE): raise NotImplementedError("LoRA is not supported for Transformer Engine (TE) expert modules.") elif isinstance(orig_module, GroupedExpertsDeepEP): - deepep_lora_cls = GroupedExpertsDeepEPLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsDeepEPLoRA - new_module = deepep_lora_cls( - 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): - lora_cls = GroupedExpertsLoRAMXFP4 if expert_weight_format == "mxfp4" else GroupedExpertsLoRA - new_module = lora_cls( - 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)}") @@ -642,7 +640,10 @@ 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, @@ -650,6 +651,7 @@ def apply_lora_to_linear_modules( 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 diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index 42d6f4a9c2..f4ec5863fc 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -362,12 +362,24 @@ class GroupedExpertsLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsLoRA): and backward (see ``MXFP4ExpertStorageMixin``). Only the LoRA adapters (and optional expert biases) remain in floating point. - When constructed from a module whose weights are still on the meta device, - packing is deferred: the module behaves exactly like ``GroupedExpertsLoRA`` - until ``pack_base_weights()`` is called (after the base checkpoint is loaded). + 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): + 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, @@ -375,7 +387,13 @@ def __init__(self, orig_module: GroupedExperts, lora_dim=8, alpha=32, lora_A_ini lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype, ) - self._init_mxfp4_storage() + 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. @@ -434,10 +452,14 @@ def _forward_grouped_mm_mxfp4(self, x, token_mask, weights, indices, n_local_exp experts_start_idx, ) - lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) - lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) - lora_down_A = _to_local(self.lora_down_A) - lora_down_B = _to_local(self.lora_down_B) + # 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) @@ -710,7 +732,13 @@ class GroupedExpertsDeepEPLoRAMXFP4(MXFP4ExpertStorageMixin, GroupedExpertsDeepE """ def __init__( - self, orig_module: GroupedExpertsDeepEP, lora_dim=8, alpha=32, lora_A_init_method="xavier", lora_dtype=None + self, + orig_module: GroupedExpertsDeepEP, + lora_dim=8, + alpha=32, + lora_A_init_method="xavier", + lora_dtype=None, + passthrough=False, ): super().__init__( orig_module, @@ -719,7 +747,13 @@ def __init__( lora_A_init_method=lora_A_init_method, lora_dtype=lora_dtype, ) - self._init_mxfp4_storage() + 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, @@ -750,10 +784,12 @@ def forward( ) permuted_probs = permuted_probs.unsqueeze(-1) - lora_gate_and_up_A = _to_local(self.lora_gate_and_up_A) - lora_gate_and_up_B = _to_local(self.lora_gate_and_up_B) - lora_down_A = _to_local(self.lora_down_A) - lora_down_B = _to_local(self.lora_down_B) + # 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) diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index 4a5d1d73da..fb0be9a4a6 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -298,7 +298,10 @@ 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 with a negligible forward change. See the + # matching guard in moe/layers.py Gate. + scores = F.softplus(scores).clamp_min(1e-12).sqrt() elif self.score_func == "sigmoid": scores = scores.sigmoid() else: diff --git a/nemo_automodel/components/moe/layers.py b/nemo_automodel/components/moe/layers.py index ea8dd83611..17a4ad1159 100644 --- a/nemo_automodel/components/moe/layers.py +++ b/nemo_automodel/components/moe/layers.py @@ -428,7 +428,11 @@ def forward( weights = original_scores.gather(1, indices) elif self.score_func == "sqrtsoftplus": # sqrt(softplus(x)) = sqrt(log(1 + exp(x))), used in DeepSeek V4. - scores = torch.sqrt(F.softplus(scores.float())) + # clamp_min keeps the sqrt argument strictly positive: softplus(x) underflows + # to exactly 0.0 in fp32 for very negative logits (x <~ -104), and sqrt'(0) = inf + # makes the backward NaN. The clamp bounds the gradient with a negligible + # (sqrt(1e-12) = 1e-6) change to the forward value. + scores = torch.sqrt(F.softplus(scores.float()).clamp_min(1e-12)) original_scores = scores if self.e_score_correction_bias is not None: From 01e71eb49744ee9789f9a467a3a8ec6abf4af9ca Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Sun, 14 Jun 2026 08:00:09 +0000 Subject: [PATCH 19/29] fix(deepseek_v4): cast lm_head weight to activation dtype for fused linear-CE Fused linear-CE for DSV4 is upstream (#2397); only the DSV4 dtype-cast remains. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- nemo_automodel/components/loss/linear_ce.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/nemo_automodel/components/loss/linear_ce.py b/nemo_automodel/components/loss/linear_ce.py index 70bfb01a08..2e55e959ec 100644 --- a/nemo_automodel/components/loss/linear_ce.py +++ b/nemo_automodel/components/loss/linear_ce.py @@ -156,6 +156,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( From 8b62ffd611b4a7688abc38db94a70ce104734bdb Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Mon, 15 Jun 2026 01:59:08 +0000 Subject: [PATCH 20/29] chore(deepseek_v4): default the mxfp4 LoRA recipes to FusedLinearCrossEntropy Both DSV4 mxfp4 LoRA recipes now use the fused linear cross-entropy so the [seq, 129280] logits are never materialized, removing the ~16 GiB fp32 logits spike (single-node context ceiling ~30k -> ~36-38k tokens on one 8xH200). Measured fit: packed 32768 at 115 GiB/rank (vs OOM ~137 GiB with MaskedCrossEntropy). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- .../deepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 12 +++++++++++- ...epseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 13 ++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) 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 index 162a19c455..7494aa39c2 100644 --- 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 @@ -65,6 +65,11 @@ model: 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: @@ -107,7 +112,12 @@ checkpoint: dequantize_base_checkpoint: true loss_fn: - _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + # 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 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 index b6f46476e4..e82a36b9ae 100644 --- 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 @@ -74,6 +74,11 @@ model: 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: @@ -116,7 +121,13 @@ checkpoint: dequantize_base_checkpoint: true loss_fn: - _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + # 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 From 391a0463fbe67533d0213db7cfa624b32c84166c Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Mon, 15 Jun 2026 02:11:18 +0000 Subject: [PATCH 21/29] chore(deepseek_v4): fp32 LoRA adapters in the mxfp4 recipes Set peft.lora_dtype=float32 in both DSV4 mxfp4 LoRA recipes. With FSDP2's default MixedPrecisionPolicy(param_dtype=bf16), this keeps fp32 master weights + fp32 AdamW state for the adapters (stability against small-update swamping, matching HF PEFT's autocast_adapter_dtype default) while the adapter matmuls still run in bf16 (FSDP casts the all-gathered params to param_dtype). Adapters are tiny so the fp32 optimizer-state cost is negligible. Validated 8xH200 EP=8 packed-4096 (fused CE + routed/shared expert LoRA + fp32 adapters): 2 steps, loss 8.96->8.77, finite grad_norm, no dtype mismatch in the triton/grouped-GEMM LoRA paths. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- .../deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 4 ++++ .../deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 4 ++++ 2 files changed, 8 insertions(+) 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 index 7494aa39c2..c92fb9cdde 100644 --- 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 @@ -102,6 +102,10 @@ peft: 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 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 index e82a36b9ae..d90a2f570c 100644 --- 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 @@ -111,6 +111,10 @@ peft: 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 From f0c93227af08369f876e79fad7db8fe8f8c193c2 Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Mon, 15 Jun 2026 03:23:50 +0000 Subject: [PATCH 22/29] test(moe): fix stale unit tests broken by lazy DeepEP buffer + packed-param requires_grad MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two L0 unit tests were stale relative to earlier branch code changes: - test_grouped_experts_deepep_token_dispatcher_init asserted init_token_dispatcher eagerly calls _init_deepep_buffer, but buffer allocation is now lazy (deferred to FusedDispatch.forward) — the revert that fixed the single-node load-time OOM. Assert it is NOT called. - ExpertParallel._partition_fn now constructs nn.Parameter(..., requires_grad=...) so non-floating packed mxfp4 params (int8 / e8m0) don't trip the default requires_grad=True. The test's stub Parameter didn't accept/store requires_grad; add it (also unblocks the requires_grad-preservation test). Both fixes verified: tests/unit_tests/moe now 450 passed, 0 failed. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Daniel Shen --- tests/unit_tests/moe/test_experts.py | 5 ++++- tests/unit_tests/moe/test_parallelizer.py | 3 ++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/moe/test_experts.py b/tests/unit_tests/moe/test_experts.py index a471c9c602..63238c11fe 100644 --- a/tests/unit_tests/moe/test_experts.py +++ b/tests/unit_tests/moe/test_experts.py @@ -759,7 +759,10 @@ def test_grouped_experts_deepep_token_dispatcher_init(self, moe_config): assert hasattr(experts, "token_dispatcher") assert experts.ep_size == 2 assert experts.ep_rank == 0 - mock_init_buffer.assert_called_once_with(mock_mesh.get_group.return_value) + # The DeepEP NVSHMEM buffer is allocated lazily (in FusedDispatch.forward), + # not eagerly in init_token_dispatcher — the revert that fixed the single-node + # load-time OOM. So init_token_dispatcher must NOT call _init_deepep_buffer. + mock_init_buffer.assert_not_called() def test_grouped_experts_deepep_apply_bias_no_bias(self, moe_config): """Test _apply_bias method with no bias.""" diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index 1ce7aa97d9..edec9c8151 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -84,8 +84,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 From 036763e4f2b2a14346cb14bc490827ca95c8051f Mon Sep 17 00:00:00 2001 From: Daniel Date: Mon, 15 Jun 2026 22:19:15 +0000 Subject: [PATCH 23/29] docs(examples): clarify dequantize_base_checkpoint for mxfp4 passthrough Co-Authored-By: Claude Fable 5 Signed-off-by: Daniel --- .../deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml | 6 ++++-- .../deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) 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 index c92fb9cdde..a5252e9b90 100644 --- 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 @@ -111,8 +111,10 @@ peft: checkpoint: enabled: false - # The DSV4-Flash checkpoint stores experts as fp4; they are dequantized to bf16 - # on load and then repacked to mxfp4 (Phase A). Required for the current flow. + # 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: 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 index d90a2f570c..0cdb9f1a7d 100644 --- 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 @@ -120,8 +120,10 @@ peft: checkpoint: enabled: false - # The DSV4-Flash checkpoint stores experts as fp4; they are dequantized to bf16 - # on load and then repacked to mxfp4 (Phase A). Required for the current flow. + # 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: From 877f9dc786baa9b8aafaecf9be26ec3cb6b1399a Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sun, 6 Sep 2026 21:10:09 -0700 Subject: [PATCH 24/29] fix(deepseek_v4): run mxfp4 LoRA with HybridEP Signed-off-by: HuiyingLi --- ..._v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) 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 index 0cdb9f1a7d..fbb48276fc 100644 --- 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 @@ -14,15 +14,16 @@ # 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 DeepEP. +# dequantized on the fly in the grouped-GEMM forward/backward, dispatched with HybridEP. # -# This is the DeepEP counterpart of deepseek_v4_flash_hellaswag_lora_mxfp4.yaml. -# mxfp4 quantizes only the expert *weights* (local to each EP rank); DeepEP only +# 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), DeepEP routes each rank only the tokens for its local experts, +# 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 @@ -30,7 +31,7 @@ # 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=deepep with +# 'gmm' path has no packed variant), so this recipe pairs dispatcher=hybridep with # experts=torch_mm. recipe: TrainFinetuneRecipeForNextTokenPrediction @@ -87,8 +88,8 @@ model: linear: torch rms_norm: torch_fp32 rope_fusion: false - # mxfp4 experts run on the torch_mm grouped-GEMM path, dispatched with DeepEP. - dispatcher: deepep + # 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 @@ -96,11 +97,12 @@ model: peft: _target_: nemo_automodel.components._peft.lora.PeftConfig target_modules: - # attention - - "*wq_a" - - "*wq_b" - - "*wkv" - - "*wo_b" + # 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" @@ -111,10 +113,10 @@ peft: 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 + # 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 From aafa6e30fc8258496402937c31a2a741038193a7 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sun, 6 Sep 2026 21:46:31 -0700 Subject: [PATCH 25/29] feat(deepseek_v4): use TE FusedAdam master weights Signed-off-by: HuiyingLi --- ..._v4_flash_hellaswag_lora_mxfp4_deepep.yaml | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) 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 index fbb48276fc..1134fdbc57 100644 --- 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 @@ -172,9 +172,25 @@ validation_dataloader: drop_last: true optimizer: - _target_: torch.optim.AdamW - fused: true + # 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 From 598b1a856b5caa374dac2aa98b1c993a1d3ccd6c Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 7 Sep 2026 00:23:20 -0700 Subject: [PATCH 26/29] refactor(deepseek_v4): own MXFP4 PEFT lifecycle Signed-off-by: HuiyingLi --- .../_transformers/infrastructure.py | 57 ++++-------------- .../components/models/deepseek_v4/model.py | 12 ++++ .../models/deepseek_v4/mxfp4_peft.py | 60 +++++++++++++++++++ .../_transformers/test_auto_model.py | 51 +++++++++++++++- 4 files changed, 134 insertions(+), 46 deletions(-) create mode 100644 nemo_automodel/components/models/deepseek_v4/mxfp4_peft.py diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 706e531722..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,43 +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) - - # Convert frozen (non-LoRA-targeted) routed experts to mxfp4-resident storage. - # LoRA-targeted experts are already GroupedExpertsLoRAMXFP4 from the call above. - # Passthrough mode (packed-at-init + adapter emits packed keys) loads the - # fp4 checkpoint straight into packed params, never materializing bf16 - # experts — capping the load-time peak. - if getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": - from nemo_automodel.components._peft.lora import convert_frozen_experts_to_mxfp4 - - # mxfp4-resident experts require expert parallelism: the packed scales - # are only loaded/applied correctly when the MoE parallelizer shards - # the experts (world_size>1, ep_size>1). At world_size=1 parallelization - # is skipped and the packed scales are not applied — the experts decode - # to unscaled fp4 (~100x too large), silently corrupting results. Fail - # loudly rather than train on garbage. - 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'." - ) - - # Put the state-dict adapter(s) in passthrough mode BEFORE the checkpoint - # load so both to_hf (destination keys) and from_hf (aggregation) keep - # experts packed. Both frozen (convert_frozen_experts_to_mxfp4 passthrough) - # and LoRA-targeted experts (GroupedExperts*LoRAMXFP4 built passthrough in - # apply_lora_to_linear_modules) load the packed fp4 keys straight into packed - # params — no bf16 expert materialization, so the load-time _aggregate_experts - # bf16 re-stack OOM is avoided. - for part in getattr(model, "parts", [model]): - adapter = getattr(part, "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) + _call_model_hook(model, "prepare_peft_checkpoint_load", peft_config) # FP8 if fp8_config is not None: @@ -827,14 +800,8 @@ def apply_model_infrastructure( "check freeze_config and the PEFT configuration." ) - # Pack deferred mxfp4-resident expert base weights now that the checkpoint is loaded. - if peft_config is not None and getattr(peft_config, "expert_weight_format", "bf16") == "mxfp4": - from nemo_automodel.components._peft.lora import pack_mxfp4_expert_base_weights - - for mp in model.parts if hasattr(model, "parts") else [model]: - num_packed = pack_mxfp4_expert_base_weights(mp) - if num_packed: - logger.info("Packed %d MoE expert modules to mxfp4-resident storage", num_packed) + 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 diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index 417aa19343..387297973a 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -1207,6 +1207,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/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() From d554c6f7925c086c36c385aa6a3334df5f22c974 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 7 Sep 2026 00:28:43 -0700 Subject: [PATCH 27/29] refactor(peft): isolate MXFP4 expert LoRA Signed-off-by: HuiyingLi --- nemo_automodel/components/_peft/lora.py | 4 +- .../components/_peft/lora_experts.py | 274 +--------------- .../components/_peft/lora_experts_mxfp4.py | 302 ++++++++++++++++++ .../_peft/test_lora_experts_mxfp4.py | 3 +- .../_peft/test_lora_experts_mxfp4_deepep.py | 6 +- 5 files changed, 310 insertions(+), 279 deletions(-) create mode 100644 nemo_automodel/components/_peft/lora_experts_mxfp4.py diff --git a/nemo_automodel/components/_peft/lora.py b/nemo_automodel/components/_peft/lora.py index 666e7a57d9..fe516794ed 100644 --- a/nemo_automodel/components/_peft/lora.py +++ b/nemo_automodel/components/_peft/lora.py @@ -24,8 +24,10 @@ from nemo_automodel.components._peft.lora_experts import ( GroupedExpertsDeepEPLoRA, - GroupedExpertsDeepEPLoRAMXFP4, GroupedExpertsLoRA, +) +from nemo_automodel.components._peft.lora_experts_mxfp4 import ( + GroupedExpertsDeepEPLoRAMXFP4, GroupedExpertsLoRAMXFP4, ) from nemo_automodel.components._peft.lora_kernel import ( diff --git a/nemo_automodel/components/_peft/lora_experts.py b/nemo_automodel/components/_peft/lora_experts.py index bae57bf573..04b267d72e 100644 --- a/nemo_automodel/components/_peft/lora_experts.py +++ b/nemo_automodel/components/_peft/lora_experts.py @@ -19,7 +19,7 @@ import torch.distributed.nn.functional as dist_nn_f import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Partial, Shard +from torch.distributed.tensor import DTensor from nemo_automodel.components.moe.experts import ( GroupedExperts, @@ -28,7 +28,6 @@ _apply_bias, _permute_tokens_for_grouped_mm, ) -from nemo_automodel.components.moe.quantized_experts import MXFP4ExpertStorageMixin from nemo_automodel.shared.utils import dtype_from_str try: @@ -399,163 +398,6 @@ def _forward_grouped_mm( return y -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 GroupedExpertsDeepEPLoRA(GroupedExpertsDeepEP): """ GroupedExpertsDeepEP + LoRA. @@ -773,117 +615,3 @@ def forward( y = self.token_dispatcher.token_unpermutation(output2) 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/_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/tests/unit_tests/_peft/test_lora_experts_mxfp4.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py index d3f17189a1..f59ba6521e 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4.py @@ -16,7 +16,8 @@ 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, GroupedExpertsLoRAMXFP4 +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 diff --git a/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py b/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py index ffa3632235..0ea9171276 100644 --- a/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py +++ b/tests/unit_tests/_peft/test_lora_experts_mxfp4_deepep.py @@ -25,10 +25,8 @@ 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, - GroupedExpertsDeepEPLoRAMXFP4, -) +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 7dacd5798069c0e70eb5da7667984ce8a40fa7b9 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 7 Sep 2026 00:30:08 -0700 Subject: [PATCH 28/29] refactor(moe): restore eager DeepEP buffer initialization Signed-off-by: HuiyingLi --- nemo_automodel/components/moe/experts.py | 11 ++--------- tests/unit_tests/moe/test_experts.py | 5 +---- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 68095ca988..9c96e90680 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -1008,15 +1008,8 @@ def init_token_dispatcher(self, ep_mesh: DeviceMesh): config=config, ep_group=ep_group, ) - # NOTE: previously called `self._init_deepep_buffer(ep_group)` here to - # eagerly allocate the DeepEP NVSHMEM buffer at model construction - # (introduced in #2076, e42584e3). On single-node EP=8 ep_shard=1 - # DSv4-Flash, the eager allocation collides with the ~135 GB load-time - # peak and OOMs at `_aggregate_experts` torch.stack. Revert to the - # original lazy allocation in FusedDispatch.forward (fused_a2a.py:136 - # via the global `_buffer` cache). Both code paths produce the same - # buffer; only the *timing* differs. - # _init_deepep_buffer remains defined below for explicit callers. + if self.dispatcher_backend == "deepep": + self._init_deepep_buffer(ep_group) def _init_deepep_buffer(self, ep_group: dist.ProcessGroup) -> None: """Initialize DeepEP communication buffers before activation checkpointing.""" diff --git a/tests/unit_tests/moe/test_experts.py b/tests/unit_tests/moe/test_experts.py index 75cac1db68..74b0f8824f 100644 --- a/tests/unit_tests/moe/test_experts.py +++ b/tests/unit_tests/moe/test_experts.py @@ -867,10 +867,7 @@ def test_grouped_experts_deepep_token_dispatcher_init(self, moe_config): assert hasattr(experts, "token_dispatcher") assert experts.ep_size == 2 assert experts.ep_rank == 0 - # The DeepEP NVSHMEM buffer is allocated lazily (in FusedDispatch.forward), - # not eagerly in init_token_dispatcher — the revert that fixed the single-node - # load-time OOM. So init_token_dispatcher must NOT call _init_deepep_buffer. - mock_init_buffer.assert_not_called() + mock_init_buffer.assert_called_once_with(mock_mesh.get_group.return_value) def test_grouped_experts_deepep_apply_bias_no_bias(self, moe_config): """Test _apply_bias method with no bias.""" From d1719376d8a798285d4d9b8eef939d6fea139127 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 7 Sep 2026 00:32:28 -0700 Subject: [PATCH 29/29] refactor(moe): keep sqrtsoftplus guard model-local Signed-off-by: HuiyingLi --- nemo_automodel/components/models/deepseek_v4/model.py | 3 +-- nemo_automodel/components/moe/layers.py | 6 +----- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index 387297973a..79ed0c2767 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -430,8 +430,7 @@ def forward( scores = F.linear(x.float(), self.weight.float()) if self.score_func == "sqrtsoftplus": # clamp_min: softplus underflows to 0.0 for very negative logits and sqrt'(0)=inf - # makes the backward NaN; bound it with a negligible forward change. See the - # matching guard in moe/layers.py Gate. + # 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() diff --git a/nemo_automodel/components/moe/layers.py b/nemo_automodel/components/moe/layers.py index 692bf0cf48..2afe823ec9 100644 --- a/nemo_automodel/components/moe/layers.py +++ b/nemo_automodel/components/moe/layers.py @@ -456,11 +456,7 @@ def _route_scores(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso weights = original_scores.gather(1, indices) elif self.score_func == "sqrtsoftplus": # sqrt(softplus(x)) = sqrt(log(1 + exp(x))), used in DeepSeek V4. - # clamp_min keeps the sqrt argument strictly positive: softplus(x) underflows - # to exactly 0.0 in fp32 for very negative logits (x <~ -104), and sqrt'(0) = inf - # makes the backward NaN. The clamp bounds the gradient with a negligible - # (sqrt(1e-12) = 1e-6) change to the forward value. - scores = torch.sqrt(F.softplus(scores.float()).clamp_min(1e-12)) + scores = torch.sqrt(F.softplus(scores.float())) original_scores = scores if correction_bias is not None: