feat: DeepSeek V4 Flash support

examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml

@@ -0,0 +1,126 @@
+# 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.
+
+# To run this recipe with 16 H100 nodes (128 GPUs)
+#   torchrun --nproc-per-node 8 examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag.yaml --nproc-per-node 8
+
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+seed: 1234
+
+step_scheduler:
+  global_batch_size: 256
+  local_batch_size: 8
+  ckpt_every_steps: 500
+  val_every_steps: 500
+  num_epochs: 1
+  max_steps: 100
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1
+  pp_size: 4
+  ep_size: 32
+
+  sequence_parallel: false
+  activation_checkpointing: false
+
+  pipeline:
+    pp_schedule: interleaved1f1b
+    pp_microbatch_size: 1
+    layers_per_stage: 2
+    round_virtual_stages_to_pp_multiple: down
+    scale_grads_in_schedule: false
+    patch_inner_model: false
+    patch_causal_lm_model: false
+
+  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: sdpa
+    linear: torch
+    rms_norm: torch_fp32
+    rope_fusion: false
+    dispatcher: torch
+    experts: torch_mm
+    enable_hf_state_dict_adapter: true
+    enable_fsdp_optimizations: true
+
+checkpoint:
+  enabled: false
+  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
+  betas:
+  - 0.9
+  - 0.95
+  eps: 1e-8
+  lr: 1e-5
+  weight_decay: 0.1

nemo_automodel/_transformers/registry.py

@@ -43,6 +43,10 @@
             "DeepseekV32ForCausalLM",
             ("nemo_automodel.components.models.deepseek_v32.model", "DeepseekV32ForCausalLM"),
         ),
+        (
+            "DeepseekV4ForCausalLM",
+            ("nemo_automodel.components.models.deepseek_v4.model", "DeepseekV4ForCausalLM"),
+        ),
         (
             "Glm4MoeForCausalLM",
             ("nemo_automodel.components.models.glm4_moe.model", "Glm4MoeForCausalLM"),
@@ -162,6 +166,7 @@
 # AutoConfig.from_pretrained can resolve them without trust_remote_code.
 _CUSTOM_CONFIG_REGISTRATIONS: Dict[str, Tuple[str, str]] = {
     "baichuan": ("nemo_automodel.components.models.baichuan.configuration", "BaichuanConfig"),
+    "deepseek_v4": ("nemo_automodel.components.models.deepseek_v4.config", "DeepseekV4Config"),
     "kimi_k25": ("nemo_automodel.components.models.kimi_k25_vl.model", "KimiK25VLConfig"),
     "kimi_vl": ("nemo_automodel.components.models.kimivl.model", "KimiVLConfig"),
     "llava_onevision": ("nemo_automodel.components.models.llava_onevision.model", "LlavaOneVisionConfig"),

nemo_automodel/components/checkpoint/_backports/hf_utils.py

@@ -48,6 +48,8 @@
     "I64": torch.int64,
     "BF16": torch.bfloat16,
     "F8_E4M3": torch.float8_e4m3fn,
+    "F8_E5M2": torch.float8_e5m2,
+    "F8_E8M0": torch.float8_e8m0fnu,
 }
 
 HF_DCP_VERSION: float = 1.0

nemo_automodel/components/checkpoint/checkpointing.py

@@ -986,7 +986,15 @@ def _get_storage_reader(
             Configured storage reader or None for other formats.
         """
         if self.config.model_save_format == SerializationFormat.SAFETENSORS or is_init_step:
-            if key_mapping is None:
+            # The upstream HuggingFaceStorageReader delegates dtype decoding to
+            # safetensors.torch._TYPES, which does not yet recognize the FP8
+            # scale dtypes emitted by some quantized HF checkpoints (e.g.
+            # DeepSeek V4's F8_E8M0 scales → KeyError('F8_E8M0') inside
+            # read_metadata → DCP ends up with metadata=None on every rank).
+            # The in-tree backport's DTYPE_MAP was extended for F8_E8M0/F8_E5M2,
+            # so prefer it for base-model HF loads. Mid-training DCP loads may
+            # still use the faster upstream reader.
+            if key_mapping is None and not is_init_step:
                 try:
                     from torch.distributed.checkpoint.hf_storage import (
                         HuggingFaceStorageReader as _UpstreamHFReader,

nemo_automodel/components/datasets/llm/hellaswag.py

@@ -34,7 +34,6 @@ def __init__(
         tokenizer,
         split="train",
         num_samples_limit=None,
-        trust_remote_code=True,
         pad_to_max_length=True,
     ):
         """Initialize the HellaSwag dataset wrapper.
@@ -44,7 +43,6 @@ def __init__(
             tokenizer (PreTrainedTokenizer): The tokenizer used to process text.
             split (str, optional): Dataset split to use (e.g., 'train', 'validation'). Defaults to 'train'.
             num_samples_limit (int, optional): Maximum number of samples to load. Defaults to None.
-            trust_remote_code (bool, optional): Whether to trust remote code. Defaults to True.
             pad_to_max_length (bool, optional): Whether to pad sequences to max length in the dataset.
                 If False, sequences will have variable lengths and padding will be handled by the collate function.
                 Defaults to True.
@@ -54,7 +52,7 @@ def __init__(
         """
         if isinstance(num_samples_limit, int):
             split = f"{split}[:{num_samples_limit}]"
-        raw_datasets = load_dataset(path_or_dataset, split=split, trust_remote_code=trust_remote_code)
+        raw_datasets = load_dataset(path_or_dataset, split=split)
         processor = SFTSingleTurnPreprocessor(tokenizer)
         processor.pad_to_max_length = pad_to_max_length
         self.dataset = processor.process(raw_datasets, self)

nemo_automodel/components/distributed/pipelining/functional.py

@@ -284,28 +284,51 @@ def _precompute_stage_shapes(
     """
     hidden_size, vocab_size = _get_hidden_and_vocab_size(model_config)
 
+    # DeepSeek V4 preserves an extra hc_mult axis between blocks, so inter-stage
+    # hidden state is [mb, seq, hc_mult, dim] until the last (norm) stage folds
+    # it back to [mb, seq, dim].
+    is_v4 = (
+        getattr(model_config, "model_type", None) == "deepseek_v4"
+        or getattr(getattr(model_config, "text_config", None), "model_type", None) == "deepseek_v4"
+    )
+    hc_mult = int(getattr(model_config, "hc_mult", 1) or 1) if is_v4 else 1
+
     for stage in stages:
         # Infer the computation dtype from the stage's parameters
         try:
             model_dtype = next(stage.submod.parameters()).dtype
         except StopIteration:
             model_dtype = torch.bfloat16
 
+        inner_submod = getattr(stage.submod, "model", stage.submod)
+        stage_has_norm = getattr(inner_submod, "norm", None) is not None
+
         # --- inputs_meta ---
         if stage.is_first:
             # First stage receives input_ids: [mb, seq_len] int64
             stage.inputs_meta = (torch.empty(microbatch_size, seq_len, device="meta", dtype=torch.long),)
         else:
-            # Non-first stages receive hidden_states: [mb, seq_len, hidden_size]
-            stage.inputs_meta = (torch.empty(microbatch_size, seq_len, hidden_size, device="meta", dtype=model_dtype),)
+            if hc_mult > 1:
+                stage.inputs_meta = (
+                    torch.empty(microbatch_size, seq_len, hc_mult, hidden_size, device="meta", dtype=model_dtype),
+                )
+            else:
+                stage.inputs_meta = (
+                    torch.empty(microbatch_size, seq_len, hidden_size, device="meta", dtype=model_dtype),
+                )
 
         # --- outputs_meta ---
         has_lm_head = hasattr(stage.submod, "lm_head") and stage.submod.lm_head is not None
         if has_lm_head:
             # Last stage with lm_head produces logits: [mb, seq_len, vocab_size]
             outputs_meta = (torch.empty(microbatch_size, seq_len, vocab_size, device="meta", dtype=model_dtype),)
+        elif hc_mult > 1 and not stage_has_norm:
+            # V4 mid-pipeline: tensor still carries the hc_mult axis.
+            outputs_meta = (
+                torch.empty(microbatch_size, seq_len, hc_mult, hidden_size, device="meta", dtype=model_dtype),
+            )
         else:
-            # Intermediate stages produce hidden_states: [mb, seq_len, hidden_size]
+            # Standard intermediate stage (or V4 final-norm stage without lm_head).
             outputs_meta = (torch.empty(microbatch_size, seq_len, hidden_size, device="meta", dtype=model_dtype),)
         stage._configure_outputs_meta(outputs_meta)
 
@@ -452,6 +475,12 @@ def split_model_into_stages(
     else:
         lm_head_fqn = "lm_head"
 
+    # DeepSeek V4: model carries an extra compressor-rotary module on every stage
+    # and an HC head on the last stage; both must survive PP module pruning.
+    is_v4_keep = getattr(getattr(model, "config", None), "model_type", None) == "deepseek_v4"
+    has_rotary_emb_compress = is_v4_keep and hasattr(text_model, "rotary_emb_compress")
+    has_hc_head = is_v4_keep and hasattr(text_model, "hc_head")
+
     # Auto-generate module split if not provided
     if module_names_per_stage is None:
         module_names_per_stage = generate_hf_model_fqn_per_model_part(
@@ -466,6 +495,14 @@ def split_model_into_stages(
             lm_head_fqn=lm_head_fqn,
         )
 
+        # V4 post-processing: keep the compressor rotary on every stage and the
+        # HC head on the last stage so the V4 PP forward can run end-to-end.
+        if has_rotary_emb_compress:
+            for stage_modules in module_names_per_stage:
+                stage_modules.append(f"{layers_prefix}rotary_emb_compress")
+        if has_hc_head:
+            module_names_per_stage[-1].append(f"{layers_prefix}hc_head")
+
     def _build_stage_from_modules(
         stage_idx: int, module_names: list[str], num_stages: int
     ) -> tuple[PipelineStage, nn.Module]:

nemo_automodel/components/models/deepseek_v4/__init__.py

@@ -0,0 +1,13 @@
+# 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.