refactor(engine): training Engine API

README.md

@@ -167,6 +167,7 @@ High-throughput scalable training
 - ✅ **DCP** - Supports PyTorch DCP and SafeTensors, sharded and consolidated layouts, merge/reshard utilities, and Hugging Face-compatible outputs.
 - ✅ **Async checkpointing** - Can write checkpoints in the background to reduce training stalls caused by I/O.
 - ✅ **Dion and Muon optimizers** - Distributed optimizer integrations with typed recipe configuration.
+- ✅ **Training Engine** - DeepSpeed-style eager `forward`, `backward`, and `step` over already-distributed models.
 - ✅ **Environment Support** - SLURM, interactive, SkyPilot, and Kubernetes (via SkyPilot) launchers.
 
 SOTA algorithms
@@ -190,7 +191,7 @@ Agentic Development and UX
 - ✅ **Agent-friendly skills** - Curated [`skills/`](https://github.com/NVIDIA-NeMo/Automodel/tree/main/skills) for common dev tasks (recipe runs, model onboarding, CI).
 
 Planned for 26.08
-- 🔜 **Unified Engine API and recipes** - Introduce a common engine and consolidate the LLM and VLM recipe paths.
+- 🔜 **Pipeline and recipe consolidation** - Consolidate the LLM and VLM pipeline-training lifecycle behind one clear recipe boundary.
 - 🔜 **Composable component configuration** - Complete the typed config and `.build()` refactor across data and remaining components.
 - 🔜 **Packed long-context training with CP** - Combine THD sequence packing with context parallelism, including DeepSeek V4 coverage.
 - 🔜 **Kernel and runtime upgrades** - Add partial CUDA graphs, evaluate FlashAttention 3/4, and upgrade to DeepEP v2.

docs/api-reference/index.mdx

@@ -19,5 +19,6 @@ This reference is built from docstrings in the [source code](https://github.com/
 | `nemo_automodel.components.checkpoint` | Async DCP and SafeTensors checkpointing |
 | `nemo_automodel.components.quantization` | FP8, QAT, calibration |
 | `nemo_automodel.components._peft` | LoRA, QLoRA adapters |
+| `nemo_automodel.engine` | DeepSpeed-style eager model wrapper for backward, gradient accumulation, clipping, optimizer updates, and scheduler steps |
 | `nemo_automodel.components.launcher` | Interactive, SkyPilot, and NeMo-Run job launchers; Slurm jobs use the repository's `slurm.sub` script |
 | `nemo_automodel.cli` | `uv run automodel <config.yaml> [--nproc-per-node N] [--key.subkey=override ...]` entry point |

nemo_automodel/__init__.py

@@ -37,9 +37,12 @@
 # Heavy dependencies (e.g., torch/transformers) are intentionally imported lazily
 # via __getattr__ so importing tokenizers doesn't pull in the full training stack.
 
-_SUBMODULES = {"recipes", "shared", "components", "models"}
+_SUBMODULES = {"recipes", "shared", "components", "models", "engine"}
 
 _LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    "AutoMFU": ("nemo_automodel._transformers.mfu", "AutoMFU"),
+    "Engine": ("nemo_automodel.engine", "Engine"),
+    "get_is_hf_model": ("nemo_automodel._transformers.model_init", "get_is_hf_model"),
     "NeMoAutoModelForCausalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForCausalLM"),
     "NeMoAutoModelForImageTextToText": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForImageTextToText"),
     "NeMoAutoModelForMultimodalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForMultimodalLM"),

nemo_automodel/_transformers/model_init.py

@@ -1252,6 +1252,11 @@ def __init_model(
                 _download_model_weights(hf_config, pretrained_model_name_or_path, process_group=process_group)
             logger.info(f"Using custom model implementation for {architectures[0]}")
             kwargs.pop("trust_remote_code", None)
+            # ``hf_config`` is passed positionally below.  Keep an in-memory
+            # config supplied to ``from_pretrained`` out of ``**kwargs`` so
+            # custom constructors that accept arbitrary keywords do not see
+            # the same argument twice.
+            kwargs.pop("config", None)
             # Treat config-related kwargs as config overrides (HF behavior) and
             # avoid forwarding them into model __init__.
             init_param_names = _get_init_param_names(model_cls)

nemo_automodel/components/datasets/vlm/collate_fns.py

@@ -56,7 +56,7 @@
     mask_fake_vision_tokens_batch,
 )
 from nemo_automodel.components.datasets.vlm.samplers import _smart_resize_image
-from nemo_automodel.components.datasets.vlm.utils import default_stop_tokens
+from nemo_automodel.components.datasets.vlm.utils import default_stop_tokens, merge_media_values
 
 # ---------------------------------------------------------------------------
 # Patch BaseVideoProcessor.fetch_videos to use decord (decord2) instead of
@@ -1362,22 +1362,6 @@ def default_collate_fn(
     return batch
 
 
-def _merge_media_values(values: list[Any]) -> torch.Tensor | list[Any]:
-    """Merge fixed-shape patch tensors or preserve variable-resolution media lists."""
-    if not values:
-        raise ValueError("Media merge requires at least one value.")
-    if all(isinstance(value, torch.Tensor) for value in values):
-        return torch.cat(values, dim=0).to(torch.bfloat16)
-    if all(isinstance(value, (list, tuple)) for value in values):
-        return [
-            item.to(torch.bfloat16) if isinstance(item, torch.Tensor) else item for value in values for item in value
-        ]
-    raise TypeError(
-        "VLM media values must be consistently tensors or variable-resolution lists, "
-        f"got {[type(value).__name__ for value in values]}."
-    )
-
-
 def pad_collate_fn(
     examples: Sequence[Dict[str, Any]],
     processor,
@@ -1434,8 +1418,8 @@ def pad_collate_fn(
         "attention_mask": torch.stack(padded_attention_mask),
     }
 
-    # Pad sequence-length tensors that mirror input_ids (e.g. mm_token_type_ids)
-    for seq_key in ("mm_token_type_ids",):
+    # Pad sequence-length tensors that mirror input_ids.
+    for seq_key in ("mm_token_type_ids", "token_type_ids"):
         if any(seq_key in ex for ex in examples):
             padded = []
             for ex in examples:
@@ -1465,7 +1449,7 @@ def pad_collate_fn(
     for key in ("pixel_values", "pixel_values_videos"):
         tensors = [ex[key] for ex in examples if key in ex and ex[key] is not None]
         if tensors:
-            batch[key] = _merge_media_values(tensors)
+            batch[key] = merge_media_values(tensors, field_name=key)
 
     # Per-sample image counts from image_grid_thw shapes (before concat)
     image_grid_per_sample = [
@@ -1621,7 +1605,7 @@ def _pad_mrope(pos, target_len):
     for key in ("pixel_values", "pixel_values_videos"):
         tensors = [x[key] for x in batch if key in x and x[key] is not None]
         if tensors:
-            result[key] = _merge_media_values(tensors)
+            result[key] = merge_media_values(tensors, field_name=key)
 
     for key in ("image_grid_thw", "image_position_ids", "video_grid_thw", "second_per_grid_ts"):
         tensors = [x[key] for x in batch if key in x and x[key] is not None]
@@ -1670,8 +1654,9 @@ def packed_sequence_thd_vlm_collater(
             (default -1000); filtered downstream in ``process_input_for_thd``.
 
     Returns:
-        Dict with ``input_ids``/``labels`` ``[batch, seq]``, ``position_ids``
-        ``[batch, seq]`` or ``[3, batch, seq]``, ``seq_lens``/``seq_lens_padded``
+        Dict with ``input_ids``/``labels`` ``[batch, seq]``, optional token-type
+        fields ``[batch, seq]``, ``position_ids`` ``[batch, seq]`` or
+        ``[3, batch, seq]``, ``seq_lens``/``seq_lens_padded``
         ``[batch, max_packs]``, ``qkv_format='thd'``, and concatenated media tensors.
     """
     if not batch:
@@ -1757,10 +1742,18 @@ def _pad_seq(tensor, pad_value, target_len, seq_dim=-1):
         "qkv_format": "thd",
     }
 
+    for key in ("mm_token_type_ids", "token_type_ids"):
+        if any(key in item and item[key] is not None for item in batch):
+            values = [
+                item[key] if item.get(key) is not None else torch.zeros_like(torch.as_tensor(item["input_ids"]))
+                for item in batch
+            ]
+            result[key] = torch.stack([_pad_seq(value, 0, max_len) for value in values])
+
     for key in ("pixel_values", "pixel_values_videos"):
         tensors = [x[key] for x in batch if key in x and x[key] is not None]
         if tensors:
-            result[key] = _merge_media_values(tensors)
+            result[key] = merge_media_values(tensors, field_name=key)
 
     for key in ("image_grid_thw", "image_position_ids", "video_grid_thw", "second_per_grid_ts"):
         tensors = [x[key] for x in batch if key in x and x[key] is not None]

nemo_automodel/components/datasets/vlm/neat_packing_vlm.py

@@ -53,6 +53,7 @@
     _smart_resize_image,
     _smart_resize_video,
 )
+from nemo_automodel.components.datasets.vlm.utils import merge_media_values
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +66,8 @@
     "second_per_grid_ts",
 )
 
+_TOKEN_TYPE_KEYS = ("mm_token_type_ids", "token_type_ids")
+
 # ---------------------------------------------------------------------------
 # Visual-token-balanced greedy knapsack
 # ---------------------------------------------------------------------------
@@ -345,9 +348,10 @@ def _shift_sample(sample: dict, has_mrope: bool = False) -> dict:
     out["labels"] = sample["labels"][1:]
     out["attention_mask"] = sample["attention_mask"][:-1]
 
-    if (mm_ttids := sample.get("mm_token_type_ids")) is not None:
-        mm_ttids = torch.as_tensor(mm_ttids)
-        out["mm_token_type_ids"] = mm_ttids[0, :-1] if mm_ttids.ndim == 2 else mm_ttids[:-1]
+    for key in _TOKEN_TYPE_KEYS:
+        if (token_types := sample.get(key)) is not None:
+            token_types = torch.as_tensor(token_types)
+            out[key] = token_types[0, :-1] if token_types.ndim == 2 else token_types[:-1]
 
     if has_mrope and "position_ids" in sample and sample["position_ids"] is not None:
         out["position_ids"] = sample["position_ids"][:, :-1]
@@ -367,7 +371,6 @@ def _aligned_length(length: int, alignment: int) -> int:
 
 def _build_packed_vlm_sample(
     samples: list[dict],
-    pack_size: int,
     padding_idx: int,
     has_mrope: bool = False,
     sequence_alignment: int = 1,
@@ -379,7 +382,10 @@ def _build_packed_vlm_sample(
     all_input_ids: list[int] = []
     all_labels: list[int] = []
     all_attention_mask: list[int] = []
-    all_mm_token_type_ids: list[int] = []
+    all_token_type_ids: dict[str, list[int]] = {key: [] for key in _TOKEN_TYPE_KEYS}
+    present_token_type_keys = tuple(
+        key for key in _TOKEN_TYPE_KEYS if any(sample.get(key) is not None for sample in samples)
+    )
     all_position_ids_1d: list[int] = []
     mrope_position_ids_list: list[torch.Tensor] = []
     seq_lens: list[int] = []
@@ -411,12 +417,13 @@ def _build_packed_vlm_sample(
         all_labels.extend(labs + [-100] * pad)
         all_attention_mask.extend([seq_idx] * padded_seq_len)
 
-        mm_ttids = sample.get("mm_token_type_ids")
-        if mm_ttids is not None:
-            mm_ttids = mm_ttids.tolist() if isinstance(mm_ttids, torch.Tensor) else list(mm_ttids)
-            all_mm_token_type_ids.extend(mm_ttids + [0] * pad)
-        else:
-            all_mm_token_type_ids.extend([0] * padded_seq_len)
+        for key in present_token_type_keys:
+            token_types = sample.get(key)
+            if token_types is None:
+                all_token_type_ids[key].extend([0] * padded_seq_len)
+                continue
+            token_types = torch.as_tensor(token_types).reshape(-1).tolist()
+            all_token_type_ids[key].extend(token_types + [0] * pad)
 
         if has_mrope and "position_ids" in sample:
             mrope_position_ids_list.append(sample["position_ids"])
@@ -443,26 +450,22 @@ def _build_packed_vlm_sample(
         "input_ids": torch.tensor(all_input_ids, dtype=torch.long),
         "labels": torch.tensor(all_labels, dtype=torch.long),
         "attention_mask": torch.tensor(all_attention_mask, dtype=torch.long),
-        "mm_token_type_ids": torch.tensor(all_mm_token_type_ids, dtype=torch.long),
         "seq_lens": seq_lens,
         "seq_lens_padded": seq_lens_padded,
         "n_images": n_images,
         "n_videos": n_videos,
     }
+    for key in present_token_type_keys:
+        packed[key] = torch.tensor(all_token_type_ids[key], dtype=torch.long)
 
     if has_mrope and mrope_position_ids_list:
         packed["position_ids"] = torch.cat(mrope_position_ids_list, dim=1)
     else:
         packed["position_ids"] = torch.tensor(all_position_ids_1d, dtype=torch.long)
 
-    if pixel_values_list and all(isinstance(value, torch.Tensor) for value in pixel_values_list):
-        packed["pixel_values"] = torch.cat(pixel_values_list, dim=0)
-    elif pixel_values_list and all(isinstance(value, (list, tuple)) for value in pixel_values_list):
-        packed["pixel_values"] = [item for value in pixel_values_list for item in value]
-    elif pixel_values_list:
-        raise TypeError("Packed VLM pixel_values must be consistently tensors or variable-resolution lists.")
-    else:
-        packed["pixel_values"] = None
+    packed["pixel_values"] = (
+        merge_media_values(pixel_values_list, field_name="pixel_values") if pixel_values_list else None
+    )
     packed["image_grid_thw"] = torch.cat(image_grid_thw_list, dim=0) if image_grid_thw_list else None
     packed["image_position_ids"] = torch.cat(image_position_ids_list, dim=0) if image_position_ids_list else None
     packed["pixel_values_videos"] = torch.cat(pixel_values_videos_list, dim=0) if pixel_values_videos_list else None
@@ -472,6 +475,43 @@ def _build_packed_vlm_sample(
     return packed
 
 
+def pack_vlm_samples(
+    samples: Sequence[dict[str, object]],
+    *,
+    padding_idx: int,
+    get_rope_index: Callable[..., object] | None = None,
+    sequence_alignment: int = 1,
+) -> dict[str, object]:
+    """Shift and concatenate pretokenized VLM samples into one packed sample.
+
+    Args:
+        samples: Processor outputs before the autoregressive input/label shift.
+        padding_idx: Token ID used for per-document alignment padding.
+        get_rope_index: Optional model callback that builds multi-axis position IDs.
+        sequence_alignment: Alignment applied independently to every shifted sample.
+
+    Returns:
+        One packed sample containing token, position, sequence, and media fields.
+    """
+    has_mrope = get_rope_index is not None
+    shifted_samples = []
+    for sample in samples:
+        prepared = dict(sample)
+        if get_rope_index is not None:
+            position_ids = _compute_mrope_position_ids(prepared, get_rope_index)
+            if position_ids is None:
+                raise ValueError("get_rope_index must accept input_ids and return VLM position_ids")
+            prepared["position_ids"] = position_ids
+        shifted_samples.append(_shift_sample(prepared, has_mrope=has_mrope))
+
+    return _build_packed_vlm_sample(
+        shifted_samples,
+        padding_idx,
+        has_mrope=has_mrope,
+        sequence_alignment=sequence_alignment,
+    )
+
+
 # ---------------------------------------------------------------------------
 # PackedDatasetWrapper — lazy packing via __getitem__
 # ---------------------------------------------------------------------------
@@ -555,8 +595,7 @@ def __init__(
             raise ValueError(f"sequence_alignment must be at least 1, got {sequence_alignment}.")
         self.sequence_alignment = sequence_alignment
         self.get_rope_index = get_rope_index
-        self.has_mrope = get_rope_index is not None
-        if self.has_mrope and self.sequence_alignment > 1:
+        if get_rope_index is not None and self.sequence_alignment > 1:
             raise NotImplementedError("Context-parallel THD packing for multi-axis mRoPE VLMs is not yet implemented.")
         self.max_retries = max_retries
 
@@ -566,18 +605,11 @@ def __len__(self):
     def __getitem__(self, pack_idx: int) -> dict:
         """Materialize one pack: tokenize + shift + concat all samples in the bin."""
         bin_indices = self.bins[pack_idx]
-        shifted_samples: list[dict] = []
+        samples: list[dict] = []
 
         for sample_idx in bin_indices:
             sample = self.inner[sample_idx]  # tokenize + load media
-
-            if self.has_mrope and self.get_rope_index is not None:
-                mrope_pos = _compute_mrope_position_ids(sample, self.get_rope_index)
-                if mrope_pos is not None:
-                    sample["position_ids"] = mrope_pos
-
-            shifted = _shift_sample(sample, has_mrope=self.has_mrope)
-            seq_len = shifted["input_ids"].shape[0]
+            seq_len = len(sample["input_ids"]) - 1
             aligned_seq_len = _aligned_length(seq_len, self.sequence_alignment)
 
             # The aligned length is the actual capacity consumed by THD CP.
@@ -592,16 +624,16 @@ def __getitem__(self, pack_idx: int) -> dict:
                 )
                 continue
 
-            shifted_samples.append(shifted)
+            samples.append(sample)
 
         # Truncate if total exceeds pack_size (estimation was wrong)
         total = 0
         kept: list[dict] = []
-        for s in shifted_samples:
-            slen = s["input_ids"].shape[0]
+        for sample in samples:
+            slen = len(sample["input_ids"]) - 1
             aligned_slen = _aligned_length(slen, self.sequence_alignment)
             if total + aligned_slen <= self.pack_size:
-                kept.append(s)
+                kept.append(sample)
                 total += aligned_slen
             else:
                 logger.debug(
@@ -615,13 +647,18 @@ def __getitem__(self, pack_idx: int) -> dict:
 
         if not kept:
             # Fallback: return a padding-only pack
-            kept = [{"input_ids": torch.tensor([], dtype=torch.long), "labels": torch.tensor([], dtype=torch.long)}]
-
-        return _build_packed_vlm_sample(
+            kept = [
+                {
+                    "input_ids": torch.tensor([self.padding_idx], dtype=torch.long),
+                    "labels": torch.tensor([self.padding_idx], dtype=torch.long),
+                    "attention_mask": torch.ones(1, dtype=torch.long),
+                }
+            ]
+
+        return pack_vlm_samples(
             kept,
-            self.pack_size,
-            self.padding_idx,
-            has_mrope=self.has_mrope,
+            padding_idx=self.padding_idx,
+            get_rope_index=self.get_rope_index,
             sequence_alignment=self.sequence_alignment,
         )