feat(moe): mxfp4-resident MoE experts for DeepSeek-V4-Flash LoRA

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

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

examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4.yaml

@@ -0,0 +1,168 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# LoRA fine-tuning of deepseek-ai/DeepSeek-V4-Flash on HellaSwag with the frozen
+# routed experts kept resident in mxfp4 (fp4-e2m1 + e8m0 block scales) and
+# dequantized on the fly in the grouped-GEMM forward/backward.
+#
+# Only the experts are quantized; every other weight (MLA attention, dense MLP,
+# embeddings, MTP, lm_head) stays bf16. The routed experts are ~90%+ of the
+# parameters, so packing them to ~4 bits cuts steady-state base-weight memory by
+# roughly 4x relative to the bf16 LoRA recipe.
+#
+# v1 constraint: mxfp4 experts require the torch_mm GroupedExperts backend, so
+# this recipe uses dispatcher=torch (NOT deepep). DeepEP + mxfp4 is a follow-up.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+seed: 1234
+
+step_scheduler:
+  global_batch_size: 128
+  local_batch_size: 1
+  ckpt_every_steps: 500
+  val_every_steps: 500
+  gc_every_steps: 10
+  num_epochs: 1
+  max_steps: 100
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  # Single node, 8x H200: ep_size must divide dp_size*cp_size (= world_size here).
+  ep_size: 8
+
+  sequence_parallel: false
+  # Recommended on: the 43-layer activation footprint, not the (now-packed) expert
+  # weights, is the steady-state pressure once experts are mxfp4.
+  activation_checkpointing: true
+
+  moe:
+    reshard_after_forward: false
+    wrap_outer_model: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config
+  config:
+    _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+    name_or_path: deepseek-ai/DeepSeek-V4-Flash
+    num_nextn_predict_layers: 0
+    # Required by FusedLinearCrossEntropy: the trainer calls the model with
+    # logits_to_keep=1 and reads out.hidden_states to apply the fused lm_head + CE
+    # without materializing full [seq, vocab] logits. The forward only populates
+    # hidden_states when config.output_hidden_states is set.
+    output_hidden_states: true
+  trust_remote_code: false
+  load_base_model: true
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: tilelang
+    linear: torch
+    rms_norm: torch_fp32
+    rope_fusion: false
+    # mxfp4 experts run on the torch_mm grouped-GEMM path (DeepEP not yet supported).
+    dispatcher: torch
+    experts: torch_mm
+    enable_hf_state_dict_adapter: true
+    enable_fsdp_optimizations: true
+
+peft:
+  _target_: nemo_automodel.components._peft.lora.PeftConfig
+  target_modules:
+    # attention
+    - "*wq_a"
+    - "*wq_b"
+    - "*wkv"
+    - "*wo_b"
+    # routed experts: GroupedExperts -> GroupedExpertsLoRAMXFP4
+    # (frozen base stays mxfp4-packed; only the LoRA adapters train)
+    - "*mlp.experts"
+    # shared experts: dense MLP nn.Linear (base stays bf16, LoRA adapters bf16)
+    - "*shared_experts.gate_proj"
+    - "*shared_experts.up_proj"
+    - "*shared_experts.down_proj"
+  dim: 8
+  alpha: 32
+  use_triton: True
+  # fp32 LoRA adapters: fp32 master weights + fp32 AdamW state for stability (small-update
+  # swamping), while FSDP2's param_dtype=bf16 still runs the adapter matmuls in bf16. Adapters
+  # are tiny so the fp32 optimizer-state cost is negligible. Matches HF PEFT's default.
+  lora_dtype: float32
+  # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly.
+  expert_weight_format: mxfp4
+
+checkpoint:
+  enabled: false
+  # The DSV4-Flash checkpoint is quantized: fp4 routed experts + fp8 non-expert
+  # projections. This enables quant-aware loading -- the fp8 non-expert weights are
+  # dequantized to bf16, while the fp4 experts are passed through packed (mxfp4,
+  # never materialized in bf16). Required to load this checkpoint.
+  dequantize_base_checkpoint: true
+
+loss_fn:
+  # Fused linear cross-entropy (Apple cut_cross_entropy): applies the lm_head + CE without
+  # materializing the [seq, vocab=129280] logits, removing the ~16 GiB fp32 logits spike and
+  # raising the single-node context ceiling. Uses the DSV4 forward's logits_to_keep /
+  # hidden_states path; loss stays fp32-quality (cut_cross_entropy accumulates the logsumexp in
+  # fp32 internally). Swap back to masked_ce.MaskedCrossEntropy if cut_cross_entropy is unavailable.
+  _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: rowan/hellaswag
+  split: train
+  tokenizer:
+    _target_: transformers.AutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+
+packed_sequence:
+  packed_sequence_size: 0
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn:
+    _target_: nemo_automodel.components.datasets.utils.default_collater
+    pad_seq_len_divisible: 64
+  shuffle: true
+
+validation_dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: rowan/hellaswag
+  split: validation
+  tokenizer:
+    _target_: transformers.AutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+
+validation_dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn:
+    _target_: nemo_automodel.components.datasets.utils.default_collater
+    pad_seq_len_divisible: 64
+  shuffle: false
+  drop_last: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  fused: true
+  betas: [0.9, 0.95]
+  eps: 1e-8
+  lr: 1e-5
+  weight_decay: 0.1

examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora_mxfp4_deepep.yaml

@@ -0,0 +1,196 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# LoRA fine-tuning of deepseek-ai/DeepSeek-V4-Flash on HellaSwag with the frozen
+# routed experts kept resident in mxfp4 (fp4-e2m1 + e8m0 block scales) and
+# dequantized on the fly in the grouped-GEMM forward/backward, dispatched with HybridEP.
+#
+# This is the HybridEP counterpart of deepseek_v4_flash_hellaswag_lora_mxfp4.yaml.
+# mxfp4 quantizes only the expert *weights* (local to each EP rank); HybridEP only
+# governs token *dispatch/combine* (fused all-to-all of bf16 activations + fp32 probs).
+# They intersect only at the two post-dispatch grouped GEMMs, which read the packed
+# weights via MXFP4GroupedMM. Unlike the torch dispatcher (which all-gathers every
+# token to every rank), HybridEP routes each rank only the tokens for its local experts,
+# so this is the scalable EP path for DSV4 at EP=8/64.
+# On an 8-GPU NVLink node, set NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8.
+#
+# Only the experts are quantized; every other weight (MLA attention, dense MLP,
+# embeddings, MTP, lm_head) stays bf16. The routed experts are ~90%+ of the
+# parameters, so packing them to ~4 bits cuts steady-state base-weight memory by
+# roughly 4x relative to the bf16 LoRA recipe.
+#
+# Constraint: mxfp4 experts require the torch_mm experts backend (the grouped_gemm
+# 'gmm' path has no packed variant), so this recipe pairs dispatcher=hybridep with
+# experts=torch_mm.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+seed: 1234
+
+step_scheduler:
+  global_batch_size: 128
+  local_batch_size: 1
+  ckpt_every_steps: 500
+  val_every_steps: 500
+  gc_every_steps: 10
+  num_epochs: 1
+  max_steps: 100
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  # Single node, 8x H200: ep_size must divide dp_size*cp_size (= world_size here).
+  ep_size: 8
+
+  sequence_parallel: false
+  # Recommended on: the 43-layer activation footprint, not the (now-packed) expert
+  # weights, is the steady-state pressure once experts are mxfp4.
+  activation_checkpointing: true
+
+  moe:
+    reshard_after_forward: false
+    wrap_outer_model: false
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config
+  config:
+    _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+    name_or_path: deepseek-ai/DeepSeek-V4-Flash
+    num_nextn_predict_layers: 0
+    # Required by FusedLinearCrossEntropy: the trainer calls the model with
+    # logits_to_keep=1 and reads out.hidden_states to apply the fused lm_head + CE
+    # without materializing full [seq, vocab] logits. The forward only populates
+    # hidden_states when config.output_hidden_states is set.
+    output_hidden_states: true
+  trust_remote_code: false
+  load_base_model: true
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: tilelang
+    linear: torch
+    rms_norm: torch_fp32
+    rope_fusion: false
+    # mxfp4 experts run on the torch_mm grouped-GEMM path, dispatched with HybridEP.
+    dispatcher: hybridep
+    experts: torch_mm
+    enable_hf_state_dict_adapter: true
+    enable_fsdp_optimizations: true
+
+peft:
+  _target_: nemo_automodel.components._peft.lora.PeftConfig
+  target_modules:
+    # Direct attention projections only. The more permissive "*wkv"/"*wq_b"
+    # patterns also match the FP32 HCA compressor/indexer projections.
+    - "*.self_attn.wq_a"
+    - "*.self_attn.wq_b"
+    - "*.self_attn.wkv"
+    - "*.self_attn.wo_b"
+    # routed experts: GroupedExpertsDeepEP -> GroupedExpertsDeepEPLoRAMXFP4
+    # (frozen base stays mxfp4-packed; only the LoRA adapters train)
+    - "*mlp.experts"
+    # shared experts: dense MLP nn.Linear (base stays bf16, LoRA adapters bf16)
+    - "*shared_experts.gate_proj"
+    - "*shared_experts.up_proj"
+    - "*shared_experts.down_proj"
+  dim: 8
+  alpha: 32
+  use_triton: True
+  # BF16 adapters match the base compute dtype and keep every patched linear's
+  # FSDP2 storage dtype uniform. FP32 adapters would require separate parameter-owning
+  # modules for the adapter weights under the current FSDP2 dtype isolation rules.
+  lora_dtype: bfloat16
+  # Keep the frozen routed experts packed as mxfp4 and dequantize on the fly.
+  expert_weight_format: mxfp4
+
+checkpoint:
+  enabled: false
+  # The DSV4-Flash checkpoint is quantized: fp4 routed experts + fp8 non-expert
+  # projections. This enables quant-aware loading -- the fp8 non-expert weights are
+  # dequantized to bf16, while the fp4 experts are passed through packed (mxfp4,
+  # never materialized in bf16). Required to load this checkpoint.
+  dequantize_base_checkpoint: true
+
+loss_fn:
+  # Fused linear cross-entropy (Apple cut_cross_entropy): applies the lm_head + CE without
+  # materializing the [seq, vocab=129280] logits, removing the ~16 GiB fp32 logits spike and
+  # raising the single-node context ceiling (~30k -> ~36-38k tokens on one 8xH200). Uses the
+  # DSV4 forward's logits_to_keep / hidden_states path; loss stays fp32-quality (cut_cross_entropy
+  # accumulates the logsumexp in fp32 internally). Swap back to masked_ce.MaskedCrossEntropy if
+  # cut_cross_entropy is unavailable.
+  _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: rowan/hellaswag
+  split: train
+  tokenizer:
+    _target_: transformers.AutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+
+packed_sequence:
+  packed_sequence_size: 0
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn:
+    _target_: nemo_automodel.components.datasets.utils.default_collater
+    pad_seq_len_divisible: 64
+  shuffle: true
+
+validation_dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: rowan/hellaswag
+  split: validation
+  tokenizer:
+    _target_: transformers.AutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: deepseek-ai/DeepSeek-V4-Flash
+
+validation_dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn:
+    _target_: nemo_automodel.components.datasets.utils.default_collater
+    pad_seq_len_divisible: 64
+  shuffle: false
+  drop_last: true
+
+optimizer:
+  # Keep the BF16 adapter parameters/checkpoint while TE FusedAdam reconstructs
+  # FP32 master weights from the BF16 values plus their stored 16-bit remainders.
+  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
+  betas: [0.9, 0.95]
+  eps: 1e-8
+  lr: 1e-5
+  weight_decay: 0.1
+  adam_w_mode: true
+  bias_correction: true
+  master_weights: true
+  master_weight_dtype: torch.float32
+  store_param_remainders: true
+  exp_avg_dtype: torch.float32
+  exp_avg_sq_dtype: torch.float32
+
+wandb:
+  enable: false
+  project: automodel-dsv4
+  name: deepseek-v4-flash-mxfp4-lora-hybridep-te-fusedadam
+  group: pr-2548
+  tags: [deepseek-v4-flash, mxfp4, lora, hybridep, te-fusedadam]
+  mode: online