feat(engine): add Datum-based forward, backward, and optimizer APIs

examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml

@@ -71,14 +71,12 @@ dataset:
   split: train
 
 packed_sequence:
-  # Set packed_sequence_size > 0 to run with packed sequences
-  packed_sequence_size: 1024
-  packing_strategy: thd
+  packed_sequence_size: 0
 
-# StatefulDataLoader with packed-sequence collate
+# StatefulDataLoader with default padded collate
 dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
-  collate_fn: nemo_automodel.components.datasets.utils.packed_sequence_thd_collater
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
   shuffle: true
 
 validation_dataset:
@@ -88,7 +86,7 @@ validation_dataset:
 
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
-  collate_fn: nemo_automodel.components.datasets.utils.packed_sequence_thd_collater
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
 
 optimizer:
   _target_: torch.optim.Adam

examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml

@@ -93,6 +93,7 @@ validation_dataset:
 
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.gemma4_prefix_collate_fn
 

examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml

@@ -98,6 +98,7 @@ validation_dataset:
 
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.gemma4_prefix_collate_fn
 

examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml

@@ -139,6 +139,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml

@@ -117,6 +117,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml

@@ -121,6 +121,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml

@@ -108,6 +108,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml

@@ -106,6 +106,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral4/mistral4_medpix.yaml

@@ -100,6 +100,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
 

examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml

@@ -59,7 +59,6 @@ model:
     rope_fusion: false
   attn_implementation: sdpa
   torch_dtype: bfloat16
-
 processor:
   _target_: transformers.AutoProcessor.from_pretrained
   pretrained_model_name_or_path: Qwen/Qwen3.5-4B

examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml

@@ -118,6 +118,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml

@@ -114,6 +114,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml

@@ -123,6 +123,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

nemo_automodel/__init__.py

@@ -40,6 +40,11 @@
 _SUBMODULES = {"recipes", "shared", "components", "models"}
 
 _LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    "CollatedLossInputs": ("nemo_automodel.components.datasets.datum", "CollatedLossInputs"),
+    "Datum": ("nemo_automodel.components.datasets.datum", "Datum"),
+    "Engine": ("nemo_automodel.engine", "Engine"),
+    "LossInputLayout": ("nemo_automodel.components.datasets.datum", "LossInputLayout"),
+    "LossFnOutputBatch": ("nemo_automodel.engine.outputs", "LossFnOutputBatch"),
     "NeMoAutoModelForCausalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForCausalLM"),
     "NeMoAutoModelForImageTextToText": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForImageTextToText"),
     "NeMoAutoModelForMultimodalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForMultimodalLM"),
@@ -56,6 +61,7 @@
     "NeMoAutoModelBiEncoder": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelBiEncoder"),
     "NeMoAutoModelCrossEncoder": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelCrossEncoder"),
     "NeMoAutoTokenizer": ("nemo_automodel._transformers.auto_tokenizer", "NeMoAutoTokenizer"),
+    "PerTokenOutput": ("nemo_automodel.engine.outputs", "PerTokenOutput"),
     "NeMoAutoDiffusionPipeline": ("nemo_automodel._diffusers.auto_diffusion_pipeline", "NeMoAutoDiffusionPipeline"),
     "ModelCapabilities": ("nemo_automodel._transformers.model_capabilities", "ModelCapabilities"),
     "query_capabilities": ("nemo_automodel._transformers.model_capabilities", "query_capabilities"),

nemo_automodel/_transformers/auto_model.py

@@ -28,6 +28,7 @@
 import inspect
 import logging
 import os
+from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from typing import TYPE_CHECKING, List, Optional, Union
 
@@ -399,6 +400,8 @@ def _build_model(
         fp8_config,
         compile_config,
         load_base_model,
+        pre_fsdp_hook: Callable[[torch.nn.Module], None] | None = None,
+        skip_task_head_prefixes_for_base_model: Sequence[str] | None = None,
         _retry_depth=0,
         **kwargs,
     ):
@@ -448,6 +451,8 @@ def _retry(**override):
                 peft_config=peft_config,
                 fp8_config=fp8_config,
                 compile_config=compile_config,
+                pre_fsdp_hook=pre_fsdp_hook,
+                skip_task_head_prefixes_for_base_model=skip_task_head_prefixes_for_base_model,
                 load_base_model=load_base_model,
                 _retry_depth=_retry_depth + 1,
                 **retry_kwargs,
@@ -636,6 +641,8 @@ def _retry(**override):
             freeze_config=freeze_config,
             weights_already_loaded=weights_already_loaded,
             inject_te_attention=inject_te_attention,
+            pre_fsdp_hook=pre_fsdp_hook,
+            skip_task_head_prefixes_for_base_model=skip_task_head_prefixes_for_base_model,
         )
 
         return model
@@ -658,6 +665,8 @@ def from_pretrained(
         peft_config: dict | None = None,
         fp8_config: Optional["FP8Config"] = None,
         compile_config: Optional["CompileConfig"] = None,
+        pre_fsdp_hook: Callable[[torch.nn.Module], None] | None = None,
+        skip_task_head_prefixes_for_base_model: Sequence[str] | None = None,
         **kwargs,
     ) -> PreTrainedModel:
         """
@@ -711,6 +720,15 @@ def from_pretrained(
                 If provided, FP8 quantization will be applied. Default: None.
             compile_config (CompileConfig | None, optional): Configuration for torch.compile.
                 If provided, the model will be compiled. Default: None.
+            pre_fsdp_hook: Optional in-place model-structure hook invoked after model and
+                kernel setup but before FSDP wrapping. The hook must return ``None``.
+                It must make the same deterministic change on every rank and must not
+                run collectives. Initially supported only without model parallelism,
+                PEFT, or quantization. Added parameters use the model's standard
+                initialization path and configured FSDP precision policy.
+            skip_task_head_prefixes_for_base_model: Native model parameter FQN
+                prefixes to omit from the pretrained base-checkpoint load.
+                Training-checkpoint restore is unaffected.
             **kwargs: Additional keyword arguments. Notable ones include:
                 - has_packed_sequence (bool): Whether using packed sequences. Default: False.
                 - cache_dir (str): Cache directory for model weights.
@@ -778,6 +796,8 @@ def from_pretrained(
             peft_config=peft_config,
             fp8_config=fp8_config,
             compile_config=compile_config,
+            pre_fsdp_hook=pre_fsdp_hook,
+            skip_task_head_prefixes_for_base_model=skip_task_head_prefixes_for_base_model,
             load_base_model=True,
             **kwargs,
         )
@@ -800,6 +820,8 @@ def from_config(
         peft_config: dict | None = None,
         fp8_config: Optional["FP8Config"] = None,
         compile_config: Optional["CompileConfig"] = None,
+        pre_fsdp_hook: Callable[[torch.nn.Module], None] | None = None,
+        skip_task_head_prefixes_for_base_model: Sequence[str] | None = None,
         **kwargs,
     ) -> PreTrainedModel:
         """
@@ -883,6 +905,8 @@ def from_config(
             peft_config=peft_config,
             fp8_config=fp8_config,
             compile_config=compile_config,
+            pre_fsdp_hook=pre_fsdp_hook,
+            skip_task_head_prefixes_for_base_model=skip_task_head_prefixes_for_base_model,
             load_base_model=kwargs.pop("load_base_model", False),
             **kwargs,
         )

nemo_automodel/_transformers/infrastructure.py

@@ -24,6 +24,7 @@
 """
 
 import logging
+from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from dataclasses import is_dataclass, replace
 from functools import partial
@@ -483,6 +484,8 @@ def apply_model_infrastructure(
     pretrained_model_name_or_path="",
     weights_already_loaded=False,
     inject_te_attention: bool = False,
+    pre_fsdp_hook: Callable[[torch.nn.Module], None] | None = None,
+    skip_task_head_prefixes_for_base_model: Sequence[str] | None = None,
     **_kwargs,
 ):
     """Apply sharding, PEFT, quantization, and checkpoint loading to a model.
@@ -521,6 +524,18 @@ def apply_model_infrastructure(
         inject_te_attention: When True, inject TransformerEngine DotProductAttention
             into all ``self_attn`` modules of HF models (has no effect on custom
             models that already use TE via BackendConfig). Default: False.
+        pre_fsdp_hook: Optional in-place model-structure hook invoked after lower-
+            precision and attention transforms and any load-before-shard base-model
+            restore, but before state-key capture and FSDP wrapping. The hook must
+            return ``None``, make the same deterministic change on every rank, and
+            avoid collectives. Fresh parameters follow whichever constructor or
+            post-shard model-initialization path applies and the configured FSDP
+            mixed-precision policy; no task-specific initializer or precision
+            override is added. The hook is responsible for keeping model
+            configuration such as weight tying consistent with its changes.
+        skip_task_head_prefixes_for_base_model: Native model parameter FQN prefixes
+            to omit from the pretrained base-checkpoint load. Full training-checkpoint
+            restores still load these parameters.
         **_kwargs: Additional keyword arguments (ignored, allows passing extra kwargs)
 
     Returns:
@@ -529,6 +544,46 @@ def apply_model_infrastructure(
     if mesh is None:
         mesh = MeshContext()
 
+    if pre_fsdp_hook is not None and not callable(pre_fsdp_hook):
+        raise TypeError("pre_fsdp_hook must be callable or None.")
+
+    if pre_fsdp_hook is not None:
+        unsupported = [
+            name
+            for name, enabled in (
+                ("tensor parallelism", mesh.tp_size != 1),
+                ("context parallelism", mesh.cp_size != 1),
+                ("expert parallelism", mesh.ep_size != 1),
+                ("pipeline parallelism", autopipeline is not None or mesh.pp_size != 1),
+                ("PEFT", peft_config is not None),
+                (
+                    "quantization",
+                    quantization_config is not None
+                    or getattr(getattr(model, "config", None), "quantization_config", None) is not None,
+                ),
+                ("FP8", fp8_config is not None),
+                ("QAT", qat_quantizer is not None),
+            )
+            if enabled
+        ]
+        if unsupported:
+            raise NotImplementedError(
+                "pre_fsdp_hook currently supports only unquantized, non-PEFT models with "
+                f"tp_size=cp_size=ep_size=pp_size=1; unsupported: {', '.join(unsupported)}."
+            )
+
+    if isinstance(skip_task_head_prefixes_for_base_model, str):
+        raise TypeError("skip_task_head_prefixes_for_base_model must be a sequence of non-empty strings, not a string.")
+    skip_task_head_prefixes = (
+        list(skip_task_head_prefixes_for_base_model) if skip_task_head_prefixes_for_base_model is not None else None
+    )
+    if skip_task_head_prefixes is not None:
+        if any(not isinstance(prefix, str) for prefix in skip_task_head_prefixes):
+            raise TypeError("skip_task_head_prefixes_for_base_model must contain only strings.")
+        if any(not prefix for prefix in skip_task_head_prefixes):
+            raise ValueError("skip_task_head_prefixes_for_base_model must not contain empty prefixes.")
+        skip_task_head_prefixes = list(dict.fromkeys(skip_task_head_prefixes))
+
     # Create a checkpointer for loading base weights only. Keep consolidation disabled
     # so load-only infrastructure does not emit save/export warnings.
     ckpt_config = CheckpointingConfig(
@@ -539,6 +594,7 @@ def apply_model_infrastructure(
         model_repo_id=pretrained_model_name_or_path,
         save_consolidated=False,
         is_peft=peft_config is not None,
+        skip_task_head_prefixes_for_base_model=skip_task_head_prefixes,
     )
     checkpointer = Checkpointer(
         ckpt_config,
@@ -610,6 +666,24 @@ def apply_model_infrastructure(
             checkpointer.load_base_model(model, device, cache_dir, pretrained_model_name_or_path, load_base_model=False)
         checkpoint_already_loaded = True
 
+    if pre_fsdp_hook is not None:
+        # A single-device meta model has already been materialized by the
+        # load-before-shard path above, while multi-rank FSDP models remain on
+        # meta until after sharding. Inspect the current tensors rather than
+        # the original construction flag so newly added modules land on the
+        # same device as the model in both cases.
+        has_meta_tensors = any(parameter.device.type == "meta" for parameter in model.parameters()) or any(
+            buffer.device.type == "meta" for buffer in model.buffers()
+        )
+        hook_context = init_empty_weights() if has_meta_tensors else nullcontext()
+        with hook_context:
+            hook_result = pre_fsdp_hook(model)
+        if hook_result is not None:
+            raise TypeError(
+                "pre_fsdp_hook must mutate the existing model in place and return None; "
+                f"got {type(hook_result).__name__}."
+            )
+
     # hold a list copy of the model state dict keys before any parallelization. To be used during checkpoint saving in safetensors format.
     state_dict_adapter = getattr(model, "state_dict_adapter", None)
     get_hf_state_dict_keys = getattr(state_dict_adapter, "get_hf_state_dict_keys", None)

nemo_automodel/components/checkpoint/checkpointing.py

@@ -100,6 +100,27 @@
 logger = logging.getLogger(__name__)
 
 
+def _remove_base_checkpoint_task_heads(
+    state_dict: dict[str, torch.Tensor],
+    prefixes: list[str],
+) -> dict[str, torch.Tensor]:
+    """Remove task-head tensors intentionally omitted from a base-model load.
+
+    Args:
+        state_dict: Full base-checkpoint state dictionary. Tensor shapes are
+            model-dependent and keys are native model parameter FQNs. The mapping
+            is mutated in place.
+        prefixes: Native model parameter FQN prefixes to remove.
+
+    Returns:
+        The same state-dict mapping with matching task-head tensors removed.
+    """
+    for key in tuple(state_dict):
+        if any(key.startswith(prefix) for prefix in prefixes):
+            state_dict.pop(key)
+    return state_dict
+
+
 def _format_restricted_load_error(f: FileLike) -> str:
     return (
         f"Refusing to load torch artifact from {f!r} with pickle-based torch.load. "
@@ -818,6 +839,10 @@ def load_model(
                         "Materialized missing tied lm_head.weight from embedding weights for %s during init load.",
                         type(model_state.model[0]).__name__,
                     )
+                _remove_base_checkpoint_task_heads(
+                    converted_state_dict,
+                    model_state.skip_task_head_prefixes,
+                )
                 # Load using full_state_dict=True to properly convert tensors to DTensors for FSDP
                 _load_full_state_dict_into_model(model_state.model, converted_state_dict)
                 return
@@ -889,6 +914,11 @@ def load_model(
                     type(model_state.model[0]).__name__,
                 )
 
+            _remove_base_checkpoint_task_heads(
+                state_dict_from_disk,
+                model_state.skip_task_head_prefixes,
+            )
+
             total_bytes = sum(
                 t.nelement() * t.element_size() for t in state_dict_from_disk.values() if isinstance(t, torch.Tensor)
             )