From 43ff2bd10e19fe270a9db75d5dd728f97f40a392 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 15 Aug 2026 17:40:00 -0700 Subject: [PATCH 01/34] feat(engine): add Datum forward-backward windows Signed-off-by: HuiyingLi --- nemo_automodel/__init__.py | 2 + nemo_automodel/components/datasets/datum.py | 284 +++++++++------- nemo_automodel/engine.py | 288 ++++++++++++++++ tests/unit_tests/datasets/test_datum.py | 135 ++++++-- tests/unit_tests/test_engine.py | 348 ++++++++++++++++++++ 5 files changed, 897 insertions(+), 160 deletions(-) create mode 100644 nemo_automodel/engine.py create mode 100644 tests/unit_tests/test_engine.py diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 740728f570..143aae8a5d 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -40,6 +40,8 @@ _SUBMODULES = {"recipes", "shared", "components", "models"} _LAZY_ATTRS: dict[str, tuple[str, str]] = { + "Datum": ("nemo_automodel.components.datasets.datum", "Datum"), + "Engine": ("nemo_automodel.engine", "Engine"), "NeMoAutoModelForCausalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForCausalLM"), "NeMoAutoModelForImageTextToText": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForImageTextToText"), "NeMoAutoModelForMultimodalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForMultimodalLM"), diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index 99db330092..e9fdae368c 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -12,42 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Typed input contract for training: :class:`Datum` and :func:`collate_datums`. - -A ``Datum`` is the single-example input boundary between user/algorithm code -(SFT, RL post-training) and the training loop. It lives in ``components.datasets`` -because feeding and collating examples is a data concern — and, crucially, -because that lets :func:`collate_datums` **reuse the canonical collaters** -(``default_collater`` for padded ``[B, T]`` and ``packed_sequence_thd_collater`` -for THD) instead of forking a second padding/packing implementation that could -drift from them. - -The companion output contract (``ModelOutput`` and the per-token extraction -helpers) lives in ``components.training`` — that side touches model logits, so -it is a forward concern, not a dataset one. - -Conventions ------------ -* A ``Datum`` holds **one** sequence. ``input_ids`` is 1-D, shape ``[T]``. -* ``loss_fn_inputs`` carries everything the loss needs, aligned to ``input_ids`` - token positions (length ``T``) for per-token entries: - - =============== ======================================================= - key meaning - =============== ======================================================= - ``target_tokens`` next-token targets, shape ``[T]`` (becomes ``labels``) - ``weights`` per-token loss mask / weight (0 disables a position) - ``logprobs`` old/behavior-policy logprobs (importance sampling) - ``advantages`` advantage signal (PPO/GRPO), per-token or per-sample - =============== ======================================================= - -* Masking convention matches the codebase: a target position with - ``weights == 0`` becomes ``ignore_index`` (default ``-100``) in ``labels``. -""" +"""The model-input and loss-input boundary used by training engines.""" from __future__ import annotations from dataclasses import dataclass, field +from typing import Any import torch import torch.nn.functional as F @@ -63,57 +33,116 @@ __all__ = ["Datum", "collate_datums"] -@dataclass +@dataclass(init=False) class Datum: - """A single training example. + """One processor-ready training example. + + ``model_inputs`` contains exactly the keyword arguments needed by the + model for one example. Text examples usually contain ``input_ids``; + multimodal examples may additionally contain fields such as + ``pixel_values`` and ``image_grid_thw``. ``loss_fn_inputs`` is kept + separate so algorithm data such as targets, weights, old log-probabilities, + and advantages is never passed to the model. Args: - input_ids: 1-D ``LongTensor`` of token ids, shape ``[T]``. - loss_fn_inputs: per-key tensors the loss consumes. Per-token entries are - 1-D and length ``T``; per-sample entries are scalar or shape ``[1]``. - See the module docstring for the well-known keys. + model_inputs: Model-ready values for one example. Values are + model-specific because LLM and VLM processors emit different + fields. + loss_fn_inputs: Tensor values consumed by the loss function. + input_ids: Deprecated convenience spelling for the old text-only API. + It cannot be combined with ``model_inputs``. """ - input_ids: torch.Tensor + model_inputs: dict[str, Any] loss_fn_inputs: dict[str, torch.Tensor] = field(default_factory=dict) + def __init__( + self, + model_inputs: dict[str, Any] | torch.Tensor | list[int] | None = None, + loss_fn_inputs: dict[str, torch.Tensor] | None = None, + *, + input_ids: torch.Tensor | list[int] | None = None, + ) -> None: + # Preserve the old positional ``Datum(input_ids, loss_fn_inputs)`` form + # while downstream users move to the model-ready mapping. + if model_inputs is not None and not isinstance(model_inputs, dict): + if input_ids is not None: + raise ValueError("pass either model_inputs or input_ids, not both") + input_ids = model_inputs + model_inputs = None + if model_inputs is not None and input_ids is not None: + raise ValueError("pass either model_inputs or input_ids, not both") + if model_inputs is None: + if input_ids is None: + raise ValueError("Datum requires model_inputs") + model_inputs = {"input_ids": input_ids} + + self.model_inputs = dict(model_inputs) + self.loss_fn_inputs = dict(loss_fn_inputs or {}) + self.__post_init__() + def __post_init__(self) -> None: - if not isinstance(self.input_ids, torch.Tensor): - self.input_ids = torch.as_tensor(self.input_ids, dtype=torch.long) - if self.input_ids.dim() != 1: - raise ValueError(f"Datum.input_ids must be 1-D [T]; got shape {tuple(self.input_ids.shape)}") + if not self.model_inputs: + raise ValueError("Datum.model_inputs cannot be empty") + + input_ids = self.model_inputs.get("input_ids") + if input_ids is not None: + if not isinstance(input_ids, torch.Tensor): + input_ids = torch.as_tensor(input_ids, dtype=torch.long) + self.model_inputs["input_ids"] = input_ids + if input_ids.ndim != 1: + raise ValueError( + "Datum.model_inputs['input_ids'] must be 1-D [T] for one token sequence; " + f"got shape {tuple(input_ids.shape)}" + ) + for key, value in self.loss_fn_inputs.items(): if not isinstance(value, torch.Tensor): self.loss_fn_inputs[key] = torch.as_tensor(value) + @property + def input_ids(self) -> torch.Tensor: + """The text token sequence, retained for source compatibility.""" + input_ids = self.model_inputs.get("input_ids") + if not isinstance(input_ids, torch.Tensor): + raise AttributeError("this Datum has no tensor model input named 'input_ids'") + return input_ids + @property def seq_len(self) -> int: - """Number of tokens in this example.""" - return int(self.input_ids.shape[0]) - - def to_features(self, *, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX) -> dict[str, list[int]]: - """Emit the per-example dict the canonical collaters expect. - - Every position of a ``Datum`` is a real token, so ``attention_mask`` is - all ones: it tells the padded collater exactly which positions it added, - instead of leaving it to infer them from the pad token *value* — which - misreads a real token that happens to equal the pad id (commonly - ``pad_token_id == eos_token_id``) as padding. - - ``labels`` is included only when ``loss_fn_inputs["target_tokens"]`` is - present, with positions where ``loss_fn_inputs["weights"] == 0`` set to - ``ignore_index``. Only integer token fields are emitted here — the - collaters cast to ``LongTensor``; float side-inputs are batched - separately by :func:`collate_datums`. - - Returns: - ``{"input_ids": [...], "attention_mask": [...], "labels": [...]}`` - as plain ``list[int]``. + """Token length used by the default text collater.""" + if isinstance(self.model_inputs.get("input_ids"), torch.Tensor): + return int(self.model_inputs["input_ids"].shape[0]) + for key in ("target_tokens", "weights"): + value = self.loss_fn_inputs.get(key) + if isinstance(value, torch.Tensor) and value.ndim == 1: + return int(value.shape[0]) + raise ValueError("cannot infer token length; use a model-specific collate_fn for this Datum") + + def to_features(self, *, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX) -> dict[str, Any]: + """Convert one text Datum for the repository's canonical collaters. + + Token-aligned 1-D model inputs become Python lists so + :func:`default_collater` can pad ragged examples. Other model inputs are + passed through unchanged. A model-specific collater should be used for + layouts such as multimodal mRoPE positions or variable media tensors. """ - features: dict[str, list[int]] = { - "input_ids": self.input_ids.tolist(), - "attention_mask": [1] * self.seq_len, - } + if "input_ids" not in self.model_inputs: + raise ValueError("the default collater requires model_inputs['input_ids']") + + features: dict[str, Any] = {} + for key, value in self.model_inputs.items(): + if isinstance(value, torch.Tensor) and value.ndim == 1 and value.shape[0] == self.seq_len: + if value.is_floating_point(): + raise ValueError( + f"the default text collater cannot preserve floating-point token field {key!r}; " + "use a model-specific collate_fn" + ) + features[key] = value.tolist() + else: + features[key] = value + features.setdefault("attention_mask", [1] * self.seq_len) + if "target_tokens" in self.loss_fn_inputs: labels = self.loss_fn_inputs["target_tokens"].clone() weights = self.loss_fn_inputs.get("weights") @@ -129,73 +158,74 @@ def collate_datums( packed: bool = False, pad_seq_len_divisible: int | None = None, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX, -) -> dict[str, torch.Tensor]: - """Collate a list of :class:`Datum` into a model-ready batch dict. - - Token fields are delegated to the **existing** canonical collaters so the - padded / THD schema (``attention_mask`` / ``qkv_format`` / ``seq_lens``) is - produced by the same code paths the dataset pipeline uses — no fork: - - * ``packed=False`` → ``default_collater`` (padded ``[B, T]``). - * ``packed=True`` → :func:`pack_features_for_thd` concatenates all datums - into one pre-packed record, then ``packed_sequence_thd_collater`` emits - the flat ``[1, total_tokens]`` THD schema (``qkv_format="thd"``, - per-sequence ``seq_lens`` for splitting outputs back per datum). - - Float per-token side-inputs (every ``loss_fn_inputs`` key shared by all datums - except ``target_tokens``, e.g. ``weights`` / ``logprobs`` / ``advantages``) - are batched under their own key — this is the part the token collaters - cannot carry (they cast to ``LongTensor``). Padded mode right-pads them to - the collated width and stacks to ``[B, T]``; packed mode concatenates them - in datum order to ``[1, total_tokens]``, aligned with ``input_ids``. - Per-sample (scalar / length-1) entries are stacked into a ``[num_datums]`` - tensor without padding in both modes. A length-1 entry on a single-token - sequence matches both shapes; it is read as per-token. +) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + """Collate text Datums into separate model and loss inputs. + + This is the default text collater. Callers with model-specific VLM + batching rules pass a thin callable around their existing collater to + ``Engine`` instead. Args: - datums: examples for this microbatch. Must be non-empty. One ``Datum`` - is treated as one sequence. - packed: pack all datums into one flat ``[1, total_tokens]`` THD row - instead of the padded ``[B, T]`` layout. - pad_seq_len_divisible: pad sequence length to a multiple of this value - (padded mode only; TP/CP/FP8 alignment). - ignore_index: label value for masked positions. + datums: Non-empty examples for one microbatch. + packed: Produce one flat THD token row instead of a padded batch. + pad_seq_len_divisible: Round the padded token width to this multiple. + ignore_index: Label fill value used internally by the THD collater. Returns: - The collater output dict, augmented with the float side-input tensors. + ``(model_inputs, loss_fn_inputs)``. Per-token loss inputs have shape + ``[B, T]`` in padded mode and ``[1, total_tokens]`` in packed mode; + scalar loss inputs have shape ``[B]``. """ - if len(datums) == 0: + if not datums: raise ValueError("collate_datums requires at least one Datum") - features = [datums[i].to_features(ignore_index=ignore_index) for i in range(len(datums))] - if packed: - # Concatenate all datums into one pre-packed record and let the - # canonical THD collater produce the [1, total] schema. - batch = packed_sequence_thd_collater([pack_features_for_thd(features, ignore_index=ignore_index)]) - else: - batch = default_collater([dict(f) for f in features], pad_seq_len_divisible) + model_keys = set(datums[0].model_inputs) + loss_keys = set(datums[0].loss_fn_inputs) + for datum in datums[1:]: + if set(datum.model_inputs) != model_keys: + raise ValueError("every Datum in a microbatch must carry the same model_inputs keys") + if set(datum.loss_fn_inputs) != loss_keys: + raise ValueError("every Datum in a microbatch must carry the same loss_fn_inputs keys") - width = int(batch["input_ids"].shape[-1]) - keys = set(datums[0].loss_fn_inputs) - for d in datums[1:]: - if set(d.loss_fn_inputs) != keys: + features = [datum.to_features(ignore_index=ignore_index) for datum in datums] + if packed: + unsupported = model_keys - {"input_ids", "attention_mask"} + if unsupported: raise ValueError( - "every Datum in a batch must carry the same loss_fn_inputs keys " - f"(got {sorted(keys)} and {sorted(d.loss_fn_inputs)}); a missing key would " - "silently drop it -- for `weights` that means losing the loss mask" + "the default packed collater only supports text input_ids; " + f"use a model-specific collate_fn for {sorted(unsupported)}" ) - for key in sorted(keys - {"target_tokens"}): - rows = [d.loss_fn_inputs[key].to(torch.float).flatten() for d in datums] - if all(r.shape[0] == d.seq_len for r, d in zip(rows, datums)): + for datum in datums: + attention_mask = datum.model_inputs.get("attention_mask") + if attention_mask is not None and not bool(torch.as_tensor(attention_mask).bool().all()): + raise ValueError( + "packed Datums must contain only real tokens; explicit attention_mask padding is unsupported" + ) + model_inputs = packed_sequence_thd_collater([pack_features_for_thd(features, ignore_index=ignore_index)]) + else: + model_inputs = default_collater([dict(feature) for feature in features], pad_seq_len_divisible) + + # Labels are loss data, not model input. The canonical collaters only see + # them so their padding/THD machinery can be reused unchanged. + model_inputs.pop("labels", None) + width = int(model_inputs["input_ids"].shape[-1]) + + loss_inputs: dict[str, torch.Tensor] = {} + for key in sorted(loss_keys): + values = [datum.loss_fn_inputs[key] for datum in datums] + per_token = all(value.ndim == 1 and value.shape[0] == datum.seq_len for value, datum in zip(values, datums)) + if per_token: if packed: - # Per-token field, packed: concatenate in datum order so the - # values ride the same flat [1, total] axis as input_ids. - batch[key] = torch.cat(rows).unsqueeze(0) + loss_inputs[key] = torch.cat(values).unsqueeze(0) else: - # Per-token field, padded: right-pad to the collated width and stack. - batch[key] = torch.stack([F.pad(r, (0, width - r.shape[0])) for r in rows]) - else: - # Per-sample field: one value per datum (shape [num_datums], which in - # packed mode is deliberately NOT the batch dim of the [1, total] rows). - batch[key] = torch.stack([r.reshape(-1)[0] for r in rows]) - return batch + loss_inputs[key] = torch.stack([F.pad(value, (0, width - value.shape[0])) for value in values]) + continue + + if not all(value.numel() == 1 for value in values): + shapes = [tuple(value.shape) for value in values] + raise ValueError( + f"the default collater only supports scalar or 1-D token-aligned loss inputs; {key!r} has {shapes}" + ) + loss_inputs[key] = torch.stack([value.reshape(()) for value in values]) + + return model_inputs, loss_inputs diff --git a/nemo_automodel/engine.py b/nemo_automodel/engine.py new file mode 100644 index 0000000000..9710e11288 --- /dev/null +++ b/nemo_automodel/engine.py @@ -0,0 +1,288 @@ +# 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. + +"""A small forward/backward engine for Datum accumulation windows.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn + +from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.distributed.mesh import MeshContext +from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh +from nemo_automodel.components.distributed.utils import get_sync_ctx +from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler +from nemo_automodel.components.training.utils import ( + prepare_after_first_microbatch, + prepare_for_final_backward, + prepare_for_grad_accumulation, +) + +CollateFn = Callable[[list[Datum]], tuple[dict[str, Any], dict[str, torch.Tensor]]] +LossFn = Callable[ + [Any, dict[str, torch.Tensor], Sequence[Datum]], + torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], +] + +__all__ = ["Engine"] + + +class Engine: + """Run model forward/backward over one optimizer accumulation window. + + The model and distributed topology are already constructed when they are + passed here. The Engine owns batching, global weight normalization, + gradient-accumulation synchronization, and backward. It deliberately does + not zero, clip, finalize expert gradients, or step them; callers choose the + optimizer boundary and retain the repository's existing distributed + gradient-finalization path. + + Args: + model: An already configured and distributed model. + device: Device on which model inputs and losses are evaluated. + mesh_context: Runtime topology. When omitted, an initialized default + process group is treated as pure data parallelism. + collate_fn: Batches one microbatch of Datums into separate model and + loss inputs. The default supports padded or packed text; + VLMs pass a thin callable around their model-specific collater. + The callable must keep model inputs and loss inputs aligned and + preserve the sum of ``weights``. + defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the + final microbatch. + + Note: + This first execution backend is eager and weight-normalized. Pipeline + and context parallel schedules and pre-reduced scalar losses are + intentionally deferred. + """ + + def __init__( + self, + model: nn.Module, + *, + device: torch.device | str, + mesh_context: MeshContext | None = None, + collate_fn: CollateFn = collate_datums, + defer_fsdp_grad_sync: bool = True, + ) -> None: + self.model = model + self.device = torch.device(device) + self.mesh_context = mesh_context + self.collate_fn = collate_fn + self.defer_fsdp_grad_sync = defer_fsdp_grad_sync + + def forward_backward( + self, + window: Sequence[Sequence[Datum]], + loss_fn: LossFn, + ) -> tuple[torch.Tensor, list[dict[str, Any]]]: + """Accumulate gradients for a complete optimizer window. + + ``window`` is explicit: each inner sequence is one eager microbatch. + ``loss_fn`` receives the raw model output, collated + ``loss_fn_inputs``, and the original Datums for that microbatch. It + returns unreduced losses with exactly the same shape as + ``loss_fn_inputs["weights"]``. It may also return one output mapping + per Datum. Those mappings are detached and preserved in input order; + the Engine deliberately does not interpret or reduce them. + + Returns: + ``(loss, loss_fn_outputs)``. ``loss`` is a detached, DP-reduced + scalar. ``loss_fn_outputs`` contains local-rank, per-Datum mappings + in window order. Model parameters are unchanged, but their + gradients contain the complete window's globally normalized + backward result. + """ + microbatches = self._validate_window(window) + self._validate_parallelism() + dp_group, dp_size = self._dp_group_and_size() + self._validate_window_size_across_dp(len(microbatches), dp_group, dp_size) + denominator = self._global_weight_sum(microbatches, dp_group, dp_size) + + self.model.train() + prepare_for_grad_accumulation([self.model], pp_enabled=False) + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(1.0 / len(microbatches)) + + local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) + loss_fn_outputs: list[dict[str, Any]] = [] + returns_outputs: bool | None = None + + for index, datums in enumerate(microbatches): + is_last = index == len(microbatches) - 1 + if is_last: + prepare_for_final_backward([self.model], pp_enabled=False) + + model_inputs, loss_inputs = self.collate_fn(datums) + model_inputs = _to_device(model_inputs, self.device) + loss_inputs = _to_device(loss_inputs, self.device) + weights = self._validate_collated_weights(datums, loss_inputs) + + with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync): + output = self.model(**model_inputs) + result = loss_fn(output, loss_inputs, datums) + has_outputs = isinstance(result, tuple) + if returns_outputs is None: + returns_outputs = has_outputs + elif returns_outputs != has_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + if isinstance(result, tuple): + losses, outputs = result + if ( + not isinstance(outputs, Sequence) + or isinstance(outputs, (str, bytes)) + or len(outputs) != len(datums) + or not all(isinstance(item, Mapping) for item in outputs) + ): + raise ValueError("loss_fn outputs must contain one mapping per Datum") + loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) + else: + losses = result + if not isinstance(losses, torch.Tensor): + raise TypeError("loss_fn must return a Tensor, optionally followed by per-Datum outputs") + if losses.shape != weights.shape: + raise ValueError( + "loss_fn losses must have exactly the same shape as weights; " + f"got losses={tuple(losses.shape)}, weights={tuple(weights.shape)}" + ) + if losses.device != weights.device: + raise ValueError("loss_fn losses and weights must be on the same device") + + numerator = (losses * weights.to(losses)).sum() + (numerator * (dp_size / denominator)).backward() + + local_loss_sum.add_(numerator.detach().to(torch.float64)) + if index == 0: + prepare_after_first_microbatch() + + if dp_size > 1: + dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=dp_group) + + loss = (local_loss_sum / denominator).detach() + return loss, loss_fn_outputs + + @staticmethod + def _validate_window(window: Sequence[Sequence[Datum]]) -> list[list[Datum]]: + if not isinstance(window, Sequence) or isinstance(window, (str, bytes)) or not window: + raise ValueError("forward_backward requires a non-empty accumulation window") + microbatches: list[list[Datum]] = [] + for index, microbatch in enumerate(window): + if not isinstance(microbatch, Sequence) or isinstance(microbatch, (str, bytes)) or not microbatch: + raise ValueError(f"microbatch {index} must be a non-empty sequence of Datum") + if not all(isinstance(datum, Datum) for datum in microbatch): + raise TypeError(f"microbatch {index} contains a value that is not a Datum") + microbatches.append(list(microbatch)) + return microbatches + + def _validate_parallelism(self) -> None: + if self.mesh_context is None: + return + if self.mesh_context.pp_size > 1: + raise NotImplementedError("Engine.forward_backward does not yet support pipeline parallelism") + if self.mesh_context.cp_size > 1: + raise NotImplementedError("Engine.forward_backward does not yet support context parallelism") + + def _dp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: + if self.mesh_context is not None and self.mesh_context.device_mesh is not None: + dp_mesh = get_flat_mesh(self.mesh_context.device_mesh, "dp") + size = int(dp_mesh.size()) + return (dp_mesh.get_group() if size > 1 else None), size + + group = self.mesh_context.process_group if self.mesh_context is not None else None + if dist.is_available() and dist.is_initialized(): + return group, dist.get_world_size(group=group) + return None, 1 + + def _global_weight_sum( + self, + microbatches: list[list[Datum]], + dp_group: dist.ProcessGroup | None, + dp_size: int, + ) -> torch.Tensor: + local_sum = 0.0 + for datum in (datum for microbatch in microbatches for datum in microbatch): + weights = datum.loss_fn_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("every Datum must contain a Tensor loss_fn_inputs['weights']") + if weights.numel() == 0 or not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()): + raise ValueError("Datum weights must be non-empty, finite, and non-negative") + local_sum += float(weights.to(torch.float64).sum()) + + denominator = torch.tensor(local_sum, dtype=torch.float64, device=self.device) + if dp_size > 1: + dist.all_reduce(denominator, op=dist.ReduceOp.SUM, group=dp_group) + if float(denominator) <= 0: + raise ValueError("forward_backward requires a positive global weight sum") + return denominator + + def _validate_window_size_across_dp( + self, + size: int, + dp_group: dist.ProcessGroup | None, + dp_size: int, + ) -> None: + if dp_size <= 1: + return + local_size = torch.tensor([size], dtype=torch.int64, device=self.device) + sizes = torch.empty(dp_size, dtype=torch.int64, device=self.device) + dist.all_gather_into_tensor(sizes, local_size, group=dp_group) + if not bool((sizes == sizes[0]).all()): + raise ValueError(f"every data-parallel rank must use the same number of microbatches; got {sizes.tolist()}") + + @staticmethod + def _validate_collated_weights( + datums: list[Datum], + loss_inputs: dict[str, torch.Tensor], + ) -> torch.Tensor: + weights = loss_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("collate_fn must return a Tensor loss input named 'weights'") + if not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()): + raise ValueError("collated weights must be finite and non-negative") + + expected_sum = sum(float(datum.loss_fn_inputs["weights"].to(torch.float64).sum()) for datum in datums) + actual_sum = weights.to(torch.float64).sum() + if not torch.isclose(actual_sum, actual_sum.new_tensor(expected_sum)): + raise ValueError("collate_fn changed the sum of Datum weights") + return weights + + +def _to_device(value: Any, device: torch.device) -> Any: + """Move tensors in common model-input containers without changing layout.""" + if isinstance(value, torch.Tensor): + return value.to(device) + if isinstance(value, dict): + return {key: _to_device(item, device) for key, item in value.items()} + if isinstance(value, list): + return [_to_device(item, device) for item in value] + if isinstance(value, tuple): + return tuple(_to_device(item, device) for item in value) + return value + + +def _detach(value: Any) -> Any: + """Detach tensor leaves without changing an output record's structure.""" + if isinstance(value, torch.Tensor): + return value.detach() + if isinstance(value, Mapping): + return {key: _detach(item) for key, item in value.items()} + if isinstance(value, list): + return [_detach(item) for item in value] + if isinstance(value, tuple): + return tuple(_detach(item) for item in value) + return value diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index 53ab2ccf63..3f8b7266e8 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -25,7 +25,7 @@ def _toy_datums(): return [ Datum( - input_ids=torch.tensor([10, 11, 12]), + model_inputs={"input_ids": torch.tensor([10, 11, 12])}, loss_fn_inputs={ "target_tokens": torch.tensor([11, 12, 13]), "weights": torch.tensor([1.0, 1.0, 0.0]), @@ -33,7 +33,7 @@ def _toy_datums(): }, ), Datum( - input_ids=torch.tensor([20, 21]), + model_inputs={"input_ids": torch.tensor([20, 21])}, loss_fn_inputs={ "target_tokens": torch.tensor([21, 22]), "weights": torch.tensor([1.0, 1.0]), @@ -46,17 +46,33 @@ def _toy_datums(): # ── Datum ───────────────────────────────────────────────────────────────── -def test_datum_coerces_and_validates(): +def test_datum_keeps_old_input_ids_convenience(): d = Datum(input_ids=[1, 2, 3], loss_fn_inputs={"weights": [1, 1, 0]}) assert isinstance(d.input_ids, torch.Tensor) assert d.input_ids.dtype == torch.long assert d.seq_len == 3 assert isinstance(d.loss_fn_inputs["weights"], torch.Tensor) + positional = Datum([4, 5], {"weights": [1, 0]}) + assert positional.input_ids.tolist() == [4, 5] + def test_datum_rejects_non_1d_input_ids(): with pytest.raises(ValueError, match="must be 1-D"): - Datum(input_ids=torch.zeros(2, 3, dtype=torch.long)) + Datum(model_inputs={"input_ids": torch.zeros(2, 3, dtype=torch.long)}) + + +def test_datum_accepts_model_specific_inputs(): + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([1, 2]), + "pixel_values": torch.randn(1, 3, 4, 4), + "image_grid_thw": torch.tensor([[1, 2, 2]]), + }, + loss_fn_inputs={"weights": torch.ones(2)}, + ) + assert datum.model_inputs["pixel_values"].shape == (1, 3, 4, 4) + assert datum.seq_len == 2 def test_to_features_applies_masking_convention(): @@ -82,22 +98,24 @@ def test_to_features_native_python_ints(): def test_collate_padded_uses_default_collater_schema(): - batch = collate_datums(_toy_datums()) + batch, loss_inputs = collate_datums(_toy_datums()) assert batch["input_ids"].shape == (2, 3) assert batch["input_ids"][1].tolist() == [20, 21, 0] # right-pad - assert batch["labels"][0].tolist() == [11, 12, CROSS_ENTROPY_IGNORE_IDX] - assert batch["labels"][1].tolist() == [21, 22, CROSS_ENTROPY_IGNORE_IDX] + assert "labels" not in batch + assert loss_inputs["target_tokens"][0].tolist() == [11, 12, 13] + assert loss_inputs["target_tokens"][1].tolist() == [21, 22, 0] # padding_mask is produced by default_collater, not by us. assert "padding_mask" in batch def test_collate_packed_concatenates_into_one_thd_row(): - batch = collate_datums(_toy_datums(), packed=True) + batch, loss_inputs = collate_datums(_toy_datums(), packed=True) # All datums share one flat [1, total_tokens] pack. assert batch["qkv_format"] == "thd" assert batch["input_ids"].shape == (1, 5) assert batch["input_ids"][0].tolist() == [10, 11, 12, 20, 21] - assert batch["labels"][0].tolist() == [11, 12, CROSS_ENTROPY_IGNORE_IDX, 21, 22] + assert "labels" not in batch + assert loss_inputs["target_tokens"][0].tolist() == [11, 12, 13, 21, 22] # position_ids restart at every sequence boundary (RoPE resets). assert batch["position_ids"][0].tolist() == [0, 1, 2, 0, 1] assert batch["seq_lens"][0].tolist() == [3, 2] @@ -105,52 +123,59 @@ def test_collate_packed_concatenates_into_one_thd_row(): def test_collate_packed_side_inputs_ride_the_flat_axis(): - batch = collate_datums(_toy_datums(), packed=True) + batch, loss_inputs = collate_datums(_toy_datums(), packed=True) # Per-token floats are concatenated in datum order, aligned with input_ids. - assert batch["advantages"].shape == (1, 5) - assert batch["advantages"][0].tolist() == pytest.approx([0.5, 0.5, 0.5, 0.9, 0.9]) - assert batch["weights"][0].tolist() == pytest.approx([1.0, 1.0, 0.0, 1.0, 1.0]) + assert set(batch).isdisjoint({"advantages", "weights", "target_tokens"}) + assert loss_inputs["advantages"].shape == (1, 5) + assert loss_inputs["advantages"][0].tolist() == pytest.approx([0.5, 0.5, 0.5, 0.9, 0.9]) + assert loss_inputs["weights"][0].tolist() == pytest.approx([1.0, 1.0, 0.0, 1.0, 1.0]) # seq_lens allows splitting flat outputs back per datum. lens = [n for n in batch["seq_lens"][0].tolist() if n > 0] - split = torch.split(batch["advantages"][0], lens) + split = torch.split(loss_inputs["advantages"][0], lens) assert [t.tolist() for t in split] == [pytest.approx([0.5, 0.5, 0.5]), pytest.approx([0.9, 0.9])] def test_collate_packed_per_sample_side_input_is_one_per_datum(): datums = [ - Datum(input_ids=torch.tensor([1, 2]), loss_fn_inputs={"advantages": torch.tensor([0.5])}), - Datum(input_ids=torch.tensor([3, 4, 5]), loss_fn_inputs={"advantages": torch.tensor([0.9])}), + Datum(model_inputs={"input_ids": torch.tensor([1, 2])}, loss_fn_inputs={"advantages": torch.tensor([0.5])}), + Datum( + model_inputs={"input_ids": torch.tensor([3, 4, 5])}, + loss_fn_inputs={"advantages": torch.tensor([0.9])}, + ), ] - batch = collate_datums(datums, packed=True) + batch, loss_inputs = collate_datums(datums, packed=True) assert batch["input_ids"].shape == (1, 5) - assert batch["advantages"].tolist() == pytest.approx([0.5, 0.9]) + assert loss_inputs["advantages"].tolist() == pytest.approx([0.5, 0.9]) def test_collate_carries_per_token_float_side_inputs(): # advantages is float per-token -> padded to collated width and stacked. - batch = collate_datums(_toy_datums()) - assert "advantages" in batch - assert batch["advantages"].dtype == torch.float - assert batch["advantages"].shape == (2, 3) - assert batch["advantages"][1].tolist() == pytest.approx([0.9, 0.9, 0.0]) # right-pad with 0 + batch, loss_inputs = collate_datums(_toy_datums()) + assert "advantages" not in batch + assert loss_inputs["advantages"].dtype == torch.float + assert loss_inputs["advantages"].shape == (2, 3) + assert loss_inputs["advantages"][1].tolist() == pytest.approx([0.9, 0.9, 0.0]) # right-pad with 0 def test_collate_per_sample_scalar_side_input(): datums = [ - Datum(input_ids=torch.tensor([1, 2]), loss_fn_inputs={"advantages": torch.tensor([0.5])}), - Datum(input_ids=torch.tensor([3, 4, 5]), loss_fn_inputs={"advantages": torch.tensor([0.9])}), + Datum(model_inputs={"input_ids": torch.tensor([1, 2])}, loss_fn_inputs={"advantages": torch.tensor([0.5])}), + Datum( + model_inputs={"input_ids": torch.tensor([3, 4, 5])}, + loss_fn_inputs={"advantages": torch.tensor([0.9])}, + ), ] - batch = collate_datums(datums) + _, loss_inputs = collate_datums(datums) # length-1 != seq_len -> treated as per-sample, one value per datum. - assert batch["advantages"].shape == (2,) - assert batch["advantages"].tolist() == pytest.approx([0.5, 0.9]) + assert loss_inputs["advantages"].shape == (2,) + assert loss_inputs["advantages"].tolist() == pytest.approx([0.5, 0.9]) def test_collate_pad_seq_len_divisible(): - batch = collate_datums(_toy_datums(), pad_seq_len_divisible=8) + batch, loss_inputs = collate_datums(_toy_datums(), pad_seq_len_divisible=8) assert batch["input_ids"].shape == (2, 8) # float side-inputs follow the collated width. - assert batch["advantages"].shape == (2, 8) + assert loss_inputs["advantages"].shape == (2, 8) def test_collate_empty_raises(): @@ -173,9 +198,53 @@ def test_collate_rejects_inconsistent_loss_input_keys(): collate_datums(datums) +def test_collate_preserves_additional_token_model_inputs(): + datums = [ + Datum( + model_inputs={"input_ids": torch.tensor([1, 2]), "token_type_ids": torch.tensor([0, 1])}, + loss_fn_inputs={"weights": torch.ones(2)}, + ), + Datum( + model_inputs={"input_ids": torch.tensor([3]), "token_type_ids": torch.tensor([1])}, + loss_fn_inputs={"weights": torch.ones(1)}, + ), + ] + model_inputs, loss_inputs = collate_datums(datums) + assert model_inputs["token_type_ids"].tolist() == [[0, 1], [1, 1]] + assert "weights" not in model_inputs + assert loss_inputs["weights"].tolist() == [[1.0, 1.0], [1.0, 0.0]] + + +def test_default_packed_collater_rejects_model_specific_inputs(): + datum = Datum( + model_inputs={"input_ids": torch.tensor([1]), "pixel_values": torch.randn(1, 3, 2, 2)}, + loss_fn_inputs={"weights": torch.ones(1)}, + ) + with pytest.raises(ValueError, match="model-specific collate_fn"): + collate_datums([datum], packed=True) + + +def test_default_packed_collater_rejects_explicitly_padded_datum(): + datum = Datum( + model_inputs={"input_ids": torch.tensor([1, 2, 0]), "attention_mask": torch.tensor([1, 1, 0])}, + loss_fn_inputs={"weights": torch.tensor([1.0, 1.0, 0.0])}, + ) + with pytest.raises(ValueError, match="only real tokens"): + collate_datums([datum], packed=True) + + +def test_default_collater_rejects_float_token_model_inputs_instead_of_casting_them(): + datum = Datum( + model_inputs={"input_ids": torch.tensor([1, 2]), "token_scores": torch.tensor([0.2, 1.7])}, + loss_fn_inputs={"weights": torch.ones(2)}, + ) + with pytest.raises(ValueError, match="cannot preserve floating-point"): + collate_datums([datum]) + + def test_collate_reads_a_length_one_entry_on_a_single_token_sequence_as_per_token(): datums = [Datum(input_ids=torch.tensor([7]), loss_fn_inputs={"advantages": torch.tensor([0.5])})] - assert collate_datums(datums)["advantages"].shape == (1, 1) + assert collate_datums(datums)[1]["advantages"].shape == (1, 1) def test_collate_padding_mask_does_not_misread_a_real_pad_valued_token(): @@ -183,6 +252,6 @@ def test_collate_padding_mask_does_not_misread_a_real_pad_valued_token(): # from matching the pad id -- id 0 is a real token here (and pad_token_id == # eos_token_id is a common config). datums = [Datum(input_ids=torch.tensor([5, 0, 7])), Datum(input_ids=torch.tensor([9, 9]))] - batch = collate_datums(datums) + batch, _ = collate_datums(datums) assert batch["padding_mask"].tolist() == [[False, False, False], [False, False, True]] assert batch["attention_mask"].tolist() == [[1, 1, 1], [1, 1, 0]] diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py new file mode 100644 index 0000000000..d1a216e2ba --- /dev/null +++ b/tests/unit_tests/test_engine.py @@ -0,0 +1,348 @@ +# 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. + +from __future__ import annotations + +from contextlib import contextmanager +from functools import partial +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F +from torch import nn + +import nemo_automodel.engine as engine_module +from nemo_automodel import Datum as PublicDatum +from nemo_automodel import Engine as PublicEngine +from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler +from nemo_automodel.engine import Engine + + +class ScaleModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.tensor(1.0)) + self.forward_calls = 0 + + def forward(self, input_ids: torch.Tensor, **_) -> torch.Tensor: + self.forward_calls += 1 + return input_ids.to(torch.float32) * self.weight + + +def _datum(values, weights=None) -> Datum: + values = torch.tensor(values, dtype=torch.long) + weights = torch.ones_like(values, dtype=torch.float32) if weights is None else torch.tensor(weights) + return Datum(model_inputs={"input_ids": values}, loss_fn_inputs={"weights": weights}) + + +def _identity_loss(output, _loss_inputs, _datums): + return output + + +def test_engine_and_datum_are_lazy_top_level_exports(): + assert PublicEngine is Engine + assert PublicDatum is Datum + + +def test_forward_backward_uses_one_denominator_for_the_window(): + model = ScaleModel() + initial_weight = model.weight.detach().clone() + engine = Engine(model, device="cpu") + + loss, outputs = engine.forward_backward( + [[_datum([1, 2])], [_datum([3])]], + _identity_loss, + ) + + assert loss.item() == pytest.approx(2.0) + assert outputs == [] + assert model.weight.grad.item() == pytest.approx(2.0) + assert torch.equal(model.weight, initial_weight) + assert model.forward_calls == 2 + + +def test_padded_and_packed_windows_have_the_same_loss_and_gradient(): + padded_model = ScaleModel() + packed_model = ScaleModel() + window = [[_datum([1, 2]), _datum([3])]] + + padded_loss, _ = Engine(padded_model, device="cpu").forward_backward(window, _identity_loss) + + def packed_identity_loss(output, loss_inputs, datums): + assert [datum.seq_len for datum in datums] == [2, 1] + return _identity_loss(output, loss_inputs, datums) + + packed_loss, _ = Engine( + packed_model, + device="cpu", + collate_fn=partial(collate_datums, packed=True), + ).forward_backward(window, packed_identity_loss) + + torch.testing.assert_close(padded_loss, packed_loss) + torch.testing.assert_close(padded_model.weight.grad, packed_model.weight.grad) + + +def test_weights_mask_loss_and_denominator(): + model = ScaleModel() + loss, _ = Engine(model, device="cpu").forward_backward( + [[_datum([1, 100], [1.0, 0.0])], [_datum([3, 5], [0.5, 1.0])]], + _identity_loss, + ) + + assert loss.item() == pytest.approx(3.0) + assert model.weight.grad.item() == pytest.approx(3.0) + + +def test_loss_fn_outputs_follow_datum_order_and_are_detached(): + model = ScaleModel() + + def loss_with_outputs(output, _loss_inputs, datums): + return output, [ + {"first_token": datum.input_ids[0], "model_value": output[index].sum()} + for index, datum in enumerate(datums) + ] + + _, outputs = Engine(model, device="cpu").forward_backward( + [[_datum([1, 2]), _datum([3])], [_datum([4])]], + loss_with_outputs, + ) + + assert [item["first_token"].item() for item in outputs] == [1, 3, 4] + assert all(not item["model_value"].requires_grad for item in outputs) + + +def test_loss_fn_outputs_must_align_with_datums(): + model = ScaleModel() + with pytest.raises(ValueError, match="one mapping per Datum"): + Engine(model, device="cpu").forward_backward( + [[_datum([1]), _datum([2])]], + lambda output, _inputs, _datums: (output, [{"only": "one"}]), + ) + assert model.weight.grad is None + + +def test_loss_fn_outputs_must_be_consistent_across_the_window(): + def inconsistent_outputs(output, _inputs, datums): + if datums[0].input_ids[0].item() == 1: + return output, [{"value": output.sum()}] + return output + + with pytest.raises(ValueError, match="every microbatch or none"): + Engine(ScaleModel(), device="cpu").forward_backward( + [[_datum([1])], [_datum([2])]], + inconsistent_outputs, + ) + + +class TinyLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = nn.Embedding(8, 4) + self.output = nn.Linear(4, 8) + + def forward(self, input_ids, **_): + return self.output(self.embedding(input_ids)) + + +def test_raw_output_and_loss_inputs_support_an_rl_loss_callback(): + datum = Datum( + model_inputs={"input_ids": torch.tensor([1, 2, 3])}, + loss_fn_inputs={ + "target_tokens": torch.tensor([2, 3, 4]), + "weights": torch.tensor([1.0, 1.0, 0.0]), + "logprobs": torch.tensor([-1.0, -1.0, 0.0]), + "advantages": torch.tensor([0.5, -0.25, 0.0]), + }, + ) + model = TinyLM() + + def policy_loss(logits, inputs, datums): + assert len(datums) == 1 + new_logprobs = -F.cross_entropy( + logits.flatten(0, 1), + inputs["target_tokens"].flatten(), + reduction="none", + ).view_as(inputs["weights"]) + ratio = torch.exp(new_logprobs - inputs["logprobs"]) + losses = -(ratio * inputs["advantages"]) + return losses, [{"policy_sum": (losses * inputs["weights"]).sum()}] + + loss, outputs = Engine(model, device="cpu").forward_backward([[datum]], policy_loss) + + assert torch.isfinite(loss) + assert torch.isfinite(outputs[0]["policy_sum"]) + assert not outputs[0]["policy_sum"].requires_grad + assert model.embedding.weight.grad is not None + assert model.output.weight.grad is not None + + +def test_lifecycle_marks_only_the_last_microbatch_for_sync(monkeypatch): + events = [] + + monkeypatch.setattr( + engine_module, "prepare_for_grad_accumulation", lambda *_args, **_kwargs: events.append("prepare") + ) + monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after_first")) + monkeypatch.setattr(engine_module, "prepare_for_final_backward", lambda *_args, **_kwargs: events.append("final")) + + @contextmanager + def sync_context(_model, is_last, _defer): + events.append(f"sync:{is_last}") + yield + + monkeypatch.setattr(engine_module, "get_sync_ctx", sync_context) + + Engine(ScaleModel(), device="cpu").forward_backward( + [[_datum([1])], [_datum([2])]], + _identity_loss, + ) + + assert events == ["prepare", "sync:False", "after_first", "final", "sync:True"] + + +def test_window_sets_the_same_moe_aux_scale_as_the_recipes(monkeypatch): + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) + + Engine(ScaleModel(), device="cpu").forward_backward( + [[_datum([1])], [_datum([2])]], + _identity_loss, + ) + + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.5) + + +class TinyVLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.text = nn.Embedding(8, 1) + self.vision = nn.Linear(1, 1, bias=False) + + def forward(self, input_ids, pixel_values): + pixels = torch.stack(pixel_values) + return self.text(input_ids).mean(dim=1).squeeze(-1) + self.vision(pixels).squeeze(-1) + + +def _vlm_collate(datums): + return ( + { + "input_ids": torch.stack([datum.model_inputs["input_ids"] for datum in datums]), + "pixel_values": [datum.model_inputs["pixel_values"] for datum in datums], + }, + {"weights": torch.stack([datum.loss_fn_inputs["weights"] for datum in datums])}, + ) + + +def test_model_specific_collater_keeps_multimodal_inputs_and_gradients(): + datums = [ + Datum( + model_inputs={"input_ids": torch.tensor([1, 2]), "pixel_values": torch.tensor([0.5])}, + loss_fn_inputs={"weights": torch.tensor(1.0)}, + ), + Datum( + model_inputs={"input_ids": torch.tensor([3, 4]), "pixel_values": torch.tensor([1.5])}, + loss_fn_inputs={"weights": torch.tensor(1.0)}, + ), + ] + model = TinyVLM() + Engine(model, device="cpu", collate_fn=_vlm_collate).forward_backward( + [datums], + lambda output, _inputs, _datums: output, + ) + + assert model.text.weight.grad is not None + assert model.text.weight.grad.abs().sum() > 0 + assert model.vision.weight.grad is not None + assert model.vision.weight.grad.abs().sum() > 0 + + +def test_zero_weights_fail_before_forward(): + model = ScaleModel() + with pytest.raises(ValueError, match="positive global weight sum"): + Engine(model, device="cpu").forward_backward( + [[_datum([1, 2], [0.0, 0.0])]], + _identity_loss, + ) + assert model.forward_calls == 0 + assert model.weight.grad is None + + +@pytest.mark.parametrize(("pp_size", "cp_size", "name"), [(2, 1, "pipeline"), (1, 2, "context")]) +def test_unsupported_parallelism_fails_before_forward(pp_size, cp_size, name): + model = ScaleModel() + mesh_context = SimpleNamespace(pp_size=pp_size, cp_size=cp_size) + + with pytest.raises(NotImplementedError, match=name): + Engine(model, device="cpu", mesh_context=mesh_context).forward_backward( + [[_datum([1])]], + _identity_loss, + ) + + assert model.forward_calls == 0 + + +def test_loss_shape_must_exactly_match_weights(): + model = ScaleModel() + with pytest.raises(ValueError, match="exactly the same shape"): + Engine(model, device="cpu").forward_backward( + [[_datum([1, 2])]], + lambda output, _inputs, _datums: output.sum(), + ) + assert model.weight.grad is None + + +def test_collater_cannot_change_weight_sum(): + model = ScaleModel() + + def bad_collate(datums): + model_inputs, loss_inputs = collate_datums(datums) + loss_inputs["weights"].zero_() + return model_inputs, loss_inputs + + with pytest.raises(ValueError, match="collate_fn changed"): + Engine(model, device="cpu", collate_fn=bad_collate).forward_backward( + [[_datum([1, 2])]], + _identity_loss, + ) + assert model.forward_calls == 0 + + +def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + model = nn.parallel.DistributedDataParallel(ScaleModel()) + bad_window = [[_datum([1])]] if rank == 0 else [[_datum([1])], [_datum([2])]] + with pytest.raises(ValueError, match="same number of microbatches"): + Engine(model, device="cpu").forward_backward(bad_window, _identity_loss) + assert model.module.forward_calls == 0 + + window = [[_datum([1, 2])], [_datum([3])]] if rank == 0 else [[_datum([4])], [_datum([5, 6])]] + loss, outputs = Engine(model, device="cpu").forward_backward(window, _identity_loss) + assert loss.item() == pytest.approx(3.5) + assert outputs == [] + assert model.module.weight.grad.item() == pytest.approx(3.5) + finally: + dist.destroy_process_group() + + +def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): + mp.spawn( + _distributed_worker, + args=(2, str(tmp_path / "engine_dist_init")), + nprocs=2, + join=True, + ) From 8790f38fdb3c012160c9a161459d769b3dfa171b Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 15 Aug 2026 18:41:26 -0700 Subject: [PATCH 02/34] feat(engine): route eager finetuning through Datum windows Signed-off-by: HuiyingLi --- nemo_automodel/components/datasets/datum.py | 30 +-- nemo_automodel/components/loss/causal_lm.py | 107 +++++++++ .../{engine.py => engine/__init__.py} | 85 +++++-- nemo_automodel/recipes/llm/train_ft.py | 181 ++++++++------ nemo_automodel/recipes/vlm/finetune.py | 225 +++++++++++------- tests/unit_tests/datasets/test_datum.py | 8 +- .../nemotron_v3/test_nemotron_v3_mtp.py | 2 +- .../recipes/test_finetune_vlm_helpers.py | 29 ++- tests/unit_tests/recipes/test_train_ft.py | 34 ++- .../recipes/test_vlm_drafter_helpers.py | 22 ++ tests/unit_tests/test_engine.py | 194 +++++++++++++-- .../test_engine_recipe_integration.py | 140 +++++++++++ 12 files changed, 842 insertions(+), 215 deletions(-) create mode 100644 nemo_automodel/components/loss/causal_lm.py rename nemo_automodel/{engine.py => engine/__init__.py} (77%) create mode 100644 tests/unit_tests/test_engine_recipe_integration.py diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index e9fdae368c..96bcf18075 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -35,17 +35,17 @@ @dataclass(init=False) class Datum: - """One processor-ready training example. + """One processor-ready training item. ``model_inputs`` contains exactly the keyword arguments needed by the - model for one example. Text examples usually contain ``input_ids``; - multimodal examples may additionally contain fields such as - ``pixel_values`` and ``image_grid_thw``. ``loss_fn_inputs`` is kept - separate so algorithm data such as targets, weights, old log-probabilities, - and advantages is never passed to the model. + model. The default collater treats each Datum as one text sequence. A + custom collater may also consume processor-ready multimodal fields or an + already-collated batch. ``loss_fn_inputs`` is kept separate so algorithm + data such as targets, weights, old log-probabilities, and advantages is + never passed to the model. Args: - model_inputs: Model-ready values for one example. Values are + model_inputs: Model-ready values for one training item. Values are model-specific because LLM and VLM processors emit different fields. loss_fn_inputs: Tensor values consumed by the loss function. @@ -90,11 +90,8 @@ def __post_init__(self) -> None: if not isinstance(input_ids, torch.Tensor): input_ids = torch.as_tensor(input_ids, dtype=torch.long) self.model_inputs["input_ids"] = input_ids - if input_ids.ndim != 1: - raise ValueError( - "Datum.model_inputs['input_ids'] must be 1-D [T] for one token sequence; " - f"got shape {tuple(input_ids.shape)}" - ) + if input_ids.ndim == 0: + raise ValueError("Datum.model_inputs['input_ids'] must have at least one dimension") for key, value in self.loss_fn_inputs.items(): if not isinstance(value, torch.Tensor): @@ -112,7 +109,10 @@ def input_ids(self) -> torch.Tensor: def seq_len(self) -> int: """Token length used by the default text collater.""" if isinstance(self.model_inputs.get("input_ids"), torch.Tensor): - return int(self.model_inputs["input_ids"].shape[0]) + input_ids = self.model_inputs["input_ids"] + if input_ids.ndim == 1: + return int(input_ids.shape[0]) + raise ValueError("the default collater requires 1-D input_ids; use a custom collate_fn") for key in ("target_tokens", "weights"): value = self.loss_fn_inputs.get(key) if isinstance(value, torch.Tensor) and value.ndim == 1: @@ -129,6 +129,8 @@ def to_features(self, *, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX) -> dict[s """ if "input_ids" not in self.model_inputs: raise ValueError("the default collater requires model_inputs['input_ids']") + if self.input_ids.ndim != 1: + raise ValueError("the default collater requires 1-D input_ids; use a custom collate_fn") features: dict[str, Any] = {} for key, value in self.model_inputs.items(): @@ -166,7 +168,7 @@ def collate_datums( ``Engine`` instead. Args: - datums: Non-empty examples for one microbatch. + datums: Non-empty training items for one microbatch. packed: Produce one flat THD token row instead of a padded batch. pad_seq_len_divisible: Round the padded token width to this multiple. ignore_index: Label fill value used internally by the THD collater. diff --git a/nemo_automodel/components/loss/causal_lm.py b/nemo_automodel/components/loss/causal_lm.py new file mode 100644 index 0000000000..2acd5bc141 --- /dev/null +++ b/nemo_automodel/components/loss/causal_lm.py @@ -0,0 +1,107 @@ +# 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. + +"""Shared causal-LM loss calculation for the LLM and VLM recipes.""" + +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn + +from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy +from nemo_automodel.components.loss.mtp import MTPLossConfig, calculate_mtp_loss +from nemo_automodel.components.loss.utils import _get_final_hidden_states, _get_lm_head_weight, calculate_loss + + +def causal_lm_loss( + loss_fn: nn.Module, + model: nn.Module, + output: Any, + labels: torch.Tensor, + mtp_config: MTPLossConfig | None, + *, + num_label_tokens: int | None, + grad_reduce_group: dist.ProcessGroup | None, + cu_seqlens: torch.Tensor | None = None, +) -> torch.Tensor: + """Return the main causal-LM loss plus an optional MTP loss. + + Args: + loss_fn: Configured causal-LM loss module. + model: Model that owns the LM head used by fused loss implementations. + output: Model output containing logits of shape ``[batch, sequence, + vocab]`` or final hidden states of shape ``[batch, sequence, + hidden]``. Optional MTP fields use the same batch and sequence + axes per prediction depth. + labels: Target token ids of shape ``[batch, sequence]`` or ``[tokens]`` + for a flattened THD stream. + mtp_config: MTP loss settings, required only when ``output`` contains + MTP predictions. + num_label_tokens: Global supervised-token denominator. ``None`` keeps + each loss as an unnormalized local sum for Engine normalization. + grad_reduce_group: Group that contributes independent fused-loss + shards, or ``None`` for an unsharded LM head. + cu_seqlens: Optional THD cumulative sequence offsets of shape + ``[num_sequences + 1]``. + + Returns: + Scalar causal-LM loss retaining its autograd graph. + """ + hidden_states = _get_final_hidden_states(output) + if isinstance(loss_fn, FusedLinearCrossEntropy) and hidden_states is None: + raise ValueError("FusedLinearCrossEntropy requires the model to output hidden states") + + lm_weight = ( + loss_fn.materialize_lm_weight( + _get_lm_head_weight(model), + grad_reduce_group=grad_reduce_group, + ) + if isinstance(loss_fn, FusedLinearCrossEntropy) + else None + ) + loss = calculate_loss( + loss_fn, + logits=getattr(output, "logits", output), + labels=labels, + model=model, + hidden_states=hidden_states, + lm_weight=lm_weight, + grad_reduce_group=grad_reduce_group, + num_label_tokens=num_label_tokens, + ) + + mtp_hidden = getattr(output, "mtp_per_depth_h", None) + mtp_logits = getattr(output, "mtp_per_depth_logits", None) + if mtp_hidden is None and mtp_logits is None: + return loss + if mtp_config is None: + raise ValueError("MTP model output requires an MTP loss config") + + scaling_factor = ( + mtp_config.scaling_factor if mtp_config.scaling_factor is not None else output.mtp_loss_scaling_factor + ) + return loss + calculate_mtp_loss( + loss_fn, + mtp_per_depth_h=mtp_hidden, + mtp_per_depth_logits=mtp_logits, + labels=labels, + model=model, + scaling_factor=scaling_factor, + num_label_tokens=num_label_tokens, + ignore_index=mtp_config.ignore_index, + cu_seqlens=cu_seqlens, + lm_weight=lm_weight, + grad_reduce_group=grad_reduce_group, + ) diff --git a/nemo_automodel/engine.py b/nemo_automodel/engine/__init__.py similarity index 77% rename from nemo_automodel/engine.py rename to nemo_automodel/engine/__init__.py index 9710e11288..2ddd50d39e 100644 --- a/nemo_automodel/engine.py +++ b/nemo_automodel/engine/__init__.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext from typing import Any import torch @@ -40,7 +41,28 @@ torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], ] -__all__ = ["Engine"] +__all__ = ["Engine", "collate_prebatched"] + + +def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + """Return one already-collated Datum without changing its layout. + + The Datum represents the whole prebatched item. Consequently, one + optional ``loss_fn_output`` mapping also describes that whole batch, not + each sample inside it. + + Args: + datums: A one-item list whose model and loss tensor fields already have + the batch layout expected by the model and loss callback. + + Returns: + Separate shallow copies of the Datum's model-input and loss-input + mappings. Tensor shapes, dtypes, devices, and storage are unchanged. + """ + if len(datums) != 1: + raise ValueError("collate_prebatched expects exactly one Datum per microbatch") + datum = datums[0] + return dict(datum.model_inputs), dict(datum.loss_fn_inputs) class Engine: @@ -59,17 +81,19 @@ class Engine: mesh_context: Runtime topology. When omitted, an initialized default process group is treated as pure data parallelism. collate_fn: Batches one microbatch of Datums into separate model and - loss inputs. The default supports padded or packed text; - VLMs pass a thin callable around their model-specific collater. - The callable must keep model inputs and loss inputs aligned and - preserve the sum of ``weights``. + loss inputs. The default supports padded text. Packed text and + VLMs pass a model-specific collater that returns final model-ready + inputs. Existing recipes whose dataloaders already collate can use + :func:`collate_prebatched`. The callable must keep model inputs and + loss inputs aligned and preserve the sum of ``weights``. + context_fn: Creates an optional context around model forward, loss, + and backward. Recipes use this for runtime contexts such as FP8. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the final microbatch. Note: This first execution backend is eager and weight-normalized. Pipeline - and context parallel schedules and pre-reduced scalar losses are - intentionally deferred. + and context parallel schedules are intentionally deferred. """ def __init__( @@ -79,12 +103,14 @@ def __init__( device: torch.device | str, mesh_context: MeshContext | None = None, collate_fn: CollateFn = collate_datums, + context_fn: Callable[[], AbstractContextManager[Any]] = nullcontext, defer_fsdp_grad_sync: bool = True, ) -> None: self.model = model self.device = torch.device(device) self.mesh_context = mesh_context self.collate_fn = collate_fn + self.context_fn = context_fn self.defer_fsdp_grad_sync = defer_fsdp_grad_sync def forward_backward( @@ -97,10 +123,22 @@ def forward_backward( ``window`` is explicit: each inner sequence is one eager microbatch. ``loss_fn`` receives the raw model output, collated ``loss_fn_inputs``, and the original Datums for that microbatch. It - returns unreduced losses with exactly the same shape as - ``loss_fn_inputs["weights"]``. It may also return one output mapping - per Datum. Those mappings are detached and preserved in input order; - the Engine deliberately does not interpret or reduce them. + returns either per-element losses with exactly the same shape as + ``loss_fn_inputs["weights"]``, or a scalar local weighted-sum + numerator. For a scalar, the callback must apply weights and masks; + the Engine will only apply global normalization. The callback may + also return one output mapping per Datum. + Those mappings are detached and preserved in input order; the Engine + deliberately does not interpret or reduce them. + + Args: + window: The complete optimizer accumulation window. Each inner + sequence is one microbatch of Datums. A Datum's token weights + may have shape ``[tokens]`` or the custom collater's batched + token layout; the loss tensor must use the identical shape. + loss_fn: Computes either that per-token loss tensor or a scalar + local weighted-sum numerator from the raw model output and + collated loss inputs. Returns: ``(loss, loss_fn_outputs)``. ``loss`` is a detached, DP-reduced @@ -129,11 +167,18 @@ def forward_backward( prepare_for_final_backward([self.model], pp_enabled=False) model_inputs, loss_inputs = self.collate_fn(datums) + if model_inputs.get("qkv_format") == "thd" and ( + "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs + ): + raise ValueError( + "packed collate_fn must return final model-ready THD inputs, not seq_lens packing metadata" + ) + self._validate_collated_weights(datums, loss_inputs) model_inputs = _to_device(model_inputs, self.device) loss_inputs = _to_device(loss_inputs, self.device) - weights = self._validate_collated_weights(datums, loss_inputs) + weights = loss_inputs["weights"] - with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync): + with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(): output = self.model(**model_inputs) result = loss_fn(output, loss_inputs, datums) has_outputs = isinstance(result, tuple) @@ -155,15 +200,18 @@ def forward_backward( losses = result if not isinstance(losses, torch.Tensor): raise TypeError("loss_fn must return a Tensor, optionally followed by per-Datum outputs") - if losses.shape != weights.shape: + if losses.ndim == 0: + numerator = losses + elif losses.shape == weights.shape: + numerator = (losses * weights.to(losses)).sum() + else: raise ValueError( - "loss_fn losses must have exactly the same shape as weights; " + "loss_fn must return a scalar local weighted sum or losses with exactly the same shape as weights; " f"got losses={tuple(losses.shape)}, weights={tuple(weights.shape)}" ) if losses.device != weights.device: raise ValueError("loss_fn losses and weights must be on the same device") - numerator = (losses * weights.to(losses)).sum() (numerator * (dp_size / denominator)).backward() local_loss_sum.add_(numerator.detach().to(torch.float64)) @@ -248,7 +296,7 @@ def _validate_window_size_across_dp( def _validate_collated_weights( datums: list[Datum], loss_inputs: dict[str, torch.Tensor], - ) -> torch.Tensor: + ) -> None: weights = loss_inputs.get("weights") if not isinstance(weights, torch.Tensor): raise ValueError("collate_fn must return a Tensor loss input named 'weights'") @@ -259,13 +307,12 @@ def _validate_collated_weights( actual_sum = weights.to(torch.float64).sum() if not torch.isclose(actual_sum, actual_sum.new_tensor(expected_sum)): raise ValueError("collate_fn changed the sum of Datum weights") - return weights def _to_device(value: Any, device: torch.device) -> Any: """Move tensors in common model-input containers without changing layout.""" if isinstance(value, torch.Tensor): - return value.to(device) + return value.to(device, non_blocking=True) if isinstance(value, dict): return {key: _to_device(item, device) for key, item in value.items()} if isinstance(value, list): diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 8674879c93..7face8ee7e 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -56,6 +56,7 @@ from nemo_automodel._transformers.utils import apply_cache_compatibility_patches from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.components.cuda_graphs import PartialCudaGraphManager +from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.loader import DataloaderConfig from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder @@ -71,12 +72,10 @@ to_float_metrics, ) from nemo_automodel.components.loggers.wandb_utils import suppress_wandb_log_messages +from nemo_automodel.components.loss.causal_lm import causal_lm_loss from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy -from nemo_automodel.components.loss.mtp import calculate_mtp_loss -from nemo_automodel.components.loss.utils import _get_lm_head_weight, calculate_loss from nemo_automodel.components.quantization.fp8 import build_fp8_config -from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG from nemo_automodel.components.training.utils import ( count_tail_padding, @@ -96,6 +95,7 @@ filter_forward_kwargs, resolve_trust_remote_code, ) +from nemo_automodel.engine import Engine, collate_prebatched from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config, shard_optimizers_for_megatron_fsdp from nemo_automodel.recipes._typed_config import RecipeConfig from nemo_automodel.recipes.base_recipe import BaseRecipe @@ -658,6 +658,26 @@ def setup(self): # Extract TE FP8 config from model backend (set after model construction) self.te_fp8 = self.model_parts[0].backend.te_fp8 if hasattr(self.model_parts[0], "backend") else None + # Packed and CP/Magi batches need recipe-owned input preparation, so + # they keep using the existing forward/backward path for now. + self.engine = None + if ( + not self.pp_enabled + and self.mesh_context.cp_size == 1 + and getattr(self.cfg.dataloader, "packing", None) is None + and not getattr(self.cfg.dataloader, "emits_thd", False) + and not self.magi.enabled + and getattr(self.loss_fn, "reduction", None) == "sum" + ): + self.engine = Engine( + self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + collate_fn=collate_prebatched, + context_fn=self.te_fp8.maybe_te_autocast if self.te_fp8 is not None else nullcontext, + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), + ) + if self.pp_enabled: self._configure_pipeline_loss_fn() @@ -1006,17 +1026,8 @@ def run_train_validation_loop(self): self._partial_cuda_graph_capture_pending = False # ------------------ helpers ------------------ - def _forward_backward_step( - self, - idx, - batch, - *, - loss_buffer, - num_label_tokens, - num_batches, - is_train: bool = True, - ): - # Move batch to device (handle both tensors and dicts of tensors like causal_mask_mapping) + def _prepare_eager_batch(self, batch): + """Move and CP-prepare one batch before an eager model call.""" batch = { k: ( {dk: dv.to(self.dist_env.device, non_blocking=True) for dk, dv in v.items() if dv is not None} @@ -1033,7 +1044,53 @@ def _forward_backward_step( num_chunks=self.pp.pp_batch_size // self.pp.pp_microbatch_size if self.pp_enabled else 1, ) train_ctx, batch = cp_sharder.shard(batch) - labels = batch.pop("labels") + return train_ctx, batch, batch.pop("labels") + + def _calculate_eager_loss(self, output, labels, model_inputs, *, num_label_tokens, is_train): + """Compute the shared CE and optional MTP loss for eager execution.""" + return causal_lm_loss( + self.loss_fn, + self.model_parts[0], + output, + labels, + getattr(self.cfg, "mtp", None), + num_label_tokens=num_label_tokens, + grad_reduce_group=self._get_dp_group(include_cp=True) if is_train else None, + cu_seqlens=model_inputs.get("cu_seqlens"), + ) + + def _make_engine_datum(self, batch): + labels = batch["labels"] + model_inputs = filter_forward_kwargs( + self.model_parts[0], {key: value for key, value in batch.items() if key != "labels"} + ) + if isinstance(self.loss_fn, FusedLinearCrossEntropy): + model_inputs["logits_to_keep"] = 1 + return Datum( + model_inputs=model_inputs, + loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, + ) + + def _engine_loss(self, output, loss_inputs, datums): + return self._calculate_eager_loss( + output, + loss_inputs["labels"], + datums[0].model_inputs, + num_label_tokens=None, + is_train=True, + ) + + def _forward_backward_step( + self, + idx, + batch, + *, + loss_buffer, + num_label_tokens, + num_batches, + is_train: bool = True, + ): + train_ctx, batch, labels = self._prepare_eager_batch(batch) fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext() if self.pp_enabled: @@ -1100,56 +1157,15 @@ def _forward_backward_step( if isinstance(self.loss_fn, FusedLinearCrossEntropy): # use num_logits_to_keep to avoid full logits matrix in memory out = model(logits_to_keep=1, **batch) - if "hidden_states" not in out: - raise ValueError( - "FusedLinearCrossEntropy requires the model to output hidden states. Set `model.output_hidden_states=True` in the config." - ) else: out = model(**batch) - - # Gather the LM head once and share it across the main loss and - # all MTP depths (FusedLinearCrossEntropy path) to avoid redundant - # full_tensor() gathers that accumulate on-device and OOM. - loss_distributed_kwargs = {} - shared_lm_weight = None - if isinstance(self.loss_fn, FusedLinearCrossEntropy): - grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None - shared_lm_weight = self.loss_fn.materialize_lm_weight( - _get_lm_head_weight(model), - grad_reduce_group=grad_reduce_group, - ) - loss_distributed_kwargs["grad_reduce_group"] = grad_reduce_group - local_loss = calculate_loss( - self.loss_fn, - logits=getattr(out, "logits", out), - labels=labels, - model=model, - hidden_states=get_final_hidden_states(out), - lm_weight=shared_lm_weight, + local_loss = self._calculate_eager_loss( + out, + labels, + batch, num_label_tokens=num_label_tokens, - **loss_distributed_kwargs, + is_train=is_train, ) - mtp_per_depth_h = getattr(out, "mtp_per_depth_h", None) - mtp_per_depth_logits = getattr(out, "mtp_per_depth_logits", None) - if mtp_per_depth_h is not None or mtp_per_depth_logits is not None: - mtp_cfg = self.cfg.mtp - scaling_factor = ( - mtp_cfg.scaling_factor if mtp_cfg.scaling_factor is not None else out.mtp_loss_scaling_factor - ) - local_loss = local_loss + calculate_mtp_loss( - self.loss_fn, - mtp_per_depth_h=mtp_per_depth_h, - mtp_per_depth_logits=mtp_per_depth_logits, - labels=labels, - model=model, - scaling_factor=scaling_factor, - num_label_tokens=num_label_tokens, - ignore_index=mtp_cfg.ignore_index, - # mask cross-boundary MTP label rolls in THD packing (matches the PP path) - cu_seqlens=batch.get("cu_seqlens"), - lm_weight=shared_lm_weight, - **loss_distributed_kwargs, - ) loss_buffer.append(local_loss.clone().detach()) if is_train: (local_loss * self._get_dp_group_size(include_cp=True)).backward() @@ -1175,7 +1191,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): num_label_tokens = self._dp_allreduce(num_label_tokens).item() num_batches = len(batches) - self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) loss_buffer = [] @@ -1186,18 +1201,29 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) + engine = getattr(self, "engine", None) + use_engine = engine is not None and num_label_tokens > 0 + if use_engine: + reporting_loss, _ = engine.forward_backward( + [[self._make_engine_datum(batch)] for batch in batches], + self._engine_loss, + ) + else: + # Engine requires a positive global weight sum. Keep the existing + # zero-label behavior on the legacy path. + self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) + prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) - for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) + for i, batch in enumerate(batches): + if i == num_batches - 1: + prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) - self._forward_backward_step( - i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches - ) + self._forward_backward_step( + i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches + ) - if i == 0: - prepare_after_first_microbatch() + if i == 0: + prepare_after_first_microbatch() grad_norm = scale_grads_and_clip_grad_norm( max_grad_norm, @@ -1273,12 +1299,13 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ).item() mfu = calculate_mfu(step_flops / 1e12, self.dist_env.world_size, time_delta) - reporting_loss = torch.sum(torch.stack(loss_buffer)) - reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) - if self.pp_enabled: - reporting_loss = reporting_loss / num_label_tokens - reporting_loss = reporting_loss.to(self.dist_env.device) - reporting_loss = self._broadcast_from_last_pp_stage(reporting_loss) + if not use_engine: + reporting_loss = torch.sum(torch.stack(loss_buffer)) + reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) + if self.pp_enabled: + reporting_loss = reporting_loss / num_label_tokens + reporting_loss = reporting_loss.to(self.dist_env.device) + reporting_loss = self._broadcast_from_last_pp_stage(reporting_loss) reporting_loss = reporting_loss.cpu().item() # fix reporting_loss, tps across ranks diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index c967066c8b..f21719a2fd 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -28,6 +28,7 @@ import logging import pathlib import time +from collections.abc import Sequence from contextlib import contextmanager, nullcontext from typing import TYPE_CHECKING, Any, Optional, Protocol @@ -46,6 +47,7 @@ ) from nemo_automodel._transformers.utils import apply_cache_compatibility_patches, resolve_get_rope_index from nemo_automodel.components.config._arg_parser import parse_args_and_load_config +from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.vlm.pp_media import stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder @@ -65,10 +67,10 @@ to_float_metrics, ) from nemo_automodel.components.loggers.wandb_utils import suppress_wandb_log_messages +from nemo_automodel.components.loss.causal_lm import causal_lm_loss from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy -from nemo_automodel.components.loss.mtp import calculate_mtp_loss -from nemo_automodel.components.loss.utils import _get_lm_head_weight, calculate_loss +from nemo_automodel.components.loss.utils import calculate_loss from nemo_automodel.components.quantization.fp8 import build_fp8_config from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG @@ -82,6 +84,7 @@ ) from nemo_automodel.components.utils.compile_utils import build_compile_config from nemo_automodel.components.utils.model_utils import VLM_INPUT_KEYS, _supports_logits_to_keep, filter_forward_kwargs +from nemo_automodel.engine import Engine, collate_prebatched from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config, shard_optimizers_for_megatron_fsdp from nemo_automodel.recipes._typed_config import RecipeConfig from nemo_automodel.recipes.base_recipe import BaseRecipe @@ -625,6 +628,26 @@ def setup(self): self.dataloader = dataloader_build.dataloader self.processor = dataloader_build.processor + # Start with the eager path whose inputs need no recipe-owned sharding or + # packing preparation. PP, CP, packed, and Magi batches keep using the + # existing forward/backward path below. + self.engine = None + if ( + not self.pp_enabled + and self.mesh_context.cp_size == 1 + and dataloader_config.packing is None + and not self.magi.enabled + and getattr(self.loss_fn, "reduction", None) == "sum" + ): + self.engine = Engine( + self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + collate_fn=collate_prebatched, + context_fn=self._cp_vision_frame_sharding_context, + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), + ) + # Build validation dataloader if the config provides it self.val_dataloader = None validation_config = self.cfg.vlm_validation_dataloader @@ -741,8 +764,9 @@ def _maybe_add_drafter_loss( base_loss: torch.Tensor, labels: torch.Tensor, model: nn.Module, - num_label_tokens: int, + num_label_tokens: int | None, log: bool = False, + log_denominator: int | float | None = None, ) -> torch.Tensor: """Return ``base_loss + lambda * sum_k CE(drafter_logits[k], shifted_labels_k)``. @@ -752,7 +776,9 @@ def _maybe_add_drafter_loss( For drafter step ``k``, labels are shifted left by ``k`` positions to match the VLM collate's pre-shifted convention (``labels[t] == input_ids[t+1]``). ``log=True`` emits a one-line breakdown on rank 0; callers should gate this - on the appropriate step / microbatch index to avoid log spam. + on the appropriate step / microbatch index to avoid log spam. A caller + that supplies unnormalized sums can set ``log_denominator`` to keep the + reported breakdown in mean-loss units. """ drafter_logits = getattr(out, "drafter_logits", None) if drafter_logits is None or len(drafter_logits) == 0: @@ -774,15 +800,89 @@ def _maybe_add_drafter_loss( total_loss = base_loss + drafter_loss_weight * drafter_loss_total if log and self.dist_env.is_main: + log_scale = 1.0 if log_denominator is None else 1.0 / log_denominator logger.info( "[joint-drafter] L_base=%.4f L_drafter=%.4f L_total=%.4f (lambda=%.3f)", - base_loss.detach().item(), - drafter_loss_total.detach().item(), - total_loss.detach().item(), + base_loss.detach().item() * log_scale, + drafter_loss_total.detach().item() * log_scale, + total_loss.detach().item() * log_scale, drafter_loss_weight, ) return total_loss + def _calculate_eager_loss( + self, + *, + out: Any, + labels: torch.Tensor, + model: nn.Module, + num_label_tokens: int | None, + is_train: bool, + log_drafter: bool = False, + log_denominator: int | float | None = None, + ) -> torch.Tensor: + """Compute base, MTP, and optional joint-drafter losses.""" + grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None + loss = causal_lm_loss( + self.loss_fn, + model, + out, + labels, + getattr(getattr(self, "cfg", None), "mtp", None), + num_label_tokens=num_label_tokens, + grad_reduce_group=grad_reduce_group, + ) + + return self._maybe_add_drafter_loss( + out=out, + base_loss=loss, + labels=labels, + model=model, + num_label_tokens=num_label_tokens, + log=log_drafter, + log_denominator=log_denominator, + ) + + def _make_engine_datum( + self, + batch: dict[str, Any], + *, + log_drafter: bool = False, + log_denominator: int | float = 1, + ) -> Datum: + """Wrap one processor-collated VLM batch as a Datum.""" + labels = batch["labels"] + model = self.model_parts[0] + model_inputs = filter_forward_kwargs(model, {key: value for key, value in batch.items() if key != "labels"}) + if isinstance(self.loss_fn, FusedLinearCrossEntropy): + model_inputs["logits_to_keep"] = 1 + return Datum( + model_inputs=model_inputs, + loss_fn_inputs={ + "labels": labels, + "weights": labels.ne(-100), + "log_drafter": torch.tensor(log_drafter), + "log_denominator": torch.tensor(log_denominator), + }, + ) + + def _engine_loss( + self, + out: Any, + loss_inputs: dict[str, torch.Tensor], + datums: Sequence[Datum], + ) -> torch.Tensor: + """Return the local loss sum; Engine owns global normalization.""" + return self._calculate_eager_loss( + out=out, + labels=loss_inputs["labels"], + model=self.model_parts[0], + num_label_tokens=None, + is_train=True, + log_drafter=bool(datums[0].loss_fn_inputs["log_drafter"].item()), + log_denominator=float(datums[0].loss_fn_inputs["log_denominator"].item()), + ) + def _maybe_set_pp_first_stage_embed_input_meta(self, model_input: torch.Tensor) -> None: if ( not self.pp_enabled @@ -928,71 +1028,17 @@ def _forward_backward_step( if isinstance(self.loss_fn, FusedLinearCrossEntropy): # use num_logits_to_keep to avoid full logits matrix in memory out = model(logits_to_keep=1, **batch) - if "hidden_states" not in out: - raise ValueError( - "FusedLinearCrossEntropy requires the model to output hidden states. " - "Set `model.text_config.output_hidden_states=True` in the config." - ) else: out = model(**batch) - grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None - shared_lm_weight = ( - self.loss_fn.materialize_lm_weight( - _get_lm_head_weight(model), - grad_reduce_group=grad_reduce_group, - ) - if isinstance(self.loss_fn, FusedLinearCrossEntropy) - else None - ) - local_loss = calculate_loss( - self.loss_fn, - logits=getattr(out, "logits", out), - labels=labels, - model=model, - hidden_states=get_final_hidden_states(out), - lm_weight=shared_lm_weight, - grad_reduce_group=grad_reduce_group, - num_label_tokens=num_label_tokens, - ) - # DSV4-style MTP loss (from main): triggers when the model emits - # ``mtp_per_depth_h`` / ``mtp_per_depth_logits``. - mtp_per_depth_h = getattr(out, "mtp_per_depth_h", None) - mtp_per_depth_logits = getattr(out, "mtp_per_depth_logits", None) - if mtp_per_depth_h is not None or mtp_per_depth_logits is not None: - mtp_cfg = self.cfg.mtp - scaling_factor = ( - mtp_cfg.scaling_factor if mtp_cfg.scaling_factor is not None else out.mtp_loss_scaling_factor - ) - local_loss = local_loss + calculate_mtp_loss( - self.loss_fn, - mtp_per_depth_h=mtp_per_depth_h, - mtp_per_depth_logits=mtp_per_depth_logits, - labels=labels, - model=model, - scaling_factor=scaling_factor, - num_label_tokens=num_label_tokens, - ignore_index=mtp_cfg.ignore_index, - lm_weight=shared_lm_weight, - grad_reduce_group=grad_reduce_group, - ) - - # Joint base + drafter co-training (Gemma4WithDrafter and - # similar): detect by presence of ``drafter_logits`` on the - # model output and add - # ``drafter_loss_weight * sum_k CE(drafter_logits[k], shifted_labels_k)`` - # to the base loss. See ``_shift_labels_left`` for the shift - # convention. Mutually exclusive with the DSV4-style MTP path - # above -- only one of ``drafter_logits`` / - # ``mtp_per_depth_*`` is set per model. - local_loss = self._maybe_add_drafter_loss( + local_loss = self._calculate_eager_loss( out=out, - base_loss=local_loss, labels=labels, model=model, num_label_tokens=num_label_tokens, + is_train=is_train, # Log once per remote-logging step on the first microbatch. - log=(idx == 0 and self.step_scheduler.is_remote_logging_step), + log_drafter=(idx == 0 and self.step_scheduler.is_remote_logging_step), ) loss_buffer.append(local_loss.clone().detach()) @@ -1030,9 +1076,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): num_label_tokens = self._dp_allreduce(num_label_tokens).item() num_batches = len(batches) - self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) - - loss_buffer = [] # number of tokens in the batch, excluding any tail padding. num_tokens_in_batch = torch.tensor( @@ -1041,18 +1084,39 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) - - for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) - - self._forward_backward_step( - i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches + engine = getattr(self, "engine", None) + use_engine = engine is not None and num_label_tokens > 0 + if use_engine: + reporting_loss, _ = engine.forward_backward( + [ + [ + self._make_engine_datum( + batch, + log_drafter=(index == 0 and self.step_scheduler.is_remote_logging_step), + log_denominator=num_label_tokens, + ) + ] + for index, batch in enumerate(batches) + ], + self._engine_loss, ) + else: + # The eager Engine requires a positive global weight sum. Preserve + # the established zero-label behavior by using the legacy path. + self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) + loss_buffer = [] + prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) + + for i, batch in enumerate(batches): + if i == num_batches - 1: + prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) + + self._forward_backward_step( + i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches + ) - if i == 0: - prepare_after_first_microbatch() + if i == 0: + prepare_after_first_microbatch() grad_norm = scale_grads_and_clip_grad_norm( max_grad_norm=max_grad_norm, @@ -1104,9 +1168,10 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): time_delta = t - self.timestamp self.timestamp = t tps = num_tokens_in_batch / time_delta - reporting_loss = torch.sum(torch.stack(loss_buffer)) - reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) - if self.pp_enabled: + if not use_engine: + reporting_loss = torch.sum(torch.stack(loss_buffer)) + reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) + if not use_engine and self.pp_enabled: # PP uses sum reduction per microbatch (no internal normalization). # Divide by num_label_tokens to get the mean loss, same as non-PP. reporting_loss = reporting_loss / num_label_tokens if num_label_tokens > 0 else reporting_loss * 0.0 diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index 3f8b7266e8..34eb78c06f 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -57,9 +57,11 @@ def test_datum_keeps_old_input_ids_convenience(): assert positional.input_ids.tolist() == [4, 5] -def test_datum_rejects_non_1d_input_ids(): - with pytest.raises(ValueError, match="must be 1-D"): - Datum(model_inputs={"input_ids": torch.zeros(2, 3, dtype=torch.long)}) +def test_datum_allows_prebatched_inputs_only_with_a_custom_collater(): + datum = Datum(model_inputs={"input_ids": torch.zeros(2, 3, dtype=torch.long)}) + assert datum.model_inputs["input_ids"].shape == (2, 3) + with pytest.raises(ValueError, match="default collater requires 1-D"): + collate_datums([datum]) def test_datum_accepts_model_specific_inputs(): diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py index b76bb2ad79..dc516927c3 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_mtp.py @@ -475,7 +475,7 @@ def test_fused_linear_ce_branch_dispatches(self, backend, monkeypatch): by stubbing ``linear_cross_entropy`` and asserting it is called once per MTP depth with the expected kwargs.""" from nemo_automodel.components.loss import linear_ce as linear_ce_mod - from nemo_automodel.recipes.llm.train_ft import calculate_mtp_loss + from nemo_automodel.components.loss.mtp import calculate_mtp_loss if not linear_ce_mod.HAVE_CUT_CROSS_ENTROPY: pytest.skip("cut_cross_entropy not installed") diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index aacc936b2d..1eb35fde68 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -432,7 +432,7 @@ def fake_calculate_loss(*args, **kwargs): ) calculate_mock = MagicMock(side_effect=fake_calculate_loss) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_loss", calculate_mock) + monkeypatch.setattr("nemo_automodel.components.loss.causal_lm.calculate_loss", calculate_mock) grad_clip_mock = MagicMock(return_value=2.5) monkeypatch.setattr( @@ -485,7 +485,7 @@ def make_thd_batch(model, device_mesh, batch, **kwargs): monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ContextParallelSharder", make_thd_batch) monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.get_sync_ctx", lambda *args, **kwargs: nullcontext()) monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.calculate_loss", + "nemo_automodel.components.loss.causal_lm.calculate_loss", lambda *args, **kwargs: torch.tensor(1.0, requires_grad=True), ) @@ -2996,6 +2996,31 @@ def test_vlm_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): assert cfg.model.backend.rope_fusion is True +def test_vlm_setup_keeps_engine_disabled_for_loss_without_sum_contract(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep", lambda _model: True) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert trainer.loss_fn == "loss_fn" + assert trainer.engine is None + + +def test_vlm_setup_builds_engine_for_eager_sum_loss(monkeypatch): + from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy + + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert isinstance(trainer.loss_fn, MaskedCrossEntropy) + assert trainer.engine is not None + + def test_vlm_setup_does_not_change_storage_dtype_for_non_kd_recipe(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True, optimizer_target="torch.optim.AdamW") _patch_vlm_setup_minimals(monkeypatch, cp_size=1) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 6c236fe00a..30aeb62580 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1080,6 +1080,31 @@ def patch_fn(model, name=None, add_backward_hooks=True): assert patch_calls == [] +def test_setup_keeps_engine_disabled_for_loss_without_sum_contract(monkeypatch): + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._supports_logits_to_keep", lambda _model: True) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert trainer.loss_fn == "loss_fn" + assert trainer.engine is None + + +def test_setup_builds_engine_for_eager_sum_loss(monkeypatch): + from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy + + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert isinstance(trainer.loss_fn, MaskedCrossEntropy) + assert trainer.engine is not None + + def test_setup_does_not_change_storage_dtype_for_non_kd_recipe(monkeypatch): cfg = _minimal_cfg_with_nvtx(nvtx_value=False, optimizer_target="torch.optim.AdamW") @@ -2549,11 +2574,14 @@ def get_local_rank(self): object.__setattr__(recipe, "model_parts", [model]) object.__setattr__(recipe, "distributed_config", SimpleNamespace(defer_fsdp_grad_sync=True)) object.__setattr__(recipe, "loss_fn", object()) # not FusedLinearCrossEntropy + object.__setattr__(recipe, "_get_dp_group", lambda include_cp=False: None) object.__setattr__(recipe, "_get_dp_group_size", lambda include_cp=False: 1) captured = {} - def _fake_calc_loss(loss_fn, *, logits, labels, model, hidden_states, lm_weight, num_label_tokens): + def _fake_calc_loss( + loss_fn, *, logits, labels, model, hidden_states, lm_weight, num_label_tokens, grad_reduce_group + ): captured["logits_is_tensor"] = isinstance(logits, torch.Tensor) assert lm_weight is None return logits.mean() @@ -2562,8 +2590,8 @@ def _fake_calc_loss(loss_fn, *, logits, labels, model, hidden_states, lm_weight, "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", lambda device_mesh, batch, *a, **k: (nullcontext, batch, None), ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calc_loss) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_final_hidden_states", lambda out: None) + monkeypatch.setattr("nemo_automodel.components.loss.causal_lm.calculate_loss", _fake_calc_loss) + monkeypatch.setattr("nemo_automodel.components.loss.causal_lm._get_final_hidden_states", lambda out: None) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_sync_ctx", lambda *a, **k: nullcontext()) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.filter_forward_kwargs", lambda model, batch: batch) diff --git a/tests/unit_tests/recipes/test_vlm_drafter_helpers.py b/tests/unit_tests/recipes/test_vlm_drafter_helpers.py index 85701122ad..3420b7d646 100644 --- a/tests/unit_tests/recipes/test_vlm_drafter_helpers.py +++ b/tests/unit_tests/recipes/test_vlm_drafter_helpers.py @@ -456,6 +456,28 @@ def test_log_main_rank_does_not_raise(self, caplog): # pin the exact floats since they depend on the toy loss arithmetic. assert any("[joint-drafter]" in r.getMessage() for r in caplog.records) + def test_log_denominator_reports_means_for_engine_sums(self, caplog): + import logging + + recipe = self._make_recipe() + recipe.loss_fn = lambda *, logits, **_kwargs: logits.sum() + out = _FakeJointOut(drafter_logits=[torch.zeros((1, 2, 3))], drafter_loss_weight=0.01) + with caplog.at_level(logging.INFO, logger="nemo_automodel.recipes.vlm.finetune"): + FinetuneRecipeForVLM._maybe_add_drafter_loss( + recipe, + out=out, + base_loss=torch.tensor(4.0), + labels=torch.tensor([[1, 2]]), + model=MagicMock(), + num_label_tokens=None, + log=True, + log_denominator=2, + ) + + message = next(record.getMessage() for record in caplog.records if "[joint-drafter]" in record.getMessage()) + assert "L_base=2.0000" in message + assert "L_total=2.0000" in message + def test_log_off_main_rank_emits_nothing(self, caplog): """``log=True`` but ``dist_env.is_main=False`` (off-main rank): no log record from the joint-drafter path.""" diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index d1a216e2ba..dcee9b9712 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -29,8 +29,10 @@ from nemo_automodel import Datum as PublicDatum from nemo_automodel import Engine as PublicEngine from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.loss.causal_lm import causal_lm_loss +from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler -from nemo_automodel.engine import Engine +from nemo_automodel.engine import Engine, collate_prebatched class ScaleModel(nn.Module): @@ -76,25 +78,60 @@ def test_forward_backward_uses_one_denominator_for_the_window(): assert model.forward_calls == 2 -def test_padded_and_packed_windows_have_the_same_loss_and_gradient(): - padded_model = ScaleModel() - packed_model = ScaleModel() - window = [[_datum([1, 2]), _datum([3])]] +def test_pre_thd_packed_collater_requires_model_ready_inputs(): + model = ScaleModel() + + with pytest.raises(ValueError, match="final model-ready THD"): + Engine( + model, + device="cpu", + collate_fn=partial(collate_datums, packed=True), + ).forward_backward([[_datum([1, 2]), _datum([3])]], _identity_loss) + + assert model.forward_calls == 0 - padded_loss, _ = Engine(padded_model, device="cpu").forward_backward(window, _identity_loss) - def packed_identity_loss(output, loss_inputs, datums): - assert [datum.seq_len for datum in datums] == [2, 1] - return _identity_loss(output, loss_inputs, datums) +def _model_ready_packed_collate(datums): + model_inputs, loss_inputs = collate_datums(datums, packed=True) + lengths = torch.tensor([datum.seq_len for datum in datums], dtype=torch.int32) + model_inputs = { + "input_ids": model_inputs["input_ids"].flatten(), + "position_ids": model_inputs["position_ids"].flatten(), + "cu_seqlens": F.pad(lengths.cumsum(0), (1, 0)), + "max_seqlen": lengths.max(), + "qkv_format": "thd", + } + loss_inputs = { + key: value.flatten() if value.ndim == 2 and value.shape[0] == 1 else value for key, value in loss_inputs.items() + } + return model_inputs, loss_inputs + + +def test_packed_rl_callback_keeps_per_datum_sequence_boundaries(): + model = ScaleModel() + first = _datum([1, 2]) + second = _datum([3]) + first.loss_fn_inputs["sequence_scale"] = torch.tensor(2.0) + second.loss_fn_inputs["sequence_scale"] = torch.tensor(0.5) + + def sequence_loss(output, _loss_inputs, datums): + chunks = output.squeeze(0).split([datum.seq_len for datum in datums]) + losses = torch.cat([chunk * datum.loss_fn_inputs["sequence_scale"] for chunk, datum in zip(chunks, datums)]) + outputs = [ + {"sequence_sum": chunk.sum(), "sequence_length": datum.seq_len} for chunk, datum in zip(chunks, datums) + ] + return losses, outputs - packed_loss, _ = Engine( - packed_model, + loss, outputs = Engine( + model, device="cpu", - collate_fn=partial(collate_datums, packed=True), - ).forward_backward(window, packed_identity_loss) + collate_fn=_model_ready_packed_collate, + ).forward_backward([[first, second]], sequence_loss) - torch.testing.assert_close(padded_loss, packed_loss) - torch.testing.assert_close(padded_model.weight.grad, packed_model.weight.grad) + assert loss.item() == pytest.approx(2.5) + assert model.weight.grad.item() == pytest.approx(2.5) + assert [item["sequence_length"] for item in outputs] == [2, 1] + assert [item["sequence_sum"].item() for item in outputs] == pytest.approx([3.0, 3.0]) def test_weights_mask_loss_and_denominator(): @@ -159,6 +196,11 @@ def forward(self, input_ids, **_): return self.output(self.embedding(input_ids)) +class TinyCausalLM(TinyLM): + def forward(self, input_ids, **_): + return SimpleNamespace(logits=super().forward(input_ids)) + + def test_raw_output_and_loss_inputs_support_an_rl_loss_callback(): datum = Datum( model_inputs={"input_ids": torch.tensor([1, 2, 3])}, @@ -191,6 +233,56 @@ def policy_loss(logits, inputs, datums): assert model.output.weight.grad is not None +def test_causal_lm_loss_matches_a_manual_accumulation_window(): + torch.manual_seed(7) + model = TinyCausalLM() + reference = TinyCausalLM() + reference.load_state_dict(model.state_dict()) + batches = [ + (torch.tensor([[1, 2, 3]]), torch.tensor([[2, 3, -100]])), + (torch.tensor([[4, 5]]), torch.tensor([[5, 6]])), + ] + loss_fn = MaskedCrossEntropy() + mtp_config = SimpleNamespace(scaling_factor=None, ignore_index=-100) + + window = [ + [ + Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, + ) + ] + for input_ids, labels in batches + ] + + def engine_loss(output, inputs, _datums): + return causal_lm_loss( + loss_fn, + model, + output, + inputs["labels"], + mtp_config, + num_label_tokens=None, + grad_reduce_group=None, + ) + + actual_loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward(window, engine_loss) + + denominator = sum((labels != -100).sum() for _, labels in batches) + reference_loss = ( + sum( + F.cross_entropy(reference(input_ids).logits.flatten(0, 1), labels.flatten(), reduction="sum") + for input_ids, labels in batches + ) + / denominator + ) + reference_loss.backward() + + torch.testing.assert_close(actual_loss, reference_loss.detach().to(actual_loss)) + for parameter, expected in zip(model.parameters(), reference.parameters()): + torch.testing.assert_close(parameter.grad, expected.grad) + + def test_lifecycle_marks_only_the_last_microbatch_for_sync(monkeypatch): events = [] @@ -215,6 +307,34 @@ def sync_context(_model, is_last, _defer): assert events == ["prepare", "sync:False", "after_first", "final", "sync:True"] +def test_forward_context_covers_forward_loss_and_backward(): + active = False + + @contextmanager + def forward_context(): + nonlocal active + active = True + try: + yield + finally: + active = False + + class ContextModel(ScaleModel): + def forward(self, input_ids, **kwargs): + assert active + return super().forward(input_ids, **kwargs) + + model = ContextModel() + model.weight.register_hook(lambda grad: grad if active else pytest.fail("context ended before backward")) + + def loss_fn(output, _inputs, _datums): + assert active + return output + + Engine(model, device="cpu", context_fn=forward_context).forward_backward([[_datum([1, 2])]], loss_fn) + assert not active + + def test_window_sets_the_same_moe_aux_scale_as_the_recipes(monkeypatch): monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) @@ -270,6 +390,48 @@ def test_model_specific_collater_keeps_multimodal_inputs_and_gradients(): assert model.vision.weight.grad.abs().sum() > 0 +def test_prebatched_datum_keeps_existing_recipe_batch_layout(): + model = ScaleModel() + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2], [3, 4]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0], [1.0, 0.0]])}, + ) + + loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([[datum]], _identity_loss) + + assert loss.item() == pytest.approx(2.0) + assert model.weight.grad.item() == pytest.approx(2.0) + + +def test_prebatched_datum_keeps_vlm_media_layout(): + model = TinyVLM() + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2], [3, 4]]), + "pixel_values": [torch.tensor([0.5]), torch.tensor([1.5])], + }, + loss_fn_inputs={"weights": torch.ones(2)}, + ) + + loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([[datum]], _identity_loss) + + assert torch.isfinite(loss) + assert model.text.weight.grad is not None + assert model.vision.weight.grad is not None + + +def test_scalar_loss_is_a_local_weighted_sum_numerator(): + model = ScaleModel() + + loss, _ = Engine(model, device="cpu").forward_backward( + [[_datum([1, 100], [1.0, 0.0])], [_datum([3, 5], [0.5, 1.0])]], + lambda output, inputs, _datums: (output * inputs["weights"]).sum(), + ) + + assert loss.item() == pytest.approx(3.0) + assert model.weight.grad.item() == pytest.approx(3.0) + + def test_zero_weights_fail_before_forward(): model = ScaleModel() with pytest.raises(ValueError, match="positive global weight sum"): @@ -300,7 +462,7 @@ def test_loss_shape_must_exactly_match_weights(): with pytest.raises(ValueError, match="exactly the same shape"): Engine(model, device="cpu").forward_backward( [[_datum([1, 2])]], - lambda output, _inputs, _datums: output.sum(), + lambda output, _inputs, _datums: output[:, :1], ) assert model.weight.grad is None diff --git a/tests/unit_tests/test_engine_recipe_integration.py b/tests/unit_tests/test_engine_recipe_integration.py new file mode 100644 index 0000000000..533ced3574 --- /dev/null +++ b/tests/unit_tests/test_engine_recipe_integration.py @@ -0,0 +1,140 @@ +# 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. + +from __future__ import annotations + +import time +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy +from nemo_automodel.engine import Engine, collate_prebatched +from nemo_automodel.recipes.llm.train_ft import TrainFinetuneRecipeForNextTokenPrediction +from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM + + +class _Config(dict): + def __init__(self): + super().__init__() + self.mtp = SimpleNamespace(scaling_factor=None, ignore_index=-100) + + +class _TinyLM(nn.Module): + def __init__(self, *, vlm: bool = False): + super().__init__() + self.embedding = nn.Embedding(8, 4) + self.vision = nn.Linear(1, 4) if vlm else None + self.output = nn.Linear(4, 8) + self.forward_calls = 0 + + def forward(self, input_ids, pixel_values=None): + self.forward_calls += 1 + hidden = self.embedding(input_ids) + if pixel_values is not None: + hidden = hidden + self.vision(pixel_values).unsqueeze(1) + return SimpleNamespace(logits=self.output(hidden)) + + +class _CountingSGD(torch.optim.SGD): + def __init__(self, parameters): + super().__init__(parameters, lr=0.1) + self.step_calls = 0 + self.zero_calls = 0 + + def step(self, closure=None): + self.step_calls += 1 + return super().step(closure) + + def zero_grad(self, *args, **kwargs): + self.zero_calls += 1 + return super().zero_grad(*args, **kwargs) + + +@pytest.mark.parametrize( + ("recipe_cls", "vlm"), + [ + (TrainFinetuneRecipeForNextTokenPrediction, False), + (FinetuneRecipeForVLM, True), + ], +) +def test_recipes_run_one_datum_engine_window_then_one_optimizer_step(recipe_cls, vlm): + model = _TinyLM(vlm=vlm) + reference = _TinyLM(vlm=vlm) + reference.load_state_dict(model.state_dict()) + recipe = object.__new__(recipe_cls) + recipe.cfg = _Config() + recipe.loss_fn = MaskedCrossEntropy() + recipe.model_parts = [model] + recipe.device_mesh = None + recipe.moe_mesh = None + recipe.pp_enabled = False + recipe.dist_env = SimpleNamespace(device=torch.device("cpu"), world_size=1, is_main=True) + recipe.distributed_config = SimpleNamespace(defer_fsdp_grad_sync=True) + recipe.engine = Engine(model, device="cpu", collate_fn=collate_prebatched) + optimizer = _CountingSGD(model.parameters()) + recipe.optimizer = [optimizer] + recipe.lr_scheduler = None + recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=lambda: None) + recipe.step_scheduler = SimpleNamespace(step=1, epoch=0, is_remote_logging_step=False) + recipe.timestamp = time.perf_counter() - 1.0 + + reductions = 0 + + def local_reduce(value, include_cp=False): + nonlocal reductions + reductions += 1 + return value + + recipe._dp_allreduce = local_reduce + batches = [ + {"input_ids": torch.tensor([[1, 2, 3]]), "labels": torch.tensor([[2, 3, -100]])}, + {"input_ids": torch.tensor([[4, 5]]), "labels": torch.tensor([[5, 6]])}, + ] + if vlm: + batches[0]["pixel_values"] = torch.tensor([[0.5]]) + batches[1]["pixel_values"] = torch.tensor([[1.5]]) + + reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.1) + denominator = sum((batch["labels"] != -100).sum() for batch in batches) + reference_loss = ( + sum( + F.cross_entropy( + reference( + batch["input_ids"], + pixel_values=batch.get("pixel_values"), + ).logits.flatten(0, 1), + batch["labels"].flatten(), + ignore_index=-100, + reduction="sum", + ) + for batch in batches + ) + / denominator + ) + reference_loss.backward() + reference_optimizer.step() + + metrics = recipe._run_train_optim_step(batches, max_grad_norm=None) + + assert metrics.metrics["loss"] == pytest.approx(reference_loss.item()) + assert model.forward_calls == 2 + assert optimizer.step_calls == 1 + assert optimizer.zero_calls == 1 + assert reductions == 2 # token counters only; Engine already reduced the loss + for actual, expected in zip(model.parameters(), reference.parameters()): + torch.testing.assert_close(actual, expected) From 9aae427ed17e590d7058c2319fd8be466640b7cd Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 15 Aug 2026 21:44:59 -0700 Subject: [PATCH 03/34] feat(engine): shard Datum windows with context parallelism Signed-off-by: HuiyingLi --- .../distributed/context_parallel/magi.py | 51 ++-- .../distributed/context_parallel/sharder.py | 6 +- .../distributed/context_parallel/utils.py | 29 +- .../components/distributed/utils.py | 10 +- nemo_automodel/engine/__init__.py | 231 +++++++++++--- nemo_automodel/recipes/llm/train_ft.py | 48 ++- nemo_automodel/recipes/vlm/finetune.py | 30 +- tests/unit_tests/distributed/test_cp_utils.py | 77 ++++- .../distributed/test_magi_attn_utils.py | 45 ++- tests/unit_tests/distributed/test_utils.py | 20 ++ .../recipes/test_finetune_vlm_helpers.py | 64 ++++ tests/unit_tests/recipes/test_train_ft.py | 28 ++ tests/unit_tests/test_engine.py | 288 ++++++++++++++++-- 13 files changed, 803 insertions(+), 124 deletions(-) diff --git a/nemo_automodel/components/distributed/context_parallel/magi.py b/nemo_automodel/components/distributed/context_parallel/magi.py index b9872aca69..3ca3f1b9d1 100644 --- a/nemo_automodel/components/distributed/context_parallel/magi.py +++ b/nemo_automodel/components/distributed/context_parallel/magi.py @@ -510,6 +510,8 @@ def magi_prepare_batch( # pragma: no cover - requires GPU + magi_attention batch: dict, cp_group: Optional[dist.ProcessGroup], chunk_size: int = DEFAULT_CHUNK_SIZE, + *, + return_local_indices: bool = False, ): """Dispatch a (batch_size==1) sequence for MagiAttention on the HF path. @@ -526,10 +528,11 @@ def magi_prepare_batch( # pragma: no cover - requires GPU + magi_attention batch: dict with at least ``input_ids`` of shape ``[1, S]``. cp_group: CP process group (size 1 -> identity shard). chunk_size: dispatch solver chunk size. + return_local_indices: Also return the dispatched global token indices. Returns: - (new_batch, key): ``new_batch`` has dispatched ``input_ids``/``position_ids``/ - ``labels`` (each ``[1, local_S]``); ``key`` is the dist-attn runtime key. + ``(new_batch, key)`` by default. With ``return_local_indices=True``, + also returns the map from the local token stream to the global input. """ from magi_attention.api import dispatch, get_position_ids, magi_attn_varlen_key from magi_attention.api.functools import compute_pad_size @@ -567,6 +570,11 @@ def magi_prepare_batch( # pragma: no cover - requires GPU + magi_attention local_input = dispatch(input_ids.squeeze(0), key=key).unsqueeze(0) # [1, local_S] position_ids = get_position_ids(key).unsqueeze(0).to(device) # [1, local_S] + local_indices = None + if return_local_indices: + local_indices = dispatch( + torch.arange(seqlen, device=device, dtype=torch.long), key=key, pad_value=seqlen + ).unsqueeze(0) _set_cp_group_on_attention(model.module if hasattr(model, "module") else model, cp_group) @@ -579,7 +587,7 @@ def magi_prepare_batch( # pragma: no cover - requires GPU + magi_attention new_batch["labels"] = dispatch(batch["labels"].squeeze(0), key=key, pad_value=-100).unsqueeze(0) # Remove anything that no longer matches the dispatched layout. new_batch.pop("attention_mask", None) - return new_batch, key + return (new_batch, key, local_indices) if return_local_indices else (new_batch, key) def _packed_cp_doc_seqlens(batch: dict, total_len: int) -> list: @@ -612,7 +620,9 @@ def _packed_cp_doc_seqlens(batch: dict, total_len: int) -> list: return seqlens -def magi_prepare_packed_cp(model, batch: dict, cp_group): # pragma: no cover - requires GPU + magi_attention +def magi_prepare_packed_cp( # pragma: no cover - requires GPU + magi_attention + model, batch: dict, cp_group, *, return_local_indices: bool = False +): """Context-parallel prep for a packed (THD) batch on the custom-model path. Takes a *global* THD batch (flat ``input_ids``/``labels``/``position_ids`` plus @@ -624,9 +634,12 @@ def magi_prepare_packed_cp(model, batch: dict, cp_group): # pragma: no cover - each rank computes a per-shard loss that the recipe's cross-CP reduction sums into the global loss (like TE-CP). + Args: + return_local_indices: Also return the dispatched global token indices. + Returns: - (new_batch, key): ``new_batch`` has the local ``input_ids``/``position_ids`` - and local ``labels``; ``key`` is the dist-attn runtime key. + ``(new_batch, key)`` by default. With ``return_local_indices=True``, + also returns the map from the local token stream to the global input. """ from magi_attention.api import dispatch, get_position_ids @@ -643,12 +656,18 @@ def magi_prepare_packed_cp(model, batch: dict, cp_group): # pragma: no cover - ) local_input = dispatch(input_ids, key=key) local_pos = get_position_ids(key).to(local_input.device) + total_len = input_ids.numel() + local_indices = None + if return_local_indices: + local_indices = dispatch( + torch.arange(total_len, device=input_ids.device, dtype=torch.long), key=key, pad_value=total_len + ) # Shard labels the same way as the input (like TE-CP): each rank computes the # loss on its own shard and the recipe's cross-CP reduction sums the shards # into the global loss. (Undispatching logits to global instead would make # every CP rank compute the full loss redundantly -> the reduction would # double-count it by a factor of cp_size.) - local_labels = dispatch(batch["labels"].reshape(-1), key=key) + local_labels = dispatch(batch["labels"].reshape(-1), key=key, pad_value=-100) new_batch = { "input_ids": local_input, "position_ids": local_pos, @@ -659,7 +678,7 @@ def magi_prepare_packed_cp(model, batch: dict, cp_group): # pragma: no cover - # magi dispatch permutation. "qkv_format": "thd", } - return new_batch, key + return (new_batch, key, local_indices) if return_local_indices else (new_batch, key) def _iter_language_model_attention(model): @@ -750,10 +769,10 @@ def prepare_llm_batch( Returns ``(train_ctx, batch, local_indices)``. magi does its own CP, so ``train_ctx`` is always ``nullcontext`` (no torch-native DTensor CP - context). ``local_indices`` is the global stream position of every - local token on the paths that dispatch the sequence (magi's - ``get_position_ids``), None otherwise; the framework installs it on - the magi ContextParallelSharder for the token-tensor verbs. + context). ``local_indices`` is the explicitly dispatched global stream + index of every local token on paths that dispatch the sequence, None + otherwise; the framework installs it on the magi + ContextParallelSharder for the token-tensor verbs. """ # cp=1 prefix-tree mask: the datasets layer cannot import this module (component # independence), so the collate attaches the tree structure and the spec is built @@ -780,16 +799,12 @@ def prepare_llm_batch( local_indices = None if self.hf_dispatch: # HF path: dispatch the (single causal) sequence across the CP group. - batch, _ = magi_prepare_batch(model, batch, self.cp_group) - # The dispatched position_ids ARE magi's get_position_ids(key): the - # global stream position of every local token. - local_indices = batch["position_ids"] + batch, _, local_indices = magi_prepare_batch(model, batch, self.cp_group, return_local_indices=True) elif self.custom and self.cp_size > 1 and is_thd: # Custom-model CP packed path: build the *global* THD layout (no TE # sharding) then dispatch it with magi's own load-balancing solver. batch = make_cp_batch_for_te(None, batch, qkv_format="thd", padding_token_id=pad_id, num_chunks=1) - batch, _ = magi_prepare_packed_cp(model, batch, self.cp_group) - local_indices = batch["position_ids"] + batch, _, local_indices = magi_prepare_packed_cp(model, batch, self.cp_group, return_local_indices=True) elif is_thd: # cp=1 packing: THD conversion (no sharding) so the batch carries # cu_seqlens -> the magi attn_func builds the per-document mask. diff --git a/nemo_automodel/components/distributed/context_parallel/sharder.py b/nemo_automodel/components/distributed/context_parallel/sharder.py index 956fde0923..1c3b09c125 100644 --- a/nemo_automodel/components/distributed/context_parallel/sharder.py +++ b/nemo_automodel/components/distributed/context_parallel/sharder.py @@ -217,8 +217,8 @@ class ShardLayout: Attributes: local_token_global_indices: The partition actually computed, for data-dependent layouts (TE's ``thd_get_partitioned_indices`` - result, magi's ``get_position_ids``); None for layouts whose index - map is a closed-form function already on the sharder. + result, magi's dispatched global token indices); None for layouts + whose index map is a closed-form function already on the sharder. original_seq_len: Pre-pad sequence length; None when the layout has no single original length (packed streams). padded_seq_len: Post-pad global sequence length; the token verbs @@ -476,6 +476,8 @@ def gather_token_tensor( out = full.gather(1, positions.clamp(min=0).to(torch.long)) return out.masked_fill(positions < 0, fill) if layout.input_row_shape is not None: + if layout.original_seq_len is not None: + full = full.narrow(seq_dim, 0, layout.original_seq_len) return full.reshape(*layout.input_row_shape, *full.shape[seq_dim + 1 :]) if layout.original_seq_len is not None: return full.narrow(seq_dim, 0, layout.original_seq_len) diff --git a/nemo_automodel/components/distributed/context_parallel/utils.py b/nemo_automodel/components/distributed/context_parallel/utils.py index 4972894f46..6be7e4c556 100644 --- a/nemo_automodel/components/distributed/context_parallel/utils.py +++ b/nemo_automodel/components/distributed/context_parallel/utils.py @@ -499,11 +499,11 @@ def _resolve_cp_sharder( # so like the TE path shard_batch returns (nullcontext, prepped_batch). # All magi internals (HF-vs-custom, recipe domain, cp group) stay in # context_parallel.magi. The dispatch-solver partition is data-dependent, so - # shard_batch installs the index map it just computed (magi's - # get_position_ids) on the sharder for the token verbs. + # shard_batch installs the token-stream permutation it just computed on the + # sharder for the token verbs. def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id=0): - input_ids = batch.get("input_ids") - row_shape = tuple(input_ids.shape[:2]) if input_ids is not None and input_ids.dim() >= 2 else None + primary = batch.get("inputs_embeds", batch.get("input_ids")) + row_shape = tuple(primary.shape[:2]) if primary is not None and primary.dim() >= 2 else None prepped, local_indices = magi.make_cp_batch( cp_mesh, batch, @@ -513,15 +513,28 @@ def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_ model=model, return_local_indices=True, ) + prepped_primary = prepped.get("inputs_embeds", prepped.get("input_ids")) + flattened = ( + row_shape is not None + and isinstance(prepped_primary, torch.Tensor) + and prepped_primary.dim() == primary.dim() - 1 + ) + cp_size = cp_mesh.size() if cp_mesh is not None else 1 + if local_indices is None and cp_size == 1 and isinstance(prepped_primary, torch.Tensor): + token_dim = 0 if flattened or prepped_primary.dim() == 1 else 1 + local_indices = torch.arange( + prepped_primary.shape[token_dim], device=prepped_primary.device, dtype=torch.long + ) layout = None if local_indices is not None: - padded = local_indices.numel() * max(getattr(magi, "cp_size", 1) or 1, 1) + padded = local_indices.numel() * cp_size original, in_rows = None, None if row_shape is not None: - if padded == row_shape[0] * row_shape[1]: - # Flatten moved no tokens and dispatch added no pad: the - # pre-flatten rows are the caller's coordinate system. + if flattened: + # THD flattens the caller's rows before magi dispatches + # (and may add dispatch padding after that). in_rows = row_shape + original = row_shape[0] * row_shape[1] elif row_shape[0] == 1 and padded >= row_shape[1]: # Single-sequence HF path: dispatch pads at the tail of # the global order, so trim restores the original length. diff --git a/nemo_automodel/components/distributed/utils.py b/nemo_automodel/components/distributed/utils.py index f41d2b92e6..e6f89cb50a 100644 --- a/nemo_automodel/components/distributed/utils.py +++ b/nemo_automodel/components/distributed/utils.py @@ -247,17 +247,17 @@ def get_sync_ctx(model, is_optim_step, defer_fsdp_grad_sync: bool): Returns: A context manager that synchronizes the model. """ - # Use `no_sync` on DDP models when we are *not* on the final micro-batch for - # this gradient update (i.e., when `is_grad` is False). This avoids an - # all-reduce for every micro-batch and greatly improves throughput. + # Use `no_sync` on wrappers that expose it when we are not on the final + # microbatch. This covers DDP and optional wrappers such as MegatronFSDP + # without importing their optional packages here. sync_ctx = nullcontext() if isinstance(model, dist.fsdp._fully_shard._fully_shard.FSDPModule): if defer_fsdp_grad_sync: model.set_requires_gradient_sync(is_optim_step) else: model.set_requires_gradient_sync(True) - elif isinstance(model, torch.nn.parallel.DistributedDataParallel) and not is_optim_step: - sync_ctx = model.no_sync() + elif not is_optim_step and callable(no_sync := getattr(model, "no_sync", None)): + sync_ctx = no_sync() return sync_ctx diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 2ddd50d39e..d122b9bd42 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -25,6 +25,11 @@ from torch import nn from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder +from nemo_automodel.components.distributed.context_parallel.sharder import ( + identity_local_indices, + shard_batch_identity, +) from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh from nemo_automodel.components.distributed.utils import get_sync_ctx @@ -34,10 +39,11 @@ prepare_for_final_backward, prepare_for_grad_accumulation, ) +from nemo_automodel.components.utils.model_utils import filter_forward_kwargs CollateFn = Callable[[list[Datum]], tuple[dict[str, Any], dict[str, torch.Tensor]]] LossFn = Callable[ - [Any, dict[str, torch.Tensor], Sequence[Datum]], + [Any, dict[str, torch.Tensor], Sequence[Datum], dict[str, Any]], torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], ] @@ -81,11 +87,13 @@ class Engine: mesh_context: Runtime topology. When omitted, an initialized default process group is treated as pure data parallelism. collate_fn: Batches one microbatch of Datums into separate model and - loss inputs. The default supports padded text. Packed text and - VLMs pass a model-specific collater that returns final model-ready - inputs. Existing recipes whose dataloaders already collate can use - :func:`collate_prebatched`. The callable must keep model inputs and - loss inputs aligned and preserve the sum of ``weights``. + loss inputs. The default supports padded and packed text. VLMs pass + a model-specific collater. Existing recipes whose dataloaders + already collate can use :func:`collate_prebatched`; the Engine then + applies any remaining CP/THD preparation. The callable must keep + model inputs and loss inputs aligned and preserve the sum of + ``weights``. + padding_token_id: Token used when the CP sharder pads ``input_ids``. context_fn: Creates an optional context around model forward, loss, and backward. Recipes use this for runtime contexts such as FP8. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the @@ -93,7 +101,8 @@ class Engine: Note: This first execution backend is eager and weight-normalized. Pipeline - and context parallel schedules are intentionally deferred. + schedules are intentionally deferred; context-parallel input layout + and transport are delegated to :class:`ContextParallelSharder`. """ def __init__( @@ -103,6 +112,7 @@ def __init__( device: torch.device | str, mesh_context: MeshContext | None = None, collate_fn: CollateFn = collate_datums, + padding_token_id: int = 0, context_fn: Callable[[], AbstractContextManager[Any]] = nullcontext, defer_fsdp_grad_sync: bool = True, ) -> None: @@ -110,6 +120,7 @@ def __init__( self.device = torch.device(device) self.mesh_context = mesh_context self.collate_fn = collate_fn + self.padding_token_id = padding_token_id self.context_fn = context_fn self.defer_fsdp_grad_sync = defer_fsdp_grad_sync @@ -121,9 +132,10 @@ def forward_backward( """Accumulate gradients for a complete optimizer window. ``window`` is explicit: each inner sequence is one eager microbatch. - ``loss_fn`` receives the raw model output, collated - ``loss_fn_inputs``, and the original Datums for that microbatch. It - returns either per-element losses with exactly the same shape as + ``loss_fn`` receives the raw model output, CP-local + ``loss_fn_inputs``, the original Datums, and the final CP-local model + inputs produced by the sharder. It returns either per-element losses + with exactly the same shape as ``loss_fn_inputs["weights"]``, or a scalar local weighted-sum numerator. For a scalar, the callback must apply weights and masks; the Engine will only apply global normalization. The callback may @@ -141,21 +153,22 @@ def forward_backward( collated loss inputs. Returns: - ``(loss, loss_fn_outputs)``. ``loss`` is a detached, DP-reduced - scalar. ``loss_fn_outputs`` contains local-rank, per-Datum mappings - in window order. Model parameters are unchanged, but their - gradients contain the complete window's globally normalized - backward result. + ``(loss, loss_fn_outputs)``. ``loss`` is a detached scalar reduced + over the DP-CP gradient group. ``loss_fn_outputs`` contains + local-rank, per-Datum mappings in window order. Model parameters + are unchanged, but their gradients contain the complete window's + globally normalized backward result. """ microbatches = self._validate_window(window) self._validate_parallelism() dp_group, dp_size = self._dp_group_and_size() - self._validate_window_size_across_dp(len(microbatches), dp_group, dp_size) + grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) + self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) denominator = self._global_weight_sum(microbatches, dp_group, dp_size) self.model.train() prepare_for_grad_accumulation([self.model], pp_enabled=False) - MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(1.0 / len(microbatches)) + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(self._cp_size() / len(microbatches)) local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) loss_fn_outputs: list[dict[str, Any]] = [] @@ -167,20 +180,59 @@ def forward_backward( prepare_for_final_backward([self.model], pp_enabled=False) model_inputs, loss_inputs = self.collate_fn(datums) + self._validate_collated_weights(datums, loss_inputs) + model_inputs = _to_device(model_inputs, self.device) + loss_inputs = _to_device(loss_inputs, self.device) + full_weights = loss_inputs["weights"] + loss_seq_dim = _loss_sequence_dim(model_inputs, full_weights) + + # ContextParallelSharder is the single owner of padded, THD, Magi, + # and model-specific CP layouts. Labels are temporarily present in + # its batch because each backend historically shards them with the + # model inputs; all other loss tensors use the sharder token verb. + cp_batch = dict(model_inputs) + labels = loss_inputs.get("labels") + cp_batch["labels"] = ( + labels.clone() if isinstance(labels, torch.Tensor) else torch.zeros_like(full_weights, dtype=torch.long) + ) + device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None + final_thd = _is_final_thd(cp_batch) + if final_thd and self._cp_size() > 1: + raise ValueError( + "context parallelism requires raw THD inputs so ContextParallelSharder can partition them" + ) + if final_thd: + sharder = ContextParallelSharder( + device_mesh=device_mesh, + shard_batch=shard_batch_identity, + local_token_global_indices=identity_local_indices, + padding_token_id=self.padding_token_id, + ) + else: + sharder = ContextParallelSharder( + self.model, + device_mesh, + cp_batch, + padding_token_id=self.padding_token_id, + ) + cp_context, model_inputs = sharder.shard(cp_batch) if model_inputs.get("qkv_format") == "thd" and ( "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs ): raise ValueError( - "packed collate_fn must return final model-ready THD inputs, not seq_lens packing metadata" + "ContextParallelSharder could not prepare raw THD inputs for this model; " + "use a THD-capable attention backend or provide final THD inputs at cp_size=1" ) - self._validate_collated_weights(datums, loss_inputs) - model_inputs = _to_device(model_inputs, self.device) - loss_inputs = _to_device(loss_inputs, self.device) + local_labels = model_inputs.pop("labels") + loss_inputs = self._shard_loss_inputs(sharder, loss_inputs, loss_seq_dim) + if labels is not None: + loss_inputs["labels"] = local_labels weights = loss_inputs["weights"] + forward_inputs = filter_forward_kwargs(self.model, model_inputs) - with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(): - output = self.model(**model_inputs) - result = loss_fn(output, loss_inputs, datums) + with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(), cp_context(): + output = self.model(**forward_inputs) + result = loss_fn(output, loss_inputs, datums, model_inputs) has_outputs = isinstance(result, tuple) if returns_outputs is None: returns_outputs = has_outputs @@ -212,14 +264,14 @@ def forward_backward( if losses.device != weights.device: raise ValueError("loss_fn losses and weights must be on the same device") - (numerator * (dp_size / denominator)).backward() + (numerator * (grad_group_size / denominator)).backward() local_loss_sum.add_(numerator.detach().to(torch.float64)) if index == 0: prepare_after_first_microbatch() - if dp_size > 1: - dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=dp_group) + if grad_group_size > 1: + dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=grad_group) loss = (local_loss_sum / denominator).detach() return loss, loss_fn_outputs @@ -238,12 +290,20 @@ def _validate_window(window: Sequence[Sequence[Datum]]) -> list[list[Datum]]: return microbatches def _validate_parallelism(self) -> None: + if any(bool(getattr(module, "calculate_per_token_loss", False)) for module in self.model.modules()): + raise NotImplementedError( + "Engine.forward_backward requires averaged distributed gradients; " + "MegatronFSDP calculate_per_token_loss=True uses summed gradients" + ) if self.mesh_context is None: return if self.mesh_context.pp_size > 1: raise NotImplementedError("Engine.forward_backward does not yet support pipeline parallelism") - if self.mesh_context.cp_size > 1: - raise NotImplementedError("Engine.forward_backward does not yet support context parallelism") + if self.mesh_context.cp_size > 1 and self.mesh_context.device_mesh is None: + raise ValueError("context parallelism requires a device mesh") + + def _cp_size(self) -> int: + return self.mesh_context.cp_size if self.mesh_context is not None else 1 def _dp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: if self.mesh_context is not None and self.mesh_context.device_mesh is not None: @@ -256,6 +316,17 @@ def _dp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: return group, dist.get_world_size(group=group) return None, 1 + def _gradient_group_and_size( + self, + dp_group: dist.ProcessGroup | None, + dp_size: int, + ) -> tuple[dist.ProcessGroup | None, int]: + if self.mesh_context is None or self.mesh_context.device_mesh is None or self._cp_size() == 1: + return dp_group, dp_size + dp_cp_mesh = get_flat_mesh(self.mesh_context.device_mesh, "dp_cp") + size = int(dp_cp_mesh.size()) + return (dp_cp_mesh.get_group() if size > 1 else None), size + def _global_weight_sum( self, microbatches: list[list[Datum]], @@ -272,25 +343,81 @@ def _global_weight_sum( local_sum += float(weights.to(torch.float64).sum()) denominator = torch.tensor(local_sum, dtype=torch.float64, device=self.device) + self._validate_weight_sum_across_cp(denominator) if dp_size > 1: dist.all_reduce(denominator, op=dist.ReduceOp.SUM, group=dp_group) if float(denominator) <= 0: raise ValueError("forward_backward requires a positive global weight sum") return denominator - def _validate_window_size_across_dp( + def _validate_window_size_across_group( self, size: int, - dp_group: dist.ProcessGroup | None, - dp_size: int, + group: dist.ProcessGroup | None, + group_size: int, ) -> None: - if dp_size <= 1: + if group_size <= 1: return local_size = torch.tensor([size], dtype=torch.int64, device=self.device) - sizes = torch.empty(dp_size, dtype=torch.int64, device=self.device) - dist.all_gather_into_tensor(sizes, local_size, group=dp_group) + sizes = torch.empty(group_size, dtype=torch.int64, device=self.device) + dist.all_gather_into_tensor(sizes, local_size, group=group) if not bool((sizes == sizes[0]).all()): - raise ValueError(f"every data-parallel rank must use the same number of microbatches; got {sizes.tolist()}") + raise ValueError(f"every gradient rank must use the same number of microbatches; got {sizes.tolist()}") + + def _validate_weight_sum_across_cp(self, local_weight_sum: torch.Tensor) -> None: + """Verify that CP replicas started from the same full-sequence weights.""" + if self._cp_size() <= 1 or not dist.is_available() or not dist.is_initialized(): + return + cp_group = self.mesh_context.device_mesh["cp"].get_group() + values = torch.empty(self._cp_size(), dtype=local_weight_sum.dtype, device=local_weight_sum.device) + dist.all_gather_into_tensor(values, local_weight_sum.reshape(1), group=cp_group) + if not torch.allclose(values, values[0].expand_as(values), rtol=1e-8, atol=1e-12): + raise ValueError(f"context-parallel ranks must use identical full-sequence weights; got {values.tolist()}") + + def _shard_loss_inputs( + self, + sharder: ContextParallelSharder, + loss_inputs: dict[str, torch.Tensor], + seq_dim: int | None, + ) -> dict[str, torch.Tensor]: + """Apply the model batch's CP token layout to loss-only tensors. + + Args: + sharder: Sharder after it has prepared the current model batch. + loss_inputs: Tensors in the pre-CP layout. ``weights`` has shape + ``[batch, sequence]`` or ``[tokens]``; any tensor with those + leading token axes is sharded identically. + seq_dim: Sequence axis in the pre-CP loss tensors, or ``None`` when + the weights do not follow the model's token axes. + + Returns: + Loss tensors in the model output's CP-local token layout. Non-token + tensors are returned unchanged; the input mapping is not mutated. + """ + weights = loss_inputs["weights"] + layout = sharder.shard_layout + if self._cp_size() == 1 and layout is None: + return {name: value for name, value in loss_inputs.items() if name != "labels"} + layout_changed = self._cp_size() > 1 or ( + layout is not None + and ( + layout.input_row_shape is not None + or layout.input_token_stream_positions is not None + or layout.original_seq_len != layout.padded_seq_len + ) + ) + if seq_dim is None: + if layout_changed: + raise ValueError("context-parallel loss weights must match the model's token axes") + return dict(loss_inputs) + + local: dict[str, torch.Tensor] = {} + for name, value in loss_inputs.items(): + if name == "labels": + continue + token_aligned = value.ndim >= weights.ndim and tuple(value.shape[: weights.ndim]) == tuple(weights.shape) + local[name] = sharder.shard_token_tensor(value, seq_dim=seq_dim, fill=0) if token_aligned else value + return local @staticmethod def _validate_collated_weights( @@ -322,6 +449,38 @@ def _to_device(value: Any, device: torch.device) -> Any: return value +def _loss_sequence_dim(model_inputs: dict[str, Any], weights: torch.Tensor) -> int | None: + """Find the sequence axis shared by primary model tokens and loss weights. + + Args: + model_inputs: Pre-CP model mapping whose ``input_ids`` has shape + ``[batch, sequence]`` or ``[tokens]``, or whose ``inputs_embeds`` + has shape ``[batch, sequence, hidden]``. + weights: Loss weights of shape ``[batch, sequence]`` or ``[tokens]``. + + Returns: + The sequence axis in ``weights``, or ``None`` when the layouts do not + describe the same token stream. + """ + primary = model_inputs.get("inputs_embeds", model_inputs.get("input_ids")) + if not isinstance(primary, torch.Tensor): + return None + if primary.ndim >= 2 and weights.ndim >= 2 and tuple(weights.shape[:2]) == tuple(primary.shape[:2]): + return 1 + if primary.ndim == 1 and weights.ndim == 1 and weights.shape == primary.shape: + return 0 + return None + + +def _is_final_thd(model_inputs: dict[str, Any]) -> bool: + """Return whether a CP1 caller already supplied the final flat THD layout.""" + if model_inputs.get("qkv_format") != "thd" or "cu_seqlens" not in model_inputs: + return False + if "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs: + raise ValueError("THD inputs cannot contain both raw seq_lens and final cu_seqlens metadata") + return True + + def _detach(value: Any) -> Any: """Detach tensor leaves without changing an output record's structure.""" if isinstance(value, torch.Tensor): diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 7face8ee7e..7e29d61bd6 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -658,26 +658,6 @@ def setup(self): # Extract TE FP8 config from model backend (set after model construction) self.te_fp8 = self.model_parts[0].backend.te_fp8 if hasattr(self.model_parts[0], "backend") else None - # Packed and CP/Magi batches need recipe-owned input preparation, so - # they keep using the existing forward/backward path for now. - self.engine = None - if ( - not self.pp_enabled - and self.mesh_context.cp_size == 1 - and getattr(self.cfg.dataloader, "packing", None) is None - and not getattr(self.cfg.dataloader, "emits_thd", False) - and not self.magi.enabled - and getattr(self.loss_fn, "reduction", None) == "sum" - ): - self.engine = Engine( - self.model_parts[0], - device=self.dist_env.device, - mesh_context=self.mesh_context, - collate_fn=collate_prebatched, - context_fn=self.te_fp8.maybe_te_autocast if self.te_fp8 is not None else nullcontext, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) - if self.pp_enabled: self._configure_pipeline_loss_fn() @@ -715,6 +695,26 @@ def setup(self): # Tokenizer + model-derived values are runtime concerns: build them here and pass them to # each DataloaderConfig.build(); the configs themselves are resolved at the RecipeConfig boundary. _, self.tokenizer = _build_tokenizer(self.cfg.model, self.cfg.dataset) + model_has_mtp = not self.pp_enabled and any( + getattr(module, "mtp", None) is not None for module in self.model_parts[0].modules() + ) + self.engine = None + if ( + not self.pp_enabled + and not self.magi.enabled + and not (self.mesh_context.cp_size > 1 and model_has_mtp) + and not getattr(self.distributed_config, "calculate_per_token_loss", False) + and getattr(self.loss_fn, "reduction", None) == "sum" + ): + self.engine = Engine( + self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + collate_fn=collate_prebatched, + padding_token_id=(self.tokenizer.pad_token_id if self.tokenizer is not None else 0) or 0, + context_fn=self.te_fp8.maybe_te_autocast if self.te_fp8 is not None else nullcontext, + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), + ) attn_implementation = None if ( self.cfg.get("packed_sequence.packed_sequence_size", 0) > 0 @@ -1061,9 +1061,7 @@ def _calculate_eager_loss(self, output, labels, model_inputs, *, num_label_token def _make_engine_datum(self, batch): labels = batch["labels"] - model_inputs = filter_forward_kwargs( - self.model_parts[0], {key: value for key, value in batch.items() if key != "labels"} - ) + model_inputs = {key: value for key, value in batch.items() if key != "labels"} if isinstance(self.loss_fn, FusedLinearCrossEntropy): model_inputs["logits_to_keep"] = 1 return Datum( @@ -1071,11 +1069,11 @@ def _make_engine_datum(self, batch): loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, ) - def _engine_loss(self, output, loss_inputs, datums): + def _engine_loss(self, output, loss_inputs, _datums, model_inputs): return self._calculate_eager_loss( output, loss_inputs["labels"], - datums[0].model_inputs, + model_inputs, num_label_tokens=None, is_train=True, ) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index f21719a2fd..b536e6cd5f 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -628,22 +628,26 @@ def setup(self): self.dataloader = dataloader_build.dataloader self.processor = dataloader_build.processor - # Start with the eager path whose inputs need no recipe-owned sharding or - # packing preparation. PP, CP, packed, and Magi batches keep using the - # existing forward/backward path below. + model_has_mtp = not self.pp_enabled and any( + getattr(module, "mtp", None) is not None for module in self.model_parts[0].modules() + ) self.engine = None if ( not self.pp_enabled - and self.mesh_context.cp_size == 1 - and dataloader_config.packing is None and not self.magi.enabled + and not (self.mesh_context.cp_size > 1 and model_has_mtp) + and not getattr(self.distributed_config, "calculate_per_token_loss", False) and getattr(self.loss_fn, "reduction", None) == "sum" ): + padding_token_id = ( + getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0 + ) self.engine = Engine( self.model_parts[0], device=self.dist_env.device, mesh_context=self.mesh_context, collate_fn=collate_prebatched, + padding_token_id=padding_token_id, context_fn=self._cp_vision_frame_sharding_context, defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), ) @@ -818,6 +822,7 @@ def _calculate_eager_loss( model: nn.Module, num_label_tokens: int | None, is_train: bool, + cu_seqlens: torch.Tensor | None = None, log_drafter: bool = False, log_denominator: int | float | None = None, ) -> torch.Tensor: @@ -831,6 +836,7 @@ def _calculate_eager_loss( getattr(getattr(self, "cfg", None), "mtp", None), num_label_tokens=num_label_tokens, grad_reduce_group=grad_reduce_group, + cu_seqlens=cu_seqlens, ) return self._maybe_add_drafter_loss( @@ -852,8 +858,7 @@ def _make_engine_datum( ) -> Datum: """Wrap one processor-collated VLM batch as a Datum.""" labels = batch["labels"] - model = self.model_parts[0] - model_inputs = filter_forward_kwargs(model, {key: value for key, value in batch.items() if key != "labels"}) + model_inputs = {key: value for key, value in batch.items() if key != "labels"} if isinstance(self.loss_fn, FusedLinearCrossEntropy): model_inputs["logits_to_keep"] = 1 return Datum( @@ -871,6 +876,7 @@ def _engine_loss( out: Any, loss_inputs: dict[str, torch.Tensor], datums: Sequence[Datum], + model_inputs: dict[str, Any], ) -> torch.Tensor: """Return the local loss sum; Engine owns global normalization.""" return self._calculate_eager_loss( @@ -879,6 +885,7 @@ def _engine_loss( model=self.model_parts[0], num_label_tokens=None, is_train=True, + cu_seqlens=model_inputs.get("cu_seqlens"), log_drafter=bool(datums[0].loss_fn_inputs["log_drafter"].item()), log_denominator=float(datums[0].loss_fn_inputs["log_denominator"].item()), ) @@ -1037,6 +1044,7 @@ def _forward_backward_step( model=model, num_label_tokens=num_label_tokens, is_train=is_train, + cu_seqlens=batch.get("cu_seqlens"), # Log once per remote-logging step on the first microbatch. log_drafter=(idx == 0 and self.step_scheduler.is_remote_logging_step), ) @@ -1085,7 +1093,13 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() engine = getattr(self, "engine", None) - use_engine = engine is not None and num_label_tokens > 0 + unsupported_thd_mrope = self._get_cp_group_size() > 1 and any( + batch.get("qkv_format") == "thd" + and isinstance(batch.get("position_ids"), torch.Tensor) + and batch["position_ids"].ndim == 3 + for batch in batches + ) + use_engine = engine is not None and num_label_tokens > 0 and not unsupported_thd_mrope if use_engine: reporting_loss, _ = engine.forward_backward( [ diff --git a/tests/unit_tests/distributed/test_cp_utils.py b/tests/unit_tests/distributed/test_cp_utils.py index c633337cb8..29702cc5d5 100644 --- a/tests/unit_tests/distributed/test_cp_utils.py +++ b/tests/unit_tests/distributed/test_cp_utils.py @@ -916,13 +916,19 @@ class _FakeMagiState: enabled = True domain = "llm" - cp_size = 2 - def __init__(self, local_indices): + def __init__(self, local_indices, *, cp_size=2, flatten=False): self._local_indices = local_indices + self.cp_size = cp_size + self._flatten = flatten def make_cp_batch(self, cp_mesh, batch, *, return_local_indices=False, **kwargs): - prepped = {"prepared": True} + prepped = dict(batch) + if self._flatten: + flat = batch["input_ids"].reshape(-1) + prepped["input_ids"] = flat if self._local_indices is None else flat[: self._local_indices.numel()] + elif self._local_indices is not None: + prepped["input_ids"] = batch["input_ids"][:, : self._local_indices.numel()] return (prepped, self._local_indices) if return_local_indices else prepped @@ -955,7 +961,7 @@ def test_magi_sharder_captures_packed_row_shape(): strategy = _cu._resolve_cp_sharder( cp2, None, - magi=_FakeMagiState(torch.tensor([[0, 3]])), + magi=_FakeMagiState(torch.tensor([0, 3]), flatten=True), is_thd=True, num_chunks=1, seq_lens_padding_value=-1000, @@ -969,6 +975,69 @@ def test_magi_sharder_captures_packed_row_shape(): assert torch.equal(sharder.shard_token_tensor(rows), torch.tensor([10.0, 40.0])) +def test_magi_cp1_vlm_no_dispatch_has_identity_token_layout(): + """VLM Magi does not dispatch at CP1, but its token verbs remain identity.""" + cp1 = _DummySubMesh(1) + strategy = _cu._resolve_cp_sharder( + cp1, + None, + magi=_FakeMagiState(None, cp_size=1), + is_thd=False, + num_chunks=1, + seq_lens_padding_value=-1000, + model=None, + ) + sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=1, tp_size=1)) + sharder.shard({"input_ids": torch.tensor([[1, 2, 3]])}) + + assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (3, 3) + weights = torch.tensor([[10.0, 20.0, 30.0]]) + assert torch.equal(sharder.shard_token_tensor(weights), weights) + + +def test_magi_cp1_custom_thd_flattens_with_identity_token_layout(): + """A CP1 custom Magi BxS -> T conversion exposes an identity THD map.""" + cp1 = _DummySubMesh(1) + strategy = _cu._resolve_cp_sharder( + cp1, + None, + magi=_FakeMagiState(None, cp_size=1, flatten=True), + is_thd=True, + num_chunks=1, + seq_lens_padding_value=-1000, + model=None, + ) + sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=1, tp_size=1)) + _, model_inputs = sharder.shard({"input_ids": torch.tensor([[1, 2], [3, 4]])}) + + assert model_inputs["input_ids"].shape == (4,) + assert sharder.shard_layout.input_row_shape == (2, 2) + assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (4, 4) + rows = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + assert torch.equal(sharder.shard_token_tensor(rows), rows.reshape(-1)) + + +def test_magi_packed_dispatch_padding_uses_zero_weight_sentinel(): + """Packed Magi can flatten BxS, add dispatch padding, and select zero-filled weights.""" + cp2 = _DummySubMesh(2) + strategy = _cu._resolve_cp_sharder( + cp2, + None, + magi=_FakeMagiState(torch.tensor([0, 3, 4]), flatten=True), + is_thd=True, + num_chunks=1, + seq_lens_padding_value=-1000, + model=None, + ) + sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=2, tp_size=1)) + sharder.shard({"input_ids": torch.tensor([[1, 2], [3, 4]])}) + + layout = sharder.shard_layout + assert (layout.original_seq_len, layout.padded_seq_len, layout.input_row_shape) == (4, 6, (2, 2)) + rows = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + assert torch.equal(sharder.shard_token_tensor(rows, fill=0), torch.tensor([10.0, 40.0, 0.0])) + + def test_make_cp_batch_for_te_identity_indices_without_cp(): """At cp<=1 the THD stream is unsharded, so the index map is the identity.""" batch = { diff --git a/tests/unit_tests/distributed/test_magi_attn_utils.py b/tests/unit_tests/distributed/test_magi_attn_utils.py index c7206377ff..d805710692 100644 --- a/tests/unit_tests/distributed/test_magi_attn_utils.py +++ b/tests/unit_tests/distributed/test_magi_attn_utils.py @@ -22,7 +22,8 @@ from __future__ import annotations -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch @@ -325,6 +326,48 @@ def test_raises_when_layout_mismatches_input(self): with pytest.raises(ValueError, match="!= flat input length 1024"): mu._packed_cp_doc_seqlens(batch, 1024) + def test_packed_dispatch_returns_token_indices_not_rope_positions(self, monkeypatch): + """The sharder map follows dispatch itself, not per-document RoPE ids.""" + expected_key = object() + order = torch.tensor([2, 5, 6, 7]) + dispatch_calls = [] + + def dispatch(value, *, key, pad_value=0): + assert key is expected_key + dispatch_calls.append((value.clone(), pad_value)) + padding = value.new_full((2,), pad_value) + return torch.cat((value, padding)).index_select(0, order) + + api = ModuleType("magi_attention.api") + api.dispatch = dispatch + api.get_position_ids = lambda key: torch.tensor([2, 0, 1, 0]) + package = ModuleType("magi_attention") + package.api = api + monkeypatch.setitem(sys.modules, "magi_attention", package) + monkeypatch.setitem(sys.modules, "magi_attention.api", api) + monkeypatch.setattr(mu, "build_flex_key", lambda *args, **kwargs: expected_key) + + model = SimpleNamespace(config=SimpleNamespace(num_attention_heads=2, num_key_value_heads=2, head_dim=4)) + batch = { + "input_ids": torch.tensor([10, 11, 12, 13, 14, 15]), + "labels": torch.tensor([20, 21, 22, 23, 24, 25]), + "cu_seqlens_padded": torch.tensor([0, 3, 6]), + } + default_result = mu.magi_prepare_packed_cp(model, batch, _FakeGroup(2)) + assert len(default_result) == 2 + assert not any(torch.equal(value, torch.arange(6)) and pad_value == 6 for value, pad_value in dispatch_calls) + + dispatch_calls.clear() + out, returned_key, local_indices = mu.magi_prepare_packed_cp( + model, batch, _FakeGroup(2), return_local_indices=True + ) + + assert returned_key is expected_key + assert torch.equal(out["position_ids"], torch.tensor([2, 0, 1, 0])) + assert torch.equal(local_indices, torch.tensor([2, 5, 6, 6])) + assert torch.equal(out["labels"], torch.tensor([22, 25, -100, -100])) + assert any(torch.equal(value, torch.arange(6)) and pad_value == 6 for value, pad_value in dispatch_calls) + class TestActiveStateAccessors: def test_attn_spec_roundtrip(self): diff --git a/tests/unit_tests/distributed/test_utils.py b/tests/unit_tests/distributed/test_utils.py index bb59ab5d04..928e5728e9 100644 --- a/tests/unit_tests/distributed/test_utils.py +++ b/tests/unit_tests/distributed/test_utils.py @@ -146,3 +146,23 @@ class Plain(torch.nn.Linear): # entering/exiting the context must be a no-op with ctx: pass + + +def test_get_sync_ctx_uses_no_sync_capability_for_non_final_microbatch(patch_dist): + class NoSyncModel: + def __init__(self): + self.no_sync_calls = 0 + + def no_sync(self): + self.no_sync_calls += 1 + return du.nullcontext() + + model = NoSyncModel() + + with du.get_sync_ctx(model, is_optim_step=False, defer_fsdp_grad_sync=False): + pass + assert model.no_sync_calls == 1 + + with du.get_sync_ctx(model, is_optim_step=True, defer_fsdp_grad_sync=False): + pass + assert model.no_sync_calls == 1 diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 1eb35fde68..0e7458b861 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -22,6 +22,7 @@ import torch.nn as nn from nemo_automodel.components.config.loader import ConfigNode +from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.vlm.pp_media import ( VLM_PP_MEDIA_KEY, chunk_step3_media, @@ -2983,6 +2984,69 @@ def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch): trainer.setup() assert cfg.model.backend.rope_fusion is False + assert trainer.engine is not None + + +def test_vlm_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=2, rope_fusion=True) + _patch_vlm_setup_minimals(monkeypatch, cp_size=2) + model = DummyModel() + model.mtp = nn.Identity() + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.build_model", lambda *args, **kwargs: model) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert trainer.engine is None + + +def test_vlm_setup_keeps_engine_disabled_for_magi(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.setup_magi", + lambda *args, **kwargs: SimpleNamespace(enabled=True), + ) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert trainer.engine is None + + +def test_vlm_engine_loss_uses_final_thd_sequence_boundaries(monkeypatch): + recipe = object.__new__(FinetuneRecipeForVLM) + recipe.model_parts = [nn.Identity()] + recipe.loss_fn = object() + recipe.cfg = SimpleNamespace(mtp=None) + recipe.dist_env = SimpleNamespace(is_main=False) + recipe._get_dp_group = lambda include_cp=False: None + recipe._maybe_add_drafter_loss = lambda **kwargs: kwargs["base_loss"] + seen = {} + + def fake_causal_lm_loss(*args, **kwargs): + seen.update(kwargs) + return torch.tensor(1.0) + + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.causal_lm_loss", fake_causal_lm_loss) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + datum = Datum( + model_inputs={"input_ids": torch.arange(5)}, + loss_fn_inputs={ + "weights": torch.ones(5), + "log_drafter": torch.tensor(False), + "log_denominator": torch.tensor(5.0), + }, + ) + + recipe._engine_loss( + object(), + {"labels": torch.arange(5)}, + [datum], + {"cu_seqlens": cu_seqlens}, + ) + + assert seen["cu_seqlens"] is cu_seqlens def test_vlm_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 30aeb62580..3a4ad9e9ce 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -2253,6 +2253,34 @@ def test_rope_fusion_disabled_when_cp_gt_1(monkeypatch): trainer.setup() assert cfg.model.backend.rope_fusion is False + assert trainer.engine is not None + + +def test_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): + cfg = _minimal_cfg_with_rope_fusion(cp_size=2, rope_fusion=True) + _patch_setup_minimals_with_cp(monkeypatch, cp_size=2) + model = DummyModel() + model.mtp = nn.Identity() + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.build_model", lambda *args, **kwargs: model) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert trainer.engine is None + + +def test_setup_keeps_engine_disabled_for_magi(monkeypatch): + cfg = _minimal_cfg_with_rope_fusion(cp_size=1, rope_fusion=True) + _patch_setup_minimals_with_cp(monkeypatch, cp_size=1) + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.setup_magi", + lambda *args, **kwargs: SimpleNamespace(enabled=True, hf_dispatch=False), + ) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert trainer.engine is None def test_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index dcee9b9712..e565cdad18 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -29,6 +29,14 @@ from nemo_automodel import Datum as PublicDatum from nemo_automodel import Engine as PublicEngine from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.distributed.config import MegatronFSDPConfig +from nemo_automodel.components.distributed.context_parallel.sharder import ( + ContextParallelSharder, + contiguous_local_indices, + shard_batch_contiguous, +) +from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes +from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh from nemo_automodel.components.loss.causal_lm import causal_lm_loss from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler @@ -46,13 +54,50 @@ def forward(self, input_ids: torch.Tensor, **_) -> torch.Tensor: return input_ids.to(torch.float32) * self.weight +class _SubMesh: + def __init__(self, size, rank=0): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def get_local_rank(self): + return self._rank + + def get_group(self): + return None + + +class _CPMesh(dict): + def __init__(self, size, rank): + super().__init__(cp=_SubMesh(size, rank), tp=_SubMesh(1)) + self.mesh_dim_names = ("cp", "tp") + + +class _DDPWithCP(nn.parallel.DistributedDataParallel): + def prepare_model_inputs_for_cp(self, batch, *, num_chunks): + return self.module.prepare_model_inputs_for_cp(batch, num_chunks=num_chunks) + + +class _DistributedCPModel(ScaleModel): + def prepare_model_inputs_for_cp(self, _batch, *, num_chunks): + assert num_chunks == 1 + return { + "cp_sharder": ContextParallelSharder( + shard_batch=partial(shard_batch_contiguous, pad_multiple=1), + local_token_global_indices=contiguous_local_indices, + ) + } + + def _datum(values, weights=None) -> Datum: values = torch.tensor(values, dtype=torch.long) weights = torch.ones_like(values, dtype=torch.float32) if weights is None else torch.tensor(weights) return Datum(model_inputs={"input_ids": values}, loss_fn_inputs={"weights": weights}) -def _identity_loss(output, _loss_inputs, _datums): +def _identity_loss(output, _loss_inputs, _datums, _model_inputs): return output @@ -78,10 +123,33 @@ def test_forward_backward_uses_one_denominator_for_the_window(): assert model.forward_calls == 2 -def test_pre_thd_packed_collater_requires_model_ready_inputs(): +def test_raw_thd_packed_collater_is_prepared_by_context_parallel_sharder(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + seen = {} + + def loss_fn(output, inputs, _datums, model_inputs): + seen.update(model_inputs) + assert inputs["weights"].shape == output.shape == (3,) + return output + + loss, _ = Engine( + model, + device="cpu", + collate_fn=partial(collate_datums, packed=True), + ).forward_backward([[_datum([1, 2]), _datum([3])]], loss_fn) + + assert loss.item() == pytest.approx(2.0) + assert "seq_lens" not in seen + assert "seq_lens_padded" not in seen + assert seen["qkv_format"] == "thd" + assert seen["cu_seqlens"].tolist() == [0, 2, 3] + + +def test_raw_thd_requires_a_thd_capable_context_parallel_sharder(): model = ScaleModel() - with pytest.raises(ValueError, match="final model-ready THD"): + with pytest.raises(ValueError, match="could not prepare raw THD inputs"): Engine( model, device="cpu", @@ -91,6 +159,78 @@ def test_pre_thd_packed_collater_requires_model_ready_inputs(): assert model.forward_calls == 0 +def test_context_parallel_shards_model_and_rl_loss_inputs_together(): + cp_context_active = False + + @contextmanager + def cp_context(): + nonlocal cp_context_active + cp_context_active = True + try: + yield + finally: + cp_context_active = False + + def shard_batch(*args, **kwargs): + _, batch, layout = shard_batch_contiguous(*args, pad_multiple=1, **kwargs) + return cp_context, batch, layout + + class CPModel(ScaleModel): + def prepare_model_inputs_for_cp(self, batch, *, num_chunks): + assert num_chunks == 1 + return { + "cp_sharder": ContextParallelSharder( + shard_batch=shard_batch, + local_token_global_indices=contiguous_local_indices, + ) + } + + def forward(self, *args, **kwargs): + assert cp_context_active + return super().forward(*args, **kwargs) + + model = CPModel() + mesh = _CPMesh(size=2, rank=1) + mesh_context = SimpleNamespace(pp_size=1, cp_size=2, device_mesh=mesh) + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4, 5, 6]])}, + loss_fn_inputs={ + "target_tokens": torch.tensor([[11, 12, 13, 14, 15, 16]]), + "weights": torch.ones(1, 6), + "advantages": torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]]), + }, + ) + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + padding_token_id=9, + ) + # This is a layout-only CPU test with a fake mesh. Distributed CP loss and + # gradient scaling are covered separately with a real process group. + engine._dp_group_and_size = lambda: (None, 1) + engine._gradient_group_and_size = lambda _group, _size: (None, 1) + model.weight.register_hook( + lambda grad: grad if cp_context_active else pytest.fail("CP context ended before backward") + ) + + def loss_fn(output, inputs, _datums, model_inputs): + assert cp_context_active + assert model_inputs["input_ids"].tolist() == [[5, 6, 9, 9]] + assert inputs["target_tokens"].tolist() == [[15, 16, 0, 0]] + assert inputs["weights"].tolist() == [[1.0, 1.0, 0.0, 0.0]] + torch.testing.assert_close(inputs["advantages"], torch.tensor([[0.5, 0.6, 0.0, 0.0]])) + return output + + loss, _ = engine.forward_backward([[datum]], loss_fn) + + assert loss.item() == pytest.approx(11 / 6) + assert model.weight.grad.item() == pytest.approx(11 / 6) + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(2.0) + assert not cp_context_active + + def _model_ready_packed_collate(datums): model_inputs, loss_inputs = collate_datums(datums, packed=True) lengths = torch.tensor([datum.seq_len for datum in datums], dtype=torch.int32) @@ -114,7 +254,7 @@ def test_packed_rl_callback_keeps_per_datum_sequence_boundaries(): first.loss_fn_inputs["sequence_scale"] = torch.tensor(2.0) second.loss_fn_inputs["sequence_scale"] = torch.tensor(0.5) - def sequence_loss(output, _loss_inputs, datums): + def sequence_loss(output, _loss_inputs, datums, _model_inputs): chunks = output.squeeze(0).split([datum.seq_len for datum in datums]) losses = torch.cat([chunk * datum.loss_fn_inputs["sequence_scale"] for chunk, datum in zip(chunks, datums)]) outputs = [ @@ -148,7 +288,7 @@ def test_weights_mask_loss_and_denominator(): def test_loss_fn_outputs_follow_datum_order_and_are_detached(): model = ScaleModel() - def loss_with_outputs(output, _loss_inputs, datums): + def loss_with_outputs(output, _loss_inputs, datums, _model_inputs): return output, [ {"first_token": datum.input_ids[0], "model_value": output[index].sum()} for index, datum in enumerate(datums) @@ -168,13 +308,13 @@ def test_loss_fn_outputs_must_align_with_datums(): with pytest.raises(ValueError, match="one mapping per Datum"): Engine(model, device="cpu").forward_backward( [[_datum([1]), _datum([2])]], - lambda output, _inputs, _datums: (output, [{"only": "one"}]), + lambda output, _inputs, _datums, _model_inputs: (output, [{"only": "one"}]), ) assert model.weight.grad is None def test_loss_fn_outputs_must_be_consistent_across_the_window(): - def inconsistent_outputs(output, _inputs, datums): + def inconsistent_outputs(output, _inputs, datums, _model_inputs): if datums[0].input_ids[0].item() == 1: return output, [{"value": output.sum()}] return output @@ -213,7 +353,7 @@ def test_raw_output_and_loss_inputs_support_an_rl_loss_callback(): ) model = TinyLM() - def policy_loss(logits, inputs, datums): + def policy_loss(logits, inputs, datums, _model_inputs): assert len(datums) == 1 new_logprobs = -F.cross_entropy( logits.flatten(0, 1), @@ -255,7 +395,7 @@ def test_causal_lm_loss_matches_a_manual_accumulation_window(): for input_ids, labels in batches ] - def engine_loss(output, inputs, _datums): + def engine_loss(output, inputs, _datums, _model_inputs): return causal_lm_loss( loss_fn, model, @@ -327,7 +467,7 @@ def forward(self, input_ids, **kwargs): model = ContextModel() model.weight.register_hook(lambda grad: grad if active else pytest.fail("context ended before backward")) - def loss_fn(output, _inputs, _datums): + def loss_fn(output, _inputs, _datums, _model_inputs): assert active return output @@ -381,7 +521,7 @@ def test_model_specific_collater_keeps_multimodal_inputs_and_gradients(): model = TinyVLM() Engine(model, device="cpu", collate_fn=_vlm_collate).forward_backward( [datums], - lambda output, _inputs, _datums: output, + lambda output, _inputs, _datums, _model_inputs: output, ) assert model.text.weight.grad is not None @@ -425,7 +565,7 @@ def test_scalar_loss_is_a_local_weighted_sum_numerator(): loss, _ = Engine(model, device="cpu").forward_backward( [[_datum([1, 100], [1.0, 0.0])], [_datum([3, 5], [0.5, 1.0])]], - lambda output, inputs, _datums: (output * inputs["weights"]).sum(), + lambda output, inputs, _datums, _model_inputs: (output * inputs["weights"]).sum(), ) assert loss.item() == pytest.approx(3.0) @@ -443,12 +583,11 @@ def test_zero_weights_fail_before_forward(): assert model.weight.grad is None -@pytest.mark.parametrize(("pp_size", "cp_size", "name"), [(2, 1, "pipeline"), (1, 2, "context")]) -def test_unsupported_parallelism_fails_before_forward(pp_size, cp_size, name): +def test_pipeline_parallelism_fails_before_forward(): model = ScaleModel() - mesh_context = SimpleNamespace(pp_size=pp_size, cp_size=cp_size) + mesh_context = SimpleNamespace(pp_size=2, cp_size=1) - with pytest.raises(NotImplementedError, match=name): + with pytest.raises(NotImplementedError, match="pipeline"): Engine(model, device="cpu", mesh_context=mesh_context).forward_backward( [[_datum([1])]], _identity_loss, @@ -457,12 +596,23 @@ def test_unsupported_parallelism_fails_before_forward(pp_size, cp_size, name): assert model.forward_calls == 0 +def test_megatron_fsdp_per_token_loss_mode_fails_before_forward(): + model = ScaleModel() + model.calculate_per_token_loss = True + + with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): + Engine(model, device="cpu").forward_backward([[_datum([1])]], _identity_loss) + + assert model.forward_calls == 0 + assert model.weight.grad is None + + def test_loss_shape_must_exactly_match_weights(): model = ScaleModel() with pytest.raises(ValueError, match="exactly the same shape"): Engine(model, device="cpu").forward_backward( [[_datum([1, 2])]], - lambda output, _inputs, _datums: output[:, :1], + lambda output, _inputs, _datums, _model_inputs: output[:, :1], ) assert model.weight.grad is None @@ -501,6 +651,83 @@ def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: dist.destroy_process_group() +def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size: int) -> None: + dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + cp_size = world_size // dp_size + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=dp_size, cp_size=cp_size), + world_size=world_size, + ) + model = _DDPWithCP(_DistributedCPModel()) + if dp_size == 1: + window = [ + [ + Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0, 0.0, 0.0]])}, + ) + ], + [ + Datum( + model_inputs={"input_ids": torch.tensor([[5, 6, 7, 8]])}, + loss_fn_inputs={"weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]])}, + ) + ], + ] + else: + dp_rank = get_flat_mesh(mesh_context.device_mesh, "dp").get_local_rank() + first = dp_rank * 4 + 1 + window = [ + [ + Datum( + model_inputs={"input_ids": torch.arange(first, first + 4).unsqueeze(0)}, + loss_fn_inputs={"weights": torch.ones(1, 4)}, + ) + ] + ] + + loss, _ = Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ).forward_backward(window, _identity_loss) + + assert loss.item() == pytest.approx(4.5) + assert model.module.weight.grad.item() == pytest.approx(4.5) + finally: + dist.destroy_process_group() + + +def _mismatched_context_parallel_weights_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=1, cp_size=world_size), + world_size=world_size, + ) + model = _DDPWithCP(_DistributedCPModel()) + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={"weights": torch.full((1, 4), float(rank + 1))}, + ) + + with pytest.raises(ValueError, match="identical full-sequence weights"): + Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ).forward_backward([[datum]], _identity_loss) + + assert model.module.forward_calls == 0 + finally: + dist.destroy_process_group() + + def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): mp.spawn( _distributed_worker, @@ -508,3 +735,30 @@ def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): nprocs=2, join=True, ) + + +def test_context_parallel_window_uses_dp_denominator_and_dp_cp_gradient_sum(tmp_path): + mp.spawn( + _context_parallel_worker, + args=(2, str(tmp_path / "engine_cp_init"), 1), + nprocs=2, + join=True, + ) + + +def test_data_and_context_parallel_composition_matches_global_reference(tmp_path): + mp.spawn( + _context_parallel_worker, + args=(4, str(tmp_path / "engine_dp_cp_init"), 2), + nprocs=4, + join=True, + ) + + +def test_context_parallel_replicas_require_the_same_full_sequence_weights(tmp_path): + mp.spawn( + _mismatched_context_parallel_weights_worker, + args=(2, str(tmp_path / "engine_cp_mismatch_init")), + nprocs=2, + join=True, + ) From 3939afa23b18f4df7bd04fb82a069869c03a7e72 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 15 Aug 2026 23:49:15 -0700 Subject: [PATCH 04/34] refactor(recipes): keep causal losses recipe-local Signed-off-by: HuiyingLi --- nemo_automodel/components/loss/causal_lm.py | 107 ------------------ nemo_automodel/recipes/llm/train_ft.py | 73 +++++++++--- nemo_automodel/recipes/vlm/finetune.py | 65 ++++++++--- .../recipes/test_finetune_vlm_helpers.py | 42 +++---- tests/unit_tests/recipes/test_train_ft.py | 4 +- tests/unit_tests/test_engine.py | 57 ---------- 6 files changed, 130 insertions(+), 218 deletions(-) delete mode 100644 nemo_automodel/components/loss/causal_lm.py diff --git a/nemo_automodel/components/loss/causal_lm.py b/nemo_automodel/components/loss/causal_lm.py deleted file mode 100644 index 2acd5bc141..0000000000 --- a/nemo_automodel/components/loss/causal_lm.py +++ /dev/null @@ -1,107 +0,0 @@ -# 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. - -"""Shared causal-LM loss calculation for the LLM and VLM recipes.""" - -from typing import Any - -import torch -import torch.distributed as dist -from torch import nn - -from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy -from nemo_automodel.components.loss.mtp import MTPLossConfig, calculate_mtp_loss -from nemo_automodel.components.loss.utils import _get_final_hidden_states, _get_lm_head_weight, calculate_loss - - -def causal_lm_loss( - loss_fn: nn.Module, - model: nn.Module, - output: Any, - labels: torch.Tensor, - mtp_config: MTPLossConfig | None, - *, - num_label_tokens: int | None, - grad_reduce_group: dist.ProcessGroup | None, - cu_seqlens: torch.Tensor | None = None, -) -> torch.Tensor: - """Return the main causal-LM loss plus an optional MTP loss. - - Args: - loss_fn: Configured causal-LM loss module. - model: Model that owns the LM head used by fused loss implementations. - output: Model output containing logits of shape ``[batch, sequence, - vocab]`` or final hidden states of shape ``[batch, sequence, - hidden]``. Optional MTP fields use the same batch and sequence - axes per prediction depth. - labels: Target token ids of shape ``[batch, sequence]`` or ``[tokens]`` - for a flattened THD stream. - mtp_config: MTP loss settings, required only when ``output`` contains - MTP predictions. - num_label_tokens: Global supervised-token denominator. ``None`` keeps - each loss as an unnormalized local sum for Engine normalization. - grad_reduce_group: Group that contributes independent fused-loss - shards, or ``None`` for an unsharded LM head. - cu_seqlens: Optional THD cumulative sequence offsets of shape - ``[num_sequences + 1]``. - - Returns: - Scalar causal-LM loss retaining its autograd graph. - """ - hidden_states = _get_final_hidden_states(output) - if isinstance(loss_fn, FusedLinearCrossEntropy) and hidden_states is None: - raise ValueError("FusedLinearCrossEntropy requires the model to output hidden states") - - lm_weight = ( - loss_fn.materialize_lm_weight( - _get_lm_head_weight(model), - grad_reduce_group=grad_reduce_group, - ) - if isinstance(loss_fn, FusedLinearCrossEntropy) - else None - ) - loss = calculate_loss( - loss_fn, - logits=getattr(output, "logits", output), - labels=labels, - model=model, - hidden_states=hidden_states, - lm_weight=lm_weight, - grad_reduce_group=grad_reduce_group, - num_label_tokens=num_label_tokens, - ) - - mtp_hidden = getattr(output, "mtp_per_depth_h", None) - mtp_logits = getattr(output, "mtp_per_depth_logits", None) - if mtp_hidden is None and mtp_logits is None: - return loss - if mtp_config is None: - raise ValueError("MTP model output requires an MTP loss config") - - scaling_factor = ( - mtp_config.scaling_factor if mtp_config.scaling_factor is not None else output.mtp_loss_scaling_factor - ) - return loss + calculate_mtp_loss( - loss_fn, - mtp_per_depth_h=mtp_hidden, - mtp_per_depth_logits=mtp_logits, - labels=labels, - model=model, - scaling_factor=scaling_factor, - num_label_tokens=num_label_tokens, - ignore_index=mtp_config.ignore_index, - cu_seqlens=cu_seqlens, - lm_weight=lm_weight, - grad_reduce_group=grad_reduce_group, - ) diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 7e29d61bd6..862aeec9b4 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -72,10 +72,12 @@ to_float_metrics, ) from nemo_automodel.components.loggers.wandb_utils import suppress_wandb_log_messages -from nemo_automodel.components.loss.causal_lm import causal_lm_loss from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy +from nemo_automodel.components.loss.mtp import calculate_mtp_loss +from nemo_automodel.components.loss.utils import _get_lm_head_weight, calculate_loss from nemo_automodel.components.quantization.fp8 import build_fp8_config +from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG from nemo_automodel.components.training.utils import ( count_tail_padding, @@ -1026,8 +1028,8 @@ def run_train_validation_loop(self): self._partial_cuda_graph_capture_pending = False # ------------------ helpers ------------------ - def _prepare_eager_batch(self, batch): - """Move and CP-prepare one batch before an eager model call.""" + def _prepare_microbatch(self, batch): + """Move and CP-prepare one batch before model execution.""" batch = { k: ( {dk: dv.to(self.dist_env.device, non_blocking=True) for dk, dv in v.items() if dv is not None} @@ -1046,17 +1048,56 @@ def _prepare_eager_batch(self, batch): train_ctx, batch = cp_sharder.shard(batch) return train_ctx, batch, batch.pop("labels") - def _calculate_eager_loss(self, output, labels, model_inputs, *, num_label_tokens, is_train): - """Compute the shared CE and optional MTP loss for eager execution.""" - return causal_lm_loss( + def _compute_causal_lm_loss(self, output, labels, model_inputs, *, num_label_tokens, is_train): + """Compute the recipe's causal-LM and optional MTP loss.""" + model = self.model_parts[0] + grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None + hidden_states = get_final_hidden_states(output) + if isinstance(self.loss_fn, FusedLinearCrossEntropy) and hidden_states is None: + raise ValueError("FusedLinearCrossEntropy requires the model to output hidden states") + + lm_weight = ( + self.loss_fn.materialize_lm_weight( + _get_lm_head_weight(model), + grad_reduce_group=grad_reduce_group, + ) + if isinstance(self.loss_fn, FusedLinearCrossEntropy) + else None + ) + loss = calculate_loss( self.loss_fn, - self.model_parts[0], - output, - labels, - getattr(self.cfg, "mtp", None), + logits=getattr(output, "logits", output), + labels=labels, + model=model, + hidden_states=hidden_states, + lm_weight=lm_weight, + num_label_tokens=num_label_tokens, + grad_reduce_group=grad_reduce_group, + ) + + mtp_hidden = getattr(output, "mtp_per_depth_h", None) + mtp_logits = getattr(output, "mtp_per_depth_logits", None) + if mtp_hidden is None and mtp_logits is None: + return loss + + mtp_config = getattr(self.cfg, "mtp", None) + if mtp_config is None: + raise ValueError("MTP model output requires an MTP loss config") + scaling_factor = ( + mtp_config.scaling_factor if mtp_config.scaling_factor is not None else output.mtp_loss_scaling_factor + ) + return loss + calculate_mtp_loss( + self.loss_fn, + mtp_per_depth_h=mtp_hidden, + mtp_per_depth_logits=mtp_logits, + labels=labels, + model=model, + scaling_factor=scaling_factor, num_label_tokens=num_label_tokens, - grad_reduce_group=self._get_dp_group(include_cp=True) if is_train else None, + ignore_index=mtp_config.ignore_index, cu_seqlens=model_inputs.get("cu_seqlens"), + lm_weight=lm_weight, + grad_reduce_group=grad_reduce_group, ) def _make_engine_datum(self, batch): @@ -1069,8 +1110,8 @@ def _make_engine_datum(self, batch): loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, ) - def _engine_loss(self, output, loss_inputs, _datums, model_inputs): - return self._calculate_eager_loss( + def _engine_loss_fn(self, output, loss_inputs, _datums, model_inputs): + return self._compute_causal_lm_loss( output, loss_inputs["labels"], model_inputs, @@ -1088,7 +1129,7 @@ def _forward_backward_step( num_batches, is_train: bool = True, ): - train_ctx, batch, labels = self._prepare_eager_batch(batch) + train_ctx, batch, labels = self._prepare_microbatch(batch) fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext() if self.pp_enabled: @@ -1157,7 +1198,7 @@ def _forward_backward_step( out = model(logits_to_keep=1, **batch) else: out = model(**batch) - local_loss = self._calculate_eager_loss( + local_loss = self._compute_causal_lm_loss( out, labels, batch, @@ -1204,7 +1245,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): if use_engine: reporting_loss, _ = engine.forward_backward( [[self._make_engine_datum(batch)] for batch in batches], - self._engine_loss, + self._engine_loss_fn, ) else: # Engine requires a positive global weight sum. Keep the existing diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index b536e6cd5f..626cc92925 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -67,10 +67,10 @@ to_float_metrics, ) from nemo_automodel.components.loggers.wandb_utils import suppress_wandb_log_messages -from nemo_automodel.components.loss.causal_lm import causal_lm_loss from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy -from nemo_automodel.components.loss.utils import calculate_loss +from nemo_automodel.components.loss.mtp import calculate_mtp_loss +from nemo_automodel.components.loss.utils import _get_lm_head_weight, calculate_loss from nemo_automodel.components.quantization.fp8 import build_fp8_config from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG @@ -814,12 +814,11 @@ def _maybe_add_drafter_loss( ) return total_loss - def _calculate_eager_loss( + def _compute_vlm_loss( self, *, out: Any, labels: torch.Tensor, - model: nn.Module, num_label_tokens: int | None, is_train: bool, cu_seqlens: torch.Tensor | None = None, @@ -827,18 +826,54 @@ def _calculate_eager_loss( log_denominator: int | float | None = None, ) -> torch.Tensor: """Compute base, MTP, and optional joint-drafter losses.""" + model = self.model_parts[0] grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None - loss = causal_lm_loss( + hidden_states = get_final_hidden_states(out) + if isinstance(self.loss_fn, FusedLinearCrossEntropy) and hidden_states is None: + raise ValueError("FusedLinearCrossEntropy requires the model to output hidden states") + + lm_weight = ( + self.loss_fn.materialize_lm_weight( + _get_lm_head_weight(model), + grad_reduce_group=grad_reduce_group, + ) + if isinstance(self.loss_fn, FusedLinearCrossEntropy) + else None + ) + loss = calculate_loss( self.loss_fn, - model, - out, - labels, - getattr(getattr(self, "cfg", None), "mtp", None), + logits=getattr(out, "logits", out), + labels=labels, + model=model, + hidden_states=hidden_states, + lm_weight=lm_weight, num_label_tokens=num_label_tokens, grad_reduce_group=grad_reduce_group, - cu_seqlens=cu_seqlens, ) + mtp_hidden = getattr(out, "mtp_per_depth_h", None) + mtp_logits = getattr(out, "mtp_per_depth_logits", None) + if mtp_hidden is not None or mtp_logits is not None: + mtp_config = getattr(getattr(self, "cfg", None), "mtp", None) + if mtp_config is None: + raise ValueError("MTP model output requires an MTP loss config") + scaling_factor = ( + mtp_config.scaling_factor if mtp_config.scaling_factor is not None else out.mtp_loss_scaling_factor + ) + loss = loss + calculate_mtp_loss( + self.loss_fn, + mtp_per_depth_h=mtp_hidden, + mtp_per_depth_logits=mtp_logits, + labels=labels, + model=model, + scaling_factor=scaling_factor, + num_label_tokens=num_label_tokens, + ignore_index=mtp_config.ignore_index, + cu_seqlens=cu_seqlens, + lm_weight=lm_weight, + grad_reduce_group=grad_reduce_group, + ) + return self._maybe_add_drafter_loss( out=out, base_loss=loss, @@ -871,7 +906,7 @@ def _make_engine_datum( }, ) - def _engine_loss( + def _engine_loss_fn( self, out: Any, loss_inputs: dict[str, torch.Tensor], @@ -879,10 +914,9 @@ def _engine_loss( model_inputs: dict[str, Any], ) -> torch.Tensor: """Return the local loss sum; Engine owns global normalization.""" - return self._calculate_eager_loss( + return self._compute_vlm_loss( out=out, labels=loss_inputs["labels"], - model=self.model_parts[0], num_label_tokens=None, is_train=True, cu_seqlens=model_inputs.get("cu_seqlens"), @@ -1038,10 +1072,9 @@ def _forward_backward_step( else: out = model(**batch) - local_loss = self._calculate_eager_loss( + local_loss = self._compute_vlm_loss( out=out, labels=labels, - model=model, num_label_tokens=num_label_tokens, is_train=is_train, cu_seqlens=batch.get("cu_seqlens"), @@ -1112,7 +1145,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ] for index, batch in enumerate(batches) ], - self._engine_loss, + self._engine_loss_fn, ) else: # The eager Engine requires a positive global weight sum. Preserve diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 0e7458b861..154625fd64 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -22,7 +22,6 @@ import torch.nn as nn from nemo_automodel.components.config.loader import ConfigNode -from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.vlm.pp_media import ( VLM_PP_MEDIA_KEY, chunk_step3_media, @@ -433,7 +432,7 @@ def fake_calculate_loss(*args, **kwargs): ) calculate_mock = MagicMock(side_effect=fake_calculate_loss) - monkeypatch.setattr("nemo_automodel.components.loss.causal_lm.calculate_loss", calculate_mock) + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_loss", calculate_mock) grad_clip_mock = MagicMock(return_value=2.5) monkeypatch.setattr( @@ -486,7 +485,7 @@ def make_thd_batch(model, device_mesh, batch, **kwargs): monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ContextParallelSharder", make_thd_batch) monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.get_sync_ctx", lambda *args, **kwargs: nullcontext()) monkeypatch.setattr( - "nemo_automodel.components.loss.causal_lm.calculate_loss", + "nemo_automodel.recipes.vlm.finetune.calculate_loss", lambda *args, **kwargs: torch.tensor(1.0, requires_grad=True), ) @@ -3014,38 +3013,41 @@ def test_vlm_setup_keeps_engine_disabled_for_magi(monkeypatch): assert trainer.engine is None -def test_vlm_engine_loss_uses_final_thd_sequence_boundaries(monkeypatch): +def test_vlm_compute_loss_uses_final_thd_sequence_boundaries(monkeypatch): recipe = object.__new__(FinetuneRecipeForVLM) recipe.model_parts = [nn.Identity()] recipe.loss_fn = object() - recipe.cfg = SimpleNamespace(mtp=None) + recipe.cfg = SimpleNamespace(mtp=SimpleNamespace(scaling_factor=0.5, ignore_index=-100)) recipe.dist_env = SimpleNamespace(is_main=False) recipe._get_dp_group = lambda include_cp=False: None recipe._maybe_add_drafter_loss = lambda **kwargs: kwargs["base_loss"] seen = {} - def fake_causal_lm_loss(*args, **kwargs): + def fake_mtp_loss(*args, **kwargs): seen.update(kwargs) - return torch.tensor(1.0) + return torch.tensor(2.0) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.causal_lm_loss", fake_causal_lm_loss) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.calculate_loss", + lambda *args, **kwargs: torch.tensor(1.0), + ) + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_mtp_loss", fake_mtp_loss) cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) - datum = Datum( - model_inputs={"input_ids": torch.arange(5)}, - loss_fn_inputs={ - "weights": torch.ones(5), - "log_drafter": torch.tensor(False), - "log_denominator": torch.tensor(5.0), - }, + out = SimpleNamespace( + logits=torch.zeros(5, 8), + mtp_per_depth_logits=[torch.zeros(5, 8)], + mtp_loss_scaling_factor=0.25, ) - recipe._engine_loss( - object(), - {"labels": torch.arange(5)}, - [datum], - {"cu_seqlens": cu_seqlens}, + loss = recipe._compute_vlm_loss( + out=out, + labels=torch.arange(5), + num_label_tokens=None, + is_train=True, + cu_seqlens=cu_seqlens, ) + assert loss.item() == pytest.approx(3.0) assert seen["cu_seqlens"] is cu_seqlens diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 3a4ad9e9ce..4f297e105c 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -2618,8 +2618,8 @@ def _fake_calc_loss( "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", lambda device_mesh, batch, *a, **k: (nullcontext, batch, None), ) - monkeypatch.setattr("nemo_automodel.components.loss.causal_lm.calculate_loss", _fake_calc_loss) - monkeypatch.setattr("nemo_automodel.components.loss.causal_lm._get_final_hidden_states", lambda out: None) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calc_loss) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_final_hidden_states", lambda out: None) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_sync_ctx", lambda *a, **k: nullcontext()) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.filter_forward_kwargs", lambda model, batch: batch) diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index e565cdad18..9959f01ee7 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -37,8 +37,6 @@ ) from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh -from nemo_automodel.components.loss.causal_lm import causal_lm_loss -from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler from nemo_automodel.engine import Engine, collate_prebatched @@ -336,11 +334,6 @@ def forward(self, input_ids, **_): return self.output(self.embedding(input_ids)) -class TinyCausalLM(TinyLM): - def forward(self, input_ids, **_): - return SimpleNamespace(logits=super().forward(input_ids)) - - def test_raw_output_and_loss_inputs_support_an_rl_loss_callback(): datum = Datum( model_inputs={"input_ids": torch.tensor([1, 2, 3])}, @@ -373,56 +366,6 @@ def policy_loss(logits, inputs, datums, _model_inputs): assert model.output.weight.grad is not None -def test_causal_lm_loss_matches_a_manual_accumulation_window(): - torch.manual_seed(7) - model = TinyCausalLM() - reference = TinyCausalLM() - reference.load_state_dict(model.state_dict()) - batches = [ - (torch.tensor([[1, 2, 3]]), torch.tensor([[2, 3, -100]])), - (torch.tensor([[4, 5]]), torch.tensor([[5, 6]])), - ] - loss_fn = MaskedCrossEntropy() - mtp_config = SimpleNamespace(scaling_factor=None, ignore_index=-100) - - window = [ - [ - Datum( - model_inputs={"input_ids": input_ids}, - loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, - ) - ] - for input_ids, labels in batches - ] - - def engine_loss(output, inputs, _datums, _model_inputs): - return causal_lm_loss( - loss_fn, - model, - output, - inputs["labels"], - mtp_config, - num_label_tokens=None, - grad_reduce_group=None, - ) - - actual_loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward(window, engine_loss) - - denominator = sum((labels != -100).sum() for _, labels in batches) - reference_loss = ( - sum( - F.cross_entropy(reference(input_ids).logits.flatten(0, 1), labels.flatten(), reduction="sum") - for input_ids, labels in batches - ) - / denominator - ) - reference_loss.backward() - - torch.testing.assert_close(actual_loss, reference_loss.detach().to(actual_loss)) - for parameter, expected in zip(model.parameters(), reference.parameters()): - torch.testing.assert_close(parameter.grad, expected.grad) - - def test_lifecycle_marks_only_the_last_microbatch_for_sync(monkeypatch): events = [] From 7adf81166f37e37a14ce76d041103f1ab710f0c9 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sun, 16 Aug 2026 18:03:43 -0700 Subject: [PATCH 05/34] feat(engine): add pipeline forward-backward support Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 112 +++++++- nemo_automodel/engine/__init__.py | 242 ++++++++++++----- nemo_automodel/recipes/llm/train_ft.py | 38 ++- .../context_parallel/run_dense_packed_cp.py | 113 ++++---- .../pipelining/test_autopipeline.py | 186 ++++++++++++- tests/unit_tests/recipes/test_train_ft.py | 250 +++++++++++++++++- tests/unit_tests/test_engine.py | 249 +++++++++++++++++ 7 files changed, 1065 insertions(+), 125 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index 59fe248029..4bee514535 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect import logging from dataclasses import dataclass from typing import Any, Callable, Literal, Optional @@ -19,7 +20,7 @@ import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.pipelining.microbatch import BlockMask, TensorChunkSpec +from torch.distributed.pipelining.microbatch import BlockMask, TensorChunkSpec, split_args_kwargs_into_chunks from torch.distributed.pipelining.microbatch import _Replicate as ReplicateChunkSpec from torch.distributed.pipelining.schedules import _PipelineSchedule from torch.distributed.pipelining.stage import PipelineStage @@ -268,6 +269,9 @@ def step( *, target: torch.Tensor | None = None, losses: list[torch.Tensor] | None = None, + loss_inputs: dict[str, Any] | None = None, + loss_fn: Callable | None = None, + return_outputs: bool = True, **kwargs: Any, ) -> Any: """Run one pipeline schedule step with model-owned input chunking. @@ -279,6 +283,17 @@ def step( ranks without the last pipeline stage. losses: Mutable list populated with scalar loss tensors, or ``None`` on ranks without the last pipeline stage. + loss_inputs: Structured loss inputs. Tensor fields of shape + [batch, ...] are split on the batch axis; scalar tensors and + non-tensor fields are replicated. Must be provided together + with ``loss_fn`` and without ``target``. + loss_fn: Callback invoked as ``loss_fn(output, loss_inputs_mb, + model_args_mb, model_kwargs_mb)`` for each microbatch. Model + tensor arguments use shape [microbatch, ...], while keyword + tensors retain their model-defined layouts. + return_outputs: Whether the last pipeline stage returns merged model + outputs when supported by the installed PyTorch version. Tensor + layouts are defined by the underlying model. **kwargs: Keyword schedule inputs. Tensor values may have arbitrary model-defined layouts; model-owned metadata identifies any nonstandard batch axis. @@ -290,16 +305,100 @@ def step( if schedule is None: raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") + if (loss_inputs is None) != (loss_fn is None): + raise ValueError("loss_inputs and loss_fn must be provided together") + if loss_inputs is not None and target is not None: + raise ValueError("target cannot be used together with loss_inputs and loss_fn") + schedule_args = (model_input,) if self._info.has_first_stage else () kwargs_chunk_spec = self._get_schedule_kwargs_chunk_spec(kwargs) - if kwargs_chunk_spec is None: - return schedule.step(*schedule_args, target=target, losses=losses, **kwargs) + schedule_options = ( + {"return_outputs": return_outputs} + if "return_outputs" in inspect.signature(schedule.step).parameters + else {} + ) + + if loss_inputs is None: + if kwargs_chunk_spec is None: + return schedule.step( + *schedule_args, + target=target, + losses=losses, + **schedule_options, + **kwargs, + ) + + previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec + schedule._kwargs_chunk_spec = kwargs_chunk_spec + try: + return schedule.step( + *schedule_args, + target=target, + losses=losses, + **schedule_options, + **kwargs, + ) + finally: + schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec + + model_args_chunks, model_kwargs_chunks = split_args_kwargs_into_chunks( + (model_input,), + kwargs, + self.num_microbatches, + kwargs_chunk_spec=kwargs_chunk_spec, + ) + if len(model_args_chunks) != self.num_microbatches: + raise ValueError(f"Expected {self.num_microbatches} model input microbatches, got {len(model_args_chunks)}") + + loss_inputs_chunk_spec = tree_map( + lambda value: ( + TensorChunkSpec(0) if isinstance(value, torch.Tensor) and value.ndim > 0 else ReplicateChunkSpec() + ), + loss_inputs, + is_leaf=lambda value: isinstance(value, BlockMask), + ) + _, loss_inputs_chunks = split_args_kwargs_into_chunks( + (), + loss_inputs, + self.num_microbatches, + kwargs_chunk_spec=loss_inputs_chunk_spec, + ) + if len(loss_inputs_chunks) != self.num_microbatches: + raise ValueError(f"Expected {self.num_microbatches} loss input microbatches, got {len(loss_inputs_chunks)}") + + def microbatch_loss(output: Any, microbatch_id: torch.Tensor) -> Any: + """Evaluate the structured loss for one pipeline microbatch. + + Args: + output: Model output whose tensor layouts are defined by the model. + microbatch_id: Tensor of shape [1] identifying the microbatch. + + Returns: + The loss value returned by ``loss_fn``. Tensor layout is defined + by the callback. + """ + index = int(microbatch_id.item()) + return loss_fn( + output, + loss_inputs_chunks[index], + model_args_chunks[index], + model_kwargs_chunks[index], + ) previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec + previous_loss_fn = schedule._loss_fn schedule._kwargs_chunk_spec = kwargs_chunk_spec + schedule._loss_fn = microbatch_loss try: - return schedule.step(*schedule_args, target=target, losses=losses, **kwargs) + return schedule.step( + *schedule_args, + target=torch.arange(self.num_microbatches, device=self.device), + losses=losses, + **schedule_options, + **kwargs, + ) finally: + schedule._loss_fn = previous_loss_fn schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec @property @@ -313,6 +412,11 @@ def parts(self) -> list[nn.Module]: def device(self) -> torch.device: return self._device + @property + def num_microbatches(self) -> int: + """Number of pipeline microbatches in one local batch.""" + return self.pp_batch_size // self.pp_microbatch_size + # -------------------------- Debug utilities -------------------------- def list_stage_modules(self) -> list[list[str]]: names_per_stage: list[list[str]] = [] diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index d122b9bd42..c1ea753fda 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -32,6 +32,7 @@ ) from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh +from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.distributed.utils import get_sync_ctx from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler from nemo_automodel.components.training.utils import ( @@ -82,10 +83,12 @@ class Engine: gradient-finalization path. Args: - model: An already configured and distributed model. + model: An already configured and distributed model, or a built + :class:`AutoPipeline`. device: Device on which model inputs and losses are evaluated. - mesh_context: Runtime topology. When omitted, an initialized default - process group is treated as pure data parallelism. + mesh_context: Runtime topology. Required for an ``AutoPipeline``. For + eager models, an initialized default process group is treated as + pure data parallelism when omitted. collate_fn: Batches one microbatch of Datums into separate model and loss inputs. The default supports padded and packed text. VLMs pass a model-specific collater. Existing recipes whose dataloaders @@ -100,14 +103,17 @@ class Engine: final microbatch. Note: - This first execution backend is eager and weight-normalized. Pipeline - schedules are intentionally deferred; context-parallel input layout - and transport are delegated to :class:`ContextParallelSharder`. + Context-parallel input layout and transport are delegated to + :class:`ContextParallelSharder`. With an :class:`AutoPipeline`, the + pipeline schedule owns its internal microbatching and backward calls; + the Engine still owns the complete outer accumulation window and loss + normalization. Pipeline execution currently requires context-parallel + size one. """ def __init__( self, - model: nn.Module, + model: nn.Module | AutoPipeline, *, device: torch.device | str, mesh_context: MeshContext | None = None, @@ -116,7 +122,9 @@ def __init__( context_fn: Callable[[], AbstractContextManager[Any]] = nullcontext, defer_fsdp_grad_sync: bool = True, ) -> None: - self.model = model + self.pipeline = model if isinstance(model, AutoPipeline) else None + self.model_parts = model.parts if self.pipeline is not None else [model] + self.model = self.model_parts[0] self.device = torch.device(device) self.mesh_context = mesh_context self.collate_fn = collate_fn @@ -131,7 +139,9 @@ def forward_backward( ) -> tuple[torch.Tensor, list[dict[str, Any]]]: """Accumulate gradients for a complete optimizer window. - ``window`` is explicit: each inner sequence is one eager microbatch. + ``window`` is explicit: each inner sequence is one eager microbatch, or + one outer pipeline batch that the schedule splits internally. Pipeline + batches currently contain exactly one already-batched Datum. ``loss_fn`` receives the raw model output, CP-local ``loss_fn_inputs``, the original Datums, and the final CP-local model inputs produced by the sharder. It returns either per-element losses @@ -139,13 +149,16 @@ def forward_backward( ``loss_fn_inputs["weights"]``, or a scalar local weighted-sum numerator. For a scalar, the callback must apply weights and masks; the Engine will only apply global normalization. The callback may - also return one output mapping per Datum. + also return one output mapping per Datum during eager execution. Those mappings are detached and preserved in input order; the Engine - deliberately does not interpret or reduce them. + deliberately does not interpret or reduce them. Pipeline execution + does not yet support per-Datum outputs. Args: window: The complete optimizer accumulation window. Each inner - sequence is one microbatch of Datums. A Datum's token weights + sequence is one eager microbatch or outer pipeline batch of + Datums. A pipeline batch must contain exactly one prebatched + Datum. A Datum's token weights may have shape ``[tokens]`` or the custom collater's batched token layout; the loss tensor must use the identical shape. loss_fn: Computes either that per-token loss tensor or a scalar @@ -154,21 +167,30 @@ def forward_backward( Returns: ``(loss, loss_fn_outputs)``. ``loss`` is a detached scalar reduced - over the DP-CP gradient group. ``loss_fn_outputs`` contains + over the DP-CP gradient group and, for pipeline execution, + synchronized across PP stages. ``loss_fn_outputs`` contains local-rank, per-Datum mappings in window order. Model parameters are unchanged, but their gradients contain the complete window's globally normalized backward result. """ microbatches = self._validate_window(window) self._validate_parallelism() + if self.pipeline is not None and any(len(microbatch) != 1 for microbatch in microbatches): + raise ValueError("pipeline Engine requires exactly one prebatched Datum in each outer batch") dp_group, dp_size = self._dp_group_and_size() grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) denominator = self._global_weight_sum(microbatches, dp_group, dp_size) - - self.model.train() - prepare_for_grad_accumulation([self.model], pp_enabled=False) - MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(self._cp_size() / len(microbatches)) + self._validate_pipeline_window(len(microbatches), denominator) + + pp_enabled = self.pipeline is not None + for part in self.model_parts: + part.train() + prepare_for_grad_accumulation(self.model_parts, pp_enabled=pp_enabled) + inner_microbatches = self.pipeline.num_microbatches if self.pipeline is not None else 1 + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( + self._cp_size() / (len(microbatches) * inner_microbatches) + ) local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) loss_fn_outputs: list[dict[str, Any]] = [] @@ -177,7 +199,7 @@ def forward_backward( for index, datums in enumerate(microbatches): is_last = index == len(microbatches) - 1 if is_last: - prepare_for_final_backward([self.model], pp_enabled=False) + prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) model_inputs, loss_inputs = self.collate_fn(datums) self._validate_collated_weights(datums, loss_inputs) @@ -202,6 +224,11 @@ def forward_backward( "context parallelism requires raw THD inputs so ContextParallelSharder can partition them" ) if final_thd: + if self.pipeline is not None and self.pipeline.num_microbatches > 1: + raise ValueError( + "pipeline Engine requires raw THD inputs so ContextParallelSharder can split them " + "for the schedule's internal microbatches" + ) sharder = ContextParallelSharder( device_mesh=device_mesh, shard_batch=shard_batch_identity, @@ -214,6 +241,7 @@ def forward_backward( device_mesh, cp_batch, padding_token_id=self.padding_token_id, + num_chunks=inner_microbatches, ) cp_context, model_inputs = sharder.shard(cp_batch) if model_inputs.get("qkv_format") == "thd" and ( @@ -228,50 +256,52 @@ def forward_backward( if labels is not None: loss_inputs["labels"] = local_labels weights = loss_inputs["weights"] - forward_inputs = filter_forward_kwargs(self.model, model_inputs) - - with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(), cp_context(): - output = self.model(**forward_inputs) - result = loss_fn(output, loss_inputs, datums, model_inputs) - has_outputs = isinstance(result, tuple) - if returns_outputs is None: - returns_outputs = has_outputs - elif returns_outputs != has_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - if isinstance(result, tuple): - losses, outputs = result - if ( - not isinstance(outputs, Sequence) - or isinstance(outputs, (str, bytes)) - or len(outputs) != len(datums) - or not all(isinstance(item, Mapping) for item in outputs) - ): - raise ValueError("loss_fn outputs must contain one mapping per Datum") - loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) - else: - losses = result - if not isinstance(losses, torch.Tensor): - raise TypeError("loss_fn must return a Tensor, optionally followed by per-Datum outputs") - if losses.ndim == 0: - numerator = losses - elif losses.shape == weights.shape: - numerator = (losses * weights.to(losses)).sum() - else: - raise ValueError( - "loss_fn must return a scalar local weighted sum or losses with exactly the same shape as weights; " - f"got losses={tuple(losses.shape)}, weights={tuple(weights.shape)}" - ) - if losses.device != weights.device: - raise ValueError("loss_fn losses and weights must be on the same device") - (numerator * (grad_group_size / denominator)).backward() - - local_loss_sum.add_(numerator.detach().to(torch.float64)) + if self.pipeline is not None: + self._pipeline_step( + model_inputs, + loss_inputs, + datums, + loss_fn, + denominator, + grad_group_size, + local_loss_sum, + cp_context, + ) + else: + forward_inputs = filter_forward_kwargs(self.model, model_inputs) + with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(), cp_context(): + output = self.model(**forward_inputs) + result = loss_fn(output, loss_inputs, datums, model_inputs) + has_outputs = isinstance(result, tuple) + if returns_outputs is None: + returns_outputs = has_outputs + elif returns_outputs != has_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + if isinstance(result, tuple): + losses, outputs = result + if ( + not isinstance(outputs, Sequence) + or isinstance(outputs, (str, bytes)) + or len(outputs) != len(datums) + or not all(isinstance(item, Mapping) for item in outputs) + ): + raise ValueError("loss_fn outputs must contain one mapping per Datum") + loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) + else: + losses = result + numerator = _weighted_numerator(losses, weights) + (numerator * (grad_group_size / denominator)).backward() + + local_loss_sum.add_(numerator.detach().to(torch.float64)) if index == 0: prepare_after_first_microbatch() if grad_group_size > 1: dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=grad_group) + pp_group, pp_size = self._pp_group_and_size() + if pp_size > 1: + dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=pp_group) loss = (local_loss_sum / denominator).detach() return loss, loss_fn_outputs @@ -289,16 +319,72 @@ def _validate_window(window: Sequence[Sequence[Datum]]) -> list[list[Datum]]: microbatches.append(list(microbatch)) return microbatches + def _pipeline_step( + self, + model_inputs: dict[str, Any], + loss_inputs: dict[str, torch.Tensor], + datums: Sequence[Datum], + loss_fn: LossFn, + denominator: torch.Tensor, + grad_group_size: int, + local_loss_sum: torch.Tensor, + cp_context: Callable[[], AbstractContextManager[Any]], + ) -> None: + primary_names = [name for name in ("input_ids", "inputs_embeds") if name in model_inputs] + if len(primary_names) != 1: + raise ValueError("pipeline Engine requires exactly one of input_ids or inputs_embeds") + primary_name = primary_names[0] + primary = model_inputs.pop(primary_name) + if not isinstance(primary, torch.Tensor) or primary.ndim < 2: + raise ValueError(f"pipeline Engine requires batched {primary_name} with a sequence dimension") + + self.pipeline.update_seq_len(primary.shape[1]) + pipeline_kwargs = { + name: value + for name, value in model_inputs.items() + if value is not None and not (isinstance(value, dict) and not value) + } + + def pipeline_loss(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): + if not model_args_mb: + raise RuntimeError("AutoPipeline loss callback did not receive the primary model input") + final_model_inputs = {primary_name: model_args_mb[0], **model_kwargs_mb} + result = loss_fn(output, loss_inputs_mb, datums, final_model_inputs) + if isinstance(result, tuple): + raise ValueError("pipeline Engine does not yet support per-Datum loss_fn outputs") + numerator = _weighted_numerator(result, loss_inputs_mb["weights"]) + local_loss_sum.add_(numerator.detach().to(torch.float64)) + return numerator * (grad_group_size / denominator) + + with self.context_fn(), cp_context(): + self.pipeline.step( + primary, + loss_inputs=loss_inputs, + loss_fn=pipeline_loss, + return_outputs=False, + **pipeline_kwargs, + ) + def _validate_parallelism(self) -> None: - if any(bool(getattr(module, "calculate_per_token_loss", False)) for module in self.model.modules()): + if any( + bool(getattr(module, "calculate_per_token_loss", False)) + for part in self.model_parts + for module in part.modules() + ): raise NotImplementedError( "Engine.forward_backward requires averaged distributed gradients; " "MegatronFSDP calculate_per_token_loss=True uses summed gradients" ) + if self.pipeline is not None and self.pipeline.scale_grads_in_schedule: + raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") + if self.pipeline is not None and self.mesh_context is None: + raise ValueError("pipeline Engine requires mesh_context") + if self.pipeline is not None and self._cp_size() > 1: + raise NotImplementedError("pipeline Engine does not yet support context parallelism") if self.mesh_context is None: return - if self.mesh_context.pp_size > 1: - raise NotImplementedError("Engine.forward_backward does not yet support pipeline parallelism") + if self.mesh_context.pp_size > 1 and self.pipeline is None: + raise NotImplementedError("pipeline parallelism requires an AutoPipeline") if self.mesh_context.cp_size > 1 and self.mesh_context.device_mesh is None: raise ValueError("context parallelism requires a device mesh") @@ -327,6 +413,27 @@ def _gradient_group_and_size( size = int(dp_cp_mesh.size()) return (dp_cp_mesh.get_group() if size > 1 else None), size + def _pp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: + if self.pipeline is None or not dist.is_available() or not dist.is_initialized(): + return None, 1 + size = int(self.pipeline.pp_mesh.size()) + return (self.pipeline.pp_mesh.get_group() if size > 1 else None), size + + def _validate_pipeline_window(self, size: int, denominator: torch.Tensor) -> None: + pp_group, pp_size = self._pp_group_and_size() + if pp_size <= 1: + return + local = torch.stack((denominator.new_tensor(size), denominator)) + gathered = torch.empty(pp_size * 2, dtype=denominator.dtype, device=denominator.device) + dist.all_gather_into_tensor(gathered, local, group=pp_group) + gathered = gathered.view(pp_size, 2) + if not bool((gathered[:, 0] == gathered[0, 0]).all()): + raise ValueError(f"pipeline stages must use the same outer window size; got {gathered[:, 0].tolist()}") + if not torch.allclose(gathered[:, 1], gathered[0, 1].expand(pp_size), rtol=1e-8, atol=1e-12): + raise ValueError( + f"pipeline stages must use the same DP-reduced weight denominator; got {gathered[:, 1].tolist()}" + ) + def _global_weight_sum( self, microbatches: list[list[Datum]], @@ -481,6 +588,23 @@ def _is_final_thd(model_inputs: dict[str, Any]) -> bool: return True +def _weighted_numerator(losses: Any, weights: torch.Tensor) -> torch.Tensor: + if not isinstance(losses, torch.Tensor): + raise TypeError("loss_fn must return a Tensor, optionally followed by per-Datum outputs") + if losses.ndim == 0: + numerator = losses + elif losses.shape == weights.shape: + numerator = (losses * weights.to(losses)).sum() + else: + raise ValueError( + "loss_fn must return a scalar local weighted sum or losses with exactly the same shape as weights; " + f"got losses={tuple(losses.shape)}, weights={tuple(weights.shape)}" + ) + if losses.device != weights.device: + raise ValueError("loss_fn losses and weights must be on the same device") + return numerator + + def _detach(value: Any) -> Any: """Detach tensor leaves without changing an output record's structure.""" if isinstance(value, torch.Tensor): diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 862aeec9b4..6b37ae3588 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -697,19 +697,35 @@ def setup(self): # Tokenizer + model-derived values are runtime concerns: build them here and pass them to # each DataloaderConfig.build(); the configs themselves are resolved at the RecipeConfig boundary. _, self.tokenizer = _build_tokenizer(self.cfg.model, self.cfg.dataset) - model_has_mtp = not self.pp_enabled and any( - getattr(module, "mtp", None) is not None for module in self.model_parts[0].modules() + model_has_mtp = any( + getattr(module, "mtp", None) is not None + for model_part in self.model_parts + for module in model_part.modules() + ) + pp_group = self._get_pp_group() if self.pp_enabled else None + if pp_group is not None: + # The MTP head may exist only on the last PP stage; all stages must choose the same path. + model_has_mtp_flag = torch.tensor(int(model_has_mtp), device=self.dist_env.device) + torch.distributed.all_reduce(model_has_mtp_flag, op=torch.distributed.ReduceOp.MAX, group=pp_group) + model_has_mtp = bool(model_has_mtp_flag.item()) + + dataloader_config = self.cfg.dataloader + pp_uses_packed_batches = self.pp_enabled and ( + _packed_seq_size > 0 or bool(getattr(dataloader_config, "emits_thd", False)) ) self.engine = None if ( - not self.pp_enabled - and not self.magi.enabled - and not (self.mesh_context.cp_size > 1 and model_has_mtp) + not self.magi.enabled + and not (model_has_mtp and (self.pp_enabled or self.mesh_context.cp_size > 1)) + and not pp_uses_packed_batches + and not (self.pp_enabled and self.mesh_context.cp_size > 1) + and not (self.pp_enabled and isinstance(self.loss_fn, FusedLinearCrossEntropy)) + and not (self.pp_enabled and self.pp.scale_grads_in_schedule) and not getattr(self.distributed_config, "calculate_per_token_loss", False) and getattr(self.loss_fn, "reduction", None) == "sum" ): self.engine = Engine( - self.model_parts[0], + self.pp if self.pp_enabled else self.model_parts[0], device=self.dist_env.device, mesh_context=self.mesh_context, collate_fn=collate_prebatched, @@ -1051,6 +1067,10 @@ def _prepare_microbatch(self, batch): def _compute_causal_lm_loss(self, output, labels, model_inputs, *, num_label_tokens, is_train): """Compute the recipe's causal-LM and optional MTP loss.""" model = self.model_parts[0] + if self.pp_enabled: + model = next( + model_part for model_part, stage in zip(self.model_parts, self.pp.info.stages) if stage.is_last + ) grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None hidden_states = get_final_hidden_states(output) if isinstance(self.loss_fn, FusedLinearCrossEntropy) and hidden_states is None: @@ -1241,7 +1261,9 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() engine = getattr(self, "engine", None) - use_engine = engine is not None and num_label_tokens > 0 + # A custom prepacked loader may emit THD even when setup did not advertise packing. + pp_uses_thd = self.pp_enabled and any(batch.get("qkv_format") == "thd" for batch in batches) + use_engine = engine is not None and num_label_tokens > 0 and not pp_uses_thd if use_engine: reporting_loss, _ = engine.forward_backward( [[self._make_engine_datum(batch)] for batch in batches], @@ -1274,7 +1296,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, pp_axis_name="pp" if self.pp_enabled else None, foreach=True, - num_label_tokens=num_label_tokens, + num_label_tokens=None if use_engine else num_label_tokens, dp_group_size=self._get_dp_group_size(include_cp=True), expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, self.device_mesh), ) diff --git a/tests/functional_tests/context_parallel/run_dense_packed_cp.py b/tests/functional_tests/context_parallel/run_dense_packed_cp.py index cd4185e86c..b6fe9fe5dd 100644 --- a/tests/functional_tests/context_parallel/run_dense_packed_cp.py +++ b/tests/functional_tests/context_parallel/run_dense_packed_cp.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Packed-THD CP and composed TP+CP parity for dense Llama, Qwen2, and Qwen3. +"""Engine/FSDP2 packed-THD CP and TP+CP parity for dense Llama, Qwen2, and Qwen3. Run with:: @@ -29,25 +29,23 @@ import torch import torch.distributed as dist -from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor -from torch.distributed.tensor.parallel import parallelize_module from transformer_engine.pytorch import DotProductAttention from transformers import LlamaConfig, Qwen2Config, Qwen3Config +from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.distributed.config import FSDP2Config from nemo_automodel.components.distributed.context_parallel.utils import ( attach_te_context_parallel, make_cp_batch_for_te, ) -from nemo_automodel.components.distributed.parallelizer import ( - _attention_is_head_sharded, - _get_parallel_plan, - _update_attention_head_counts_for_tp, -) +from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager +from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.llama.model import LlamaForCausalLM from nemo_automodel.components.models.qwen2.model import Qwen2ForCausalLM from nemo_automodel.components.models.qwen3.model import Qwen3ForCausalLM +from nemo_automodel.engine import Engine, collate_prebatched NUM_HIDDEN_LAYERS = 2 @@ -110,31 +108,14 @@ def _build_model( return model.to(device=device, dtype=torch.bfloat16).train() -def _apply_tensor_parallel(model: torch.nn.Module, tp_mesh) -> None: - """Apply the production dense-model TP plan without an additional FSDP axis. - - Args: - model: Replicated model whose projection parameters are plain tensors. - tp_mesh: One-dimensional tensor-parallel mesh. - """ - if tp_mesh.size() == 1: - return - plan = _get_parallel_plan(model, tp_size=tp_mesh.size()) - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message=".*could not be resolved.*", category=UserWarning) - parallelize_module(model, tp_mesh, plan) - if _attention_is_head_sharded(plan): - _update_attention_head_counts_for_tp(model, tp_mesh.size()) - - def _full_tensor(tensor: torch.Tensor) -> torch.Tensor: - """Materialize a replicated local tensor from a TP-sharded DTensor. + """Materialize a replicated global tensor from a distributed tensor. Args: - tensor: Plain tensor or DTensor with a sharded vocabulary or parameter axis. + tensor: Plain tensor or DTensor sharded by FSDP, TP, or both. Returns: - Plain tensor with the global TP shape, replicated within the TP group. + Plain tensor with the global shape, replicated over its device mesh. """ return tensor.full_tensor() if isinstance(tensor, DTensor) else tensor @@ -171,11 +152,14 @@ def _reconstruct_global_tokens( def _run_model( model_kind: str, device: torch.device, - cp_mesh, - tp_mesh, + mesh_context: MeshContext, + distributed_config: FSDP2Config, ) -> None: import transformer_engine_torch as tex + assert mesh_context.device_mesh is not None + cp_mesh = mesh_context.device_mesh["cp"] + tp_mesh = mesh_context.device_mesh["tp"] cp_rank = cp_mesh.get_local_rank() cp_group = cp_mesh.get_group() batch = { @@ -191,22 +175,48 @@ def _run_model( baseline_batch = make_cp_batch_for_te(None, _clone_batch(batch)) baseline_labels = baseline_batch.pop("labels") baseline_logits = baseline_model(**baseline_batch).logits.squeeze(0) - loss_normalizer = baseline_logits.numel() - (baseline_logits.float().square().sum() / loss_normalizer).backward() + baseline_loss = baseline_logits.float().square().mean() + baseline_loss.backward() baseline_grads = [layer.self_attn.q_proj.weight.grad.detach().float() for layer in baseline_model.model.layers] assert baseline_labels.shape == (8,) cp_model = _build_model(model_kind, device) - _apply_tensor_parallel(cp_model, tp_mesh) + cp_model = FSDP2Manager(distributed_config, device_mesh=mesh_context.device_mesh).parallelize(cp_model) attention_cp_mesh = cp_mesh if cp_mesh.size() > 1 else None configured = attach_te_context_parallel(cp_model, attention_cp_mesh, tp_mesh) assert configured == NUM_HIDDEN_LAYERS - cp_batch = make_cp_batch_for_te(attention_cp_mesh, _clone_batch(batch)) - cp_batch.pop("labels") - local_logits = _full_tensor(cp_model(**cp_batch).logits).squeeze(0) - (local_logits.float().square().sum() / loss_normalizer).backward() - cu_seqlens = cp_batch["cu_seqlens"] + raw_model_inputs = _clone_batch(batch) + raw_labels = raw_model_inputs.pop("labels") + datum = Datum( + model_inputs=raw_model_inputs, + loss_fn_inputs={ + "labels": raw_labels, + "weights": torch.ones_like(raw_labels, dtype=torch.float32), + }, + ) + assert "seq_lens" in datum.model_inputs and "cu_seqlens" not in datum.model_inputs + + observed: dict[str, torch.Tensor] = {} + + def loss_fn(output, loss_inputs, _datums, model_inputs): + local_logits = _full_tensor(output.logits).squeeze(0) + observed["logits"] = local_logits.detach() + observed["cu_seqlens"] = model_inputs["cu_seqlens"].detach() + observed["labels"] = loss_inputs["labels"].detach() + return local_logits.float().square().mean(dim=-1) + + engine_loss, _ = Engine( + cp_model, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + defer_fsdp_grad_sync=distributed_config.defer_fsdp_grad_sync, + ).forward_backward([[datum]], loss_fn) + + local_logits = observed["logits"] + assert observed["labels"].shape == (8 // cp_mesh.size(),) + cu_seqlens = observed["cu_seqlens"] indices = tex.thd_get_partitioned_indices(cu_seqlens, 8, cp_mesh.size(), cp_rank).to(torch.int32) cp_logits = _reconstruct_global_tokens( local_logits, @@ -215,22 +225,30 @@ def _run_model( group=cp_group, ) cp_grads = [_full_tensor(layer.self_attn.q_proj.weight.grad).detach().float() for layer in cp_model.model.layers] - for cp_grad in cp_grads: - dist.all_reduce(cp_grad, group=cp_group) + cp_loss = cp_logits.float().square().mean() + # The Engine must normalize the exact distributed logits tightly. The + # baseline comparison is looser because independently executed BF16 + # attention paths can differ by roughly one ULP, which doubles under x**2. + torch.testing.assert_close(engine_loss.float(), cp_loss, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(engine_loss.float(), baseline_loss.detach(), atol=1e-6, rtol=1e-2) torch.testing.assert_close(cp_logits, baseline_logits, atol=3e-2, rtol=3e-2) for cp_grad, baseline_grad in zip(cp_grads, baseline_grads): torch.testing.assert_close(cp_grad, baseline_grad, atol=1e-3, rtol=5e-2) assert torch.isfinite(cp_logits).all() assert all(torch.isfinite(cp_grad).all() for cp_grad in cp_grads) if dist.get_rank() == 0: + loss_diff = (engine_loss.float() - baseline_loss.detach()).abs().item() + normalization_diff = (engine_loss.float() - cp_loss).abs().item() output_diff = (cp_logits.float() - baseline_logits.float()).abs().max().item() grad_diff = max( (cp_grad - baseline_grad).abs().max().item() for cp_grad, baseline_grad in zip(cp_grads, baseline_grads) ) print( f"{model_kind} full TP={tp_mesh.size()} CP={cp_mesh.size()}: " - f"packed parity passed (logits max={output_diff:.6f}, grad max={grad_diff:.6f})" + f"Engine packed parity passed " + f"(loss={loss_diff:.6f}, normalization={normalization_diff:.6f}, " + f"logits max={output_diff:.6f}, grad max={grad_diff:.6f})" ) @@ -296,12 +314,17 @@ def main() -> None: if dist.get_world_size() % args.tp_size != 0: raise ValueError(f"World size {dist.get_world_size()} must be divisible by TP size {args.tp_size}.") cp_size = dist.get_world_size() // args.tp_size - mesh = init_device_mesh("cuda", (cp_size, args.tp_size), mesh_dim_names=("cp", "tp")) - cp_mesh = mesh["cp"] - tp_mesh = mesh["tp"] + distributed_config = FSDP2Config() + mesh_context = MeshContext.build( + distributed_config, + ParallelismSizes(dp_size=1, cp_size=cp_size, tp_size=args.tp_size), + world_size=dist.get_world_size(), + ) + assert mesh_context.device_mesh is not None + assert mesh_context.device_mesh["dp_shard_cp"].size() == cp_size try: for model_kind in ("llama", "qwen2", "qwen3"): - _run_model(model_kind, device, cp_mesh, tp_mesh) + _run_model(model_kind, device, mesh_context, distributed_config) dist.barrier() if args.tp_size == 1: for model_kind in ("qwen2", "qwen3"): diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 9bc08cce91..214ff2299b 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -161,6 +161,7 @@ def test_valid_autopipeline(self): assert ap.pp_schedule == "1f1b" assert ap.pp_microbatch_size == 1 assert ap.pp_batch_size == 4 + assert ap.num_microbatches == 4 assert ap._device == torch.device("cpu") def test_invalid_batch_size(self): @@ -237,31 +238,41 @@ def get_pipeline_kwargs_chunk_dims(self, kwargs): class _KwargsChunkSchedule: - def __init__(self, *, fail_on_step: bool = False): + def __init__(self, *, fail_on_step: bool = False, invoke_loss: bool = False): self._kwargs_chunk_spec = None + self._loss_fn = Mock(return_value=torch.tensor(0.0)) self.fail_on_step = fail_on_step + self.invoke_loss = invoke_loss self.args_during_step = None self.kwargs_chunk_spec_during_step = None + self.loss_fn_during_step = None self.kwargs_split = None + self.target_during_step = None + self.return_outputs_during_step = None + self.loss_results = [] - def step(self, *args, target=None, losses=None, **kwargs): + def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): """Split schedule inputs using the chunk spec active during the call. Args: *args: Positional schedule inputs. Tensor values have arbitrary model-defined layouts. - target: Optional tensor of shape [batch, sequence] containing loss - targets. + target: Optional tensor of shape [batch, ...] containing loss targets + or structured-loss microbatch IDs. losses: Optional mutable list populated with scalar loss tensors. + return_outputs: Whether to return the schedule result. **kwargs: Keyword schedule inputs. Tensor values have arbitrary model-defined layouts. Returns: A sentinel string identifying the schedule result. """ - del target, losses + del losses self.args_during_step = args self.kwargs_chunk_spec_during_step = self._kwargs_chunk_spec + self.loss_fn_during_step = self._loss_fn + self.target_during_step = target + self.return_outputs_during_step = return_outputs if self.fail_on_step: raise RuntimeError("schedule failed") _, self.kwargs_split = split_args_kwargs_into_chunks( @@ -270,9 +281,26 @@ def step(self, *args, target=None, losses=None, **kwargs): 2, kwargs_chunk_spec=self._kwargs_chunk_spec, ) + if self.invoke_loss: + assert target is not None + target_chunks = torch.tensor_split(target, 2) + for index in (1, 0): + self.loss_results.append(self._loss_fn(torch.tensor(float(index)), target_chunks[index])) return "schedule-result" +class _LegacyStepSchedule(_KwargsChunkSchedule): + """Schedule with the PyTorch 2.6-2.9 step signature.""" + + def __init__(self): + super().__init__() + self.received_return_outputs = False + + def step(self, *args, target=None, losses=None, **kwargs): + self.received_return_outputs = "return_outputs" in kwargs + return super().step(*args, target=target, losses=losses, **kwargs) + + class TestAutoPipelineKwargsChunkSpec: def _pipeline_with_parts(self, *parts: nn.Module, schedule=None, has_first_stage: bool = True): ap = AutoPipeline( @@ -318,12 +346,29 @@ def test_step_splits_mrope_position_ids_on_model_owned_batch_axis(self): def test_step_without_model_hook_uses_pytorch_default_chunking(self): ap = self._pipeline_with_parts(nn.Module()) + original_loss_fn = ap.info.schedule._loss_fn ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) assert ap.info.schedule.kwargs_chunk_spec_during_step is None assert ap.info.schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) assert ap.info.schedule._kwargs_chunk_spec is None + assert ap.info.schedule._loss_fn is original_loss_fn + assert ap.info.schedule.return_outputs_during_step is True + + @pytest.mark.parametrize("structured", [False, True]) + def test_step_does_not_forward_return_outputs_to_older_pytorch(self, structured): + schedule = _LegacyStepSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + structured_kwargs = ( + {"loss_inputs": {"weights": torch.ones(2, 8)}, "loss_fn": lambda *_: torch.tensor(0.0)} + if structured + else {} + ) + + ap.step(torch.zeros(2, 8), return_outputs=False, **structured_kwargs) + + assert schedule.received_return_outputs is False def test_only_canonical_model_part_supplies_chunk_policy(self): ap = self._pipeline_with_parts( @@ -361,6 +406,137 @@ def test_model_hook_cannot_configure_unknown_kwarg(self): with pytest.raises(ValueError, match="unknown kwarg"): ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) + def test_structured_loss_receives_aligned_microbatches(self): + schedule = _KwargsChunkSchedule(invoke_loss=True) + ap = self._pipeline_with_parts( + _KwargsChunkHookPart({"position_ids": 1}), + schedule=schedule, + has_first_stage=False, + ) + input_ids = torch.arange(16).view(2, 8) + position_ids = torch.arange(48).view(3, 2, 8) + weights = torch.arange(16).view(2, 8) + scale = torch.tensor(0.5) + seen = [] + + def loss_fn(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): + """Record one structured-loss callback. + + Args: + output: Scalar tensor produced by the fake last stage. + loss_inputs_mb: Mapping containing ``weights`` with shape + [microbatch, sequence] and scalar tensor ``scale``. + model_args_mb: Tuple containing input IDs with shape + [microbatch, sequence]. + model_kwargs_mb: Mapping containing position IDs with shape + [axes, microbatch, sequence]. + + Returns: + Scalar loss tensor. + """ + seen.append((output, loss_inputs_mb, model_args_mb, model_kwargs_mb)) + return output + + result = ap.step( + input_ids, + loss_inputs={"weights": weights, "scale": scale, "label": "train"}, + loss_fn=loss_fn, + return_outputs=False, + position_ids=position_ids, + ) + + assert result == "schedule-result" + assert schedule.args_during_step == () + torch.testing.assert_close(schedule.target_during_step, torch.arange(2)) + assert schedule.target_during_step.device == ap.device + assert schedule.return_outputs_during_step is False + assert [item[0].item() for item in seen] == [1.0, 0.0] + for index, (_, loss_inputs_mb, model_args_mb, model_kwargs_mb) in zip((1, 0), seen): + torch.testing.assert_close(loss_inputs_mb["weights"], weights[index : index + 1]) + assert loss_inputs_mb["scale"] is scale + assert loss_inputs_mb["label"] == "train" + torch.testing.assert_close(model_args_mb[0], input_ids[index : index + 1]) + torch.testing.assert_close(model_kwargs_mb["position_ids"], position_ids[:, index : index + 1]) + + @pytest.mark.parametrize( + ("loss_inputs", "loss_fn"), + [ + ({"weights": torch.ones(2, 8)}, None), + (None, Mock()), + ], + ) + def test_structured_loss_requires_inputs_and_callback(self, loss_inputs, loss_fn): + ap = self._pipeline_with_parts(nn.Module()) + + with pytest.raises(ValueError, match="must be provided together"): + ap.step(torch.zeros(2, 8), loss_inputs=loss_inputs, loss_fn=loss_fn) + + def test_structured_loss_rejects_legacy_target(self): + ap = self._pipeline_with_parts(nn.Module()) + + with pytest.raises(ValueError, match="target cannot be used"): + ap.step( + torch.zeros(2, 8), + target=torch.zeros(2, 8), + loss_inputs={"weights": torch.ones(2, 8)}, + loss_fn=Mock(), + ) + + def test_structured_loss_requires_exact_microbatch_count(self): + ap = self._pipeline_with_parts(nn.Module()) + + with pytest.raises(ValueError, match="Expected 2 model input microbatches, got 1"): + ap.step( + torch.zeros(1, 8), + loss_inputs={"weights": torch.ones(2, 8)}, + loss_fn=Mock(), + ) + + with pytest.raises(ValueError, match="Expected 2 loss input microbatches, got 1"): + ap.step( + torch.zeros(2, 8), + loss_inputs={"weights": torch.ones(1, 8)}, + loss_fn=Mock(), + ) + + def test_structured_loss_restores_schedule_state_after_callback_failure(self): + schedule = _KwargsChunkSchedule(invoke_loss=True) + original_loss_fn = schedule._loss_fn + original_chunk_spec = {"position_ids": TensorChunkSpec(0)} + schedule._kwargs_chunk_spec = original_chunk_spec + ap = self._pipeline_with_parts(_KwargsChunkHookPart({"position_ids": 1}), schedule=schedule) + + def failing_loss(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): + """Raise while evaluating a structured microbatch loss. + + Args: + output: Scalar tensor produced by the fake last stage. + loss_inputs_mb: Mapping containing weights with shape + [microbatch, sequence]. + model_args_mb: Tuple containing input IDs with shape + [microbatch, sequence]. + model_kwargs_mb: Mapping containing position IDs with shape + [axes, microbatch, sequence]. + + Raises: + RuntimeError: Always. + """ + del output, loss_inputs_mb, model_args_mb, model_kwargs_mb + raise RuntimeError("loss failed") + + with pytest.raises(RuntimeError, match="loss failed"): + ap.step( + torch.zeros(2, 8), + loss_inputs={"weights": torch.ones(2, 8)}, + loss_fn=failing_loss, + position_ids=torch.zeros(3, 2, 8), + ) + + assert schedule.loss_fn_during_step is not original_loss_fn + assert schedule.kwargs_chunk_spec_during_step["position_ids"].split_dim == 1 + assert schedule._loss_fn is original_loss_fn + assert schedule._kwargs_chunk_spec is original_chunk_spec + # ----------------------------- # Core build/materialize/step tests diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 4f297e105c..c806dc4942 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1105,6 +1105,175 @@ def test_setup_builds_engine_for_eager_sum_loss(monkeypatch): assert trainer.engine is not None +def test_setup_builds_engine_for_eager_fused_loss(monkeypatch): + from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy + + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + fused_loss = FusedLinearCrossEntropy() + monkeypatch.setattr( + RecipeConfig, + "loss_fn", + property(lambda self: SimpleNamespace(build=lambda: fused_loss)), + ) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._supports_logits_to_keep", lambda _model: True) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert trainer.loss_fn is fused_loss + assert trainer.engine is not None + + +@pytest.mark.parametrize( + ( + "local_has_mtp", + "peer_has_mtp", + "cp_size", + "packed_sequence_size", + "dataloader_emits_thd", + "fused_loss", + "scale_grads_in_schedule", + "expect_engine", + ), + [ + (False, False, 1, 0, False, False, False, True), + (True, False, 1, 0, False, False, False, False), + (False, True, 1, 0, False, False, False, False), + (False, False, 2, 0, False, False, False, False), + (False, False, 1, 8, False, False, False, False), + (False, False, 1, 0, True, False, False, False), + (False, False, 1, 0, False, True, False, False), + (False, False, 1, 0, False, False, True, False), + ], +) +def test_setup_engine_gate_for_pipeline( + monkeypatch, + local_has_mtp, + peer_has_mtp, + cp_size, + packed_sequence_size, + dataloader_emits_thd, + fused_loss, + scale_grads_in_schedule, + expect_engine, +): + from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy + + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + cfg.step_scheduler.local_batch_size = 2 + cfg.step_scheduler.global_batch_size = 2 + cfg.packed_sequence = ConfigNode({"packed_sequence_size": packed_sequence_size}) + _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + fused_loss_fn = FusedLinearCrossEntropy() if fused_loss else None + if fused_loss_fn is not None: + monkeypatch.setattr( + RecipeConfig, + "loss_fn", + property(lambda self: SimpleNamespace(build=lambda: fused_loss_fn)), + ) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._supports_logits_to_keep", lambda _model: True) + monkeypatch.setattr( + RecipeConfig, + "dataloader", + property( + lambda self: SimpleNamespace( + build=lambda **kwargs: "dl", + dataset_builds_on_all_ranks=False, + emits_thd=dataloader_emits_thd, + seed=42, + ) + ), + ) + + class DummyAutoPipeline(SimpleNamespace): + pass + + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.AutoPipeline", DummyAutoPipeline) + monkeypatch.setattr("nemo_automodel.engine.AutoPipeline", DummyAutoPipeline) + parts = [DummyModel()] + parts[0]._pp_return_hidden_states_supported = True + if local_has_mtp: + parts[0].mtp = nn.Identity() + pipeline = DummyAutoPipeline( + parts=parts, + pp_batch_size=2, + pp_microbatch_size=1, + scale_grads_in_schedule=scale_grads_in_schedule, + info=SimpleNamespace( + has_first_stage=True, + has_last_stage=False, + schedule=SimpleNamespace(), + stages=[SimpleNamespace(is_last=False)], + ), + ) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.build_model", lambda *args, **kwargs: pipeline) + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.create_distributed_setup_from_config", + lambda cfg, world_size: SimpleNamespace( + mesh_context=SimpleNamespace( + pp_enabled=True, + device_mesh=None, + moe_mesh=None, + cp_size=cp_size, + pp_size=2, + ), + strategy_config=None, + pipeline_config=SimpleNamespace(pp_seq_len=None), + moe_parallel_config=None, + activation_checkpointing=False, + ), + ) + pp_group = object() + monkeypatch.setattr(TrainFinetuneRecipeForNextTokenPrediction, "_get_pp_group", lambda self: pp_group) + reduced_mtp_flags = [] + + def all_reduce_mtp_flag(flag, *, op, group): + assert op == torch.distributed.ReduceOp.MAX + assert group is pp_group + reduced_mtp_flags.append(bool(flag.item())) + if peer_has_mtp: + flag.fill_(1) + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce_mtp_flag) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert reduced_mtp_flags == [local_has_mtp] + if fused_loss_fn is not None: + assert trainer.loss_fn is fused_loss_fn + assert (trainer.engine is not None) is expect_engine + if expect_engine: + assert trainer.engine.pipeline is pipeline + + +def test_setup_keeps_engine_disabled_for_per_token_megatron_fsdp(monkeypatch): + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.create_distributed_setup_from_config", + lambda cfg, world_size: SimpleNamespace( + mesh_context=SimpleNamespace( + pp_enabled=False, + device_mesh=None, + moe_mesh=None, + cp_size=1, + pp_size=1, + ), + strategy_config=SimpleNamespace(calculate_per_token_loss=True), + pipeline_config=None, + moe_parallel_config=None, + activation_checkpointing=False, + ), + ) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + trainer.setup() + + assert trainer.engine is None + + def test_setup_does_not_change_storage_dtype_for_non_kd_recipe(monkeypatch): cfg = _minimal_cfg_with_nvtx(nvtx_value=False, optimizer_target="torch.optim.AdamW") @@ -2167,6 +2336,78 @@ def test_pp_scale_includes_pipeline_microbatches_and_token_normalization(self, m # Base CP-aware average: 2 / 8. PP post-normalization compensation: 6 / 8. assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.1875) + def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=True) + batches = [ + {"input_ids": torch.tensor([[1, 2, 3]]), "labels": torch.tensor([[1, 2, -100]])}, + {"input_ids": torch.tensor([[4, 5, 6]]), "labels": torch.tensor([[4, 5, -100]])}, + ] + datums = [object(), object()] + make_datum = MagicMock(side_effect=datums) + monkeypatch.setattr(recipe, "_make_engine_datum", make_datum) + monkeypatch.setattr(recipe, "_forward_backward_step", MagicMock(side_effect=AssertionError("legacy path"))) + monkeypatch.setattr( + recipe, + "_broadcast_from_last_pp_stage", + MagicMock(side_effect=AssertionError("Engine loss must not be broadcast again")), + ) + + engine = MagicMock() + engine.forward_backward.return_value = (torch.tensor(0.25), []) + object.__setattr__(recipe, "engine", engine) + + finalizer_calls = [] + + def finalize_grads(*args, **kwargs): + finalizer_calls.append((args, kwargs)) + return torch.tensor(1.0) + + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", finalize_grads) + optimizer = SimpleNamespace( + step=MagicMock(), + zero_grad=MagicMock(), + param_groups=[{"lr": 0.01}], + ) + object.__setattr__(recipe, "optimizer", [optimizer]) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) + + metrics = recipe._run_train_optim_step(batches) + + engine.forward_backward.assert_called_once_with([[datums[0]], [datums[1]]], recipe._engine_loss_fn) + assert make_datum.call_count == 2 + assert len(finalizer_calls) == 1 + assert finalizer_calls[0][1]["num_label_tokens"] is None + assert finalizer_calls[0][1]["pp_enabled"] is True + assert finalizer_calls[0][1]["pp_axis_name"] == "pp" + optimizer.step.assert_called_once_with() + optimizer.zero_grad.assert_called_once_with() + assert metrics.metrics["loss"] == pytest.approx(0.25) + + def test_pp_thd_batch_uses_legacy_forward_backward(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=True) + batch = { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[1, 2, -100]]), + "qkv_format": "thd", + } + engine = MagicMock() + object.__setattr__(recipe, "engine", engine) + + def legacy_step(_idx, _batch, *, loss_buffer, **_kwargs): + loss_buffer.append(torch.tensor(0.5)) + + legacy_step = MagicMock(side_effect=legacy_step) + monkeypatch.setattr(recipe, "_forward_backward_step", legacy_step) + finalizer = MagicMock(return_value=torch.tensor(1.0)) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", finalizer) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) + + recipe._run_train_optim_step([batch]) + + engine.forward_backward.assert_not_called() + legacy_step.assert_called_once() + assert finalizer.call_args.kwargs["num_label_tokens"] == 2 + @pytest.mark.parametrize("dp_size", [1, 8]) def test_non_pp_scale_is_independent_of_dp_size(self, monkeypatch, dp_size): from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler @@ -2256,9 +2497,10 @@ def test_rope_fusion_disabled_when_cp_gt_1(monkeypatch): assert trainer.engine is not None -def test_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): - cfg = _minimal_cfg_with_rope_fusion(cp_size=2, rope_fusion=True) - _patch_setup_minimals_with_cp(monkeypatch, cp_size=2) +@pytest.mark.parametrize(("cp_size", "expect_engine"), [(1, True), (2, False)]) +def test_setup_engine_gate_for_mtp_with_context_parallelism(monkeypatch, cp_size, expect_engine): + cfg = _minimal_cfg_with_rope_fusion(cp_size=cp_size, rope_fusion=True) + _patch_setup_minimals_with_cp(monkeypatch, cp_size=cp_size) model = DummyModel() model.mtp = nn.Identity() monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.build_model", lambda *args, **kwargs: model) @@ -2266,7 +2508,7 @@ def test_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) trainer.setup() - assert trainer.engine is None + assert (trainer.engine is not None) is expect_engine def test_setup_keeps_engine_disabled_for_magi(monkeypatch): diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 9959f01ee7..62b605ddec 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -37,6 +37,7 @@ ) from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh +from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler from nemo_automodel.engine import Engine, collate_prebatched @@ -73,6 +74,58 @@ def __init__(self, size, rank): self.mesh_dim_names = ("cp", "tp") +class _FakeAutoPipeline(AutoPipeline): + def __init__(self, model, *, parts=None, num_microbatches=2, scale_grads=False, events=None): + self.compute_model = model + self._parts = parts or [model] + self._num_microbatches = num_microbatches + self.scale_grads_in_schedule = scale_grads + self.pp_mesh = _SubMesh(2) + self.events = events + self.step_calls = 0 + self.backward_calls = 0 + self.updated_seq_lens = [] + self.callback_losses = [] + + @property + def parts(self): + return self._parts + + @property + def num_microbatches(self): + return self._num_microbatches + + def update_seq_len(self, seq_len): + self.updated_seq_lens.append(seq_len) + + def step(self, model_input, *, loss_inputs, loss_fn, return_outputs, **kwargs): + assert return_outputs is False + self.step_calls += 1 + if self.events is not None: + self.events.append("step") + + def chunks(value): + if isinstance(value, torch.Tensor) and value.ndim > 0: + return value.chunk(self.num_microbatches, dim=0) + return (value,) * self.num_microbatches + + model_chunks = chunks(model_input) + kwargs_chunks = {name: chunks(value) for name, value in kwargs.items()} + loss_chunks = {name: chunks(value) for name, value in loss_inputs.items()} + for index in range(self.num_microbatches): + model_kwargs = {name: values[index] for name, values in kwargs_chunks.items()} + loss_inputs_mb = {name: values[index] for name, values in loss_chunks.items()} + output = self.compute_model(model_chunks[index], **model_kwargs) + scaled_loss = loss_fn(output, loss_inputs_mb, (model_chunks[index],), model_kwargs) + self.callback_losses.append(scaled_loss.detach()) + scaled_loss.backward() + self.backward_calls += 1 + + +def _pipeline_mesh_context(): + return SimpleNamespace(pp_size=2, cp_size=1, device_mesh=None, process_group=None) + + class _DDPWithCP(nn.parallel.DistributedDataParallel): def prepare_model_inputs_for_cp(self, batch, *, num_chunks): return self.module.prepare_model_inputs_for_cp(batch, num_chunks=num_chunks) @@ -390,6 +443,202 @@ def sync_context(_model, is_last, _defer): assert events == ["prepare", "sync:False", "after_first", "final", "sync:True"] +def test_pipeline_window_uses_schedule_microbatches_and_global_normalization(): + active = False + context_events = [] + backward_calls = 0 + + @contextmanager + def forward_context(): + nonlocal active + assert not active + active = True + context_events.append("enter") + try: + yield + finally: + active = False + context_events.append("exit") + + class PipelineModel(ScaleModel): + def forward(self, input_ids, **kwargs): + assert active + return super().forward(input_ids, **kwargs) + + model = PipelineModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + + def check_backward_context(grad): + nonlocal backward_calls + assert active + backward_calls += 1 + return grad + + model.weight.register_hook(check_backward_context) + + def loss_fn(output, inputs, datums, model_inputs): + assert active + assert len(datums) == 1 + assert output.shape == inputs["weights"].shape == model_inputs["input_ids"].shape == (1, 2) + return output + + loss, outputs = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + context_fn=forward_context, + ).forward_backward( + [ + [_datum([[1, 2], [3, 4]])], + [_datum([[5, 6], [7, 8]])], + ], + loss_fn, + ) + + assert loss.item() == pytest.approx(4.5) + assert model.weight.grad.item() == pytest.approx(4.5) + assert outputs == [] + assert pipeline.step_calls == 2 + # The fake schedule performs and counts every backward, then returns None. + # A second Engine-owned backward would either fail or change these counts. + assert pipeline.backward_calls == backward_calls == 4 + assert model.forward_calls == 4 + assert pipeline.updated_seq_lens == [2, 2] + scaled_losses = torch.stack(pipeline.callback_losses) + torch.testing.assert_close( + scaled_losses, + torch.tensor([3 / 8, 7 / 8, 11 / 8, 15 / 8], dtype=scaled_losses.dtype), + ) + assert context_events == ["enter", "exit", "enter", "exit"] + assert not active + + +def test_pipeline_lifecycle_and_moe_scale_cover_outer_and_inner_microbatches(monkeypatch): + events = [] + model = ScaleModel() + other_part = ScaleModel() + model.eval() + other_part.eval() + pipeline = _FakeAutoPipeline( + model, + parts=[model, other_part], + num_microbatches=2, + events=events, + ) + + def prepare(parts, *, pp_enabled): + assert parts == [model, other_part] + assert pp_enabled is True + events.append("prepare") + + def prepare_final(parts, *, pp_enabled): + assert parts == [model, other_part] + assert pp_enabled is True + events.append("final") + + monkeypatch.setattr(engine_module, "prepare_for_grad_accumulation", prepare) + monkeypatch.setattr(engine_module, "prepare_for_final_backward", prepare_final) + monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after_first")) + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) + + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward_backward( + [ + [_datum([[1], [2]])], + [_datum([[3], [4]])], + ], + _identity_loss, + ) + + assert events == ["prepare", "step", "after_first", "final", "step"] + assert model.training + assert other_part.training + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) + + +def test_pipeline_rejects_per_datum_outputs_before_schedule_backward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + + with pytest.raises(ValueError, match="per-Datum"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward_backward( + [[_datum([[1], [2]])]], + lambda output, _inputs, _datums, _model_inputs: (output, [{"metric": output.sum()}]), + ) + + assert pipeline.step_calls == 1 + assert pipeline.backward_calls == 0 + assert model.forward_calls == 1 + assert model.weight.grad is None + + +def test_pipeline_rejects_schedule_gradient_scaling_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, scale_grads=True) + + with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): + Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()).forward_backward( + [[_datum([1])]], _identity_loss + ) + + assert pipeline.step_calls == 0 + assert model.forward_calls == 0 + assert model.weight.grad is None + + +def test_pipeline_rejects_multiple_datums_in_one_outer_batch_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model) + + with pytest.raises(ValueError, match="exactly one prebatched Datum"): + Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()).forward_backward( + [[_datum([1]), _datum([2])]], + _identity_loss, + ) + + assert pipeline.step_calls == 0 + assert model.forward_calls == 0 + assert model.weight.grad is None + + +def test_pipeline_requires_mesh_context_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model) + + with pytest.raises(ValueError, match="requires mesh_context"): + Engine(pipeline, device="cpu").forward_backward([[_datum([1])]], _identity_loss) + + assert pipeline.step_calls == 0 + assert model.forward_calls == 0 + assert model.weight.grad is None + + +def test_pipeline_rejects_context_parallelism_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model) + mesh_context = SimpleNamespace(pp_size=2, cp_size=2, device_mesh=_CPMesh(size=2, rank=0)) + + with pytest.raises(NotImplementedError, match="does not yet support context parallelism"): + Engine(pipeline, device="cpu", mesh_context=mesh_context).forward_backward( + [[_datum([1])]], + _identity_loss, + ) + + assert pipeline.step_calls == 0 + assert model.forward_calls == 0 + assert model.weight.grad is None + + def test_forward_context_covers_forward_loss_and_backward(): active = False From a8a3655011dea861e96bad496fa2c016d12cd1a8 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 17 Aug 2026 02:02:38 -0700 Subject: [PATCH 06/34] feat(engine): unify recipe forward backward across parallelism Signed-off-by: HuiyingLi --- .../mistral/mixtral-8x7b-v0-1_squad.yaml | 10 +- .../qwen3_5_4b_cp2_vision_frame_shard.yaml | 2 + .../qwen3_5_122b_128k_ep8cp32.yaml | 2 + .../qwen3_6_35b_medpix_ep8cp2_4k.yaml | 2 + .../components/datasets/vlm/pp_media.py | 110 +- .../distributed/context_parallel/magi.py | 21 + .../distributed/context_parallel/utils.py | 85 +- .../distributed/pipelining/autopipeline.py | 211 ++- .../distributed/pipelining/functional.py | 20 +- .../distributed/pipelining/hf_utils.py | 116 +- .../components/distributed/thd_utils.py | 285 ++- nemo_automodel/components/loss/mtp.py | 11 + .../nemotron_parse/nemotron_parse_loss.py | 11 +- nemo_automodel/engine/__init__.py | 625 +++++-- nemo_automodel/recipes/base_recipe.py | 36 + nemo_automodel/recipes/llm/train_ft.py | 327 ++-- nemo_automodel/recipes/vlm/finetune.py | 388 ++-- .../L2_PP_Dense_Packed_Test.sh | 22 + .../context_parallel/run_cp_pp_image_sink.py | 99 +- .../context_parallel/run_cp_pp_layer2_sink.py | 214 ++- .../context_parallel/run_dense_packed_cp.py | 6 +- .../context_parallel/run_packed_pp.py | 318 ++++ .../context_parallel/test_context_parallel.py | 5 + .../moe/test_experts_ep_tp_grad_parity.py | 124 +- .../pipelining/test_autopipeline.py | 290 ++- .../distributed/pipelining/test_functional.py | 18 + .../distributed/pipelining/test_hf_utils.py | 60 + tests/unit_tests/distributed/test_cp_utils.py | 24 +- .../distributed/test_magi_attn_utils.py | 132 ++ .../unit_tests/distributed/test_thd_utils.py | 87 + .../loss/test_mtp_lm_head_gather.py | 23 + .../loss/test_nemotron_parse_loss.py | 35 +- tests/unit_tests/recipes/test_base_recipe.py | 23 + .../recipes/test_finetune_vlm_cp_wiring.py | 284 +-- .../recipes/test_finetune_vlm_helpers.py | 1607 ++++------------- tests/unit_tests/recipes/test_train_ft.py | 387 ++-- tests/unit_tests/test_engine.py | 554 ++++-- .../test_engine_recipe_integration.py | 18 +- 38 files changed, 3566 insertions(+), 3026 deletions(-) create mode 100755 tests/functional_tests/context_parallel/L2_PP_Dense_Packed_Test.sh create mode 100644 tests/functional_tests/context_parallel/run_packed_pp.py diff --git a/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml b/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml index 9c2ffaa1aa..f223ce8a40 100644 --- a/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml +++ b/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 diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml index d0cf627a66..47944ccdaa 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml @@ -59,6 +59,8 @@ model: rope_fusion: false attn_implementation: sdpa torch_dtype: bfloat16 + # MTP's future-token shift is not yet context-parallel aware. + num_nextn_predict_layers: 0 processor: _target_: transformers.AutoProcessor.from_pretrained diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml index d47dff48b9..3af3c425c3 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml @@ -70,6 +70,8 @@ model: text_config: # FusedLinearCrossEntropy consumes the decoder hidden states directly. output_hidden_states: true + # MTP's future-token shift is not yet context-parallel aware. + num_nextn_predict_layers: 0 # Qwen3.5 stores MTP expert projections as per-expert tensors in HF checkpoints. mtp_expert_hf_layout: split backend: diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml index 1e841b9d73..b363b07f82 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml @@ -31,6 +31,8 @@ model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.6-35B-A3B trust_remote_code: false + # MTP's future-token shift is not yet context-parallel aware. + num_nextn_predict_layers: 0 backend: _target_: nemo_automodel.components.models.common.BackendConfig attn: te diff --git a/nemo_automodel/components/datasets/vlm/pp_media.py b/nemo_automodel/components/datasets/vlm/pp_media.py index 9d20f89a4c..1afb858430 100644 --- a/nemo_automodel/components/datasets/vlm/pp_media.py +++ b/nemo_automodel/components/datasets/vlm/pp_media.py @@ -39,21 +39,96 @@ def chunk_vlm_media( - pixel_values: torch.Tensor, - image_grid: torch.Tensor, + pixel_values: torch.Tensor | list[torch.Tensor], + image_grid: torch.Tensor | None, batch_size: int, n_microbatches: int, n_images_per_sample: torch.Tensor | None = None, -) -> tuple[list[torch.Tensor], list[torch.Tensor]]: +) -> tuple[list[torch.Tensor | list[torch.Tensor]], list[torch.Tensor] | None]: """Split VLM pixel values and media metadata into PP microbatch chunks. - Handles four layouts: + Handles five layouts: 1. ``[N, C, H, W]`` with ``N == batch_size`` -- one full image per sample. 2. ``[N, max_patches, D]`` with ``N == batch_size`` -- padded patches per image. 3. Flat patches ``[total_patches, D]`` with per-sample media counts from ``n_images_per_sample``. 4. Flat patches with ``n_images == batch_size`` -- legacy one-image-per-sample. + 5. Variable-resolution media lists, split at sample boundaries using + ``n_images_per_sample`` (or one media item per sample when counts are absent). + + Args: + pixel_values: Tensor of shape [media, channels, height, width], [media, patches, hidden], or + [patches, hidden], or a list of ``media`` tensors with arbitrary processor-defined shapes. + image_grid: Optional tensor of shape [media, grid_dims]. It may be ``None`` only for + variable-resolution media lists. + batch_size: Number of text samples represented by the media. + n_microbatches: Number of pipeline microbatches to materialize. + n_images_per_sample: Optional integer tensor of shape [batch] mapping samples to media entries. + + Returns: + A pair containing the media chunks and optional grid chunks in pipeline-microbatch order. Tensor chunks + are views of ``pixel_values``; list chunks retain references to the original media tensors. """ + if isinstance(pixel_values, list): + if not all(isinstance(value, torch.Tensor) for value in pixel_values): + raise TypeError("variable-resolution pixel_values must be a list of tensors") + if n_images_per_sample is None: + if len(pixel_values) != batch_size: + raise ValueError( + "VLM PP chunking requires n_images_per_sample for variable-resolution media " + f"when len(pixel_values)={len(pixel_values)} differs from batch_size={batch_size}." + ) + media_counts = torch.ones(batch_size, dtype=torch.long) + else: + if not isinstance(n_images_per_sample, torch.Tensor) or n_images_per_sample.ndim != 1: + raise ValueError("n_images_per_sample must be a one-dimensional tensor") + if n_images_per_sample.numel() != batch_size: + raise ValueError( + f"n_images_per_sample must have length batch_size={batch_size}, " + f"got shape={tuple(n_images_per_sample.shape)}." + ) + if ( + n_images_per_sample.dtype == torch.bool + or n_images_per_sample.is_floating_point() + or n_images_per_sample.is_complex() + ): + raise ValueError("n_images_per_sample must contain integer media counts") + media_counts = n_images_per_sample.to(dtype=torch.long, device="cpu") + if bool((media_counts < 0).any()): + raise ValueError("n_images_per_sample must contain non-negative media counts") + + total_media = int(media_counts.sum().item()) + if total_media != len(pixel_values): + raise ValueError( + "VLM PP chunking cannot align variable-resolution media with sample counts: " + f"len(pixel_values)={len(pixel_values)}, sum(n_images_per_sample)={total_media}." + ) + if image_grid is not None: + if not isinstance(image_grid, torch.Tensor) or image_grid.ndim == 0: + raise ValueError("image_grid must be a non-scalar tensor when variable-resolution media uses a grid") + if image_grid.shape[0] != total_media: + raise ValueError( + "VLM PP chunking cannot align image_grid with variable-resolution media: " + f"image_grid.shape[0]={image_grid.shape[0]}, len(pixel_values)={total_media}." + ) + + media_offsets = torch.cat((torch.zeros(1, dtype=torch.long), media_counts.cumsum(dim=0))) + samples_per_mb = -(-batch_size // n_microbatches) + pixel_values_chunks: list[torch.Tensor | list[torch.Tensor]] = [] + image_grid_chunks: list[torch.Tensor] | None = [] if image_grid is not None else None + for mb_idx in range(n_microbatches): + sample_start = min(mb_idx * samples_per_mb, batch_size) + sample_end = min(sample_start + samples_per_mb, batch_size) + media_start = int(media_offsets[sample_start].item()) + media_end = int(media_offsets[sample_end].item()) + pixel_values_chunks.append(pixel_values[media_start:media_end]) + if image_grid_chunks is not None: + image_grid_chunks.append(image_grid[media_start:media_end]) + return pixel_values_chunks, image_grid_chunks + + if image_grid is None: + raise ValueError("VLM PP media prep requires image-grid metadata with tensor pixel_values.") + n_images = image_grid.shape[0] pixel_values_chunks: list[torch.Tensor] = [] image_grid_chunks: list[torch.Tensor] = [] @@ -230,6 +305,17 @@ def prepare_vlm_media_for_pp( The returned batch no longer carries raw media tensors that PyTorch PP would chunk by row incorrectly; instead it carries ``VLM_PP_MEDIA_KEY`` with per-microbatch media chunks. + + Args: + batch: Mutable processor batch containing ``input_ids`` of shape [batch, sequence] and optional media + tensors or variable-resolution media lists accepted by :func:`chunk_vlm_media`. This mapping is + mutated in place: raw media fields are removed and replaced by pre-chunked PP storage. + batch_size: Number of text samples in ``batch``. + n_microbatches: Number of pipeline microbatches to materialize. + + Returns: + The mutated ``batch`` mapping. ``VLM_PP_MEDIA_KEY`` maps media field names to lists in pipeline-microbatch + order; each entry is either a tensor chunk or a list of variable-resolution media tensors. """ if n_microbatches < 1: raise ValueError(f"n_microbatches must be >= 1, got {n_microbatches}") @@ -251,9 +337,9 @@ def prepare_vlm_media_for_pp( n_videos_per_sample = batch.pop("n_videos_per_sample", None) image_grid = _select_image_grid(image_grid_hws, image_grid_thw, image_sizes, image_position_ids) - pp_media: dict[str, list[torch.Tensor]] = {} + pp_media: dict[str, list[Any]] = {} - if pixel_values is not None and image_grid is None: + if isinstance(pixel_values, torch.Tensor) and image_grid is None: step3_media = chunk_step3_media( pixel_values, batch_size=batch_size, @@ -264,10 +350,10 @@ def prepare_vlm_media_for_pp( ) pp_media.update(step3_media) - if pixel_values_videos is not None and video_grid_thw is None: + if isinstance(pixel_values_videos, torch.Tensor) and video_grid_thw is None: raise ValueError("VLM PP media prep requires video_grid_thw with pixel_values_videos.") - if pixel_values is not None and image_grid is not None: + if pixel_values is not None and (image_grid is not None or isinstance(pixel_values, list)): pixel_values_chunks, image_grid_chunks = chunk_vlm_media( pixel_values, image_grid, @@ -276,9 +362,10 @@ def prepare_vlm_media_for_pp( n_images_per_sample=n_images_per_sample, ) pp_media["pixel_values"] = pixel_values_chunks - pp_media["image_grid_hws"] = image_grid_chunks + if image_grid_chunks is not None: + pp_media["image_grid_hws"] = image_grid_chunks - if pixel_values_videos is not None and video_grid_thw is not None: + if pixel_values_videos is not None and (video_grid_thw is not None or isinstance(pixel_values_videos, list)): pixel_values_videos_chunks, video_grid_thw_chunks = chunk_vlm_media( pixel_values_videos, video_grid_thw, @@ -287,7 +374,8 @@ def prepare_vlm_media_for_pp( n_images_per_sample=n_videos_per_sample, ) pp_media["pixel_values_videos"] = pixel_values_videos_chunks - pp_media["video_grid_thw"] = video_grid_thw_chunks + if video_grid_thw_chunks is not None: + pp_media["video_grid_thw"] = video_grid_thw_chunks if pp_media: batch[VLM_PP_MEDIA_KEY] = pp_media diff --git a/nemo_automodel/components/distributed/context_parallel/magi.py b/nemo_automodel/components/distributed/context_parallel/magi.py index 3ca3f1b9d1..e7d846825f 100644 --- a/nemo_automodel/components/distributed/context_parallel/magi.py +++ b/nemo_automodel/components/distributed/context_parallel/magi.py @@ -780,6 +780,11 @@ def prepare_llm_batch( # (a spec or None) is self-clearing, so a stale spec never leaks into the next # batch; plain batches omit "prefix_tree". prefix_tree = batch.pop("prefix_tree", None) + if prefix_tree is not None and self.cp_size > 1: + raise NotImplementedError( + "The prefix-tree attention mask currently requires cp_size=1; " + "distributed prefix-tree dispatch is not wired yet." + ) if prefix_tree is not None and self.hf_dispatch: # The prefix-tree mask is handed to the attn_func out-of-band (the HF # attention interface has a fixed signature and cannot receive a custom @@ -798,6 +803,12 @@ def prepare_llm_batch( set_active_attn_spec(spec) local_indices = None if self.hf_dispatch: + if is_thd: + raise NotImplementedError( + "The HF magi backend does not support packed THD batches because its fixed " + "attention interface cannot preserve packed document boundaries; use the custom-model " + "magi backend (model.backend.attn='magi')." + ) # HF path: dispatch the (single causal) sequence across the CP group. batch, _, local_indices = magi_prepare_batch(model, batch, self.cp_group, return_local_indices=True) elif self.custom and self.cp_size > 1 and is_thd: @@ -805,6 +816,11 @@ def prepare_llm_batch( # sharding) then dispatch it with magi's own load-balancing solver. batch = make_cp_batch_for_te(None, batch, qkv_format="thd", padding_token_id=pad_id, num_chunks=1) batch, _, local_indices = magi_prepare_packed_cp(model, batch, self.cp_group, return_local_indices=True) + elif self.custom and self.cp_size > 1: + # A plain causal custom-model batch uses the same dist key and + # dispatch as the HF path. The custom attention callable reads the + # key from the active CP group instead of the stamped modules. + batch, _, local_indices = magi_prepare_batch(model, batch, self.cp_group, return_local_indices=True) elif is_thd: # cp=1 packing: THD conversion (no sharding) so the batch carries # cu_seqlens -> the magi attn_func builds the per-document mask. @@ -866,6 +882,11 @@ def make_cp_batch( """ del cp_mesh local_indices = None + if self.custom: + # Engine callers do not have to run recipe-level setup_magi first. + # Refreshing this process-local handle also makes the active group + # explicit for every outer PP accumulation batch. + set_active_cp_group(self.cp_group) if self.domain == "vlm": _, batch = self.prepare_vlm_batch(model, batch) else: diff --git a/nemo_automodel/components/distributed/context_parallel/utils.py b/nemo_automodel/components/distributed/context_parallel/utils.py index 6be7e4c556..5e77187818 100644 --- a/nemo_automodel/components/distributed/context_parallel/utils.py +++ b/nemo_automodel/components/distributed/context_parallel/utils.py @@ -29,6 +29,8 @@ ) from nemo_automodel.components.distributed.thd_utils import ( split_batch_into_thd_chunks, + split_final_thd_batch, + stack_thd_chunks, thd_padding_mask_from_token_ids, ) @@ -419,6 +421,11 @@ def _prepare_cp_sharder( """ batch_is_thd = batch.get("qkv_format") == "thd" + if batch_is_thd and bool(getattr(model, "_te_attention_injected", False)): + raise NotImplementedError( + "THD inputs require a native THD-capable model backend; " + "Transformer Engine injected into a stock Hugging Face model supports padded BSHD only" + ) magi_state = _magi_state_from_model(model, device_mesh) magi_enabled = magi_state is not None and getattr(magi_state, "enabled", False) backend_uses_thd = batch_is_thd and (magi_enabled or _uses_te_attention(model)) @@ -770,21 +777,30 @@ def make_cp_batch_for_te( if qkv_format != "thd": raise ValueError(f"Currently only 'thd' format is supported, got: {qkv_format}") - batch = split_batch_into_thd_chunks( - batch, - num_chunks=num_chunks, - seq_lens_padding_value=seq_lens_padding_value, - padding_token_id=padding_token_id, - ) + final_thd = "cu_seqlens" in batch and "seq_lens" not in batch and "seq_lens_padded" not in batch + if final_thd and num_chunks > 1: + batch = stack_thd_chunks( + split_final_thd_batch(batch, num_chunks, seq_lens_padding_value), + seq_lens_padding_value, + ) + elif not final_thd: + batch = split_batch_into_thd_chunks( + batch, + num_chunks=num_chunks, + seq_lens_padding_value=seq_lens_padding_value, + padding_token_id=padding_token_id, + ) if cp_mesh is None or cp_mesh.size() <= 1: if not return_local_indices: return batch # Unsharded THD stream: identity index map. Chunked streams are # per-chunk token spaces with no single step-wide map -> None. - input_ids = batch["input_ids"] + primary = batch.get("inputs_embeds", batch.get("input_ids")) + if not isinstance(primary, torch.Tensor): + raise ValueError("THD batch requires tensor input_ids or inputs_embeds") local_indices = ( - torch.arange(input_ids.shape[-1], device=input_ids.device, dtype=torch.long) if num_chunks <= 1 else None + torch.arange(primary.shape[0], device=primary.device, dtype=torch.long) if num_chunks <= 1 else None ) return batch, local_indices @@ -797,22 +813,16 @@ def make_cp_batch_for_te( # Extract each chunk from the batched result and shard it chunks = [] for i in range(num_chunks): - chunk_batch = {k: v[i] if isinstance(v, torch.Tensor) else v for k, v in batch.items()} + chunk_batch = { + key: value[i] + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == num_chunks + else value + for key, value in batch.items() + } chunks.append( _shard_thd_chunk_for_te(chunk_batch, cp_mesh, qkv_format, seq_lens_padding_value, padding_token_id)[0] ) - - return_dict = { - "input_ids": torch.stack([chunk["input_ids"] for chunk in chunks]), - "labels": torch.stack([chunk["labels"] for chunk in chunks]), - "position_ids": torch.stack([chunk["position_ids"] for chunk in chunks]), - "cu_seqlens": torch.stack([chunk["cu_seqlens"] for chunk in chunks]), - "max_seqlen": torch.stack([chunk["max_seqlen"] for chunk in chunks]), - "qkv_format": qkv_format, - "padding_mask": torch.stack([chunk["padding_mask"] for chunk in chunks]), - "cp_size": cp_mesh.size() if cp_mesh is not None else 1, - "cp_rank": torch.distributed.get_rank(group=cp_mesh.get_group()) if cp_mesh is not None else 0, - } + return_dict = stack_thd_chunks(chunks, seq_lens_padding_value) # Chunked mode: each chunk is its own token space, so there is no single # step-wide local-token index map to expose. @@ -846,13 +856,19 @@ def _shard_thd_chunk_for_te( # The partition is the same for every token-aligned key; it is also this # rank's local-token global index map, returned so the caller can install # it on the THD sharder (ContextParallelSharder token verbs). - local_indices = tex.thd_get_partitioned_indices( - filtered_cu_seqlens_padded, batch["input_ids"].size(0), cp_size, cp_rank - ) - mask_keys = ["input_ids", "labels", "position_ids", "padding_mask"] - for key in mask_keys: - if key in batch: - batch[key] = batch[key].index_select(0, local_indices) + primary_name = "inputs_embeds" if "inputs_embeds" in batch else "input_ids" + primary = batch[primary_name] + total_tokens = primary.size(0) + local_indices = tex.thd_get_partitioned_indices(filtered_cu_seqlens_padded, total_tokens, cp_size, cp_rank) + token_keys = {"input_ids", "inputs_embeds", "labels", "position_ids", "padding_mask"} + for key, value in batch.items(): + if ( + (key in token_keys or key.startswith("__engine_loss__")) + and isinstance(value, torch.Tensor) + and value.ndim > 0 + and value.shape[0] == total_tokens + ): + batch[key] = value.index_select(0, local_indices) # Keep model-owned payloads (for example VLM media) by default. Only remove # source metadata that is invalid after the THD CP conversion; the update @@ -862,11 +878,16 @@ def _shard_thd_chunk_for_te( output_batch.pop("cu_seqlens_padded", None) max_seqlen = (filtered_cu_seqlens_padded[1:] - filtered_cu_seqlens_padded[:-1]).max().item() + output_batch[primary_name] = ( + batch[primary_name].to(torch.int64).contiguous() + if primary_name == "input_ids" + else batch[primary_name].contiguous() + ) + output_batch["labels"] = batch["labels"].to(torch.int64).contiguous() + if isinstance(batch.get("position_ids"), torch.Tensor): + output_batch["position_ids"] = batch["position_ids"].to(torch.int64).contiguous() output_batch.update( { - "input_ids": batch["input_ids"].to(torch.int64).contiguous(), - "labels": batch["labels"].to(torch.int64).contiguous(), - "position_ids": batch["position_ids"].to(torch.int64).contiguous(), "cu_seqlens": cu_seqlens_padded.to(torch.int32).contiguous(), "max_seqlen": torch.tensor(max_seqlen).to(torch.int32).to(device=cu_seqlens_padded.device), "qkv_format": qkv_format, @@ -881,6 +902,8 @@ def _shard_thd_chunk_for_te( if "padding_mask" in batch: output_batch["padding_mask"] = batch["padding_mask"].bool().contiguous() else: + if primary_name != "input_ids": + raise ValueError("THD inputs_embeds require an explicit padding_mask") output_batch["padding_mask"] = thd_padding_mask_from_token_ids( output_batch["input_ids"], padding_token_id ).contiguous() diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index 4bee514535..c3ebc62aa2 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -20,7 +20,7 @@ import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.pipelining.microbatch import BlockMask, TensorChunkSpec, split_args_kwargs_into_chunks +from torch.distributed.pipelining.microbatch import BlockMask, TensorChunkSpec from torch.distributed.pipelining.microbatch import _Replicate as ReplicateChunkSpec from torch.distributed.pipelining.schedules import _PipelineSchedule from torch.distributed.pipelining.stage import PipelineStage @@ -129,6 +129,8 @@ def __init__( ) self._model_config = None self._pp_current_seq_len: Optional[int] = None + self._pp_current_microbatch_size: Optional[int] = None + self._pp_current_input_signature: tuple[tuple[int, ...], torch.dtype] | None = None def build( self, @@ -190,8 +192,14 @@ def build( def info(self) -> PipelineInfo: return self._info - def update_seq_len(self, seq_len: int) -> None: - """Reset pipeline stage infrastructure for a new sequence length. + def update_seq_len( + self, + seq_len: int, + *, + microbatch_size: int | None = None, + input_tensor: torch.Tensor | None = None, + ) -> None: + """Reset pipeline stage infrastructure for a new input shape. VLM training batches can have wildly different sequence lengths across steps (image batches vs. text-only batches). PyTorch's PipelineStage locks in recv @@ -203,8 +211,25 @@ def update_seq_len(self, seq_len: int) -> None: Args: seq_len: Sequence length of the upcoming batch (``input_ids.shape[1]``). + microbatch_size: Materialized leading batch extent. Defaults to the + configured padded microbatch size. Packed THD streams use one + synthetic batch row even when they were formed from multiple + source examples. + input_tensor: Actual positional input for the first stage. When + provided, its shape and dtype define first-stage metadata; + this supports floating-point embeddings and flat THD tensors. """ - if seq_len == self._pp_current_seq_len: + effective_microbatch_size = self.pp_microbatch_size if microbatch_size is None else microbatch_size + input_signature = ( + (tuple(input_tensor.shape), input_tensor.dtype) if isinstance(input_tensor, torch.Tensor) else None + ) + if effective_microbatch_size <= 0: + raise ValueError(f"microbatch_size must be positive, got {effective_microbatch_size}") + if ( + seq_len == self._pp_current_seq_len + and effective_microbatch_size == self._pp_current_microbatch_size + and input_signature == self._pp_current_input_signature + ): return if self._model_config is None: raise RuntimeError("AutoPipeline.build() must be called before update_seq_len()") @@ -212,12 +237,23 @@ def update_seq_len(self, seq_len: int) -> None: self._info.schedule, self._info.stages, self._model_config, - self.pp_microbatch_size, + effective_microbatch_size, seq_len, tensor_dtype=self.dtype, + first_stage_input_meta=( + torch.empty(tuple(input_tensor.shape), device="meta", dtype=input_tensor.dtype) + if isinstance(input_tensor, torch.Tensor) + else None + ), ) self._pp_current_seq_len = seq_len - logger.debug(f"PP stage shapes updated for seq_len={seq_len}") + self._pp_current_microbatch_size = effective_microbatch_size + self._pp_current_input_signature = input_signature + logger.debug( + "PP stage shapes updated for seq_len=%s, microbatch_size=%s", + seq_len, + effective_microbatch_size, + ) def _get_schedule_kwargs_chunk_spec(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: """Build pipeline microbatch chunking metadata for keyword inputs. @@ -269,8 +305,6 @@ def step( *, target: torch.Tensor | None = None, losses: list[torch.Tensor] | None = None, - loss_inputs: dict[str, Any] | None = None, - loss_fn: Callable | None = None, return_outputs: bool = True, **kwargs: Any, ) -> Any: @@ -283,14 +317,6 @@ def step( ranks without the last pipeline stage. losses: Mutable list populated with scalar loss tensors, or ``None`` on ranks without the last pipeline stage. - loss_inputs: Structured loss inputs. Tensor fields of shape - [batch, ...] are split on the batch axis; scalar tensors and - non-tensor fields are replicated. Must be provided together - with ``loss_fn`` and without ``target``. - loss_fn: Callback invoked as ``loss_fn(output, loss_inputs_mb, - model_args_mb, model_kwargs_mb)`` for each microbatch. Model - tensor arguments use shape [microbatch, ...], while keyword - tensors retain their model-defined layouts. return_outputs: Whether the last pipeline stage returns merged model outputs when supported by the installed PyTorch version. Tensor layouts are defined by the underlying model. @@ -305,11 +331,6 @@ def step( if schedule is None: raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") - if (loss_inputs is None) != (loss_fn is None): - raise ValueError("loss_inputs and loss_fn must be provided together") - if loss_inputs is not None and target is not None: - raise ValueError("target cannot be used together with loss_inputs and loss_fn") - schedule_args = (model_input,) if self._info.has_first_stage else () kwargs_chunk_spec = self._get_schedule_kwargs_chunk_spec(kwargs) schedule_options = ( @@ -318,89 +339,99 @@ def step( else {} ) - if loss_inputs is None: - if kwargs_chunk_spec is None: - return schedule.step( - *schedule_args, - target=target, - losses=losses, - **schedule_options, - **kwargs, - ) - - previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec - schedule._kwargs_chunk_spec = kwargs_chunk_spec - try: - return schedule.step( - *schedule_args, - target=target, - losses=losses, - **schedule_options, - **kwargs, - ) - finally: - schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec - - model_args_chunks, model_kwargs_chunks = split_args_kwargs_into_chunks( - (model_input,), - kwargs, - self.num_microbatches, - kwargs_chunk_spec=kwargs_chunk_spec, - ) - if len(model_args_chunks) != self.num_microbatches: - raise ValueError(f"Expected {self.num_microbatches} model input microbatches, got {len(model_args_chunks)}") - - loss_inputs_chunk_spec = tree_map( - lambda value: ( - TensorChunkSpec(0) if isinstance(value, torch.Tensor) and value.ndim > 0 else ReplicateChunkSpec() - ), - loss_inputs, - is_leaf=lambda value: isinstance(value, BlockMask), - ) - _, loss_inputs_chunks = split_args_kwargs_into_chunks( - (), - loss_inputs, - self.num_microbatches, - kwargs_chunk_spec=loss_inputs_chunk_spec, - ) - if len(loss_inputs_chunks) != self.num_microbatches: - raise ValueError(f"Expected {self.num_microbatches} loss input microbatches, got {len(loss_inputs_chunks)}") - - def microbatch_loss(output: Any, microbatch_id: torch.Tensor) -> Any: - """Evaluate the structured loss for one pipeline microbatch. - - Args: - output: Model output whose tensor layouts are defined by the model. - microbatch_id: Tensor of shape [1] identifying the microbatch. - - Returns: - The loss value returned by ``loss_fn``. Tensor layout is defined - by the callback. - """ - index = int(microbatch_id.item()) - return loss_fn( - output, - loss_inputs_chunks[index], - model_args_chunks[index], - model_kwargs_chunks[index], + if kwargs_chunk_spec is None: + return schedule.step( + *schedule_args, + target=target, + losses=losses, + **schedule_options, + **kwargs, ) previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec - previous_loss_fn = schedule._loss_fn schedule._kwargs_chunk_spec = kwargs_chunk_spec - schedule._loss_fn = microbatch_loss try: return schedule.step( *schedule_args, - target=torch.arange(self.num_microbatches, device=self.device), + target=target, losses=losses, **schedule_options, **kwargs, ) finally: - schedule._loss_fn = previous_loss_fn schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec + def step_microbatches( + self, + model_inputs: list[dict[str, Any]], + *, + loss_fn: Callable[[Any, int], Any], + losses: list[torch.Tensor] | None = None, + return_outputs: bool = False, + ) -> Any: + """Run a schedule step over already prepared model microbatches. + + The caller owns microbatch preparation. This method passes the prepared + inputs to the PyTorch schedule unchanged, while the schedule continues + to own execution order and backward calls. + + Args: + model_inputs: Exactly :attr:`num_microbatches` complete model-input + mappings. Each mapping contains exactly one of ``input_ids`` or + ``inputs_embeds``; all remaining items are model keyword inputs. + loss_fn: Callback invoked as ``loss_fn(output, microbatch_index)``. + The index identifies the corresponding item in ``model_inputs`` + regardless of the schedule's execution order. + losses: Mutable list populated by the schedule on the last stage. + return_outputs: Whether the last stage returns merged model outputs + when supported by the installed PyTorch version. + + Returns: + The value returned by the underlying PyTorch pipeline schedule. + """ + schedule = self._info.schedule + if schedule is None: + raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") + if len(model_inputs) != self.num_microbatches: + raise ValueError(f"Expected {self.num_microbatches} model input microbatches, got {len(model_inputs)}") + + model_args_chunks: list[tuple[Any, ...]] = [] + model_kwargs_chunks: list[dict[str, Any]] = [] + for index, inputs in enumerate(model_inputs): + if not isinstance(inputs, dict): + raise TypeError(f"model input microbatch {index} must be a dict") + primary_names = [name for name in ("input_ids", "inputs_embeds") if name in inputs] + if len(primary_names) != 1: + raise ValueError( + f"model input microbatch {index} must contain exactly one of input_ids or inputs_embeds" + ) + kwargs = dict(inputs) + primary = kwargs.pop(primary_names[0]) + model_args_chunks.append((primary,) if self._info.has_first_stage else ()) + model_kwargs_chunks.append(kwargs) + + def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: + return loss_fn(output, int(microbatch_id.item())) + + schedule_options = ( + {"return_outputs": return_outputs} + if "return_outputs" in inspect.signature(schedule.step).parameters + else {} + ) + previous_split_inputs = schedule._split_inputs + previous_loss_fn = schedule._loss_fn + schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) + schedule._loss_fn = indexed_loss + try: + return schedule.step( + target=torch.arange(self.num_microbatches, device=self.device), + losses=losses, + **schedule_options, + ) + finally: + schedule._loss_fn = previous_loss_fn + schedule._split_inputs = previous_split_inputs + @property def parts(self) -> list[nn.Module]: if self._info.model_parts is None: diff --git a/nemo_automodel/components/distributed/pipelining/functional.py b/nemo_automodel/components/distributed/pipelining/functional.py index 859de1a6dc..2afb1bcddc 100644 --- a/nemo_automodel/components/distributed/pipelining/functional.py +++ b/nemo_automodel/components/distributed/pipelining/functional.py @@ -311,6 +311,7 @@ def _precompute_stage_shapes( microbatch_size: int, seq_len: int, tensor_dtype: torch.dtype | None = None, + first_stage_input_meta: torch.Tensor | None = None, ) -> None: """Precompute input/output meta tensors for each pipeline stage to bypass serial shape inference. @@ -354,11 +355,15 @@ def _precompute_stage_shapes( seq_len=seq_len, dtype=model_dtype, ) + if stage.is_first and first_stage_input_meta is not None: + inputs_meta = (first_stage_input_meta,) _set_stage_metas(stage, inputs_meta, outputs_meta) continue # --- inputs_meta --- - if stage.is_first: + if stage.is_first and first_stage_input_meta is not None: + inputs_meta = (first_stage_input_meta,) + elif stage.is_first: # First stage receives input_ids: [mb, seq_len] int64 inputs_meta = (torch.empty(microbatch_size, seq_len, device="meta", dtype=torch.long),) else: @@ -449,6 +454,7 @@ def reset_pp_stage_shapes( microbatch_size: int, seq_len: int, tensor_dtype: torch.dtype | None = None, + first_stage_input_meta: torch.Tensor | None = None, ) -> None: """Reset pipeline stage infrastructure and recompute shapes for a new sequence length. @@ -468,6 +474,9 @@ def reset_pp_stage_shapes( model_config: The HuggingFace model config (``model.config``). microbatch_size: Per-microbatch batch size used by the schedule. seq_len: Sequence length of the upcoming batch (e.g. ``input_ids.shape[1]``). + first_stage_input_meta: Optional exact metadata for the positional input + consumed by the first stage. This preserves floating-point + ``inputs_embeds`` and flat THD input shapes. """ for stage in stages: # PyTorch <= 2.10 stores static metadata in these fields. @@ -495,7 +504,14 @@ def reset_pp_stage_shapes( stage.grad_send_info = None # Analytically set shapes for the new seq_len (no forward pass) - _precompute_stage_shapes(stages, model_config, microbatch_size, seq_len, tensor_dtype=tensor_dtype) + _precompute_stage_shapes( + stages, + model_config, + microbatch_size, + seq_len, + tensor_dtype=tensor_dtype, + first_stage_input_meta=first_stage_input_meta, + ) # Trigger _initialize_stage(s) on the next step() call. # PipelineScheduleSingle uses singular, PipelineScheduleMulti uses plural. diff --git a/nemo_automodel/components/distributed/pipelining/hf_utils.py b/nemo_automodel/components/distributed/pipelining/hf_utils.py index 419023f3fc..f7d2e7c0fa 100644 --- a/nemo_automodel/components/distributed/pipelining/hf_utils.py +++ b/nemo_automodel/components/distributed/pipelining/hf_utils.py @@ -124,6 +124,31 @@ def pipeline_forward( causal_mask_mapping: Optional[dict] = None, **kwargs, ) -> Union[torch.Tensor, BaseModelOutputWithPast]: + """Run one generic decoder pipeline stage. + + Args: + input_ids: Token IDs ``[B, S]`` on the first stage, or floating-point + hidden states ``[B, S, H]`` on later stages. Packed pipeline + microbatches retain a singleton batch axis: ``[1, T]`` or + ``[1, T, H]``. + attention_mask: Optional padded attention mask ``[B, S]``. Packed + THD attention ignores this mask and uses document boundaries. + position_ids: Positions ``[B, S]`` or packed positions ``[1, T]``. + past_key_values: Optional padded-generation cache; unsupported for + packed THD training. + inputs_embeds: Hidden states ``[B, S, H]`` or packed ``[1, T, H]``. + use_cache: Whether to update a padded-generation cache. + cache_position: Optional padded cache positions ``[S]``. + causal_mask_mapping: Optional padded causal masks keyed by attention + type. THD layers receive ``None`` and use ``cu_seqlens`` instead. + **kwargs: Model metadata. Packed THD inputs set ``qkv_format='thd'`` + and carry ``cu_seqlens`` ``[1, N + 1]`` plus scalar + ``max_seqlen``; CP may also supply ``cp_size`` and ``cp_rank``. + + Returns: + Hidden states ``[B, S, H]`` (packed: ``[1, T, H]``), either directly + for a pipeline stage or in ``BaseModelOutputWithPast``. + """ # For VLM models the text components (embed_tokens, layers, norm) live on a # nested text module (e.g. model.language_model) rather than directly on self. # get_text_module returns self when no nesting exists (e.g. LlamaModel). @@ -145,30 +170,65 @@ def pipeline_forward( else: raise ValueError("inputs_embeds must be provided for pipeline stages without embed_tokens") - if use_cache and past_key_values is None: - from transformers.cache_utils import DynamicCache - - past_key_values = DynamicCache() + is_thd = kwargs.get("qkv_format") == "thd" + if is_thd: + if past_key_values is not None: + raise ValueError("Packed THD training does not support past_key_values.") + if position_ids is None: + raise ValueError("Packed THD input requires position_ids.") + + kwargs.pop("padding_mask", None) + if inputs_embeds.ndim > 2: + inputs_embeds = inputs_embeds.squeeze(0) + if inputs_embeds.ndim != 2: + raise ValueError( + f"Packed pipeline THD hidden states must be [1, T, H], got {tuple(inputs_embeds.shape)}." + ) + if position_ids.ndim > 1: + position_ids = position_ids.squeeze(0) + if position_ids.ndim != 1: + raise ValueError(f"Packed pipeline THD position_ids must be [1, T], got {tuple(position_ids.shape)}.") + for key, value in kwargs.items(): + if not isinstance(value, torch.Tensor): + continue + if key == "max_seqlen": + kwargs[key] = value.item() + continue + if value.ndim > 1: + value = value.squeeze(0) + if key in ("cu_seqlens", "cu_seqlens_padded"): + value = value[value != -1000].contiguous() + kwargs[key] = value + + attention_mask = None + cache_position = None + causal_mask_mapping = {"full_attention": None} + use_cache = False + else: + if use_cache and past_key_values is None: + from transformers.cache_utils import DynamicCache - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device - ) + past_key_values = DynamicCache() - if position_ids is None: - position_ids = cache_position.unsqueeze(0) + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) - # Attention mask handling (compilation-friendly): - # causal_mask_mapping is precomputed in the data pipeline (default_collater + - # add_causal_masks_to_batch) and passed to the first stage. The PP schedule - # cannot forward this dict to non-first stages, which therefore arrive with - # causal_mask_mapping=None. Build it once per stage and cache it (see - # _build_or_reuse_pp_causal_mask) instead of recomputing every microbatch. - if causal_mask_mapping is None: - causal_mask_mapping = _build_or_reuse_pp_causal_mask( - self, inputs_embeds, attention_mask, cache_position, position_ids - ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + # Attention mask handling (compilation-friendly): + # causal_mask_mapping is precomputed in the data pipeline (default_collater + + # add_causal_masks_to_batch) and passed to the first stage. The PP schedule + # cannot forward this dict to non-first stages, which therefore arrive with + # causal_mask_mapping=None. Build it once per stage and cache it (see + # _build_or_reuse_pp_causal_mask) instead of recomputing every microbatch. + if causal_mask_mapping is None: + causal_mask_mapping = _build_or_reuse_pp_causal_mask( + self, inputs_embeds, attention_mask, cache_position, position_ids + ) hidden_states = inputs_embeds @@ -176,7 +236,15 @@ def pipeline_forward( position_embeddings = None rotary_emb = get_text_module(self).rotary_emb if rotary_emb is not None: - position_embeddings = rotary_emb(hidden_states, position_ids) + if is_thd: + position_embeddings = rotary_emb( + hidden_states, + position_ids, + qkv_format="thd", + cp_size=kwargs.get("cp_size", 1), + ) + else: + position_embeddings = rotary_emb(hidden_states, position_ids) if hasattr(text_module, "layers") and text_module.layers is not None: # Works for dict-like or list-like containers @@ -196,12 +264,16 @@ def pipeline_forward( use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, + **(kwargs if is_thd else {}), ) hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs if hasattr(text_module, "norm") and text_module.norm is not None: hidden_states = text_module.norm(hidden_states) + if is_thd: + hidden_states = hidden_states.unsqueeze(0) + if model_class_name == "PipelineStage": return hidden_states else: diff --git a/nemo_automodel/components/distributed/thd_utils.py b/nemo_automodel/components/distributed/thd_utils.py index bf275c02ef..69384d3287 100644 --- a/nemo_automodel/components/distributed/thd_utils.py +++ b/nemo_automodel/components/distributed/thd_utils.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + import torch @@ -83,7 +85,7 @@ def _thd_padding_mask( def process_input_for_thd( - batch: dict[str, torch.Tensor], + batch: dict[str, Any], seq_lens_padding_value: int = -1000, padding_token_id: int = 0, ) -> dict[str, torch.Tensor]: @@ -272,19 +274,161 @@ def process_input_for_thd( if max_seqlen is not None: result["max_seqlen"] = max_seqlen - # Pass through any field this function neither transforms nor consumes (e.g. - # VLM media tensors like pixel_values / image_grid_thw), tensor or not, so - # callers don't need to pop and restore them around the THD conversion. - _consumed = {"seq_lens", "seq_lens_padded"} + # Engine-prefixed loss fields use the primary stream's [B, S] -> [T] + # transform. The prefix makes the token-layout intent explicit: model-owned + # tensors may coincidentally start with [B, S] but need their original + # layout (for example a VLM's global vision mask). + _consumed = {"input_ids", "labels", "position_ids", "seq_lens", "seq_lens_padded"} for key, value in batch.items(): - if key not in result and key not in _consumed: + if key in result or key in _consumed: + continue + if ( + key.startswith("__engine_loss__") + and isinstance(value, torch.Tensor) + and value.ndim >= 2 + and tuple(value.shape[:2]) == (batch_size, seq_len) + ): + result[key] = value.reshape(total_tokens, *value.shape[2:]) + else: result[key] = value return result +def stack_thd_chunks(chunks: list[dict[str, Any]], seq_lens_padding_value: int = -1000) -> dict[str, Any]: + """Stack independently prepared THD microbatches. + + Args: + chunks: Non-empty list of THD mappings. Token tensors have shape + [tokens, ...]; cumulative-length tensors have shape [sequences + 1] + and may differ in length across chunks. + seq_lens_padding_value: Right-padding sentinel for cumulative-length + tensors. + + Returns: + One mapping whose tensor fields have a leading [microbatches] axis. + Cumulative-length fields are padded to + [microbatches, max_sequences + 1]; non-tensor metadata is replicated. + """ + if not chunks: + raise ValueError("stack_thd_chunks requires at least one chunk") + + result: dict[str, Any] = {} + keys = set().union(*(chunk.keys() for chunk in chunks)) + for key in keys: + if key == "cu_seqlens_padded": + values = [chunk.get("cu_seqlens_padded", chunk["cu_seqlens"]) for chunk in chunks] + else: + if not all(key in chunk for chunk in chunks): + raise ValueError(f"THD chunks have inconsistent field {key!r}") + values = [chunk[key] for chunk in chunks] + tensor_fields = [isinstance(value, torch.Tensor) for value in values] + if any(tensor_fields) and not all(tensor_fields): + raise ValueError(f"THD chunks disagree whether field {key!r} is a tensor") + if not all(tensor_fields): + result[key] = values[0] + continue + tensors = values + if key in {"cu_seqlens", "cu_seqlens_padded"}: + max_length = max(tensor.numel() for tensor in tensors) + tensors = [ + torch.cat( + ( + tensor, + torch.full( + (max_length - tensor.numel(),), + seq_lens_padding_value, + dtype=tensor.dtype, + device=tensor.device, + ), + ) + ) + if tensor.numel() < max_length + else tensor + for tensor in tensors + ] + result[key] = torch.stack(tensors) + return result + + +def split_final_thd_batch( + batch: dict[str, Any], + num_chunks: int, + seq_lens_padding_value: int = -1000, +) -> list[dict[str, Any]]: + """Split an already flattened THD stream at equal token boundaries. + + Args: + batch: Final THD mapping. ``input_ids`` has shape [tokens] or + ``inputs_embeds`` has shape [tokens, hidden]. Token-aligned + auxiliary tensors use shape [tokens, ...]. ``cu_seqlens`` and + optional ``cu_seqlens_padded`` have shape [sequences + 1]. + num_chunks: Number of equal token chunks. Every boundary must also be + a packed-sequence boundary. + seq_lens_padding_value: Padding sentinel in cumulative-length tensors. + + Returns: + THD mappings in microbatch order. Token tensors have shape + [tokens / num_chunks, ...] and cumulative lengths are rebased to zero. + """ + if num_chunks <= 0: + raise ValueError(f"num_chunks must be positive, got {num_chunks}") + primary_name = "inputs_embeds" if "inputs_embeds" in batch else "input_ids" + primary = batch.get(primary_name) + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError("final THD input requires tensor input_ids or inputs_embeds") + position_ids = batch.get("position_ids") + if isinstance(position_ids, torch.Tensor) and position_ids.ndim == 3 and num_chunks > 1: + raise NotImplementedError("chunked final THD does not support three-dimensional mRoPE position_ids") + + total_tokens = primary.shape[0] + if total_tokens % num_chunks != 0: + raise ValueError(f"final THD token count {total_tokens} must be divisible by {num_chunks} chunks") + tokens_per_chunk = total_tokens // num_chunks + + cu_seqlens = batch.get("cu_seqlens") + if not isinstance(cu_seqlens, torch.Tensor) or cu_seqlens.ndim != 1: + raise ValueError("final THD input requires one-dimensional cu_seqlens") + cu_seqlens = cu_seqlens[cu_seqlens != seq_lens_padding_value] + cu_seqlens_padded = batch.get("cu_seqlens_padded", cu_seqlens) + if not isinstance(cu_seqlens_padded, torch.Tensor) or cu_seqlens_padded.ndim != 1: + raise ValueError("final THD cu_seqlens_padded must be one-dimensional") + cu_seqlens_padded = cu_seqlens_padded[cu_seqlens_padded != seq_lens_padding_value] + if cu_seqlens.numel() != cu_seqlens_padded.numel(): + raise ValueError("cu_seqlens and cu_seqlens_padded must describe the same sequences") + + boundary_indices: list[int] = [] + for offset in range(0, total_tokens + 1, tokens_per_chunk): + matches = (cu_seqlens_padded == offset).nonzero(as_tuple=False).flatten() + if matches.numel() != 1: + raise ValueError(f"final THD chunk boundary {offset} is not a unique packed-sequence boundary") + boundary_indices.append(int(matches.item())) + + chunks: list[dict[str, Any]] = [] + for index in range(num_chunks): + token_start = index * tokens_per_chunk + sequence_start = boundary_indices[index] + sequence_end = boundary_indices[index + 1] + local_cu = cu_seqlens[sequence_start : sequence_end + 1] - cu_seqlens[sequence_start] + local_cu_padded = cu_seqlens_padded[sequence_start : sequence_end + 1] - cu_seqlens_padded[sequence_start] + chunk: dict[str, Any] = {} + for key, value in batch.items(): + if key in {"cu_seqlens", "cu_seqlens_padded", "max_seqlen"}: + continue + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == total_tokens: + chunk[key] = value.narrow(0, token_start, tokens_per_chunk) + else: + chunk[key] = value + chunk["cu_seqlens"] = local_cu.to(torch.int32) + if not torch.equal(local_cu, local_cu_padded): + chunk["cu_seqlens_padded"] = local_cu_padded.to(torch.int32) + chunk["max_seqlen"] = (local_cu[1:] - local_cu[:-1]).max().to(torch.int32) + chunks.append(chunk) + return chunks + + def split_batch_into_thd_chunks( - batch: dict[str, torch.Tensor], + batch: dict[str, Any], num_chunks: int, seq_lens_padding_value: int = -1000, padding_token_id: int = 0, @@ -308,8 +452,12 @@ def split_batch_into_thd_chunks( - 'position_ids': [batch_size, seq_len] (required) - 'seq_lens': [batch_size, num_packs] - 'seq_lens_padded': [batch_size, num_packs] - num_chunks: Number of chunks to split the batch into. Must evenly divide batch_size. - If num_chunks <= 1, returns the result from process_input_for_thd directly. + num_chunks: Number of chunks to split the batch into. Normally this + must evenly divide ``batch_size``. A single already-packed row is + instead split by its sequence metadata; every chunk must receive + the same number of sequences and the same padded token width. If + ``num_chunks <= 1``, returns the result from + :func:`process_input_for_thd` directly. seq_lens_padding_value: Value used to indicate padding in seq_lens/seq_lens_padded tensors and for padding cu_seqlens to uniform length (default: -1000) padding_token_id: Filler token id. Only consulted by the metadata-free @@ -353,53 +501,78 @@ def split_batch_into_thd_chunks( >>> # result['cu_seqlens'][0]: tensor([0, 6, 12], dtype=torch.int32) >>> # result['cu_seqlens'][1]: tensor([0, 6, 12], dtype=torch.int32) """ - # NOTE: 3D mRoPE position_ids ([n_rope, batch, seq]) are only validated for the - # num_chunks<=1 path (cp_size=1). The multi-chunk stacking below has not been - # validated for mRoPE and should not be used for VLM+CP/PP THD yet. if num_chunks <= 1: return process_input_for_thd(batch, seq_lens_padding_value, padding_token_id) + position_ids = batch.get("position_ids") + if isinstance(position_ids, torch.Tensor) and position_ids.ndim == 3: + raise NotImplementedError("chunked THD does not support three-dimensional mRoPE position_ids") + + batch_size = batch["input_ids"].shape[0] + if batch_size == 1: + seq_lens = batch.get("seq_lens") + seq_lens_padded = batch.get("seq_lens_padded", seq_lens) + if not isinstance(seq_lens, torch.Tensor) or not isinstance(seq_lens_padded, torch.Tensor): + raise ValueError("a single packed THD row requires seq_lens and seq_lens_padded for chunking") + real_lengths = seq_lens.reshape(-1) + padded_lengths = seq_lens_padded.reshape(-1) + real_lengths = real_lengths[real_lengths != seq_lens_padding_value] + padded_lengths = padded_lengths[padded_lengths != seq_lens_padding_value] + if ( + real_lengths.numel() < num_chunks + or real_lengths.numel() != padded_lengths.numel() + or real_lengths.numel() % num_chunks != 0 + ): + raise ValueError( + f"packed THD sequence count {real_lengths.numel()} must divide evenly across {num_chunks} chunks" + ) + sequences_per_chunk = real_lengths.numel() // num_chunks + padded_by_chunk = padded_lengths.reshape(num_chunks, sequences_per_chunk) + chunk_widths = padded_by_chunk.sum(dim=1) + if not bool((chunk_widths == chunk_widths[0]).all()): + raise ValueError(f"packed THD pipeline chunks must have equal token widths; got {chunk_widths.tolist()}") + total_tokens = batch["input_ids"].shape[1] + if int(chunk_widths.sum().item()) != total_tokens: + raise ValueError( + f"packed THD metadata spans {int(chunk_widths.sum().item())} tokens, " + f"but input_ids has width {total_tokens}" + ) - def pad_and_stack(tensor_list, padding_value): - """Pad tensors to same length and stack them.""" - max_len = max(len(t) for t in tensor_list) - padded = [] - for t in tensor_list: - if len(t) < max_len: - pad = torch.full((max_len - len(t),), padding_value, dtype=t.dtype, device=t.device) - t = torch.cat([t, pad]) - padded.append(t) - return torch.stack(padded) - - chunk_size = batch["input_ids"].shape[0] // num_chunks - - # Process all chunks - chunk_results = [ - process_input_for_thd( - { - k: v[i * chunk_size : (i + 1) * chunk_size] if isinstance(v, torch.Tensor) else v - for k, v in batch.items() - }, - seq_lens_padding_value, - padding_token_id, - ) - for i in range(num_chunks) - ] - - stacked: dict = { - "input_ids": torch.stack([c["input_ids"] for c in chunk_results]), - "labels": torch.stack([c["labels"] for c in chunk_results]), - "position_ids": torch.stack([c["position_ids"] for c in chunk_results]), - "cu_seqlens": pad_and_stack([c["cu_seqlens"] for c in chunk_results], seq_lens_padding_value), - "padding_mask": torch.stack([c["padding_mask"] for c in chunk_results]), - } - # Emit cu_seqlens_padded whenever any chunk emits it; absorbed chunks - # fall back to their cu_seqlens (semantically equal) for rectangularity. - if any("cu_seqlens_padded" in c for c in chunk_results): - stacked["cu_seqlens_padded"] = pad_and_stack( - [c.get("cu_seqlens_padded", c["cu_seqlens"]) for c in chunk_results], - seq_lens_padding_value, - ) - if all("max_seqlen" in c for c in chunk_results): - stacked["max_seqlen"] = torch.stack([c["max_seqlen"] for c in chunk_results]) - stacked.update({k: v for k, v in chunk_results[0].items() if not isinstance(v, torch.Tensor)}) - return stacked + token_keys = {"input_ids", "labels", "position_ids", "attention_mask", "padding_mask"} + chunk_results = [] + token_start = 0 + for index in range(num_chunks): + width = int(chunk_widths[index].item()) + sequence_start = index * sequences_per_chunk + chunk: dict[str, Any] = {} + for key, value in batch.items(): + if key in {"seq_lens", "seq_lens_padded"}: + chunk[key] = value.narrow(1, sequence_start, sequences_per_chunk) + elif ( + (key in token_keys or key.startswith("__engine_loss__")) + and isinstance(value, torch.Tensor) + and value.ndim >= 2 + and tuple(value.shape[:2]) == (1, total_tokens) + ): + chunk[key] = value.narrow(1, token_start, width) + else: + chunk[key] = value + chunk_results.append(process_input_for_thd(chunk, seq_lens_padding_value, padding_token_id)) + token_start += width + return stack_thd_chunks(chunk_results, seq_lens_padding_value) + + if batch_size % num_chunks != 0: + raise ValueError(f"THD batch size {batch_size} must be divisible by {num_chunks} chunks") + chunk_size = batch_size // num_chunks + + chunk_results = [] + for index in range(num_chunks): + chunk = { + key: ( + value.narrow(0, index * chunk_size, chunk_size) + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == batch_size + else value + ) + for key, value in batch.items() + } + chunk_results.append(process_input_for_thd(chunk, seq_lens_padding_value, padding_token_id)) + return stack_thd_chunks(chunk_results, seq_lens_padding_value) diff --git a/nemo_automodel/components/loss/mtp.py b/nemo_automodel/components/loss/mtp.py index d93fe9e82b..2139d02871 100644 --- a/nemo_automodel/components/loss/mtp.py +++ b/nemo_automodel/components/loss/mtp.py @@ -358,6 +358,17 @@ def forward(self, output, labels: torch.Tensor) -> torch.Tensor: if isinstance(self.loss_fn, FusedLinearCrossEntropy) and isinstance(output, torch.Tensor): logits = None hidden_states = output + # A single THD pipeline microbatch keeps a synthetic leading + # pipeline axis on the last-stage hidden states, while its + # labels use the native flat token layout. Remove only that + # synthetic axis before invoking fused linear CE. + if ( + labels.ndim == 1 + and hidden_states.ndim == 3 + and hidden_states.shape[0] == 1 + and hidden_states.shape[1] == labels.shape[0] + ): + hidden_states = hidden_states.squeeze(0) else: logits = getattr(output, "logits", output) hidden_states = _get_final_hidden_states(output) diff --git a/nemo_automodel/components/models/nemotron_parse/nemotron_parse_loss.py b/nemo_automodel/components/models/nemotron_parse/nemotron_parse_loss.py index 6faf4190cd..766f4211a6 100644 --- a/nemo_automodel/components/models/nemotron_parse/nemotron_parse_loss.py +++ b/nemo_automodel/components/models/nemotron_parse/nemotron_parse_loss.py @@ -106,13 +106,18 @@ def forward( loss_full[coordinate_mask] *= self.coordinate_weight valid_tokens = (labels != self.ignore_index).sum() + loss_sum = loss_full.sum() if valid_tokens == 0: - return torch.tensor(0.0, device=logits.device, dtype=logits.dtype) + return loss_sum * 0 if num_label_tokens is not None: assert self.reduction == "sum", ( f"num_label_tokens is only supported when reduction='sum', got reduction='{self.reduction}'" ) - return loss_full.sum() / (num_label_tokens + 1e-6) + return loss_sum / (num_label_tokens + 1e-6) - return loss_full.sum() / (valid_tokens + 1e-6) + if self.reduction == "sum": + return loss_sum + if self.reduction == "mean": + return loss_sum / (valid_tokens + 1e-6) + raise ValueError(f"Unsupported reduction={self.reduction!r}; expected 'sum' or 'mean'") diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index c1ea753fda..359e606b32 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -26,10 +26,6 @@ from nemo_automodel.components.datasets.datum import Datum, collate_datums from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder -from nemo_automodel.components.distributed.context_parallel.sharder import ( - identity_local_indices, - shard_batch_identity, -) from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh from nemo_automodel.components.distributed.pipelining import AutoPipeline @@ -44,13 +40,20 @@ CollateFn = Callable[[list[Datum]], tuple[dict[str, Any], dict[str, torch.Tensor]]] LossFn = Callable[ - [Any, dict[str, torch.Tensor], Sequence[Datum], dict[str, Any]], + [Any, dict[str, torch.Tensor]], torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], ] +_LOSS_FIELD_PREFIX = "__engine_loss__" +_LOSS_METADATA = ("cu_seqlens", "cu_seqlens_padded", "max_seqlen", "padding_mask") + __all__ = ["Engine", "collate_prebatched"] +def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextManager[Any]: + return nullcontext() + + def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: """Return one already-collated Datum without changing its layout. @@ -89,6 +92,8 @@ class Engine: mesh_context: Runtime topology. Required for an ``AutoPipeline``. For eager models, an initialized default process group is treated as pure data parallelism when omitted. + microbatch_size: Number of Datum items collated into each outer batch. + Use one with :func:`collate_prebatched`. collate_fn: Batches one microbatch of Datums into separate model and loss inputs. The default supports padded and packed text. VLMs pass a model-specific collater. Existing recipes whose dataloaders @@ -97,18 +102,25 @@ class Engine: model inputs and loss inputs aligned and preserve the sum of ``weights``. padding_token_id: Token used when the CP sharder pads ``input_ids``. - context_fn: Creates an optional context around model forward, loss, - and backward. Recipes use this for runtime contexts such as FP8. + context_fn: Creates an optional context from the CP-prepared model-input + mapping. It covers model forward, loss, and backward, or the full + pipeline schedule. Recipes use it for FP8 and model input staging. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the final microbatch. Note: Context-parallel input layout and transport are delegated to :class:`ContextParallelSharder`. With an :class:`AutoPipeline`, the - pipeline schedule owns its internal microbatching and backward calls; - the Engine still owns the complete outer accumulation window and loss - normalization. Pipeline execution currently requires context-parallel - size one. + pipeline schedule owns execution order and backward calls; the Engine + owns exact microbatch materialization, the complete outer accumulation + window, and loss normalization. Pipeline output mappings are + synchronized to every physical PP rank. Magi's current packed contract + uses one inner pipeline microbatch; recipes enforce that configuration. + Packed pipeline batches with multiple inner microbatches must split at + sequence boundaries into equal-width token chunks. Token-aligned loss + fields follow those chunks. Per-Datum scalar loss fields do not yet + carry enough boundary metadata to be split in that layout and are + rejected instead of being replicated silently. """ def __init__( @@ -117,16 +129,20 @@ def __init__( *, device: torch.device | str, mesh_context: MeshContext | None = None, + microbatch_size: int = 1, collate_fn: CollateFn = collate_datums, padding_token_id: int = 0, - context_fn: Callable[[], AbstractContextManager[Any]] = nullcontext, + context_fn: Callable[[dict[str, Any]], AbstractContextManager[Any]] = _nullcontext_for_batch, defer_fsdp_grad_sync: bool = True, ) -> None: + if isinstance(microbatch_size, bool) or not isinstance(microbatch_size, int) or microbatch_size <= 0: + raise ValueError(f"microbatch_size must be a positive integer, got {microbatch_size!r}") self.pipeline = model if isinstance(model, AutoPipeline) else None self.model_parts = model.parts if self.pipeline is not None else [model] self.model = self.model_parts[0] self.device = torch.device(device) self.mesh_context = mesh_context + self.microbatch_size = microbatch_size self.collate_fn = collate_fn self.padding_token_id = padding_token_id self.context_fn = context_fn @@ -134,33 +150,40 @@ def __init__( def forward_backward( self, - window: Sequence[Sequence[Datum]], + datums: Sequence[Datum], loss_fn: LossFn, ) -> tuple[torch.Tensor, list[dict[str, Any]]]: """Accumulate gradients for a complete optimizer window. - ``window`` is explicit: each inner sequence is one eager microbatch, or - one outer pipeline batch that the schedule splits internally. Pipeline - batches currently contain exactly one already-batched Datum. - ``loss_fn`` receives the raw model output, CP-local - ``loss_fn_inputs``, the original Datums, and the final CP-local model - inputs produced by the sharder. It returns either per-element losses - with exactly the same shape as + ``datums`` is a flat optimizer accumulation window. The Engine groups + it into outer batches of ``microbatch_size`` and invokes ``collate_fn`` + once per group. A Datum normally represents one sample. Recipes that + retain worker-side collation instead wrap each prepared batch in one + Datum and configure ``microbatch_size=1`` with + :func:`collate_prebatched`. + + ``loss_fn`` receives the raw model output and CP-local + ``loss_fn_inputs``. It returns either per-element losses with exactly + the same shape as ``loss_fn_inputs["weights"]``, or a scalar local weighted-sum numerator. For a scalar, the callback must apply weights and masks; the Engine will only apply global normalization. The callback may - also return one output mapping per Datum during eager execution. + also return one output mapping per Datum. Those mappings are detached and preserved in input order; the Engine - deliberately does not interpret or reduce them. Pipeline execution - does not yet support per-Datum outputs. + deliberately does not interpret or reduce them. Under pipeline + parallelism the last stage computes the mappings and broadcasts them + to every stage in the pipeline group. + + A prebatched Datum intentionally hides its inner sample boundaries. + It can therefore return a single output mapping only when a pipeline + schedule has one inner microbatch. Callers that need one output per + sample use ordinary flat Datums. Args: - window: The complete optimizer accumulation window. Each inner - sequence is one eager microbatch or outer pipeline batch of - Datums. A pipeline batch must contain exactly one prebatched - Datum. A Datum's token weights - may have shape ``[tokens]`` or the custom collater's batched - token layout; the loss tensor must use the identical shape. + datums: Flat sequence of Datum items in the complete optimizer + accumulation window. A Datum's token weights may have shape + [tokens] or the custom collater's batched token layout; the + loss tensor must use the identical shape. loss_fn: Computes either that per-token loss tensor or a scalar local weighted-sum numerator from the raw model output and collated loss inputs. @@ -169,18 +192,19 @@ def forward_backward( ``(loss, loss_fn_outputs)``. ``loss`` is a detached scalar reduced over the DP-CP gradient group and, for pipeline execution, synchronized across PP stages. ``loss_fn_outputs`` contains - local-rank, per-Datum mappings in window order. Model parameters - are unchanged, but their gradients contain the complete window's + per-Datum mappings in window order; pipeline execution returns the + same mappings on every physical stage rank. Model parameters are + unchanged, but their gradients contain the complete window's globally normalized backward result. """ - microbatches = self._validate_window(window) + microbatches = self._group_datums(datums) self._validate_parallelism() - if self.pipeline is not None and any(len(microbatch) != 1 for microbatch in microbatches): - raise ValueError("pipeline Engine requires exactly one prebatched Datum in each outer batch") dp_group, dp_size = self._dp_group_and_size() grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) denominator = self._global_weight_sum(microbatches, dp_group, dp_size) + zero_denominator = bool(denominator == 0) + safe_denominator = torch.where(denominator > 0, denominator, torch.ones_like(denominator)) self._validate_pipeline_window(len(microbatches), denominator) pp_enabled = self.pipeline is not None @@ -201,78 +225,36 @@ def forward_backward( if is_last: prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) - model_inputs, loss_inputs = self.collate_fn(datums) - self._validate_collated_weights(datums, loss_inputs) - model_inputs = _to_device(model_inputs, self.device) - loss_inputs = _to_device(loss_inputs, self.device) - full_weights = loss_inputs["weights"] - loss_seq_dim = _loss_sequence_dim(model_inputs, full_weights) - - # ContextParallelSharder is the single owner of padded, THD, Magi, - # and model-specific CP layouts. Labels are temporarily present in - # its batch because each backend historically shards them with the - # model inputs; all other loss tensors use the sharder token verb. - cp_batch = dict(model_inputs) - labels = loss_inputs.get("labels") - cp_batch["labels"] = ( - labels.clone() if isinstance(labels, torch.Tensor) else torch.zeros_like(full_weights, dtype=torch.long) - ) - device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None - final_thd = _is_final_thd(cp_batch) - if final_thd and self._cp_size() > 1: - raise ValueError( - "context parallelism requires raw THD inputs so ContextParallelSharder can partition them" - ) - if final_thd: - if self.pipeline is not None and self.pipeline.num_microbatches > 1: - raise ValueError( - "pipeline Engine requires raw THD inputs so ContextParallelSharder can split them " - "for the schedule's internal microbatches" - ) - sharder = ContextParallelSharder( - device_mesh=device_mesh, - shard_batch=shard_batch_identity, - local_token_global_indices=identity_local_indices, - padding_token_id=self.padding_token_id, - ) - else: - sharder = ContextParallelSharder( - self.model, - device_mesh, - cp_batch, - padding_token_id=self.padding_token_id, - num_chunks=inner_microbatches, - ) - cp_context, model_inputs = sharder.shard(cp_batch) - if model_inputs.get("qkv_format") == "thd" and ( - "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs - ): - raise ValueError( - "ContextParallelSharder could not prepare raw THD inputs for this model; " - "use a THD-capable attention backend or provide final THD inputs at cp_size=1" - ) - local_labels = model_inputs.pop("labels") - loss_inputs = self._shard_loss_inputs(sharder, loss_inputs, loss_seq_dim) - if labels is not None: - loss_inputs["labels"] = local_labels - weights = loss_inputs["weights"] + cp_context, model_inputs, loss_inputs = self._prepare_batch(datums, inner_microbatches) if self.pipeline is not None: - self._pipeline_step( + batch_returns_outputs, batch_outputs = self._pipeline_step( model_inputs, loss_inputs, datums, loss_fn, - denominator, + safe_denominator, + zero_denominator, grad_group_size, local_loss_sum, cp_context, ) + if batch_returns_outputs is not None: + if returns_outputs is None: + returns_outputs = batch_returns_outputs + elif returns_outputs != batch_returns_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + loss_fn_outputs.extend(batch_outputs) else: - forward_inputs = filter_forward_kwargs(self.model, model_inputs) - with get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(), cp_context(): + loss_inputs = _with_loss_metadata(model_inputs, loss_inputs) + with ( + get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), + self.context_fn(model_inputs), + cp_context(), + ): + forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - result = loss_fn(output, loss_inputs, datums, model_inputs) + result = loss_fn(output, loss_inputs) has_outputs = isinstance(result, tuple) if returns_outputs is None: returns_outputs = has_outputs @@ -290,8 +272,10 @@ def forward_backward( loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) else: losses = result - numerator = _weighted_numerator(losses, weights) - (numerator * (grad_group_size / denominator)).backward() + numerator = _weighted_numerator(losses, loss_inputs["weights"]) + if zero_denominator: + numerator = numerator * 0 + (numerator * (grad_group_size / safe_denominator)).backward() local_loss_sum.add_(numerator.detach().to(torch.float64)) if index == 0: @@ -303,21 +287,121 @@ def forward_backward( if pp_size > 1: dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=pp_group) - loss = (local_loss_sum / denominator).detach() + loss = (local_loss_sum / safe_denominator).detach() return loss, loss_fn_outputs - @staticmethod - def _validate_window(window: Sequence[Sequence[Datum]]) -> list[list[Datum]]: - if not isinstance(window, Sequence) or isinstance(window, (str, bytes)) or not window: - raise ValueError("forward_backward requires a non-empty accumulation window") - microbatches: list[list[Datum]] = [] - for index, microbatch in enumerate(window): - if not isinstance(microbatch, Sequence) or isinstance(microbatch, (str, bytes)) or not microbatch: - raise ValueError(f"microbatch {index} must be a non-empty sequence of Datum") - if not all(isinstance(datum, Datum) for datum in microbatch): - raise TypeError(f"microbatch {index} contains a value that is not a Datum") - microbatches.append(list(microbatch)) - return microbatches + def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: + if not isinstance(datums, Sequence) or isinstance(datums, (str, bytes)) or not datums: + raise ValueError("forward_backward requires a non-empty flat sequence of Datum") + if not all(isinstance(datum, Datum) for datum in datums): + raise TypeError("forward_backward received a value that is not a Datum") + return [ + list(datums[start : start + self.microbatch_size]) for start in range(0, len(datums), self.microbatch_size) + ] + + def _prepare_batch( + self, + datums: list[Datum], + num_pipeline_microbatches: int, + ) -> tuple[Callable[[], AbstractContextManager[Any]], dict[str, Any], dict[str, torch.Tensor]]: + """Collate, move, and CP-shard one outer batch. + + Args: + datums: Datum items in one outer batch. Tensor layouts are defined + by the configured collater. + num_pipeline_microbatches: Number of pipeline microbatches that + the prepared outer batch must materialize. + + Returns: + The CP context factory, CP-local model inputs, and CP-local loss + inputs. Token-aligned model and loss tensors use the same padded, + packed THD, Magi, or model-owned local sequence layout. + """ + model_inputs, loss_inputs = self.collate_fn(datums) + self._validate_collated_weights(datums, loss_inputs) + model_inputs = _to_device(model_inputs, self.device) + loss_inputs = _to_device(loss_inputs, self.device) + full_weights = loss_inputs["weights"] + loss_seq_dim = _loss_sequence_dim(model_inputs, full_weights) + + cp_batch = dict(model_inputs) + labels = loss_inputs.get("labels") + cp_batch["labels"] = ( + labels.clone() if isinstance(labels, torch.Tensor) else torch.zeros_like(full_weights, dtype=torch.long) + ) + is_thd = cp_batch.get("qkv_format") == "thd" + position_ids = cp_batch.get("position_ids") + if ( + is_thd + and isinstance(position_ids, torch.Tensor) + and position_ids.ndim == 3 + and (self._cp_size() > 1 or num_pipeline_microbatches > 1) + ): + raise NotImplementedError( + "THD pipeline/context parallelism does not yet support three-dimensional mRoPE position_ids" + ) + + thd_loss_fields: list[str] = [] + if is_thd: + ambiguous_per_datum_fields = [ + name + for name, value in loss_inputs.items() + if name != "labels" + and isinstance(value, torch.Tensor) + and value.ndim > 0 + and value.shape[0] == len(datums) + and len(datums) > 1 + and not _is_token_aligned(value, full_weights) + ] + if num_pipeline_microbatches > 1 and ambiguous_per_datum_fields: + raise NotImplementedError( + "packed pipeline microbatching cannot yet split per-Datum loss fields " + f"{ambiguous_per_datum_fields}; use token-aligned fields or a prepared collater" + ) + for name, value in loss_inputs.items(): + if name == "labels" or not _is_token_aligned(value, full_weights): + continue + key = f"{_LOSS_FIELD_PREFIX}{name}" + if key in cp_batch: + raise ValueError(f"model inputs contain reserved Engine key {key!r}") + cp_batch[key] = value + thd_loss_fields.append(name) + + device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None + sharder = ContextParallelSharder( + self.model, + device_mesh, + cp_batch, + padding_token_id=self.padding_token_id, + num_chunks=num_pipeline_microbatches, + ) + cp_context, model_inputs = sharder.shard(cp_batch) + if model_inputs.get("qkv_format") == "thd" and ( + "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs + ): + raise ValueError( + "ContextParallelSharder could not prepare raw THD inputs for this model; " + "use a THD-capable attention backend" + ) + + local_labels = model_inputs.pop("labels") + if is_thd: + local_loss_inputs = dict(loss_inputs) + for name in thd_loss_fields: + key = f"{_LOSS_FIELD_PREFIX}{name}" + candidate = model_inputs.pop(key, None) + if isinstance(candidate, torch.Tensor) and _matches_primary_token_layout(candidate, model_inputs): + local_loss_inputs[name] = candidate + else: + local_loss_inputs[name] = sharder.shard_token_tensor( + loss_inputs[name], seq_dim=loss_seq_dim or 0, fill=0 + ) + loss_inputs = local_loss_inputs + else: + loss_inputs = self._shard_loss_inputs(sharder, loss_inputs, loss_seq_dim) + if labels is not None: + loss_inputs["labels"] = local_labels + return cp_context, model_inputs, loss_inputs def _pipeline_step( self, @@ -326,45 +410,182 @@ def _pipeline_step( datums: Sequence[Datum], loss_fn: LossFn, denominator: torch.Tensor, + zero_denominator: bool, grad_group_size: int, local_loss_sum: torch.Tensor, cp_context: Callable[[], AbstractContextManager[Any]], - ) -> None: - primary_names = [name for name in ("input_ids", "inputs_embeds") if name in model_inputs] - if len(primary_names) != 1: - raise ValueError("pipeline Engine requires exactly one of input_ids or inputs_embeds") - primary_name = primary_names[0] - primary = model_inputs.pop(primary_name) - if not isinstance(primary, torch.Tensor) or primary.ndim < 2: - raise ValueError(f"pipeline Engine requires batched {primary_name} with a sequence dimension") - - self.pipeline.update_seq_len(primary.shape[1]) - pipeline_kwargs = { - name: value - for name, value in model_inputs.items() - if value is not None and not (isinstance(value, dict) and not value) - } - - def pipeline_loss(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): - if not model_args_mb: - raise RuntimeError("AutoPipeline loss callback did not receive the primary model input") - final_model_inputs = {primary_name: model_args_mb[0], **model_kwargs_mb} - result = loss_fn(output, loss_inputs_mb, datums, final_model_inputs) - if isinstance(result, tuple): - raise ValueError("pipeline Engine does not yet support per-Datum loss_fn outputs") - numerator = _weighted_numerator(result, loss_inputs_mb["weights"]) - local_loss_sum.add_(numerator.detach().to(torch.float64)) - return numerator * (grad_group_size / denominator) - - with self.context_fn(), cp_context(): - self.pipeline.step( - primary, - loss_inputs=loss_inputs, + ) -> tuple[bool | None, list[dict[str, Any]]]: + outputs_by_microbatch: list[list[dict[str, Any]] | None] = [None] * self.pipeline.num_microbatches + returns_outputs: bool | None = None + + with self.context_fn(model_inputs), cp_context(): + model_microbatches, loss_microbatches = self._materialize_pipeline_microbatches(model_inputs, loss_inputs) + primary = model_microbatches[0].get("inputs_embeds", model_microbatches[0].get("input_ids")) + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError("pipeline Engine requires a tensor input_ids or inputs_embeds") + if model_microbatches[0].get("qkv_format") == "thd" and self.pipeline.num_microbatches == 1: + seq_len = primary.shape[0] + else: + seq_len = primary.shape[1] if primary.ndim >= 2 else primary.shape[0] + effective_microbatch_size = 1 if model_microbatches[0].get("qkv_format") == "thd" else primary.shape[0] + self.pipeline.update_seq_len( + seq_len, + microbatch_size=effective_microbatch_size, + input_tensor=primary, + ) + + def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: + nonlocal returns_outputs + loss_inputs_mb = loss_microbatches[microbatch_index] + result = loss_fn(output, loss_inputs_mb) + has_outputs = isinstance(result, tuple) + if returns_outputs is None: + returns_outputs = has_outputs + elif returns_outputs != has_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + if isinstance(result, tuple): + losses, batch_outputs = result + if ( + not isinstance(batch_outputs, Sequence) + or isinstance(batch_outputs, (str, bytes)) + or not all(isinstance(item, Mapping) for item in batch_outputs) + ): + raise ValueError("loss_fn outputs must be a sequence of mappings") + if len(datums) == 1 and self.pipeline.num_microbatches > 1: + raise ValueError( + "a prebatched Datum may return outputs only when num_microbatches=1 because " + "its inner sample boundaries are not part of the Datum contract" + ) + outputs_by_microbatch[microbatch_index] = [_detach(dict(item)) for item in batch_outputs] + else: + losses = result + numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) + if zero_denominator: + numerator = numerator * 0 + local_loss_sum.add_(numerator.detach().to(torch.float64)) + return numerator * (grad_group_size / denominator) + + losses = [] if self.pipeline.info.has_last_stage else None + self.pipeline.step_microbatches( + model_microbatches, loss_fn=pipeline_loss, + losses=losses, return_outputs=False, - **pipeline_kwargs, ) + outputs: list[dict[str, Any]] = [] + if self.pipeline.info.has_last_stage and returns_outputs: + if any(items is None for items in outputs_by_microbatch): + raise RuntimeError("pipeline schedule did not evaluate loss_fn for every logical microbatch") + outputs = [item for items in outputs_by_microbatch if items is not None for item in items] + if len(outputs) != len(datums): + raise ValueError( + f"pipeline loss_fn returned {len(outputs)} outputs across the outer batch, " + f"expected one for each of its {len(datums)} Datums" + ) + outputs = self._broadcast_pipeline_outputs(outputs) + return bool(outputs), outputs + + def _broadcast_pipeline_outputs(self, outputs: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Broadcast last-stage per-Datum outputs to every pipeline stage.""" + pp_group, pp_size = self._pp_group_and_size() + if pp_size <= 1: + return outputs + + has_last_stage = self.pipeline.info.has_last_stage + local_state = torch.tensor( + [int(has_last_stage), int(has_last_stage and bool(outputs))], dtype=torch.int64, device=self.device + ) + stage_states = torch.empty(pp_size * 2, dtype=torch.int64, device=self.device) + dist.all_gather_into_tensor(stage_states, local_state, group=pp_group) + stage_states = stage_states.view(pp_size, 2) + source_ranks = (stage_states[:, 0] == 1).nonzero(as_tuple=False).flatten() + if source_ranks.numel() != 1: + raise RuntimeError("pipeline output synchronization requires exactly one physical last-stage rank") + + source_group_rank = int(source_ranks.item()) + if not bool(stage_states[source_group_rank, 1]): + return [] + source_global_rank = dist.get_global_rank(pp_group, source_group_rank) + object_list: list[Any] = [ + _to_device(outputs, torch.device("cpu")) if dist.get_rank(group=pp_group) == source_group_rank else None + ] + dist.broadcast_object_list(object_list, src=source_global_rank, group=pp_group, device=self.device) + received = object_list[0] + if not isinstance(received, list) or not all(isinstance(item, dict) for item in received): + raise RuntimeError("pipeline output synchronization received invalid per-Datum outputs") + return _to_device(received, self.device) + + def _materialize_pipeline_microbatches( + self, + model_inputs: dict[str, Any], + loss_inputs: dict[str, torch.Tensor], + ) -> tuple[list[dict[str, Any]], list[dict[str, torch.Tensor]]]: + """Split one CP-prepared outer batch into exact pipeline inputs. + + Args: + model_inputs: CP-local model tensors. Padded tensors use shape + [batch, sequence, ...]. Chunked THD tensors use shape + [microbatches, tokens, ...]. + loss_inputs: CP-local loss tensors. Token-aligned fields have the + same leading token axes as the primary model tensor. + + Returns: + Parallel lists of complete model and loss mappings, each with + exactly ``pipeline.num_microbatches`` items. Tensor slicing returns + views that retain a size-one pipeline microbatch axis. + """ + num_microbatches = self.pipeline.num_microbatches + if num_microbatches == 1: + return [dict(model_inputs)], [_with_loss_metadata(model_inputs, loss_inputs)] + + primary_name = _primary_name(model_inputs) + primary = model_inputs[primary_name] + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError(f"pipeline Engine requires tensor {primary_name}") + + if model_inputs.get("qkv_format") == "thd": + if primary.shape[0] != num_microbatches: + raise ValueError( + f"THD sharder produced {primary.shape[0]} chunks, expected {num_microbatches} pipeline microbatches" + ) + model_microbatches = [ + _select_chunk(model_inputs, index, num_microbatches) for index in range(num_microbatches) + ] + loss_microbatches = [ + _select_chunk(loss_inputs, index, num_microbatches) for index in range(num_microbatches) + ] + else: + batch_size = primary.shape[0] + if batch_size % num_microbatches != 0: + raise ValueError( + f"pipeline outer batch size {batch_size} must be divisible by {num_microbatches} microbatches" + ) + materialized_batch_size = batch_size // num_microbatches + if materialized_batch_size != self.pipeline.pp_microbatch_size: + raise ValueError( + f"materialized pipeline microbatch has batch size {materialized_batch_size}, " + f"but AutoPipeline is configured for pp_microbatch_size={self.pipeline.pp_microbatch_size}" + ) + custom_dims: dict[str, int] = {} + chunk_dims = getattr(self.model, "get_pipeline_kwargs_chunk_dims", None) + if chunk_dims is not None: + custom_dims = chunk_dims(model_inputs) or {} + model_microbatches = [ + _slice_batch_mapping(model_inputs, index, num_microbatches, batch_size, custom_dims) + for index in range(num_microbatches) + ] + loss_microbatches = [ + _slice_batch_mapping(loss_inputs, index, num_microbatches, batch_size) + for index in range(num_microbatches) + ] + + loss_microbatches = [ + _with_loss_metadata(model_microbatch, loss_microbatch) + for model_microbatch, loss_microbatch in zip(model_microbatches, loss_microbatches) + ] + return model_microbatches, loss_microbatches + def _validate_parallelism(self) -> None: if any( bool(getattr(module, "calculate_per_token_loss", False)) @@ -379,8 +600,6 @@ def _validate_parallelism(self) -> None: raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") if self.pipeline is not None and self.mesh_context is None: raise ValueError("pipeline Engine requires mesh_context") - if self.pipeline is not None and self._cp_size() > 1: - raise NotImplementedError("pipeline Engine does not yet support context parallelism") if self.mesh_context is None: return if self.mesh_context.pp_size > 1 and self.pipeline is None: @@ -453,8 +672,6 @@ def _global_weight_sum( self._validate_weight_sum_across_cp(denominator) if dp_size > 1: dist.all_reduce(denominator, op=dist.ReduceOp.SUM, group=dp_group) - if float(denominator) <= 0: - raise ValueError("forward_backward requires a positive global weight sum") return denominator def _validate_window_size_across_group( @@ -579,13 +796,119 @@ def _loss_sequence_dim(model_inputs: dict[str, Any], weights: torch.Tensor) -> i return None -def _is_final_thd(model_inputs: dict[str, Any]) -> bool: - """Return whether a CP1 caller already supplied the final flat THD layout.""" - if model_inputs.get("qkv_format") != "thd" or "cu_seqlens" not in model_inputs: +def _is_token_aligned(value: Any, weights: torch.Tensor) -> bool: + return ( + isinstance(value, torch.Tensor) + and weights.ndim > 0 + and value.ndim >= weights.ndim + and tuple(value.shape[: weights.ndim]) == tuple(weights.shape) + ) + + +def _primary_name(model_inputs: Mapping[str, Any]) -> str: + names = [name for name in ("input_ids", "inputs_embeds") if name in model_inputs] + if len(names) != 1: + raise ValueError("model inputs must contain exactly one of input_ids or inputs_embeds") + return names[0] + + +def _matches_primary_token_layout(tensor: torch.Tensor, model_inputs: Mapping[str, Any]) -> bool: + primary_name = _primary_name(model_inputs) + primary = model_inputs[primary_name] + if not isinstance(primary, torch.Tensor) or tensor.ndim == 0: return False - if "seq_lens" in model_inputs or "seq_lens_padded" in model_inputs: - raise ValueError("THD inputs cannot contain both raw seq_lens and final cu_seqlens metadata") - return True + if model_inputs.get("qkv_format") == "thd": + token_dims = primary.ndim - int(primary_name == "inputs_embeds") + else: + token_dims = 1 if primary.ndim == 1 else 2 + return tensor.ndim >= token_dims and tuple(tensor.shape[:token_dims]) == tuple(primary.shape[:token_dims]) + + +def _with_loss_metadata( + model_inputs: Mapping[str, Any], loss_inputs: Mapping[str, torch.Tensor] +) -> dict[str, torch.Tensor]: + result = dict(loss_inputs) + for name in _LOSS_METADATA: + value = model_inputs.get(name) + if isinstance(value, torch.Tensor): + result[name] = value + return result + + +def _select_chunk(value: Any, index: int, num_chunks: int) -> Any: + """Select one already materialized THD chunk without dropping its batch axis. + + Args: + value: Tensor leaves use shape [microbatches, ...] when chunked; + arbitrary nested containers and replicated metadata are accepted. + index: Pipeline microbatch index. + num_chunks: Expected leading microbatch extent. + + Returns: + A matching container whose chunked tensor leaves retain shape [1, ...]. + """ + if isinstance(value, torch.Tensor): + return value.narrow(0, index, 1) if value.ndim > 0 and value.shape[0] == num_chunks else value + if isinstance(value, dict): + return {name: _select_chunk(item, index, num_chunks) for name, item in value.items()} + if isinstance(value, list): + return [_select_chunk(item, index, num_chunks) for item in value] + if isinstance(value, tuple): + return tuple(_select_chunk(item, index, num_chunks) for item in value) + return value + + +def _slice_batch_mapping( + values: Mapping[str, Any], + index: int, + num_chunks: int, + batch_size: int, + custom_dims: Mapping[str, int] | None = None, +) -> dict[str, Any]: + """Slice batch-aligned tensor leaves into one padded microbatch. + + Args: + values: Mapping whose batch-aligned tensor leaves have shape + [batch, ...], except keys listed in ``custom_dims``. + index: Pipeline microbatch index. + num_chunks: Number of equal pipeline microbatches. + batch_size: Full outer batch extent. + custom_dims: Optional top-level key to batch-axis mapping. A tensor for + key ``name`` then has shape [..., batch, ...] with the batch axis at + ``custom_dims[name]``. + + Returns: + Shallow container copy whose batch-aligned tensors are narrow views; + scalar and non-batch metadata is replicated. + """ + custom_dims = custom_dims or {} + + def slice_value(value: Any, dim: int | None = None) -> Any: + if isinstance(value, torch.Tensor): + if value.ndim == 0: + return value + resolved_dim = dim if dim is not None else (0 if value.shape[0] == batch_size else None) + if resolved_dim is None: + return value + if value.shape[resolved_dim] != batch_size: + raise ValueError( + f"pipeline batch axis {resolved_dim} has length {value.shape[resolved_dim]}, expected {batch_size}" + ) + chunk_size = batch_size // num_chunks + return value.narrow(resolved_dim, index * chunk_size, chunk_size) + if isinstance(value, dict): + return {name: slice_value(item) for name, item in value.items()} + if isinstance(value, list): + return [slice_value(item) for item in value] + if isinstance(value, tuple): + return tuple(slice_value(item) for item in value) + return value + + return { + name: slice_value(value, custom_dims.get(name)) + for name, value in values.items() + if value is not None and not (isinstance(value, dict) and not value) + } def _weighted_numerator(losses: Any, weights: torch.Tensor) -> torch.Tensor: diff --git a/nemo_automodel/recipes/base_recipe.py b/nemo_automodel/recipes/base_recipe.py index 69520f9e2e..c948b5cbfe 100644 --- a/nemo_automodel/recipes/base_recipe.py +++ b/nemo_automodel/recipes/base_recipe.py @@ -756,6 +756,42 @@ def _get_cp_group_size(self): return 1 return device_mesh["cp"].size() + def _validate_mtp_context_parallelism(self, model_parts: list[nn.Module]) -> None: + """Reject MTP until future-token shifts are CP-aware. + + MTP currently rolls inputs and labels after context-parallel sharding. + A local roll cannot recover the next token across CP rank boundaries, + and round-robin layouts introduce additional false adjacencies. The + model-wide reduction makes the decision identical on every PP stage, + including stages that do not locally own the MTP module. + """ + mesh_context = getattr(self, "mesh_context", None) + if mesh_context is None or mesh_context.cp_size <= 1: + return + + modules = ( + module + for part in model_parts + for module in (part.modules() if callable(getattr(part, "modules", None)) else (part,)) + ) + local_has_mtp = any( + getattr(module, "mtp", None) is not None + or bool(getattr(getattr(module, "mtp_config", None), "enabled", False)) + for module in modules + ) + enabled = torch.tensor( + int(local_has_mtp), + dtype=torch.int32, + device=getattr(getattr(self, "dist_env", None), "device", torch.device("cpu")), + ) + if dist.is_initialized(): + dist.all_reduce(enabled, op=dist.ReduceOp.MAX, group=getattr(mesh_context, "process_group", None)) + if bool(enabled.item()): + raise NotImplementedError( + "MTP with context parallelism is not supported because future-token shifts are not CP-aware; " + "set the model's MTP layer count to 0 or use cp_size=1" + ) + def _set_moe_aux_loss_backward_scale(self, *, num_batches: int, num_label_tokens: int) -> None: """Set the per-microbatch MoE auxiliary-loss scale for one optimizer step. diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 6b37ae3588..865b68c2b7 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -32,7 +32,7 @@ import time from contextlib import nullcontext from dataclasses import replace -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any import mlflow import torch @@ -64,7 +64,7 @@ from nemo_automodel.components.distributed.init_utils import initialize_distributed from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.distributed.pipelining import AutoPipeline -from nemo_automodel.components.distributed.utils import FirstRankPerNode, dp_eval_sample_shard, get_sync_ctx +from nemo_automodel.components.distributed.utils import FirstRankPerNode, dp_eval_sample_shard from nemo_automodel.components.loggers.log_utils import setup_logging from nemo_automodel.components.loggers.metric_logger import MetricsSample, build_metric_logger from nemo_automodel.components.loggers.mlflow_utils import ( @@ -82,9 +82,6 @@ from nemo_automodel.components.training.utils import ( count_tail_padding, get_expert_tp_replication_factor, - prepare_after_first_microbatch, - prepare_for_final_backward, - prepare_for_grad_accumulation, scale_grads_and_clip_grad_norm, ) from nemo_automodel.components.utils.compile_utils import ( @@ -151,6 +148,31 @@ def _should_pack_validation( ) +def _validate_pipeline_thd_model(model: nn.Module) -> None: + """Require an explicit model-owned THD path for pipeline training. + + Args: + model: First local pipeline model part. + + Raises: + ValueError: If the model has no native THD capability, CP preparation + hook, or model-owned TE/Magi attention backend. + """ + backend_attn = getattr(getattr(model, "backend", None), "attn", None) + if ( + bool(getattr(model, "supports_thd", False)) + or callable(getattr(model, "prepare_model_inputs_for_cp", None)) + or backend_attn in ("te", "magi") + ): + return + + raise ValueError( + f"Pipeline parallelism with THD batches is not supported for {type(model).__name__}. " + "Generic Hugging Face pipeline stages do not consume packed document boundaries. " + "Use a model with native THD support or disable THD packing." + ) + + def _should_precompute_pp_causal_masks(model_config: Any) -> bool: """Return whether the recipe should attach PP causal-mask precomputation.""" return getattr(model_config, "model_type", None) != "deepseek_v4" @@ -441,9 +463,7 @@ class TrainFinetuneRecipeForNextTokenPrediction(BaseRecipe): This class orchestrates training, from setup to main training loop. """ - # MagiAttention is disabled until setup() resolves it from config; this - # disabled default keeps _forward_backward_step working if setup() is skipped - # (e.g. unit tests that exercise the step directly). It is read-only. + # MagiAttention is disabled until setup() resolves it from config. It is read-only. magi = MagiState() def __init__(self, cfg): @@ -508,6 +528,14 @@ def setup(self): if not self._should_setup_training_components(): return + if getattr(self.distributed_config, "calculate_per_token_loss", False): + raise NotImplementedError( + "Engine-backed finetuning does not support " + "MegatronFSDP calculate_per_token_loss=True; use averaged gradients instead." + ) + if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): + raise ValueError("Engine-backed finetuning requires distributed.pipeline.scale_grads_in_schedule=False") + # MagiAttention (FFA / context-parallel) backend, enabled via # model.attn_implementation="magi" (HF) or model.backend.attn="magi" (custom). self.magi = setup_magi(self.cfg, self.device_mesh) @@ -546,9 +574,18 @@ def setup(self): pp_batch_size = self.cfg.get("step_scheduler.local_batch_size", 1) pp_microbatch_size = self.cfg.get("distributed.pipeline.pp_microbatch_size", 1) - assert pp_batch_size // pp_microbatch_size >= self.mesh_context.pp_size, ( - f"pp_batch_size {pp_batch_size} // pp_microbatch_size {pp_microbatch_size} must be >= pp_size {self.mesh_context.pp_size}" - ) + if self.magi.enabled: + if pp_batch_size != 1 or pp_microbatch_size != 1: + raise ValueError( + "Magi pipeline training requires local_batch_size=1 and pp_microbatch_size=1; " + "use outer gradient accumulation for larger optimizer windows" + ) + else: + if pp_batch_size // pp_microbatch_size < self.mesh_context.pp_size: + raise ValueError( + f"pp_batch_size {pp_batch_size} // pp_microbatch_size {pp_microbatch_size} " + f"must be >= pp_size {self.mesh_context.pp_size}" + ) # THD override logic if ( @@ -563,9 +600,8 @@ def setup(self): f"Overriding pp_batch_size: {pp_batch_size}, pp_microbatch_size: {pp_microbatch_size} for THD" ) - assert not isinstance(self.distributed_config, MegatronFSDPConfig), ( - "MegatronFSDPConfig is not supported when pipeline parallelism is enabled" - ) + if isinstance(self.distributed_config, MegatronFSDPConfig): + raise ValueError("MegatronFSDPConfig is not supported when pipeline parallelism is enabled") # Update pipeline_config runtime fields self.pipeline_config.pp_batch_size = pp_batch_size @@ -624,6 +660,9 @@ def setup(self): cfg_qat=self.cfg.get("qat", None), sdpa_method=self.cfg.get("sdpa_method", None), ) + if self.pp_enabled and self.cfg.dataloader.emits_thd: + first_model_part = model.parts[0] if isinstance(model, AutoPipeline) else model + _validate_pipeline_thd_model(first_model_part) self.embedding_row_repair_report = None embedding_row_repair = self.cfg.embedding_row_repair if embedding_row_repair is not None and embedding_row_repair.enabled: @@ -654,12 +693,15 @@ def setup(self): self.model_parts = [model] self.pp = None + self._validate_mtp_context_parallelism(self.model_parts) + # Loss-function capability check self.loss_fn = _maybe_downgrade_loss_fn(self.loss_fn, self.model_parts[0], self.pp is not None) # Extract TE FP8 config from model backend (set after model construction) self.te_fp8 = self.model_parts[0].backend.te_fp8 if hasattr(self.model_parts[0], "backend") else None + self.pipeline_loss_fn = None if self.pp_enabled: self._configure_pipeline_loss_fn() @@ -697,42 +739,22 @@ def setup(self): # Tokenizer + model-derived values are runtime concerns: build them here and pass them to # each DataloaderConfig.build(); the configs themselves are resolved at the RecipeConfig boundary. _, self.tokenizer = _build_tokenizer(self.cfg.model, self.cfg.dataset) - model_has_mtp = any( - getattr(module, "mtp", None) is not None - for model_part in self.model_parts - for module in model_part.modules() + if getattr(self.loss_fn, "reduction", None) != "sum": + raise ValueError("Engine-backed finetuning requires a loss with reduction='sum'") + self.engine = Engine( + self.pp if self.pp_enabled else self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + microbatch_size=1, + collate_fn=collate_prebatched, + padding_token_id=(self.tokenizer.pad_token_id if self.tokenizer is not None else 0) or 0, + context_fn=( + (lambda _model_inputs: self.te_fp8.maybe_te_autocast()) + if self.te_fp8 is not None + else (lambda _model_inputs: nullcontext()) + ), + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), ) - pp_group = self._get_pp_group() if self.pp_enabled else None - if pp_group is not None: - # The MTP head may exist only on the last PP stage; all stages must choose the same path. - model_has_mtp_flag = torch.tensor(int(model_has_mtp), device=self.dist_env.device) - torch.distributed.all_reduce(model_has_mtp_flag, op=torch.distributed.ReduceOp.MAX, group=pp_group) - model_has_mtp = bool(model_has_mtp_flag.item()) - - dataloader_config = self.cfg.dataloader - pp_uses_packed_batches = self.pp_enabled and ( - _packed_seq_size > 0 or bool(getattr(dataloader_config, "emits_thd", False)) - ) - self.engine = None - if ( - not self.magi.enabled - and not (model_has_mtp and (self.pp_enabled or self.mesh_context.cp_size > 1)) - and not pp_uses_packed_batches - and not (self.pp_enabled and self.mesh_context.cp_size > 1) - and not (self.pp_enabled and isinstance(self.loss_fn, FusedLinearCrossEntropy)) - and not (self.pp_enabled and self.pp.scale_grads_in_schedule) - and not getattr(self.distributed_config, "calculate_per_token_loss", False) - and getattr(self.loss_fn, "reduction", None) == "sum" - ): - self.engine = Engine( - self.pp if self.pp_enabled else self.model_parts[0], - device=self.dist_env.device, - mesh_context=self.mesh_context, - collate_fn=collate_prebatched, - padding_token_id=(self.tokenizer.pad_token_id if self.tokenizer is not None else 0) or 0, - context_fn=self.te_fp8.maybe_te_autocast if self.te_fp8 is not None else nullcontext, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) attn_implementation = None if ( self.cfg.get("packed_sequence.packed_sequence_size", 0) > 0 @@ -759,7 +781,10 @@ def materialize_loader(config): supports_seq_lens=_supports_seq_lens(self.model_parts[0]), cp_size=self.cfg.get("distributed.cp_size", 1), attn_implementation=attn_implementation, - collate_wrapper=collate_wrapper, + # Raw THD batches already encode causal document boundaries + # in cu_seqlens. Building a dense [B, 1, S, S] PP mask is + # both redundant and prohibitive for long packed streams. + collate_wrapper=None if getattr(config, "emits_thd", False) else collate_wrapper, ) self.dataloader = materialize_loader(self.cfg.dataloader) @@ -922,11 +947,14 @@ def _configure_pipeline_loss_fn(self): if isinstance(self.loss_fn, FusedLinearCrossEntropy): last_stage_model._pp_return_hidden_states = True - self.pp.info.schedule._loss_fn = self.cfg.mtp.build( + self.pipeline_loss_fn = self.cfg.mtp.build( self.loss_fn, last_stage_model, grad_reduce_group=self._get_dp_group(include_cp=True), ) + # Validation still executes the schedule directly. Training supplies the + # same loss through Engine's per-microbatch callback. + self.pp.info.schedule._loss_fn = self.pipeline_loss_fn def _setup_qat(self, cfg, model_parts: list[nn.Module]): if not cfg.get("qat.enabled", False): @@ -1044,8 +1072,17 @@ def run_train_validation_loop(self): self._partial_cuda_graph_capture_pending = False # ------------------ helpers ------------------ - def _prepare_microbatch(self, batch): - """Move and CP-prepare one batch before model execution.""" + def _prepare_validation_batch(self, batch: dict[str, Any]): + """Move and CP-prepare one validation batch. + + Args: + batch: Worker-collated inputs. Padded token tensors have shape + [batch, sequence]; packed THD token tensors have shape [tokens]. + + Returns: + The CP context factory, CP-local model-input mapping, and labels + matching the model output's local token axes. + """ batch = { k: ( {dk: dv.to(self.dist_env.device, non_blocking=True) for dk, dv in v.items() if dv is not None} @@ -1120,7 +1157,17 @@ def _compute_causal_lm_loss(self, output, labels, model_inputs, *, num_label_tok grad_reduce_group=grad_reduce_group, ) - def _make_engine_datum(self, batch): + def _make_engine_datum(self, batch: dict[str, Any]) -> Datum: + """Wrap one worker-collated text batch for ``collate_prebatched``. + + Args: + batch: Model inputs plus labels. Padded labels have shape [batch, + sequence]; packed labels have shape [tokens]. + + Returns: + A Datum whose model tensors preserve their input layout and whose + labels and weights share the same token axes. + """ labels = batch["labels"] model_inputs = {key: value for key, value in batch.items() if key != "labels"} if isinstance(self.loss_fn, FusedLinearCrossEntropy): @@ -1130,26 +1177,42 @@ def _make_engine_datum(self, batch): loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, ) - def _engine_loss_fn(self, output, loss_inputs, _datums, model_inputs): + def _engine_loss_fn(self, output: Any, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Compute a local summed causal-LM loss for Engine normalization. + + Args: + output: Model output with logits shaped [batch, sequence, vocab], + packed logits shaped [tokens, vocab], or the PP/MTP tuple contract. + loss_inputs: CP-local labels and weights with matching token axes, + plus optional packed-sequence metadata. + + Returns: + Scalar local loss-sum tensor. + """ + if self.pp_enabled: + if self.pipeline_loss_fn is None: + raise RuntimeError("The last pipeline stage has no configured causal-LM loss") + self.pipeline_loss_fn.cu_seqlens = loss_inputs.get("cu_seqlens") + return self.pipeline_loss_fn(output, loss_inputs["labels"]) return self._compute_causal_lm_loss( output, loss_inputs["labels"], - model_inputs, + loss_inputs, num_label_tokens=None, is_train=True, ) - def _forward_backward_step( - self, - idx, - batch, - *, - loss_buffer, - num_label_tokens, - num_batches, - is_train: bool = True, - ): - train_ctx, batch, labels = self._prepare_microbatch(batch) + def _forward_validation_step(self, batch: dict[str, Any]) -> torch.Tensor: + """Run one recipe-owned forward-only validation step. + + Args: + batch: Worker-collated inputs and labels. Padded token tensors have + shape [batch, sequence]; packed THD token tensors have shape [tokens]. + + Returns: + Detached scalar local loss-sum tensor. + """ + train_ctx, batch, labels = self._prepare_validation_batch(batch) fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext() if self.pp_enabled: @@ -1178,56 +1241,31 @@ def _forward_backward_step( cu_seqlens = batch_filtered.get("cu_seqlens") if isinstance(cu_seqlens, torch.Tensor) and cu_seqlens.dim() == 2: cu_seqlens = cu_seqlens.squeeze(0) # [1, T] -> [T] - pp_loss_fn = getattr(self.pp.info.schedule, "_loss_fn", None) if self.pp.info.has_last_stage else None - if pp_loss_fn is not None and hasattr(pp_loss_fn, "cu_seqlens"): - pp_loss_fn.cu_seqlens = cu_seqlens - if is_train: - # Use step for training (forward + backward) - if self.pp.info.has_first_stage: - self.pp.info.schedule.step(input_ids, target=targets, losses=losses, **batch_filtered) - else: - self.pp.info.schedule.step(target=targets, losses=losses, **batch_filtered) + if self.pipeline_loss_fn is not None: + self.pipeline_loss_fn.cu_seqlens = cu_seqlens + if self.pp.info.has_first_stage: + self.pp.info.schedule.eval(input_ids, target=targets, losses=losses, **batch_filtered) else: - # Use eval for validation (forward only, no backward) - if self.pp.info.has_first_stage: - self.pp.info.schedule.eval(input_ids, target=targets, losses=losses, **batch_filtered) - else: - self.pp.info.schedule.eval(target=targets, losses=losses, **batch_filtered) + self.pp.info.schedule.eval(target=targets, losses=losses, **batch_filtered) if self.pp.info.has_last_stage: - local_loss = torch.sum(torch.stack(losses)) - else: - local_loss = torch.tensor(0.0, device=self.dist_env.device) + return torch.sum(torch.stack(losses)).detach() + return torch.zeros((), device=self.dist_env.device) - loss_buffer.append(local_loss.clone().detach()) - else: - model = self.model_parts[0] - sync_ctx = ( - get_sync_ctx( - model, - idx == num_batches - 1, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) - if is_train - else nullcontext() - ) - with train_ctx(), sync_ctx, fp8_ctx: - batch = filter_forward_kwargs(model, batch) - if isinstance(self.loss_fn, FusedLinearCrossEntropy): - # use num_logits_to_keep to avoid full logits matrix in memory - out = model(logits_to_keep=1, **batch) - else: - out = model(**batch) - local_loss = self._compute_causal_lm_loss( - out, - labels, - batch, - num_label_tokens=num_label_tokens, - is_train=is_train, - ) - loss_buffer.append(local_loss.clone().detach()) - if is_train: - (local_loss * self._get_dp_group_size(include_cp=True)).backward() + model = self.model_parts[0] + with train_ctx(), fp8_ctx: + batch = filter_forward_kwargs(model, batch) + if isinstance(self.loss_fn, FusedLinearCrossEntropy): + out = model(logits_to_keep=1, **batch) + else: + out = model(**batch) + return self._compute_causal_lm_loss( + out, + labels, + batch, + num_label_tokens=None, + is_train=False, + ).detach() def _broadcast_from_last_pp_stage(self, tensor: torch.Tensor) -> torch.Tensor: """Broadcast a PP last-stage scalar to the other ranks in its pipeline group.""" @@ -1236,12 +1274,16 @@ def _broadcast_from_last_pp_stage(self, tensor: torch.Tensor) -> torch.Tensor: torch.distributed.broadcast(tensor, src=pp_src_rank, group=pp_group) return tensor - def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): + def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: float | None = None) -> MetricsSample: """Execute a single training step. Args: - batches: List of batches of training data. + batches: Worker-collated optimizer window. Padded token tensors use + shape [batch, sequence]; packed tensors use their THD token layout. max_grad_norm: Gradient clipping norm. Optional, if None will not clip gradients. + + Returns: + Metrics for the completed optimizer step. """ num_label_tokens = torch.tensor( @@ -1249,10 +1291,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ) num_label_tokens = self._dp_allreduce(num_label_tokens).item() - num_batches = len(batches) - - loss_buffer = [] - # number of tokens in the batch, excluding any tail padding. num_tokens_in_batch = torch.tensor( sum(batch["labels"].numel() - count_tail_padding(batch["labels"]) for batch in batches), @@ -1260,31 +1298,10 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - engine = getattr(self, "engine", None) - # A custom prepacked loader may emit THD even when setup did not advertise packing. - pp_uses_thd = self.pp_enabled and any(batch.get("qkv_format") == "thd" for batch in batches) - use_engine = engine is not None and num_label_tokens > 0 and not pp_uses_thd - if use_engine: - reporting_loss, _ = engine.forward_backward( - [[self._make_engine_datum(batch)] for batch in batches], - self._engine_loss_fn, - ) - else: - # Engine requires a positive global weight sum. Keep the existing - # zero-label behavior on the legacy path. - self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) - - for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) - - self._forward_backward_step( - i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches - ) - - if i == 0: - prepare_after_first_microbatch() + reporting_loss, _ = self.engine.forward_backward( + [self._make_engine_datum(batch) for batch in batches], + self._engine_loss_fn, + ) grad_norm = scale_grads_and_clip_grad_norm( max_grad_norm, @@ -1296,7 +1313,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, pp_axis_name="pp" if self.pp_enabled else None, foreach=True, - num_label_tokens=None if use_engine else num_label_tokens, + num_label_tokens=None, dp_group_size=self._get_dp_group_size(include_cp=True), expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, self.device_mesh), ) @@ -1360,14 +1377,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ).item() mfu = calculate_mfu(step_flops / 1e12, self.dist_env.world_size, time_delta) - if not use_engine: - reporting_loss = torch.sum(torch.stack(loss_buffer)) - reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) - if self.pp_enabled: - reporting_loss = reporting_loss / num_label_tokens - reporting_loss = reporting_loss.to(self.dist_env.device) - reporting_loss = self._broadcast_from_last_pp_stage(reporting_loss) - reporting_loss = reporting_loss.cpu().item() # fix reporting_loss, tps across ranks @@ -1403,18 +1412,8 @@ def _run_validation_epoch(self, val_dataloader): total_num_label_tokens = 0 for batch in val_dataloader: - loss_buffer = [] num_label_tokens = (batch["labels"] != -100).sum().item() - self._forward_backward_step( - 0, - batch, - loss_buffer=loss_buffer, - num_label_tokens=None, # we will normalize outside. - num_batches=1, - is_train=False, - ) - - total_loss += torch.sum(torch.stack(loss_buffer)).item() + total_loss += self._forward_validation_step(batch).item() total_num_label_tokens += num_label_tokens total_loss = self._dp_allreduce(total_loss, include_cp=True) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 626cc92925..8beb181d6e 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -28,9 +28,8 @@ import logging import pathlib import time -from collections.abc import Sequence -from contextlib import contextmanager, nullcontext -from typing import TYPE_CHECKING, Any, Optional, Protocol +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Protocol import mlflow import torch @@ -48,7 +47,7 @@ from nemo_automodel._transformers.utils import apply_cache_compatibility_patches, resolve_get_rope_index from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.components.datasets.datum import Datum -from nemo_automodel.components.datasets.vlm.pp_media import stage_vlm_media_for_pp +from nemo_automodel.components.datasets.vlm.pp_media import VLM_PP_MEDIA_KEY, stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi @@ -59,7 +58,7 @@ ) from nemo_automodel.components.distributed.init_utils import initialize_distributed from nemo_automodel.components.distributed.pipelining import AutoPipeline -from nemo_automodel.components.distributed.utils import FirstRankPerNode, get_sync_ctx +from nemo_automodel.components.distributed.utils import FirstRankPerNode from nemo_automodel.components.loggers.log_utils import setup_logging from nemo_automodel.components.loggers.metric_logger import MetricsSample, build_metric_logger from nemo_automodel.components.loggers.mlflow_utils import ( @@ -77,9 +76,6 @@ from nemo_automodel.components.training.utils import ( count_tail_padding, get_expert_tp_replication_factor, - prepare_after_first_microbatch, - prepare_for_final_backward, - prepare_for_grad_accumulation, scale_grads_and_clip_grad_norm, ) from nemo_automodel.components.utils.compile_utils import build_compile_config @@ -400,9 +396,7 @@ def build_dataloader( class FinetuneRecipeForVLM(BaseRecipe): """Recipe for fine-tuning a VLM model.""" - # MagiAttention is disabled until setup() resolves it from config; this - # disabled default keeps the train step working if setup() is skipped (e.g. - # unit tests that exercise the step directly). It is read-only. + # MagiAttention is disabled until setup() resolves it from config. It is read-only. magi = MagiState() def __init__(self, cfg): @@ -462,6 +456,14 @@ def setup(self): if not self._should_setup_training_components(): return + if getattr(self.distributed_config, "calculate_per_token_loss", False): + raise NotImplementedError( + "Engine-backed VLM finetuning does not support " + "MegatronFSDP calculate_per_token_loss=True; use averaged gradients instead." + ) + if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): + raise ValueError("Engine-backed VLM finetuning requires distributed.pipeline.scale_grads_in_schedule=False") + # MagiAttention (FFA) backend for the language backbone; the vision tower # stays on SDPA. Enabled via model.attn_implementation="magi" (HF VLMs) or # model.backend.attn="magi" (custom VLMs, e.g. qwen3_vl_moe). @@ -490,13 +492,21 @@ def setup(self): pp_batch_size = self.cfg.get("step_scheduler.local_batch_size", 1) pp_microbatch_size = self.cfg.get("distributed.pipeline.pp_microbatch_size", 1) - assert pp_batch_size // pp_microbatch_size >= self.mesh_context.pp_size, ( - f"pp_batch_size {pp_batch_size} // pp_microbatch_size {pp_microbatch_size} must be >= pp_size {self.mesh_context.pp_size}" - ) + if self.magi.enabled: + if pp_batch_size != 1 or pp_microbatch_size != 1: + raise ValueError( + "Magi pipeline training requires local_batch_size=1 and pp_microbatch_size=1; " + "use outer gradient accumulation for larger optimizer windows" + ) + else: + if pp_batch_size // pp_microbatch_size < self.mesh_context.pp_size: + raise ValueError( + f"pp_batch_size {pp_batch_size} // pp_microbatch_size {pp_microbatch_size} " + f"must be >= pp_size {self.mesh_context.pp_size}" + ) - assert not isinstance(self.distributed_config, MegatronFSDPConfig), ( - "MegatronFSDPConfig is not supported when pipeline parallelism is enabled" - ) + if isinstance(self.distributed_config, MegatronFSDPConfig): + raise ValueError("MegatronFSDPConfig is not supported when pipeline parallelism is enabled") # Update pipeline_config runtime fields self.pipeline_config.pp_batch_size = pp_batch_size @@ -564,6 +574,8 @@ def setup(self): else: self.model_parts = [model] self.pp = None + self._validate_mtp_context_parallelism(self.model_parts) + self.pipeline_loss_fn = None if self.pp_enabled: self._configure_pipeline_loss_fn() @@ -628,29 +640,19 @@ def setup(self): self.dataloader = dataloader_build.dataloader self.processor = dataloader_build.processor - model_has_mtp = not self.pp_enabled and any( - getattr(module, "mtp", None) is not None for module in self.model_parts[0].modules() + if getattr(self.loss_fn, "reduction", None) != "sum": + raise ValueError("Engine-backed VLM finetuning requires a loss with reduction='sum'") + padding_token_id = getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0 + self.engine = Engine( + self.pp if self.pp_enabled else self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + microbatch_size=1, + collate_fn=collate_prebatched, + padding_token_id=padding_token_id, + context_fn=self._engine_context, + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), ) - self.engine = None - if ( - not self.pp_enabled - and not self.magi.enabled - and not (self.mesh_context.cp_size > 1 and model_has_mtp) - and not getattr(self.distributed_config, "calculate_per_token_loss", False) - and getattr(self.loss_fn, "reduction", None) == "sum" - ): - padding_token_id = ( - getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0 - ) - self.engine = Engine( - self.model_parts[0], - device=self.dist_env.device, - mesh_context=self.mesh_context, - collate_fn=collate_prebatched, - padding_token_id=padding_token_id, - context_fn=self._cp_vision_frame_sharding_context, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) # Build validation dataloader if the config provides it self.val_dataloader = None @@ -884,67 +886,65 @@ def _compute_vlm_loss( log_denominator=log_denominator, ) - def _make_engine_datum( - self, - batch: dict[str, Any], - *, - log_drafter: bool = False, - log_denominator: int | float = 1, - ) -> Datum: - """Wrap one processor-collated VLM batch as a Datum.""" + def _make_engine_datum(self, batch: dict[str, Any]) -> Datum: + """Wrap one processor-collated VLM batch for ``collate_prebatched``. + + Args: + batch: Model and media inputs plus labels. Padded labels have shape + [batch, sequence]; packed labels have shape [tokens]. + + Returns: + A Datum preserving model/media layouts with labels and weights on + matching token axes. Non-first PP stages omit raw media tensors. + """ labels = batch["labels"] model_inputs = {key: value for key, value in batch.items() if key != "labels"} + if self.pp_enabled and not self.pp.info.has_first_stage: + for key in VLM_INPUT_KEYS: + if key != "input_ids": + model_inputs.pop(key, None) + model_inputs.pop(VLM_PP_MEDIA_KEY, None) if isinstance(self.loss_fn, FusedLinearCrossEntropy): model_inputs["logits_to_keep"] = 1 return Datum( model_inputs=model_inputs, - loss_fn_inputs={ - "labels": labels, - "weights": labels.ne(-100), - "log_drafter": torch.tensor(log_drafter), - "log_denominator": torch.tensor(log_denominator), - }, + loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, ) def _engine_loss_fn( self, out: Any, loss_inputs: dict[str, torch.Tensor], - datums: Sequence[Datum], - model_inputs: dict[str, Any], + *, + log_drafter: bool = False, + log_denominator: int | float | None = None, ) -> torch.Tensor: - """Return the local loss sum; Engine owns global normalization.""" + """Compute a local summed VLM loss for Engine normalization. + + Args: + out: Model output with logits shaped [batch, sequence, vocab], + packed logits shaped [tokens, vocab], or the PP/MTP tuple contract. + loss_inputs: CP-local labels and weights with matching token axes, + plus optional packed-sequence metadata. + + Returns: + Scalar local loss-sum tensor. + """ + if self.pp_enabled: + if self.pipeline_loss_fn is None: + raise RuntimeError("The last pipeline stage has no configured causal-LM loss") + self.pipeline_loss_fn.cu_seqlens = loss_inputs.get("cu_seqlens") + return self.pipeline_loss_fn(out, loss_inputs["labels"]) return self._compute_vlm_loss( out=out, labels=loss_inputs["labels"], num_label_tokens=None, is_train=True, - cu_seqlens=model_inputs.get("cu_seqlens"), - log_drafter=bool(datums[0].loss_fn_inputs["log_drafter"].item()), - log_denominator=float(datums[0].loss_fn_inputs["log_denominator"].item()), + cu_seqlens=loss_inputs.get("cu_seqlens"), + log_drafter=log_drafter, + log_denominator=log_denominator, ) - def _maybe_set_pp_first_stage_embed_input_meta(self, model_input: torch.Tensor) -> None: - if ( - not self.pp_enabled - or not getattr(self.pp.info, "has_first_stage", False) - or not model_input.dtype.is_floating_point - or model_input.ndim != 3 - ): - return - - for stage in self.pp.info.stages: - if stage.is_first: - stage.inputs_meta = ( - torch.empty( - self.pp.pp_microbatch_size, - model_input.shape[1], - model_input.shape[2], - device="meta", - dtype=model_input.dtype, - ), - ) - @contextmanager def _cp_vision_frame_sharding_context(self): """Publish the CP-only group while a VLM forward may run its vision tower.""" @@ -967,124 +967,22 @@ def _cp_vision_frame_sharding_context(self): finally: reset_cp_vision_group(token) - def _forward_backward_step( - self, - idx, - batch, - *, - loss_buffer, - num_label_tokens, - num_batches, - is_train: bool = True, - ): - batch = {k: _move_to_device(v, self.dist_env.device) for k, v in batch.items()} - - # Single CP dispatch (magi / model-owned / generic). The pre-embed hook is - # a plain method call (prepare_model_inputs_for_cp): sharder-only, it - # touches no weights and consumes nothing. Invoke it on EVERY pp stage so - # its aux-only sharder keeps input_ids full-length everywhere; otherwise - # non-first stages hit the generic round-robin sharder, feed an - # already-local seq_len to update_seq_len, and get_pipeline_stage_metas - # ÷cp a second time -> the inter-stage hidden truncates to S/cp² - # (text-decoder RoPE size mismatch). - _is_first_or_no_pp = not self.pp_enabled or getattr(self.pp.info, "has_first_stage", False) - _cp_active = ( - self.device_mesh is not None - and "cp" in getattr(self.device_mesh, "mesh_dim_names", ()) - and self.device_mesh["cp"].size() > 1 - ) - if _cp_active and not _is_first_or_no_pp and hasattr(self.model_parts[0], "prepare_model_inputs_for_cp"): - # Non-first PP stages don't embed; drop raw multimodal inputs so their - # forwards see only text. - for k in VLM_INPUT_KEYS: - if k != "input_ids": - batch.pop(k, None) - # THD packed VLM inputs (qkv_format='thd' from the packing collator) use TE - # sequence metadata even without context parallelism (#3052). Standard - # one-dimensional RoPE can follow the generic TE CP partition. Multi-axis - # mRoPE still needs axis-aware sharding before it can use this path. - _use_te_vlm = batch.get("qkv_format", None) == "thd" - position_ids = batch.get("position_ids") - if ( - _use_te_vlm - and self.mesh_context.cp_size > 1 - and isinstance(position_ids, torch.Tensor) - and position_ids.ndim == 3 - ): - raise NotImplementedError( - "Context-parallel THD packing for multi-axis mRoPE VLMs is not yet implemented; " - "use one-dimensional position_ids or cp_size=1." - ) - _padding_id = getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0 - cp_sharder = ContextParallelSharder( - self.model_parts[0], - self.device_mesh, - batch, - padding_token_id=_padding_id, - invoke_pre_embed=True, - ) - train_ctx, batch = cp_sharder.shard(batch) - labels = batch.pop("labels") - - if self.pp_enabled: - if not is_train: - logging.info("Skipping forward pass for validation because pipeline parallelism is enabled") - return - - with self._cp_vision_frame_sharding_context(), train_ctx(): - losses = [] if self.pp.info.has_last_stage else None - if self.pp.info.has_last_stage: - masked_labels = labels.clone() - targets = masked_labels - else: - targets = None - - model_input_key = "inputs_embeds" if "inputs_embeds" in batch else "input_ids" - model_input = batch.pop(model_input_key) - self.pp.update_seq_len(model_input.shape[1]) - self._maybe_set_pp_first_stage_embed_input_meta(model_input) - - with stage_vlm_media_for_pp(self.pp, self.model_parts, batch): - self.pp.step(model_input, target=targets, losses=losses, **batch) - - if self.pp.info.has_last_stage: - local_loss = torch.sum(torch.stack(losses)) - else: - local_loss = torch.tensor(0.0, device=self.dist_env.device) + @contextmanager + def _engine_context(self, model_inputs: dict[str, Any]): + """Install VLM runtime state around one Engine forward/backward call. - loss_buffer.append(local_loss.clone().detach()) - else: - model = self.model_parts[0] - sync_ctx = ( - get_sync_ctx( - model, - idx == num_batches - 1, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) - if is_train - else nullcontext() - ) - with sync_ctx, self._cp_vision_frame_sharding_context(), train_ctx(): - batch = filter_forward_kwargs(model, batch) - if isinstance(self.loss_fn, FusedLinearCrossEntropy): - # use num_logits_to_keep to avoid full logits matrix in memory - out = model(logits_to_keep=1, **batch) - else: - out = model(**batch) - - local_loss = self._compute_vlm_loss( - out=out, - labels=labels, - num_label_tokens=num_label_tokens, - is_train=is_train, - cu_seqlens=batch.get("cu_seqlens"), - # Log once per remote-logging step on the first microbatch. - log_drafter=(idx == 0 and self.step_scheduler.is_remote_logging_step), - ) + Args: + model_inputs: CP-local outer-batch mapping. Padded token tensors use + shape [batch, sequence]; THD tensors use the sharder-produced + packed layout. PP media is stored as per-microbatch tensor lists. + """ + if not self.pp_enabled: + with self._cp_vision_frame_sharding_context(): + yield + return - loss_buffer.append(local_loss.clone().detach()) - if is_train: - (local_loss * self._get_dp_group_size(include_cp=True)).backward() + with self._cp_vision_frame_sharding_context(), stage_vlm_media_for_pp(self.pp, self.model_parts, model_inputs): + yield def _configure_pipeline_loss_fn(self): if self.pp is None or not self.pp.info.has_last_stage: @@ -1098,26 +996,33 @@ def _configure_pipeline_loss_fn(self): if last_stage_model is None: raise RuntimeError("Pipeline reports a last stage, but no last-stage model part was found") - self.pp.info.schedule._loss_fn = self.cfg.mtp.build( + if isinstance(self.loss_fn, FusedLinearCrossEntropy): + last_stage_model._pp_return_hidden_states = True + + self.pipeline_loss_fn = self.cfg.mtp.build( self.loss_fn, last_stage_model, grad_reduce_group=self._get_dp_group(include_cp=True), ) + self.pp.info.schedule._loss_fn = self.pipeline_loss_fn - def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): + def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: float | None = None) -> MetricsSample: """Execute a single training step. Args: - batches: List of batches of training data. + batches: Processor-collated optimizer window. Padded token tensors + use shape [batch, sequence]; packed tensors use their THD token + layout, and media tensors retain model-specific layouts. max_grad_norm: Gradient clipping norm. Optional, if None will not clip gradients. + + Returns: + Metrics for the completed optimizer step. """ num_label_tokens = torch.tensor( sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long ) num_label_tokens = self._dp_allreduce(num_label_tokens).item() - num_batches = len(batches) - # number of tokens in the batch, excluding any tail padding. num_tokens_in_batch = torch.tensor( sum(batch["labels"].numel() - count_tail_padding(batch["labels"]) for batch in batches), @@ -1125,45 +1030,23 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - engine = getattr(self, "engine", None) - unsupported_thd_mrope = self._get_cp_group_size() > 1 and any( - batch.get("qkv_format") == "thd" - and isinstance(batch.get("position_ids"), torch.Tensor) - and batch["position_ids"].ndim == 3 - for batch in batches - ) - use_engine = engine is not None and num_label_tokens > 0 and not unsupported_thd_mrope - if use_engine: - reporting_loss, _ = engine.forward_backward( - [ - [ - self._make_engine_datum( - batch, - log_drafter=(index == 0 and self.step_scheduler.is_remote_logging_step), - log_denominator=num_label_tokens, - ) - ] - for index, batch in enumerate(batches) - ], - self._engine_loss_fn, + log_drafter = self.step_scheduler.is_remote_logging_step + + def engine_loss_fn(out: Any, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + nonlocal log_drafter + should_log = log_drafter + log_drafter = False + return self._engine_loss_fn( + out, + loss_inputs, + log_drafter=should_log, + log_denominator=max(num_label_tokens, 1), ) - else: - # The eager Engine requires a positive global weight sum. Preserve - # the established zero-label behavior by using the legacy path. - self._set_moe_aux_loss_backward_scale(num_batches=num_batches, num_label_tokens=num_label_tokens) - loss_buffer = [] - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) - - for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) - - self._forward_backward_step( - i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches - ) - if i == 0: - prepare_after_first_microbatch() + reporting_loss, _ = self.engine.forward_backward( + [self._make_engine_datum(batch) for batch in batches], + engine_loss_fn, + ) grad_norm = scale_grads_and_clip_grad_norm( max_grad_norm=max_grad_norm, @@ -1175,7 +1058,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, pp_axis_name="pp" if self.pp_enabled else None, foreach=True, - num_label_tokens=num_label_tokens, + num_label_tokens=None, dp_group_size=self._get_dp_group_size(include_cp=True), expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, self.device_mesh), ) @@ -1215,33 +1098,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None): time_delta = t - self.timestamp self.timestamp = t tps = num_tokens_in_batch / time_delta - if not use_engine: - reporting_loss = torch.sum(torch.stack(loss_buffer)) - reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) - if not use_engine and self.pp_enabled: - # PP uses sum reduction per microbatch (no internal normalization). - # Divide by num_label_tokens to get the mean loss, same as non-PP. - reporting_loss = reporting_loss / num_label_tokens if num_label_tokens > 0 else reporting_loss * 0.0 - reporting_loss = reporting_loss.float().to(self.dist_env.device) - # Send loss to first rank from the last PP stage of rank0's mesh coords. - # This avoids picking a global-rank sender from a different EP/PP group. - if self.device_mesh is not None and "pp" in self.device_mesh.mesh_dim_names: - dim_names = list(self.device_mesh.mesh_dim_names) - mesh = self.device_mesh.mesh - idx = [] - for name in dim_names: - if name == "pp": - idx.append(-1) - else: - idx.append(0) - src_rank = mesh[tuple(idx)].item() - else: - src_rank = self.device_mesh.mesh.reshape(-1)[-1].item() - if self.dist_env.rank == src_rank: - torch.distributed.send(reporting_loss, dst=0) - elif self.dist_env.is_main: - torch.distributed.recv(reporting_loss, src=src_rank) - reporting_loss = reporting_loss.item() # fix reporting_loss, tps across ranks diff --git a/tests/functional_tests/context_parallel/L2_PP_Dense_Packed_Test.sh b/tests/functional_tests/context_parallel/L2_PP_Dense_Packed_Test.sh new file mode 100755 index 0000000000..ea64c63dfe --- /dev/null +++ b/tests/functional_tests/context_parallel/L2_PP_Dense_Packed_Test.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# 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. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1" + +python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run \ + tests/functional_tests/context_parallel/run_packed_pp.py diff --git a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py index b63bc805a0..72e90ab4e4 100644 --- a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py +++ b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py @@ -14,8 +14,8 @@ """images x cp2xpp2 for minimax: the in-forward vision splice under CP+PP (L1, 2/4 GPU). -Regression cover for images x context-parallelism x pipeline-parallelism after the -pre-embed sink. Media rides the existing per-microbatch side channel +Regression cover for images x context-parallelism x pipeline-parallelism through +the Datum Engine after the pre-embed sink. Media rides the existing per-microbatch side channel (prepare_vlm_media_for_pp -> stage_vlm_media_for_pp -> stage-0 chunk pull); the in-forward embed + vision splice runs on the microbatch's full sequence with the CP ring dispatcher suspended around the (non-causal, unsharded) vision tower @@ -32,6 +32,7 @@ import os import sys +from contextlib import nullcontext import torch import torch.distributed as dist @@ -132,17 +133,25 @@ def main(): torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") - from torch.distributed.device_mesh import init_device_mesh - + from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.vlm.pp_media import prepare_vlm_media_for_pp, stage_vlm_media_for_pp - from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder + from nemo_automodel.components.distributed.config import FSDP2Config + from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.moe.parallelizer import apply_cp + from nemo_automodel.engine import Engine, collate_prebatched pp1 = bool(os.environ.get("NEMO_CP_PP_TEST_PP1")) pp_size = 1 if pp1 else 2 cp_size = world // pp_size - mesh = init_device_mesh("cuda", (pp_size, 1, cp_size), mesh_dim_names=("pp", "dp", "cp")) + mesh_context = MeshContext.build( + FSDP2Config(), + ParallelismSizes(dp_size=1, pp_size=pp_size, cp_size=cp_size), + world_size=world, + ) + mesh = mesh_context.device_mesh + if mesh is None: + raise RuntimeError("FSDP2 CP/PP validation requires a device mesh") torch.manual_seed(0) model = build_minimax(device) @@ -153,9 +162,22 @@ def main(): model.train() vocab = model.config.text_config.vocab_size - def loss_fn(output, labels): + def loss_fn(output, loss_inputs): + """Return image-text token losses in the Engine's CP-local layout. + + Args: + output: Model output whose logits have shape [batch, sequence, vocab]. + loss_inputs: Mapping containing labels and weights with shape + [batch, sequence] in the same CP-local layout as the logits. + + Returns: + Tensor of shape [batch, sequence] containing per-token causal-LM loss. + """ logits = output[0] if isinstance(output, tuple) else getattr(output, "logits", output) - return F.cross_entropy(logits.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100) + labels = loss_inputs["labels"] + return F.cross_entropy( + logits.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100, reduction="none" + ).reshape_as(loss_inputs["weights"]) def cp_only(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw): if cp_axis_name is not None and world_mesh[cp_axis_name].size() > 1: @@ -166,9 +188,7 @@ def cp_only(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw): pp = AutoPipeline( world_mesh=mesh, moe_mesh=None, - pp_axis_name="pp", - dp_axis_names=("dp",), - cp_axis_name="cp", + **mesh_context.pipeline_axis_kwargs(), pp_schedule="1f1b", pp_microbatch_size=1, pp_batch_size=2, @@ -176,10 +196,35 @@ def cp_only(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw): dtype=torch.bfloat16, pp_seq_len=seqlen, ).build(model, loss_fn=loss_fn, parallelize_fn=cp_only) - model_part0, has_last, has_first = pp.parts[0], pp.info.has_last_stage, pp.info.has_first_stage + model_part0 = pp.parts[0] + engine_model = pp else: - cp_only(model, mesh, None, dp_axis_names=("dp",), cp_axis_name="cp") + cp_only(model, mesh, None, **mesh_context.parallelize_axis_kwargs()) model_part0 = model + engine_model = model + + def engine_context(model_inputs): + """Stage pre-chunked media for one Engine batch. + + Args: + model_inputs: Mapping whose text tensors have shape [batch, sequence] + and whose PP media field contains per-microbatch tensor chunks. + + Returns: + Context manager that stages media on PP stage zero, or a no-op context + for the eager comparison leg. + """ + if pp_size > 1: + return stage_vlm_media_for_pp(pp, pp.parts, model_inputs) + return nullcontext() + + engine = Engine( + engine_model, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + context_fn=engine_context, + ) losses = [] for step in range(20): @@ -191,30 +236,12 @@ def cp_only(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw): batch = prepare_vlm_media_for_pp(batch, batch_size=2, n_microbatches=2) else: batch.update({"pixel_values": pv, "image_grid_thw": grid}) - cp_sharder = ContextParallelSharder(model_part0, mesh, batch) - train_ctx, batch = cp_sharder.shard(batch) labels = batch.pop("labels") - if pp_size > 1: - with train_ctx(), stage_vlm_media_for_pp(pp, pp.parts, batch): - sl = [] if has_last else None - mi = batch.pop("input_ids") - pp.update_seq_len(mi.shape[1]) - ( - pp.info.schedule.step(mi, target=labels, losses=sl, **batch) - if has_first - else pp.info.schedule.step(target=labels, losses=sl, **batch) - ) - local = torch.stack(sl).mean() if has_last else torch.tensor(0.0, device=device) - else: - with train_ctx(): - out = model( - input_ids=batch["input_ids"], - pixel_values=batch["pixel_values"], - image_grid_thw=batch["image_grid_thw"], - position_ids=batch["position_ids"], - ) - local = loss_fn(out, labels) - local.backward() + datum = Datum( + model_inputs=batch, + loss_fn_inputs={"labels": labels, "weights": torch.ones_like(labels, dtype=torch.float32)}, + ) + local, _ = engine.forward_backward([datum], loss_fn) losses.append(float(local.detach())) embed = model_part0.get_input_embeddings() diff --git a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py index 063bfdca01..1f24e015f7 100644 --- a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py +++ b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py @@ -14,11 +14,12 @@ """cp2xpp2 layer-2 verification for the sunk pre-embed path (L1, 2/4 GPUs). -Drives the REAL AutoPipeline split + schedule.step under cp2xpp2 (4 GPUs) and -cp2xpp1 (2 GPUs) with a tiny random-init text-only config, exercising the whole +Drives the real Datum Engine over an AutoPipeline under cp2xpp2 (4 GPUs) and +an eager model under cp2xpp1 (2 GPUs) with a tiny random-init text-only config, +exercising the whole sunk layer-2 contract: the sharder-only hook, the in-forward embed + shard_sequence_for_cp_round_robin, the asymmetric get_pipeline_stage_metas (full-length -first-stage ids, local sharded outputs), and per-microbatch backward. Asserts: +first-stage ids, local sharded outputs), and Engine-owned microbatch backward. Asserts: (1) 20 steps run clean -- no "backward through the graph a second time" (the double-backward the old shared pre-embed graph caused under PP*CP); @@ -53,17 +54,47 @@ def build_step3p7(device): layers = 4 cfg = Step3p7Config( - vision_config={"width": 8, "layers": 0, "heads": 2, "num_channels": 3, "image_size": 8, "patch_size": 2, - "mlp_ratio": 2.0, "hidden_act": "gelu", "use_ln_pre": False, "use_ln_post": False, - "use_abs_posemb": False, "use_rope2d": False}, - text_config={"hidden_size": 16, "intermediate_size": 32, "num_attention_heads": 4, "num_attention_groups": 2, - "num_hidden_layers": layers, "vocab_size": 32, "moe_num_experts": 2, "moe_top_k": 1, - "moe_intermediate_size": 8, "share_expert_dims": 8, "head_dim": 4, "torch_dtype": "bfloat16", - "moe_layers_enum": (), "layer_types": ["full_attention"] * layers, "num_nextn_predict_layers": 1}, + vision_config={ + "width": 8, + "layers": 0, + "heads": 2, + "num_channels": 3, + "image_size": 8, + "patch_size": 2, + "mlp_ratio": 2.0, + "hidden_act": "gelu", + "use_ln_pre": False, + "use_ln_post": False, + "use_abs_posemb": False, + "use_rope2d": False, + }, + text_config={ + "hidden_size": 16, + "intermediate_size": 32, + "num_attention_heads": 4, + "num_attention_groups": 2, + "num_hidden_layers": layers, + "vocab_size": 32, + "moe_num_experts": 2, + "moe_top_k": 1, + "moe_intermediate_size": 8, + "share_expert_dims": 8, + "head_dim": 4, + "torch_dtype": "bfloat16", + "moe_layers_enum": (), + "layer_types": ["full_attention"] * layers, + "num_nextn_predict_layers": 1, + }, image_token_id=31, ) - backend = BackendConfig(attn="sdpa", linear="torch", rms_norm="torch", dispatcher="torch", - rope_fusion=False, enable_hf_state_dict_adapter=False) + backend = BackendConfig( + attn="sdpa", + linear="torch", + rms_norm="torch", + dispatcher="torch", + rope_fusion=False, + enable_hf_state_dict_adapter=False, + ) model = Step3p7ForConditionalGeneration(cfg, backend=backend) model.initialize_weights(dtype=torch.bfloat16) return model.to(device).to(torch.bfloat16) @@ -75,23 +106,63 @@ def build_minimax(device): from nemo_automodel.components.models.minimax_m3_vl.model import MiniMaxM3SparseForConditionalGeneration tiny = dict( - hidden_size=64, intermediate_size=32, dense_intermediate_size=48, shared_intermediate_size=32, - num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, head_dim=16, rotary_dim=8, - partial_rotary_factor=0.5, vocab_size=128, max_position_embeddings=512, rms_norm_eps=1e-6, - rope_theta=10000.0, num_local_experts=4, num_experts_per_tok=2, n_shared_experts=1, - moe_layer_freq=[0, 1, 1, 1], use_gemma_norm=True, use_qk_norm=True, qk_norm_type="per_head", - scoring_func="sigmoid", use_routing_bias=True, routed_scaling_factor=2.0, swiglu_alpha=1.702, - swiglu_limit=7.0, num_mtp_modules=0, sparse_attention_config=dict(use_sparse_attention=False), + hidden_size=64, + intermediate_size=32, + dense_intermediate_size=48, + shared_intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + rotary_dim=8, + partial_rotary_factor=0.5, + vocab_size=128, + max_position_embeddings=512, + rms_norm_eps=1e-6, + rope_theta=10000.0, + num_local_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + moe_layer_freq=[0, 1, 1, 1], + use_gemma_norm=True, + use_qk_norm=True, + qk_norm_type="per_head", + scoring_func="sigmoid", + use_routing_bias=True, + routed_scaling_factor=2.0, + swiglu_alpha=1.702, + swiglu_limit=7.0, + num_mtp_modules=0, + sparse_attention_config=dict(use_sparse_attention=False), ) vision = dict( - hidden_size=32, num_attention_heads=4, num_hidden_layers=2, intermediate_size=64, patch_size=2, - num_channels=3, rope_theta=10000.0, hidden_act="gelu", layer_norm_eps=1e-5, + hidden_size=32, + num_attention_heads=4, + num_hidden_layers=2, + intermediate_size=64, + patch_size=2, + num_channels=3, + rope_theta=10000.0, + hidden_act="gelu", + layer_norm_eps=1e-5, img_token_compression_config={"spatial_merge_size": 2, "temporal_patch_size": 2}, ) - backend = BackendConfig(linear="torch", attn="sdpa", rms_norm="torch", rope_fusion=False, - dispatcher="torch", fake_balanced_gate=False, enable_hf_state_dict_adapter=False) - cfg = MiniMaxM3VLConfig(vision_config=dict(vision), text_config={**tiny, "torch_dtype": "bfloat16"}, - image_token_index=100, video_token_index=101, projector_hidden_size=tiny["hidden_size"]) + backend = BackendConfig( + linear="torch", + attn="sdpa", + rms_norm="torch", + rope_fusion=False, + dispatcher="torch", + fake_balanced_gate=False, + enable_hf_state_dict_adapter=False, + ) + cfg = MiniMaxM3VLConfig( + vision_config=dict(vision), + text_config={**tiny, "torch_dtype": "bfloat16"}, + image_token_index=100, + video_token_index=101, + projector_hidden_size=tiny["hidden_size"], + ) model = MiniMaxM3SparseForConditionalGeneration(cfg, backend=backend) model.initialize_weights(dtype=torch.bfloat16) return model.to(device).to(torch.bfloat16) @@ -103,16 +174,24 @@ def main(): torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") - from torch.distributed.device_mesh import init_device_mesh - - from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder + from nemo_automodel.components.datasets.datum import Datum + from nemo_automodel.components.distributed.config import FSDP2Config + from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.moe.parallelizer import apply_cp + from nemo_automodel.engine import Engine, collate_prebatched pp1 = bool(os.environ.get("NEMO_CP_PP_TEST_PP1")) # cp2xpp1 comparison leg pp_size = 1 if pp1 else 2 cp_size = world // pp_size - mesh = init_device_mesh("cuda", (pp_size, 1, cp_size), mesh_dim_names=("pp", "dp", "cp")) + mesh_context = MeshContext.build( + FSDP2Config(), + ParallelismSizes(dp_size=1, pp_size=pp_size, cp_size=cp_size), + world_size=world, + ) + mesh = mesh_context.device_mesh + if mesh is None: + raise RuntimeError("FSDP2 CP/PP validation requires a device mesh") which = os.environ.get("NEMO_CP_PP_MODEL", "minimax") torch.manual_seed(0) @@ -125,7 +204,20 @@ def main(): vocab = model.config.text_config.vocab_size mtp_used = {"any": False} - def loss_fn(output, labels): + def loss_fn(output, loss_inputs): + """Return token losses in the Engine's CP-local token layout. + + Args: + output: Model output whose logits have shape [batch, sequence, vocab]. + MTP models may return one additional logits tensor of the same + shape per prediction depth. + loss_inputs: Mapping containing labels and weights with shape + [batch, sequence] in the same CP-local layout as the logits. + + Returns: + Tensor of shape [batch, sequence] containing the summed causal-LM + loss over the base and optional MTP prediction depths. + """ # step3p7's last PP stage emits (logits, *mtp_per_depth_logits); minimax # emits a bare logits tensor. Handle both, threading the MTP depths. if isinstance(output, tuple): @@ -133,10 +225,16 @@ def loss_fn(output, labels): else: logits = getattr(output, "logits", output) mtp = list(getattr(output, "mtp_per_depth_logits", None) or []) - loss = F.cross_entropy(logits.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100) + labels = loss_inputs["labels"] + weights = loss_inputs["weights"] + loss = F.cross_entropy( + logits.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100, reduction="none" + ).reshape_as(weights) for m in mtp: mtp_used["any"] = True - loss = loss + F.cross_entropy(m.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100) + loss = loss + F.cross_entropy( + m.reshape(-1, vocab).float(), labels.reshape(-1), ignore_index=-100, reduction="none" + ).reshape_as(weights) return loss def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw): @@ -146,14 +244,29 @@ def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name= seqlen = 32 if pp_size > 1: pp = AutoPipeline( - world_mesh=mesh, moe_mesh=None, pp_axis_name="pp", dp_axis_names=("dp",), cp_axis_name="cp", - pp_schedule="1f1b", pp_microbatch_size=1, pp_batch_size=2, device=device, dtype=torch.bfloat16, + world_mesh=mesh, + moe_mesh=None, + **mesh_context.pipeline_axis_kwargs(), + pp_schedule="1f1b", + pp_microbatch_size=1, + pp_batch_size=2, + device=device, + dtype=torch.bfloat16, pp_seq_len=seqlen, ).build(model, loss_fn=loss_fn, parallelize_fn=cp_only_parallelize) - model_part0, has_last, has_first = pp.parts[0], pp.info.has_last_stage, pp.info.has_first_stage + model_part0 = pp.parts[0] + engine_model = pp else: - cp_only_parallelize(model, mesh, None, dp_axis_names=("dp",), cp_axis_name="cp") + cp_only_parallelize(model, mesh, None, **mesh_context.parallelize_axis_kwargs()) model_part0 = model + engine_model = model + + engine = Engine( + engine_model, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ) losses = [] for step in range(20): @@ -161,27 +274,12 @@ def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name= input_ids = torch.randint(2, vocab, (2, seqlen), device=device) dist.broadcast(input_ids, src=0) pos = torch.arange(seqlen, device=device).unsqueeze(0).expand(2, -1).contiguous() - batch = {"input_ids": input_ids.clone(), "labels": input_ids.clone(), "position_ids": pos.clone()} - cp_sharder = ContextParallelSharder(model_part0, mesh, batch) - train_ctx, batch = cp_sharder.shard(batch) - labels = batch.pop("labels") - if pp_size > 1: - with train_ctx(): - step_losses = [] if has_last else None - model_input = batch.pop("input_ids") - pp.update_seq_len(model_input.shape[1]) - if has_first: - pp.info.schedule.step(model_input, target=labels, losses=step_losses, **batch) - else: - pp.info.schedule.step(target=labels, losses=step_losses, **batch) - # Per-microbatch loss_fn returns the microbatch mean; averaging over - # microbatches gives the batch mean, comparable to the cp2xpp1 leg. - local = torch.stack(step_losses).mean() if has_last else torch.tensor(0.0, device=device) - else: - with train_ctx(): - out = model(input_ids=batch["input_ids"], position_ids=batch["position_ids"]) - local = loss_fn(out, labels) # full output so MTP depths are included - local.backward() + labels = input_ids.clone() + datum = Datum( + model_inputs={"input_ids": input_ids.clone(), "position_ids": pos.clone()}, + loss_fn_inputs={"labels": labels, "weights": torch.ones_like(labels, dtype=torch.float32)}, + ) + local, _ = engine.forward_backward([datum], loss_fn) losses.append(float(local.detach())) # (3) embeddings receive gradients -- only the first PP stage owns embed_tokens. @@ -196,7 +294,7 @@ def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name= last = torch.tensor(losses[-1], device=device) gflag = torch.tensor(1.0 if embed_grad else 0.0, device=device) mflag = torch.tensor(1.0 if mtp_used["any"] else 0.0, device=device) - dist.all_reduce(last, op=dist.ReduceOp.MAX) # last stage holds the real loss + dist.all_reduce(last, op=dist.ReduceOp.MAX) # Engine synchronizes loss; MAX is a cross-rank sanity check. dist.all_reduce(gflag, op=dist.ReduceOp.MAX) # first stage owns embeddings dist.all_reduce(mflag, op=dist.ReduceOp.MAX) # last stage owns the MTP heads diff --git a/tests/functional_tests/context_parallel/run_dense_packed_cp.py b/tests/functional_tests/context_parallel/run_dense_packed_cp.py index b6fe9fe5dd..d590574b43 100644 --- a/tests/functional_tests/context_parallel/run_dense_packed_cp.py +++ b/tests/functional_tests/context_parallel/run_dense_packed_cp.py @@ -199,10 +199,10 @@ def _run_model( observed: dict[str, torch.Tensor] = {} - def loss_fn(output, loss_inputs, _datums, model_inputs): + def loss_fn(output, loss_inputs): local_logits = _full_tensor(output.logits).squeeze(0) observed["logits"] = local_logits.detach() - observed["cu_seqlens"] = model_inputs["cu_seqlens"].detach() + observed["cu_seqlens"] = loss_inputs["cu_seqlens"].detach() observed["labels"] = loss_inputs["labels"].detach() return local_logits.float().square().mean(dim=-1) @@ -212,7 +212,7 @@ def loss_fn(output, loss_inputs, _datums, model_inputs): mesh_context=mesh_context, collate_fn=collate_prebatched, defer_fsdp_grad_sync=distributed_config.defer_fsdp_grad_sync, - ).forward_backward([[datum]], loss_fn) + ).forward_backward([datum], loss_fn) local_logits = observed["logits"] assert observed["labels"].shape == (8 // cp_mesh.size(),) diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py new file mode 100644 index 0000000000..8a3e6fe6ee --- /dev/null +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -0,0 +1,318 @@ +# 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. + +"""Two-GPU Engine/AutoPipeline parity for packed THD Llama. + +The test runs the same two-microbatch, four-document update through PP=2 from +both raw THD metadata (``seq_lens``) and final THD metadata (``cu_seqlens``). +Each path must match a native eager Llama in loss and every local stage +gradient. A final padded two-Datum update verifies that Engine broadcasts the +callback's per-Datum mappings to both pipeline ranks in logical input order. + +Run with:: + + torchrun --standalone --nproc-per-node=2 run_packed_pp.py +""" + +from __future__ import annotations + +import os +import warnings + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from transformers import LlamaConfig + +from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.distributed.config import FSDP2Config +from nemo_automodel.components.distributed.context_parallel.utils import make_cp_batch_for_te +from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes +from nemo_automodel.components.distributed.pipelining import AutoPipeline +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.llama.model import LlamaForCausalLM +from nemo_automodel.engine import Engine, collate_prebatched + +VOCAB_SIZE = 64 +SEQ_LEN = 8 + + +def _build_model(device: torch.device) -> LlamaForCausalLM: + """Build the deterministic two-layer model used by eager and PP paths.""" + torch.manual_seed(1234) + config = LlamaConfig( + vocab_size=VOCAB_SIZE, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + attention_dropout=0.0, + tie_word_embeddings=False, + use_cache=False, + torch_dtype=torch.bfloat16, + ) + config._attn_implementation = "sdpa" + backend = BackendConfig(attn="te", linear="torch", rms_norm="torch_fp32", rope_fusion=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + model = LlamaForCausalLM(config, backend=backend) + return model.to(device=device, dtype=torch.bfloat16).train() + + +def _raw_batch(device: torch.device) -> tuple[dict[str, object], torch.Tensor, torch.Tensor]: + """Return two packed rows containing two documents each. + + Token tensors use ``[B=2, S=8]``. The reset positions and lengths describe + documents of lengths ``[3, 5]`` and ``[4, 4]``. + """ + input_ids = torch.tensor( + [[1, 2, 3, 11, 12, 13, 14, 15], [21, 22, 23, 24, 31, 32, 33, 34]], + device=device, + ) + labels = torch.tensor( + [[2, 3, -100, 12, 13, 14, 15, -100], [22, 23, 24, -100, 32, 33, 34, -100]], + device=device, + ) + weights = torch.tensor( + [[1.0, 2.0, 0.0, 1.5, 0.5, 2.5, 1.0, 0.0], [0.5, 1.0, 2.0, 0.0, 1.0, 3.0, 1.5, 0.0]], + device=device, + ) + model_inputs: dict[str, object] = { + "input_ids": input_ids, + "position_ids": torch.tensor( + [[0, 1, 2, 0, 1, 2, 3, 4], [0, 1, 2, 3, 0, 1, 2, 3]], + device=device, + ), + "seq_lens": torch.tensor([[3, 5], [4, 4]], device=device), + "seq_lens_padded": torch.tensor([[3, 5], [4, 4]], device=device), + "qkv_format": "thd", + } + return model_inputs, labels, weights + + +def _clone_mapping(values: dict[str, object]) -> dict[str, object]: + return {name: value.clone() if isinstance(value, torch.Tensor) else value for name, value in values.items()} + + +def _token_losses(output, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + logits = getattr(output, "logits", output) + return F.cross_entropy( + logits.float().reshape(-1, VOCAB_SIZE), + loss_inputs["labels"].reshape(-1), + ignore_index=-100, + reduction="none", + ).reshape_as(loss_inputs["weights"]) + + +def _eager_reference( + device: torch.device, +) -> tuple[torch.Tensor, dict[str, torch.Tensor], dict[str, object], torch.Tensor, torch.Tensor]: + raw_inputs, labels, weights = _raw_batch(device) + prepared = make_cp_batch_for_te( + None, + {**_clone_mapping(raw_inputs), "labels": labels.clone()}, + ) + prepared_labels = prepared.pop("labels") + + model = _build_model(device) + output = model(**prepared) + token_losses = _token_losses(output, {"labels": prepared_labels, "weights": weights.reshape(-1)}) + loss = (token_losses * weights.reshape(-1)).sum() / weights.sum() + loss.backward() + grads = {name: parameter.grad.detach().clone() for name, parameter in model.named_parameters()} + del model + return loss.detach(), grads, raw_inputs, labels, weights + + +def _build_pipeline( + device: torch.device, + mesh_context: MeshContext, +) -> AutoPipeline: + model = _build_model(device) + pipeline = AutoPipeline( + world_mesh=mesh_context.device_mesh, + moe_mesh=None, + **mesh_context.pipeline_axis_kwargs(), + pp_schedule="1f1b", + pp_microbatch_size=1, + pp_batch_size=2, + device=device, + dtype=torch.bfloat16, + pp_seq_len=SEQ_LEN, + ).build(model, loss_fn=_token_losses) + del model + return pipeline + + +def _assert_local_grad_parity( + pipeline: AutoPipeline, + reference_grads: dict[str, torch.Tensor], +) -> float: + max_diff = 0.0 + checked = 0 + for part in pipeline.parts: + for name, parameter in part.named_parameters(): + if parameter.grad is None: + raise AssertionError(f"pipeline parameter {name} has no gradient") + torch.testing.assert_close( + parameter.grad.float(), + reference_grads[name].float(), + atol=5e-3, + rtol=5e-2, + ) + max_diff = max(max_diff, (parameter.grad.float() - reference_grads[name].float()).abs().max().item()) + checked += 1 + if checked == 0: + raise AssertionError("pipeline rank owns no checked parameters") + return max_diff + + +def _run_thd_layout( + layout: str, + device: torch.device, + mesh_context: MeshContext, + reference_loss: torch.Tensor, + reference_grads: dict[str, torch.Tensor], + raw_inputs: dict[str, object], + labels: torch.Tensor, + weights: torch.Tensor, +) -> AutoPipeline: + pipeline = _build_pipeline(device, mesh_context) + if layout == "raw": + model_inputs = _clone_mapping(raw_inputs) + loss_labels = labels.clone() + loss_weights = weights.clone() + elif layout == "final": + model_inputs = make_cp_batch_for_te( + None, + {**_clone_mapping(raw_inputs), "labels": labels.clone()}, + ) + loss_labels = model_inputs.pop("labels") + loss_weights = weights.reshape(-1).clone() + else: + raise ValueError(f"unknown THD layout: {layout}") + + datum = Datum( + model_inputs=model_inputs, + loss_fn_inputs={"labels": loss_labels, "weights": loss_weights}, + ) + loss, outputs = Engine( + pipeline, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ).forward_backward([datum], _token_losses) + + torch.testing.assert_close(loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) + assert outputs == [] + grad_diff = _assert_local_grad_parity(pipeline, reference_grads) + if dist.get_rank() == 0: + print( + f"PP2 {layout} THD parity passed " + f"(loss diff={(loss.float() - reference_loss.float()).abs().item():.6f}, grad max={grad_diff:.6f})" + ) + return pipeline + + +def _run_padded_output_broadcast(pipeline: AutoPipeline, device: torch.device, mesh_context: MeshContext) -> None: + for part in pipeline.parts: + part.zero_grad(set_to_none=True) + + datums = [] + for sample_id, offset in ((17, 0), (29, 8)): + input_ids = torch.arange(1 + offset, 1 + offset + SEQ_LEN, device=device) % VOCAB_SIZE + labels = input_ids.roll(-1) + labels[-1] = -100 + datums.append( + Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs={ + "labels": labels, + "weights": (labels != -100).to(torch.float32), + "sample_id": torch.tensor(sample_id, device=device), + }, + ) + ) + + def loss_with_output(output, loss_inputs): + logits = getattr(output, "logits", output) + losses = _token_losses(output, loss_inputs) + return losses, [{"sample_id": loss_inputs["sample_id"][0], "score": logits.float().mean()}] + + loss, outputs = Engine( + pipeline, + device=device, + mesh_context=mesh_context, + microbatch_size=2, + ).forward_backward(datums, loss_with_output) + + expected_ids = torch.tensor([17, 29], device=device) + output_ids = torch.stack([item["sample_id"] for item in outputs]).to(device=device, dtype=torch.long) + torch.testing.assert_close(output_ids, expected_ids) + gathered = [torch.empty_like(output_ids) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, output_ids) + assert all(torch.equal(ids, expected_ids) for ids in gathered) + assert torch.isfinite(loss) + assert all(torch.isfinite(item["score"]) for item in outputs) + if dist.get_rank() == 0: + print("PP2 padded per-Datum outputs passed (logical order [17, 29] synchronized on both ranks)") + + +def main() -> None: + dist.init_process_group("nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if dist.get_world_size() != 2: + raise ValueError("packed PP functional requires exactly two ranks") + + mesh_context = MeshContext.build( + FSDP2Config(), + ParallelismSizes(dp_size=1, pp_size=2), + world_size=dist.get_world_size(), + ) + try: + reference_loss, reference_grads, raw_inputs, labels, weights = _eager_reference(device) + _run_thd_layout( + "raw", + device, + mesh_context, + reference_loss, + reference_grads, + raw_inputs, + labels, + weights, + ) + dist.barrier() + final_pipeline = _run_thd_layout( + "final", + device, + mesh_context, + reference_loss, + reference_grads, + raw_inputs, + labels, + weights, + ) + dist.barrier() + _run_padded_output_broadcast(final_pipeline, device, mesh_context) + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/functional_tests/context_parallel/test_context_parallel.py b/tests/functional_tests/context_parallel/test_context_parallel.py index 1a406aa6a4..99d4e286af 100644 --- a/tests/functional_tests/context_parallel/test_context_parallel.py +++ b/tests/functional_tests/context_parallel/test_context_parallel.py @@ -32,6 +32,7 @@ CP_QWEN3_5_MOE_LINEAR_ATTN_TEST_FILENAME = "L2_CP_Qwen3_5MoE_LinearAttn_Test.sh" CP_DENSE_PACKED_TEST_FILENAME = "L2_CP_Dense_Packed_Test.sh" TP_DENSE_PACKED_TEST_FILENAME = "L2_TP_Dense_Packed_Test.sh" +PP_DENSE_PACKED_TEST_FILENAME = "L2_PP_Dense_Packed_Test.sh" TP_CP_DENSE_PACKED_TEST_FILENAME = "L2_TP_CP_Dense_Packed_Test.sh" TP_CP_DENSE_PACKED_REQUIRED_GPUS = 4 @@ -71,6 +72,10 @@ def test_tp_dense_packed(self): """Test packed THD parity for dense Llama, Qwen2, and Qwen3 with TP=2 and CP=1.""" run_test_script(TEST_FOLDER, TP_DENSE_PACKED_TEST_FILENAME) + def test_pp_dense_packed(self): + """Test Engine PP=2 raw/final THD loss and gradient parity for dense Llama.""" + run_test_script(TEST_FOLDER, PP_DENSE_PACKED_TEST_FILENAME) + @pytest.mark.skipif( torch.cuda.device_count() < TP_CP_DENSE_PACKED_REQUIRED_GPUS, reason="requires 4 GPUs for TP=2 x CP=2; remove once context_parallel CI runs on a 4-GPU runner", diff --git a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py index f41c5c9f19..7dc5a3cb6b 100644 --- a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py +++ b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py @@ -12,16 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Scheduled two-rank CPU parity test for expert gradients under composed TP x EP. +"""Scheduled two-rank CPU parity test for the Engine and finalizer under TP x EP. The custom-MoE tensor-parallel path keeps the token path (attention, router) replicated across TP ranks, so every TP rank feeds the same tokens into the expert-parallel all-gather and each expert gradient accumulates ``tp_size`` -identical contributions. This test drives the real ``GroupedExperts`` -forward/backward with tp=2 replicated tokens through a 2-rank EP mesh and -asserts that ``scale_grads_and_clip_grad_norm`` with the factor returned by -``get_expert_tp_replication_factor`` restores the single-process fp32 -reference gradients. +identical contributions. This test drives the real ``GroupedExperts`` through +the Datum Engine with tp=2 replicated tokens and a 2-rank EP mesh, then asserts +that the existing gradient finalizer restores the single-process fp32 loss, +gradients, global gradient norm, and one-step parameter update. """ from __future__ import annotations @@ -37,12 +36,15 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import Shard, distribute_tensor +from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.experts import GroupedExperts from nemo_automodel.components.training.utils import ( get_expert_tp_replication_factor, scale_grads_and_clip_grad_norm, ) +from nemo_automodel.engine import Engine, collate_prebatched _TP_SIZE = 2 _WORLD_SIZE = _TP_SIZE @@ -51,6 +53,35 @@ _DIM = 16 _MOE_INTER_DIM = 32 _NUM_TOKENS = 6 +_LEARNING_RATE = 0.05 + + +class _ExpertModel(nn.Module): + """Expose ``GroupedExperts`` through the Engine's primary-input convention.""" + + def __init__(self, experts: GroupedExperts) -> None: + super().__init__() + self.experts = experts + + def forward( + self, + input_ids: torch.Tensor, + token_mask: torch.Tensor, + router_weights: torch.Tensor, + router_indices: torch.Tensor, + ) -> torch.Tensor: + """Run the expert layer. + + Args: + input_ids: Tensor of shape [tokens, hidden] containing expert inputs. + token_mask: Boolean tensor of shape [tokens] selecting active tokens. + router_weights: Tensor of shape [tokens, top_k] containing route weights. + router_indices: Integer tensor of shape [tokens, top_k] containing expert IDs. + + Returns: + Tensor of shape [tokens, hidden] containing the combined expert outputs. + """ + return self.experts(input_ids, token_mask, router_weights, router_indices) def _free_port() -> int: @@ -93,6 +124,11 @@ def _global_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Te return x, weights, indices, token_mask +def _loss_weights() -> torch.Tensor: + """Return deterministic nonuniform positive weights of shape [tokens, hidden].""" + return torch.linspace(0.25, 1.25, steps=_NUM_TOKENS * _DIM).reshape(_NUM_TOKENS, _DIM) + + def _build_experts(config: MoEConfig) -> GroupedExperts: generator = torch.Generator().manual_seed(4321) experts = GroupedExperts(config) @@ -102,15 +138,17 @@ def _build_experts(config: MoEConfig) -> GroupedExperts: return experts -def _reference_forward_backward() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _reference_forward_backward() -> tuple[GroupedExperts, torch.Tensor, torch.Tensor]: """Single-process (tp=1, ep=1) fp32 forward/backward as ground truth.""" experts = _build_experts(_tiny_moe_config()) x, weights, indices, token_mask = _global_inputs() y = experts(x, token_mask, weights, indices) - y.sum().backward() + loss_weights = _loss_weights() + loss = (y.square() * loss_weights).sum() / loss_weights.sum() + loss.backward() assert experts.gate_and_up_projs.grad is not None assert experts.down_projs.grad is not None - return y.detach(), experts.gate_and_up_projs.grad.detach(), experts.down_projs.grad.detach() + return experts, y.detach(), loss.detach() def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: @@ -121,7 +159,12 @@ def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: os.environ["WORLD_SIZE"] = str(world_size) dist.init_process_group("gloo", rank=rank, world_size=world_size) - y_ref, gate_up_grad_ref, down_grad_ref = _reference_forward_backward() + reference_experts, y_ref, loss_ref = _reference_forward_backward() + gate_up_grad_ref = reference_experts.gate_and_up_projs.grad.detach().clone() + down_grad_ref = reference_experts.down_projs.grad.detach().clone() + reference_grad_norm = torch.linalg.vector_norm( + torch.cat((gate_up_grad_ref.reshape(-1), down_grad_ref.reshape(-1))).to(torch.float64) + ) # Composed TP x EP topology on 2 ranks: the same two ranks form the TP # replica group of the token path and the EP group of the experts, as @@ -138,13 +181,42 @@ def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: # TP path is active; get_expert_tp_replication_factor keys off it. experts._nemo_moe_tp_requires_replica_sync = True - # TP-replicated token path: every rank feeds the identical full batch - # into the EP all-gather. + # The Engine sees the TP-replicated token path and the EP experts as one + # ordinary forward/backward boundary. x, weights, indices, token_mask = _global_inputs() - y_local = experts(x, token_mask, weights, indices) - torch.testing.assert_close(y_local, y_ref, rtol=1e-4, atol=1e-5) + loss_weights = _loss_weights() + datum = Datum( + model_inputs={ + "input_ids": x, + "token_mask": token_mask, + "router_weights": weights, + "router_indices": indices, + }, + loss_fn_inputs={"weights": loss_weights}, + ) + observed: dict[str, torch.Tensor] = {} + + def loss_fn(output: torch.Tensor, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Return squared expert outputs in the Engine loss-weight layout. + + Args: + output: Tensor of shape [tokens, hidden] containing expert outputs. + loss_inputs: Mapping whose ``weights`` tensor has shape [tokens, hidden]. - y_local.sum().backward() + Returns: + Tensor of shape [tokens, hidden] containing per-element squared loss. + """ + observed["output"] = output.detach() + return output.square() + + engine_loss, _ = Engine( + _ExpertModel(experts), + device="cpu", + mesh_context=MeshContext.from_meshes(world_mesh, ep_mesh), + collate_fn=collate_prebatched, + ).forward_backward([datum], loss_fn) + torch.testing.assert_close(observed["output"], y_ref, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(engine_loss, loss_ref.to(torch.float64), rtol=1e-5, atol=1e-7) n_local_experts = _N_EXPERTS // world_size start = rank * n_local_experts @@ -168,8 +240,8 @@ def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: # make the TP replication factor the only expert divisor. replication_factor = get_expert_tp_replication_factor([experts], world_mesh) assert replication_factor == _TP_SIZE - scale_grads_and_clip_grad_norm( - max_grad_norm=None, + grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=1e6, model_parts=[experts], moe_mesh=ep_mesh, ep_axis_name="ep", @@ -180,11 +252,27 @@ def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: experts.gate_and_up_projs.grad.to_local(), gate_up_grad_ref_local, rtol=1e-4, atol=1e-5 ) torch.testing.assert_close(experts.down_projs.grad.to_local(), down_grad_ref_local, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(grad_norm, reference_grad_norm, rtol=1e-5, atol=1e-7) + + torch.optim.SGD(reference_experts.parameters(), lr=_LEARNING_RATE).step() + torch.optim.SGD(experts.parameters(), lr=_LEARNING_RATE).step() + torch.testing.assert_close( + experts.gate_and_up_projs.to_local(), + reference_experts.gate_and_up_projs[start:end], + rtol=1e-4, + atol=1e-5, + ) + torch.testing.assert_close( + experts.down_projs.to_local(), + reference_experts.down_projs[start:end], + rtol=1e-4, + atol=1e-5, + ) finally: if dist.is_initialized(): dist.destroy_process_group() @pytest.mark.skipif(not dist.is_available(), reason="torch.distributed is not available") -def test_tp_replicated_tokens_through_ep_match_reference_after_replication_scaling(): +def test_engine_tp_replicated_tokens_through_ep_matches_reference_after_finalization(): mp.spawn(_ep_tp_grad_parity_worker, args=(_WORLD_SIZE, _free_port()), nprocs=_WORLD_SIZE, join=True) diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 214ff2299b..5e4f168edb 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -244,13 +244,24 @@ def __init__(self, *, fail_on_step: bool = False, invoke_loss: bool = False): self.fail_on_step = fail_on_step self.invoke_loss = invoke_loss self.args_during_step = None + self.args_split = None self.kwargs_chunk_spec_during_step = None self.loss_fn_during_step = None + self.split_inputs_during_step = None self.kwargs_split = None self.target_during_step = None + self.losses_during_step = None self.return_outputs_during_step = None self.loss_results = [] + def _split_inputs(self, args, kwargs=None): + return split_args_kwargs_into_chunks( + args, + kwargs, + 2, + kwargs_chunk_spec=self._kwargs_chunk_spec, + ) + def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): """Split schedule inputs using the chunk spec active during the call. @@ -258,7 +269,7 @@ def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): *args: Positional schedule inputs. Tensor values have arbitrary model-defined layouts. target: Optional tensor of shape [batch, ...] containing loss targets - or structured-loss microbatch IDs. + or prepared-input microbatch IDs. losses: Optional mutable list populated with scalar loss tensors. return_outputs: Whether to return the schedule result. **kwargs: Keyword schedule inputs. Tensor values have arbitrary @@ -267,20 +278,16 @@ def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): Returns: A sentinel string identifying the schedule result. """ - del losses self.args_during_step = args self.kwargs_chunk_spec_during_step = self._kwargs_chunk_spec self.loss_fn_during_step = self._loss_fn + self.split_inputs_during_step = self._split_inputs self.target_during_step = target + self.losses_during_step = losses self.return_outputs_during_step = return_outputs if self.fail_on_step: raise RuntimeError("schedule failed") - _, self.kwargs_split = split_args_kwargs_into_chunks( - args, - kwargs, - 2, - kwargs_chunk_spec=self._kwargs_chunk_spec, - ) + self.args_split, self.kwargs_split = self._split_inputs(args, kwargs) if self.invoke_loss: assert target is not None target_chunks = torch.tensor_split(target, 2) @@ -356,20 +363,109 @@ def test_step_without_model_hook_uses_pytorch_default_chunking(self): assert ap.info.schedule._loss_fn is original_loss_fn assert ap.info.schedule.return_outputs_during_step is True - @pytest.mark.parametrize("structured", [False, True]) - def test_step_does_not_forward_return_outputs_to_older_pytorch(self, structured): + def test_step_does_not_forward_return_outputs_to_older_pytorch(self): schedule = _LegacyStepSchedule() ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - structured_kwargs = ( - {"loss_inputs": {"weights": torch.ones(2, 8)}, "loss_fn": lambda *_: torch.tensor(0.0)} - if structured - else {} - ) - ap.step(torch.zeros(2, 8), return_outputs=False, **structured_kwargs) + ap.step(torch.zeros(2, 8), return_outputs=False) + + assert schedule.received_return_outputs is False + + def test_step_microbatches_passes_prepared_inputs_without_resplitting(self): + schedule = _KwargsChunkSchedule(invoke_loss=True) + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + input_ids = [torch.full((1, 8), index, dtype=torch.long) for index in range(2)] + position_ids = [torch.full((3, 1, 8), index, dtype=torch.long) for index in range(2)] + metadata = [object(), object()] + model_inputs = [ + { + "input_ids": input_ids[index], + "position_ids": position_ids[index], + "metadata": metadata[index], + } + for index in range(2) + ] + original_split_inputs = schedule._split_inputs + original_loss_fn = schedule._loss_fn + seen = [] + losses = [] + + def loss_fn(output, index): + seen.append((output, index)) + return output + + result = ap.step_microbatches(model_inputs, loss_fn=loss_fn, losses=losses, return_outputs=False) + + assert result == "schedule-result" + assert schedule.args_during_step == () + assert schedule.args_split[0][0] is input_ids[0] + assert schedule.args_split[1][0] is input_ids[1] + assert schedule.kwargs_split[0]["position_ids"] is position_ids[0] + assert schedule.kwargs_split[1]["position_ids"] is position_ids[1] + assert schedule.kwargs_split[0]["metadata"] is metadata[0] + assert schedule.kwargs_split[1]["metadata"] is metadata[1] + assert schedule.target_during_step.tolist() == [0, 1] + assert schedule.losses_during_step is losses + assert schedule.return_outputs_during_step is False + assert [(output.item(), index) for output, index in seen] == [(1.0, 1), (0.0, 0)] + assert schedule._split_inputs == original_split_inputs + assert schedule._loss_fn is original_loss_fn + assert model_inputs[0]["input_ids"] is input_ids[0] + assert model_inputs[1]["input_ids"] is input_ids[1] + + def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): + schedule = _KwargsChunkSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule, has_first_stage=False) + model_inputs = [{"inputs_embeds": torch.zeros(1, 8, 4), "position_ids": torch.zeros(1, 8)} for _ in range(2)] + + ap.step_microbatches(model_inputs, loss_fn=Mock()) + + assert schedule.args_split == [(), ()] + assert all("inputs_embeds" not in kwargs for kwargs in schedule.kwargs_split) + + def test_step_microbatches_does_not_forward_return_outputs_to_older_pytorch(self): + schedule = _LegacyStepSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + return_outputs=False, + ) assert schedule.received_return_outputs is False + @pytest.mark.parametrize( + "model_inputs", + [ + [{"input_ids": torch.zeros(1, 8)}], + [{"attention_mask": torch.ones(1, 8)} for _ in range(2)], + [{"input_ids": torch.zeros(1, 8), "inputs_embeds": torch.zeros(1, 8, 4)} for _ in range(2)], + ], + ) + def test_step_microbatches_validates_prepared_inputs(self, model_inputs): + ap = self._pipeline_with_parts(nn.Module()) + + with pytest.raises(ValueError, match="Expected 2|exactly one"): + ap.step_microbatches(model_inputs, loss_fn=Mock()) + + def test_step_microbatches_restores_schedule_state_after_failure(self): + schedule = _KwargsChunkSchedule(fail_on_step=True) + original_split_inputs = schedule._split_inputs + original_loss_fn = schedule._loss_fn + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + with pytest.raises(RuntimeError, match="schedule failed"): + ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + ) + + assert schedule.split_inputs_during_step is not original_split_inputs + assert schedule.loss_fn_during_step is not original_loss_fn + assert schedule._split_inputs == original_split_inputs + assert schedule._loss_fn is original_loss_fn + def test_only_canonical_model_part_supplies_chunk_policy(self): ap = self._pipeline_with_parts( _KwargsChunkHookPart({"position_ids": 1}), @@ -406,137 +502,6 @@ def test_model_hook_cannot_configure_unknown_kwarg(self): with pytest.raises(ValueError, match="unknown kwarg"): ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) - def test_structured_loss_receives_aligned_microbatches(self): - schedule = _KwargsChunkSchedule(invoke_loss=True) - ap = self._pipeline_with_parts( - _KwargsChunkHookPart({"position_ids": 1}), - schedule=schedule, - has_first_stage=False, - ) - input_ids = torch.arange(16).view(2, 8) - position_ids = torch.arange(48).view(3, 2, 8) - weights = torch.arange(16).view(2, 8) - scale = torch.tensor(0.5) - seen = [] - - def loss_fn(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): - """Record one structured-loss callback. - - Args: - output: Scalar tensor produced by the fake last stage. - loss_inputs_mb: Mapping containing ``weights`` with shape - [microbatch, sequence] and scalar tensor ``scale``. - model_args_mb: Tuple containing input IDs with shape - [microbatch, sequence]. - model_kwargs_mb: Mapping containing position IDs with shape - [axes, microbatch, sequence]. - - Returns: - Scalar loss tensor. - """ - seen.append((output, loss_inputs_mb, model_args_mb, model_kwargs_mb)) - return output - - result = ap.step( - input_ids, - loss_inputs={"weights": weights, "scale": scale, "label": "train"}, - loss_fn=loss_fn, - return_outputs=False, - position_ids=position_ids, - ) - - assert result == "schedule-result" - assert schedule.args_during_step == () - torch.testing.assert_close(schedule.target_during_step, torch.arange(2)) - assert schedule.target_during_step.device == ap.device - assert schedule.return_outputs_during_step is False - assert [item[0].item() for item in seen] == [1.0, 0.0] - for index, (_, loss_inputs_mb, model_args_mb, model_kwargs_mb) in zip((1, 0), seen): - torch.testing.assert_close(loss_inputs_mb["weights"], weights[index : index + 1]) - assert loss_inputs_mb["scale"] is scale - assert loss_inputs_mb["label"] == "train" - torch.testing.assert_close(model_args_mb[0], input_ids[index : index + 1]) - torch.testing.assert_close(model_kwargs_mb["position_ids"], position_ids[:, index : index + 1]) - - @pytest.mark.parametrize( - ("loss_inputs", "loss_fn"), - [ - ({"weights": torch.ones(2, 8)}, None), - (None, Mock()), - ], - ) - def test_structured_loss_requires_inputs_and_callback(self, loss_inputs, loss_fn): - ap = self._pipeline_with_parts(nn.Module()) - - with pytest.raises(ValueError, match="must be provided together"): - ap.step(torch.zeros(2, 8), loss_inputs=loss_inputs, loss_fn=loss_fn) - - def test_structured_loss_rejects_legacy_target(self): - ap = self._pipeline_with_parts(nn.Module()) - - with pytest.raises(ValueError, match="target cannot be used"): - ap.step( - torch.zeros(2, 8), - target=torch.zeros(2, 8), - loss_inputs={"weights": torch.ones(2, 8)}, - loss_fn=Mock(), - ) - - def test_structured_loss_requires_exact_microbatch_count(self): - ap = self._pipeline_with_parts(nn.Module()) - - with pytest.raises(ValueError, match="Expected 2 model input microbatches, got 1"): - ap.step( - torch.zeros(1, 8), - loss_inputs={"weights": torch.ones(2, 8)}, - loss_fn=Mock(), - ) - - with pytest.raises(ValueError, match="Expected 2 loss input microbatches, got 1"): - ap.step( - torch.zeros(2, 8), - loss_inputs={"weights": torch.ones(1, 8)}, - loss_fn=Mock(), - ) - - def test_structured_loss_restores_schedule_state_after_callback_failure(self): - schedule = _KwargsChunkSchedule(invoke_loss=True) - original_loss_fn = schedule._loss_fn - original_chunk_spec = {"position_ids": TensorChunkSpec(0)} - schedule._kwargs_chunk_spec = original_chunk_spec - ap = self._pipeline_with_parts(_KwargsChunkHookPart({"position_ids": 1}), schedule=schedule) - - def failing_loss(output, loss_inputs_mb, model_args_mb, model_kwargs_mb): - """Raise while evaluating a structured microbatch loss. - - Args: - output: Scalar tensor produced by the fake last stage. - loss_inputs_mb: Mapping containing weights with shape - [microbatch, sequence]. - model_args_mb: Tuple containing input IDs with shape - [microbatch, sequence]. - model_kwargs_mb: Mapping containing position IDs with shape - [axes, microbatch, sequence]. - - Raises: - RuntimeError: Always. - """ - del output, loss_inputs_mb, model_args_mb, model_kwargs_mb - raise RuntimeError("loss failed") - - with pytest.raises(RuntimeError, match="loss failed"): - ap.step( - torch.zeros(2, 8), - loss_inputs={"weights": torch.ones(2, 8)}, - loss_fn=failing_loss, - position_ids=torch.zeros(3, 2, 8), - ) - - assert schedule.loss_fn_during_step is not original_loss_fn - assert schedule.kwargs_chunk_spec_during_step["position_ids"].split_dim == 1 - assert schedule._loss_fn is original_loss_fn - assert schedule._kwargs_chunk_spec is original_chunk_spec - # ----------------------------- # Core build/materialize/step tests @@ -1201,18 +1166,43 @@ def test_update_seq_len_calls_reset(self, monkeypatch): captured_args = [] - def mock_reset(schedule, stages, model_config, microbatch_size, seq_len, tensor_dtype=None): - captured_args.append((schedule, stages, model_config, microbatch_size, seq_len, tensor_dtype)) + def mock_reset( + schedule, + stages, + model_config, + microbatch_size, + seq_len, + tensor_dtype=None, + first_stage_input_meta=None, + ): + captured_args.append( + (schedule, stages, model_config, microbatch_size, seq_len, tensor_dtype, first_stage_input_meta) + ) monkeypatch.setattr(ap_mod, "reset_pp_stage_shapes", mock_reset) ap.update_seq_len(256) assert len(captured_args) == 1 - _, _, _, mb_size, sl, tensor_dtype = captured_args[0] + _, _, _, mb_size, sl, tensor_dtype, first_stage_input_meta = captured_args[0] assert mb_size == 2 # pp_microbatch_size assert sl == 256 assert tensor_dtype is ap.dtype + assert first_stage_input_meta is None + + # THD materialization collapses source examples into one synthetic + # batch row. The same sequence length must still refresh stage metadata + # when its effective batch extent changes. + ap.update_seq_len(256, microbatch_size=1) + assert len(captured_args) == 2 + assert captured_args[-1][3:5] == (1, 256) + + embeds = torch.empty(1, 256, 32, dtype=torch.bfloat16) + ap.update_seq_len(256, microbatch_size=1, input_tensor=embeds) + input_meta = captured_args[-1][-1] + assert input_meta.device.type == "meta" + assert input_meta.shape == embeds.shape + assert input_meta.dtype == embeds.dtype def test_update_seq_len_tracks_current(self, monkeypatch): """update_seq_len should track current seq_len and reset on change.""" diff --git a/tests/unit_tests/distributed/pipelining/test_functional.py b/tests/unit_tests/distributed/pipelining/test_functional.py index 8ad068dd98..1a23f2bf1d 100644 --- a/tests/unit_tests/distributed/pipelining/test_functional.py +++ b/tests/unit_tests/distributed/pipelining/test_functional.py @@ -1280,6 +1280,24 @@ def test_multi_stage_reset(self): assert schedule._stages_forward_initialized is False assert schedule._stages_backward_initialized is False + def test_exact_first_stage_input_meta_overrides_default_token_ids(self): + stage = self._make_stage(is_first=True, is_last=False, has_lm_head=False) + schedule = self._make_schedule() + config = self._make_config() + embeds_meta = torch.empty(1, 32, 64, device="meta", dtype=torch.bfloat16) + + reset_pp_stage_shapes( + schedule, + [stage], + config, + microbatch_size=1, + seq_len=32, + first_stage_input_meta=embeds_meta, + ) + + assert stage.inputs_meta[0].shape == (1, 32, 64) + assert stage.inputs_meta[0].dtype == torch.bfloat16 + def test_shapes_change_on_new_seq_len(self): """Calling reset twice with different seq_lens should produce different shapes.""" stage = self._make_stage(is_first=True, is_last=False, has_lm_head=False) diff --git a/tests/unit_tests/distributed/pipelining/test_hf_utils.py b/tests/unit_tests/distributed/pipelining/test_hf_utils.py index 20325204c8..3e505690e5 100644 --- a/tests/unit_tests/distributed/pipelining/test_hf_utils.py +++ b/tests/unit_tests/distributed/pipelining/test_hf_utils.py @@ -121,6 +121,66 @@ def test_forward_with_float_input_ids(self): assert isinstance(output, torch.Tensor) + def test_forward_packed_thd_uses_document_boundaries(self): + class TokenEmbedding(nn.Module): + def forward(self, input_ids): + return input_ids.to(torch.float32).unsqueeze(-1) + + class RecordingRotary(nn.Module): + def __init__(self): + super().__init__() + self.kwargs = None + + def forward(self, hidden_states, position_ids, **kwargs): + self.kwargs = kwargs + return hidden_states, position_ids + + class DocumentLocalCumsum(nn.Module): + def __init__(self): + super().__init__() + self.attention_mask = "not-called" + self.cu_seqlens = None + self.max_seqlen = None + + def forward(self, hidden_states, attention_mask=None, **kwargs): + self.attention_mask = attention_mask + self.cu_seqlens = kwargs["cu_seqlens"] + self.max_seqlen = kwargs["max_seqlen"] + boundaries = self.cu_seqlens.tolist() + documents = [ + hidden_states[start:end].cumsum(dim=0) for start, end in zip(boundaries[:-1], boundaries[1:]) + ] + return torch.cat(documents) + + class PackedDecoder(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = TokenEmbedding() + self.rotary_emb = RecordingRotary() + self.layers = nn.ModuleList([DocumentLocalCumsum()]) + self.norm = None + + model = PackedDecoder() + forward_fn = create_pipeline_forward_inner("PipelineStage") + output = forward_fn( + model, + input_ids=torch.tensor([[1, 2, 3, 4, 5]]), + attention_mask=torch.ones(1, 5), + position_ids=torch.tensor([[0, 1, 0, 1, 2]]), + qkv_format="thd", + cu_seqlens=torch.tensor([[0, 2, 5, -1000]], dtype=torch.int32), + max_seqlen=torch.tensor([3], dtype=torch.int32), + padding_mask=torch.zeros(1, 5, dtype=torch.bool), + ) + + # Each packed document starts a fresh causal history: [1, 2] and [3, 4, 5]. + torch.testing.assert_close(output, torch.tensor([[[1.0], [3.0], [3.0], [7.0], [12.0]]])) + layer = model.layers[0] + assert layer.attention_mask is None + torch.testing.assert_close(layer.cu_seqlens, torch.tensor([0, 2, 5], dtype=torch.int32)) + assert layer.max_seqlen == 3 + assert model.rotary_emb.kwargs == {"qkv_format": "thd", "cp_size": 1} + class TestCreatePipelineForwardCausalLM: """Test create_pipeline_forward_causal_lm function.""" diff --git a/tests/unit_tests/distributed/test_cp_utils.py b/tests/unit_tests/distributed/test_cp_utils.py index 29702cc5d5..be2a62001a 100644 --- a/tests/unit_tests/distributed/test_cp_utils.py +++ b/tests/unit_tests/distributed/test_cp_utils.py @@ -618,8 +618,8 @@ def thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_ran assert result["padding_mask"].shape == result["input_ids"].shape -def test_shard_thd_chunk_skips_missing_padding_mask(monkeypatch): - """Test that _shard_thd_chunk_for_te handles missing padding_mask gracefully.""" +def test_shard_thd_chunk_only_partitions_explicit_token_fields(monkeypatch): + """Token loss fields follow TE indices while same-length model payloads remain global.""" cp_mesh = _DummySubMesh(size=2) def mock_get_rank(group=None): @@ -628,7 +628,7 @@ def mock_get_rank(group=None): class MockTex: @staticmethod def thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_rank): - return torch.arange(total_tokens) + return torch.tensor([0, total_tokens - 1]) import sys @@ -637,12 +637,15 @@ def thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_ran monkeypatch.setattr(torch.distributed, "get_rank", mock_get_rank) # Batch without padding_mask — should not raise KeyError + media_payload = torch.arange(4.0) batch = { "input_ids": torch.tensor([1, 2, 3, 4]), "labels": torch.tensor([10, 20, 30, 40]), "position_ids": torch.tensor([0, 1, 2, 3]), "cu_seqlens": torch.tensor([0, 4], dtype=torch.int32), "cu_seqlens_padded": torch.tensor([0, 4], dtype=torch.int32), + "__engine_loss__advantages": torch.tensor([10.0, 20.0, 30.0, 40.0]), + "media_payload": media_payload, } result, local_indices = _cu._shard_thd_chunk_for_te(batch, cp_mesh, "thd", -1000, 0) @@ -650,8 +653,10 @@ def thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_ran assert "input_ids" in result assert "attention_mask" not in result assert "cu_seqlens_padded" not in result - # the partition IS the local-token global index map (mock returns arange) - assert torch.equal(local_indices, torch.arange(4)) + assert torch.equal(local_indices, torch.tensor([0, 3])) + assert torch.equal(result["input_ids"], torch.tensor([1, 4])) + assert torch.equal(result["__engine_loss__advantages"], torch.tensor([10.0, 40.0])) + assert result["media_payload"] is media_payload def test_make_cp_batch_for_te_unsupported_format(): @@ -813,6 +818,15 @@ def fake_make_cp_batch_for_te( assert (seen["pad"], seen["fmt"], seen["chunks"], seen["sent"]) == (7, "thd", 3, -1000) +def test_sharder_constructor_rejects_injected_hf_te_for_thd(): + """Injected HF TE attention is still BSHD-only and cannot consume THD metadata.""" + model = type("_Model", (), {"_te_attention_injected": True})() + batch = {"input_ids": torch.tensor([[1, 2]]), "qkv_format": "thd"} + + with pytest.raises(NotImplementedError, match="stock Hugging Face model supports padded BSHD only"): + ContextParallelSharder(model, _DummyDeviceMesh(cp_size=1, tp_size=1), batch) + + def test_sharder_constructor_does_not_infer_te_from_batch_alone(monkeypatch): """A THD-origin batch does not force TE preparation on a non-TE model.""" monkeypatch.setattr( diff --git a/tests/unit_tests/distributed/test_magi_attn_utils.py b/tests/unit_tests/distributed/test_magi_attn_utils.py index d805710692..5d923c1343 100644 --- a/tests/unit_tests/distributed/test_magi_attn_utils.py +++ b/tests/unit_tests/distributed/test_magi_attn_utils.py @@ -80,6 +80,13 @@ def __init__(self): self.visual = nn.Linear(4, 4) # vision tower: must NOT be stamped +class _CausalLM(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(num_attention_heads=2, num_key_value_heads=2, head_dim=4) + self.self_attn = _FakeAttention() + + # --------------------------------------------------------------------------- # # AttnMaskSpec builders (pure Python) # --------------------------------------------------------------------------- # @@ -166,6 +173,83 @@ def test_prepare_llm_batch_custom_prefix_tree_ok(self): finally: mu.set_active_attn_spec(None) + def test_prepare_llm_batch_refreshes_prefix_spec_for_each_outer_pipeline_batch(self, monkeypatch): + """A one-microbatch PP call cannot leak its mask into the next GA call.""" + active_groups = [] + monkeypatch.setattr(mu, "set_active_cp_group", active_groups.append) + st = MagiState(enabled=True, custom=True, cp_group=None, cp_size=1) + + first = { + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "prefix_tree": ([2, 2], [[0, 1]]), + } + try: + st.make_cp_batch(None, first, model=None, return_local_indices=True) + first_spec = mu.get_active_attn_spec() + assert first_spec is not None + assert first_spec.fingerprint() == AttnMaskSpec.prefix_tree([2, 2], [[0, 1]])[0].fingerprint() + + # With pp_batch_size == pp_microbatch_size == 1, Engine enters the + # schedule once per outer GA item. Preparing the next item must + # clear the prior prefix-tree mask before that forward starts. + st.make_cp_batch( + None, + {"input_ids": torch.zeros(1, 4, dtype=torch.long)}, + model=None, + return_local_indices=True, + ) + assert mu.get_active_attn_spec() is None + assert active_groups == [None, None] + finally: + mu.set_active_attn_spec(None) + + def test_prepare_llm_batch_custom_cp2_causal_dispatches(self, monkeypatch): + group = _FakeGroup(2) + expected_batch = {"input_ids": torch.tensor([[1, 3]]), "labels": torch.tensor([[11, 13]])} + expected_indices = torch.tensor([[0, 2]]) + calls = [] + + def prepare(model, batch, cp_group, *, return_local_indices): + calls.append((model, batch, cp_group, return_local_indices)) + return expected_batch, object(), expected_indices + + monkeypatch.setattr(mu, "magi_prepare_batch", prepare) + st = MagiState(enabled=True, custom=True, cp_group=group, cp_size=2) + model = object() + batch = {"input_ids": torch.arange(4).view(1, 4), "labels": torch.arange(10, 14).view(1, 4)} + + train_ctx, out, local_indices = st.prepare_llm_batch( + model, + batch, + device_mesh=None, + is_thd=False, + pad_id=0, + num_chunks=1, + ) + + from contextlib import nullcontext + + assert train_ctx is nullcontext + assert out is expected_batch + assert local_indices is expected_indices + assert calls == [(model, batch, group, True)] + + def test_prepare_llm_batch_hf_packed_rejects_lost_document_boundaries(self): + st = MagiState(enabled=True, custom=False, cp_group=None, cp_size=1) + batch = { + "input_ids": torch.zeros(1, 8, dtype=torch.long), + "labels": torch.zeros(1, 8, dtype=torch.long), + "seq_lens": torch.tensor([4, 4]), + } + with pytest.raises(NotImplementedError, match="cannot preserve packed document boundaries"): + st.prepare_llm_batch(model=None, batch=batch, device_mesh=None, is_thd=True, pad_id=0, num_chunks=1) + + def test_prepare_llm_batch_prefix_tree_cp2_rejects_undispatched_spec(self): + st = MagiState(enabled=True, custom=True, cp_group=_FakeGroup(2), cp_size=2) + batch = {"input_ids": torch.zeros(1, 4, dtype=torch.long), "prefix_tree": ([2, 2], [[0, 1]])} + with pytest.raises(NotImplementedError, match="requires cp_size=1"): + st.prepare_llm_batch(model=None, batch=batch, device_mesh=None, is_thd=False, pad_id=0, num_chunks=1) + # --------------------------------------------------------------------------- # # setup_magi @@ -253,6 +337,54 @@ def test_iter_language_model_attention_skips_vision(self): assert mods == [model.language_model.self_attn] +class TestMagiPrepareBatch: + def test_dispatches_inputs_labels_and_loss_indices_with_one_layout(self, monkeypatch): + """HF/custom causal CP keeps loss tokens aligned for a PP microbatch.""" + expected_key = object() + order = torch.tensor([2, 5, 6, 7]) + dispatch_calls = [] + + def dispatch(value, *, key, pad_value=0): + assert key is expected_key + dispatch_calls.append((value.clone(), pad_value)) + padded = torch.cat((value, value.new_full((2,), pad_value))) + return padded.index_select(0, order) + + api = ModuleType("magi_attention.api") + api.dispatch = dispatch + api.get_position_ids = lambda key: torch.tensor([2, 5, 0, 0]) + api.magi_attn_varlen_key = lambda **kwargs: expected_key + functools = ModuleType("magi_attention.api.functools") + functools.compute_pad_size = lambda *args, **kwargs: 2 + package = ModuleType("magi_attention") + package.api = api + monkeypatch.setitem(sys.modules, "magi_attention", package) + monkeypatch.setitem(sys.modules, "magi_attention.api", api) + monkeypatch.setitem(sys.modules, "magi_attention.api.functools", functools) + + model = _CausalLM() + batch = { + "input_ids": torch.tensor([[10, 11, 12, 13, 14, 15]]), + "labels": torch.tensor([[20, 21, 22, 23, 24, 25]]), + "attention_mask": torch.ones(1, 6), + } + out, returned_key, local_indices = mu.magi_prepare_batch( + model, + batch, + _FakeGroup(2), + return_local_indices=True, + ) + + assert returned_key is expected_key + assert torch.equal(out["input_ids"], torch.tensor([[12, 15, 0, 0]])) + assert torch.equal(out["labels"], torch.tensor([[22, 25, -100, -100]])) + assert torch.equal(local_indices, torch.tensor([[2, 5, 6, 6]])) + assert torch.equal(out["position_ids"], torch.tensor([[2, 5, 0, 0]])) + assert "attention_mask" not in out + assert model.self_attn.cp_group.size() == 2 + assert any(torch.equal(value, torch.arange(6)) and pad_value == 6 for value, pad_value in dispatch_calls) + + class TestMagiPrepareVlm: """magi_prepare_vlm is pure Python (no magi import) for the cp_size==1 path.""" diff --git a/tests/unit_tests/distributed/test_thd_utils.py b/tests/unit_tests/distributed/test_thd_utils.py index 17df4ee572..72f0fb5b2b 100644 --- a/tests/unit_tests/distributed/test_thd_utils.py +++ b/tests/unit_tests/distributed/test_thd_utils.py @@ -704,3 +704,90 @@ def test_process_input_for_thd_2d_position_ids_unchanged(): } out = process_input_for_thd(batch) assert tuple(out["position_ids"].shape) == (B * S,) + + +def test_process_input_for_thd_flattens_only_explicit_engine_loss_fields(): + vision_mask = torch.tensor([[False, True, False], [True, False, True]]) + advantages = torch.arange(6, dtype=torch.float32).view(2, 3) + batch = { + "input_ids": torch.arange(6).view(2, 3), + "labels": torch.arange(6).view(2, 3), + "position_ids": torch.arange(3).expand(2, -1), + "seq_lens": torch.tensor([[3], [3]]), + "seq_lens_padded": torch.tensor([[3], [3]]), + "_global_vision_mask": vision_mask, + "__engine_loss__advantages": advantages, + } + + out = process_input_for_thd(batch) + + assert out["_global_vision_mask"] is vision_mask + assert torch.equal(out["__engine_loss__advantages"], advantages.flatten()) + + +def test_split_batch_into_thd_chunks_keeps_engine_loss_fields_on_the_token_stream(): + advantages = torch.arange(8, dtype=torch.float32).view(2, 4) + old_logprobs = -advantages + batch = { + "input_ids": torch.arange(1, 9).view(2, 4), + "labels": torch.arange(1, 9).view(2, 4), + "position_ids": torch.arange(4).expand(2, -1), + "seq_lens": torch.tensor([[4], [4]]), + "seq_lens_padded": torch.tensor([[4], [4]]), + "__engine_loss__advantages": advantages, + "__engine_loss__old_logprobs": old_logprobs, + } + + out = split_batch_into_thd_chunks(batch, num_chunks=2) + + assert out["__engine_loss__advantages"].shape == (2, 4) + assert torch.equal(out["__engine_loss__advantages"], advantages) + assert torch.equal(out["__engine_loss__old_logprobs"], old_logprobs) + + +def test_split_single_packed_row_uses_document_boundaries_for_pipeline_chunks(): + advantages = torch.arange(8, dtype=torch.float32).unsqueeze(0) + batch = { + "input_ids": torch.arange(1, 9).unsqueeze(0), + "labels": torch.arange(11, 19).unsqueeze(0), + "position_ids": torch.tensor([[0, 1, 0, 1, 0, 1, 0, 1]]), + "seq_lens": torch.tensor([[2, 2, 2, 2]]), + "seq_lens_padded": torch.tensor([[2, 2, 2, 2]]), + "__engine_loss__advantages": advantages, + } + + out = split_batch_into_thd_chunks(batch, num_chunks=2) + + assert out["input_ids"].shape == (2, 4) + assert out["input_ids"].tolist() == [[1, 2, 3, 4], [5, 6, 7, 8]] + assert out["__engine_loss__advantages"].tolist() == [[0, 1, 2, 3], [4, 5, 6, 7]] + assert out["cu_seqlens"].tolist() == [[0, 2, 4], [0, 2, 4]] + + +def test_split_single_packed_row_rejects_variable_pipeline_token_widths(): + batch = { + "input_ids": torch.arange(8).unsqueeze(0), + "labels": torch.arange(8).unsqueeze(0), + "position_ids": torch.arange(8).unsqueeze(0), + "seq_lens": torch.tensor([[1, 1, 2, 4]]), + "seq_lens_padded": torch.tensor([[1, 1, 2, 4]]), + } + + with pytest.raises(ValueError, match="equal token widths"): + split_batch_into_thd_chunks(batch, num_chunks=2) + + +def test_stack_thd_chunks_rejects_inconsistent_fields(): + from nemo_automodel.components.distributed.thd_utils import stack_thd_chunks + + chunks = [ + {"input_ids": torch.tensor([1]), "cu_seqlens": torch.tensor([0, 1])}, + { + "input_ids": torch.tensor([2]), + "cu_seqlens": torch.tensor([0, 1]), + "unexpected": torch.tensor([3]), + }, + ] + + with pytest.raises(ValueError, match="inconsistent field 'unexpected'"): + stack_thd_chunks(chunks) diff --git a/tests/unit_tests/loss/test_mtp_lm_head_gather.py b/tests/unit_tests/loss/test_mtp_lm_head_gather.py index 4492897029..55af542e98 100644 --- a/tests/unit_tests/loss/test_mtp_lm_head_gather.py +++ b/tests/unit_tests/loss/test_mtp_lm_head_gather.py @@ -260,6 +260,29 @@ def fake_calc(loss_fn, **kw): assert captured["logits"] is None +def test_pipeline_loss_fused_ce_aligns_single_thd_hidden_states_with_flat_labels(): + """The synthetic PP axis on a single THD chunk is not a token axis.""" + from nemo_automodel.components.loss.mtp import PipelineCausalLMLoss + + m = _TinyModel() + m.training = False + loss_mod = PipelineCausalLMLoss(FusedLinearCrossEntropy(), m) + hidden = torch.randn(1, S, H) + labels = torch.randint(0, V, (S,)) + captured = {} + + def fake_calc(loss_fn, **kw): + captured.update(kw) + return torch.zeros((), requires_grad=True) + + with mock.patch.object(_mtp, "calculate_loss", side_effect=fake_calc): + loss_mod(hidden, labels) + + assert captured["hidden_states"].shape == (S, H) + torch.testing.assert_close(captured["hidden_states"], hidden.squeeze(0)) + assert captured["logits"] is None + + def test_pipeline_loss_fused_ce_with_mtp_tuple_raises(): """FusedLinearCrossEntropy cannot consume an MTP tuple output (it carries logits, not the hidden states the fused loss needs), so the last-stage loss diff --git a/tests/unit_tests/loss/test_nemotron_parse_loss.py b/tests/unit_tests/loss/test_nemotron_parse_loss.py index 1b89569ce9..00253b4476 100644 --- a/tests/unit_tests/loss/test_nemotron_parse_loss.py +++ b/tests/unit_tests/loss/test_nemotron_parse_loss.py @@ -76,7 +76,7 @@ def test_backward_compatibility(): logits = torch.randn(2, 10, 100) labels = torch.randint(0, 100, (2, 10)) - loss_fn = NemotronParseLoss(coordinate_weight=10.0, class_token_start_idx=50000) + loss_fn = NemotronParseLoss(coordinate_weight=10.0, class_token_start_idx=50000, reduction="mean") loss_new = loss_fn(logits=logits, labels=labels) loss_ref = _compute_reference_loss(logits, labels) @@ -126,7 +126,7 @@ def test_fp32_upcast(): assert torch.isfinite(loss_fp32) assert torch.isfinite(loss_bf16) - assert torch.allclose(loss_fp32, loss_bf16, rtol=1e-2) + assert torch.allclose(loss_fp32, loss_bf16.float(), rtol=1e-2) def test_invalid_logits_shape(): @@ -265,10 +265,12 @@ def test_mixed_coordinate_and_regular_tokens(): """Test loss computation with mixed coordinate and regular tokens.""" torch.manual_seed(42) logits = torch.randn(2, 10, 50150) # Increased vocab size to accommodate labels - labels = torch.tensor([ - [10, 20, 50001, 50002, 30, 40, 50003, 50, 60, 70], # Mixed - [50100, 50101, 50102, 100, 200, 300, 50103, 400, 500, 600] # Mixed - ]) + labels = torch.tensor( + [ + [10, 20, 50001, 50002, 30, 40, 50003, 50, 60, 70], # Mixed + [50100, 50101, 50102, 100, 200, 300, 50103, 400, 500, 600], # Mixed + ] + ) loss_fn = NemotronParseLoss(coordinate_weight=10.0, class_token_start_idx=50000) loss = loss_fn(logits=logits, labels=labels) @@ -346,3 +348,24 @@ def test_reduction_parameter_stored(): assert loss_fn_sum.reduction == "sum" assert loss_fn_mean.reduction == "mean" + + +def test_sum_and_mean_reductions_have_distinct_normalization(): + logits = torch.randn(2, 5, 100) + labels = torch.randint(0, 100, (2, 5)) + labels[0, :2] = -100 + loss_sum = NemotronParseLoss(reduction="sum")(logits=logits, labels=labels) + loss_mean = NemotronParseLoss(reduction="mean")(logits=logits, labels=labels) + + torch.testing.assert_close(loss_sum, loss_mean * (labels != -100).sum()) + + +def test_all_ignored_sum_keeps_a_zero_autograd_path(): + logits = torch.randn(2, 5, 100, requires_grad=True) + labels = torch.full((2, 5), -100) + + loss = NemotronParseLoss(reduction="sum")(logits=logits, labels=labels) + loss.backward() + + assert logits.grad is not None + assert torch.count_nonzero(logits.grad) == 0 diff --git a/tests/unit_tests/recipes/test_base_recipe.py b/tests/unit_tests/recipes/test_base_recipe.py index 38351971cb..93cd4df5f1 100644 --- a/tests/unit_tests/recipes/test_base_recipe.py +++ b/tests/unit_tests/recipes/test_base_recipe.py @@ -328,6 +328,29 @@ def fake_all_reduce(tensor, op=None, group=None): assert calls[0][1] is None +def test_mtp_cp_validation_reaches_consensus_across_pipeline_stages(monkeypatch): + """A stage without the MTP module must still reject when a peer owns it.""" + group = object() + recipe = SimpleNamespace( + mesh_context=SimpleNamespace(cp_size=2, process_group=group), + dist_env=SimpleNamespace(device=torch.device("cpu")), + ) + part = nn.Linear(2, 2) + calls = [] + + def fake_all_reduce(flag, op=None, group=None): + calls.append((op, group)) + flag.fill_(1) + + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + + with pytest.raises(NotImplementedError, match="MTP with context parallelism"): + BaseRecipe._validate_mtp_context_parallelism(recipe, [part]) + + assert calls == [(torch.distributed.ReduceOp.MAX, group)] + + def test_optimizer_checkpoint_part_ids_use_global_pipeline_stage_indices(tmp_path): recipe_inst = _ToyRecipe(tmp_path) recipe_inst.pp = SimpleNamespace( diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py index c1a5917201..2c4a2e1b76 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -12,20 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the VLM-CP wiring in ``recipes/vlm/finetune.py``. - -These reproduce the ``_forward_backward_step``-style and -``_run_validation_epoch``-style batch handling without instantiating the -full recipe — exercising the code shape that gets shipped: - - - Invoke the sharder-only ``prepare_model_inputs_for_cp`` directly through - ``ContextParallelSharder`` construction (a plain method call; nothing consumed, so input_ids - and multimodal inputs stay in the batch for the model's own forward) - - PP gating: the sharder-only hook is invoked on every stage (all PP-capable - VLMs are sunk — they embed + shard in their own forward); media is dropped on - non-first stages so those stage forwards see only text inputs - - Validation: count labels after ``ContextParallelSharder.shard`` and inside train_ctx - - Validation: position_ids ``.to(self.dist_env.device)`` (not model.device) +"""Tests for VLM context-parallel wiring in ``recipes/vlm/finetune.py``. + +Training forward/backward and CP sharding are owned by :class:`Engine`. These +tests cover the VLM recipe responsibilities that remain around that core: +pipeline media staging setup, vision-frame context publication, and the +recipe-owned validation forward path. """ from __future__ import annotations @@ -39,6 +31,7 @@ import nemo_automodel.recipes.vlm.finetune as vlm_finetune from nemo_automodel.components.config.loader import ConfigNode from nemo_automodel.components.distributed.cp_vision_frame_shard import CpVisionFrameShardingConfig +from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM @@ -57,59 +50,6 @@ def _identity_cp_shard(sharder, batch): return nullcontext, batch -def _make_recipe_with_pp_stages(*, pp_enabled=True, has_first_stage=True, pp_microbatch_size=2): - first_stage = SimpleNamespace(is_first=True, inputs_meta=("old-first",)) - later_stage = SimpleNamespace(is_first=False, inputs_meta=("old-later",)) - recipe = SimpleNamespace( - pp_enabled=pp_enabled, - pp=SimpleNamespace( - pp_microbatch_size=pp_microbatch_size, - info=SimpleNamespace(has_first_stage=has_first_stage, stages=[first_stage, later_stage]), - ), - ) - return recipe, first_stage, later_stage - - -def test_maybe_set_pp_first_stage_embed_input_meta_sets_first_stage_meta(): - recipe, first_stage, later_stage = _make_recipe_with_pp_stages(pp_microbatch_size=3) - model_input = torch.empty(5, 11, 13, dtype=torch.bfloat16) - - FinetuneRecipeForVLM._maybe_set_pp_first_stage_embed_input_meta(recipe, model_input) - - assert later_stage.inputs_meta == ("old-later",) - assert len(first_stage.inputs_meta) == 1 - meta = first_stage.inputs_meta[0] - assert tuple(meta.shape) == (3, 11, 13) - assert meta.dtype == torch.bfloat16 - assert meta.device.type == "meta" - - -@pytest.mark.parametrize( - ("recipe_kwargs", "model_input"), - [ - ({"pp_enabled": False}, torch.empty(5, 11, 13)), - ({"has_first_stage": False}, torch.empty(5, 11, 13)), - ({}, torch.empty(5, 11, 13, dtype=torch.int64)), - ({}, torch.empty(5, 11)), - ], -) -def test_maybe_set_pp_first_stage_embed_input_meta_guard_conditions(recipe_kwargs, model_input): - recipe, first_stage, later_stage = _make_recipe_with_pp_stages(**recipe_kwargs) - - FinetuneRecipeForVLM._maybe_set_pp_first_stage_embed_input_meta(recipe, model_input) - - assert first_stage.inputs_meta == ("old-first",) - assert later_stage.inputs_meta == ("old-later",) - - -class _FakeCPMesh: - mesh_dim_names = ("cp",) - - def __getitem__(self, key): - assert key == "cp" - return SimpleNamespace(size=lambda: 2, get_group=lambda: "cp-group") - - class _UnsupportedVisionModel: supports_cp_vision_frame_sharding = False @@ -205,212 +145,13 @@ def _reset(actual_token): assert events[2] == ("reset", token) -class _ScheduleSpy: - def __init__(self): - self.calls = [] - - def step(self, model_input=None, *, target=None, losses=None, **batch): - self.calls.append({"model_input": model_input, "target": target, "batch": batch}) - if losses is not None: - losses.append(torch.tensor(1.25)) - - -class _PPSpy(SimpleNamespace): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.step_batches = [] - - def step(self, model_input, *, target=None, losses=None, **kwargs): - """Record and forward an AutoPipeline step. - - Args: - model_input: Tensor of shape [batch, ...] containing the first - pipeline stage's input. - target: Optional tensor of shape [batch, sequence] containing loss - targets. - losses: Optional mutable list populated with scalar loss tensors. - **kwargs: Keyword schedule inputs. Tensor values have arbitrary - model-defined layouts. - - Returns: - The value returned by the schedule spy. - """ - self.step_batches.append(dict(kwargs)) - schedule_args = (model_input,) if self.info.has_first_stage else () - return self.info.schedule.step(*schedule_args, target=target, losses=losses, **kwargs) - - -def test_forward_backward_step_pp_cp_first_stage_sunk_keeps_input_ids_full(monkeypatch): - """Sunk model on the FIRST PP stage under CP: the sharder-only hook is invoked - (consumes nothing), so input_ids stays full-length, update_seq_len sees the - full seq_len, and the full-length input_ids is fed to the pipeline schedule - (the model embeds + shards inside its own forward).""" - labels = torch.arange(12, dtype=torch.long).reshape(2, 6) - model = _SunkSpyVLM() - schedule = _ScheduleSpy() - seq_lens = [] - first_stage = SimpleNamespace(is_first=True, inputs_meta=None) - recipe = object.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) - recipe.device_mesh = _FakeCPMesh() - recipe.cp_vision_frame_sharding = CpVisionFrameShardingConfig(enabled=True) - recipe.distributed_config = SimpleNamespace(defer_fsdp_grad_sync=True) - recipe.model_parts = [model] - recipe.pp_enabled = True - recipe.pp = _PPSpy( - pp_microbatch_size=2, - info=SimpleNamespace( - has_first_stage=True, - has_last_stage=True, - stages=[first_stage, SimpleNamespace(is_first=False, inputs_meta=None)], - schedule=schedule, - ), - update_seq_len=seq_lens.append, - ) - batch = { - "input_ids": torch.ones(2, 6, dtype=torch.long), - "pixel_values": torch.zeros(2, 3, 4, 4), - "labels": labels, - } - seen_cp_batch = {} - - def _shard(sharder, cp_batch): - """Capture the global model-input mapping before CP transport. - - Args: - sharder: Sharder configured by the VLM recipe. - cp_batch: Mutable model-input mapping whose tensor values have - global batch and sequence extents. - - Returns: - The null context factory and the same input mapping. - """ - del sharder - seen_cp_batch.update(cp_batch) - return nullcontext, cp_batch - - monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _shard) - monkeypatch.setattr(vlm_finetune, "stage_vlm_media_for_pp", lambda *args, **kwargs: nullcontext()) - monkeypatch.setattr(FinetuneRecipeForVLM, "_maybe_set_pp_first_stage_embed_input_meta", lambda self, mi: None) - - loss_buffer = [] - FinetuneRecipeForVLM._forward_backward_step( - recipe, - 0, - batch, - loss_buffer=loss_buffer, - num_label_tokens=labels.numel(), - num_batches=1, - ) - - assert len(model.calls) == 1 - # Sharder-only: input_ids stays full, no inputs_embeds injected. - assert "input_ids" in seen_cp_batch - assert tuple(seen_cp_batch["input_ids"].shape) == (2, 6) - assert "inputs_embeds" not in seen_cp_batch - assert seq_lens == [6] - assert [set(call.keys()) for call in recipe.pp.step_batches] == [{"pixel_values"}] - assert len(schedule.calls) == 1 - assert tuple(schedule.calls[0]["model_input"].shape) == (2, 6) - assert torch.equal(schedule.calls[0]["target"], labels) - assert torch.equal(loss_buffer[0], torch.tensor(1.25)) - - -class _SunkSpyVLM: - """Sunk VLM: sharder-only CP hook (embeds/shards in forward, consumes nothing).""" - - def __init__(self): - self.calls = [] - - def prepare_model_inputs_for_cp(self, batch, *, num_chunks=1): - # Sharder-only: nothing consumed, no inputs_embeds — input_ids stays full. - self.calls.append({"batch": dict(batch), "num_chunks": num_chunks}) - return {} - - def __call__(self, **kwargs): - raise AssertionError("CP prepare must call prepare_model_inputs_for_cp directly, not __call__") - - -def _run_nonfirst_stage_fbstep(monkeypatch, model): - """Drive _forward_backward_step for a non-first (has_first_stage=False) PP+CP stage.""" - labels = torch.arange(12, dtype=torch.long).reshape(2, 6) - schedule = _ScheduleSpy() - seq_lens = [] - recipe = object.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) - recipe.device_mesh = _FakeCPMesh() - recipe.cp_vision_frame_sharding = CpVisionFrameShardingConfig(enabled=True) - recipe.distributed_config = SimpleNamespace(defer_fsdp_grad_sync=True) - recipe.model_parts = [model] - recipe.pp_enabled = True - recipe.pp = _PPSpy( - pp_microbatch_size=2, - info=SimpleNamespace( - has_first_stage=False, - has_last_stage=True, - stages=[SimpleNamespace(is_first=False, inputs_meta=None)], - schedule=schedule, - ), - update_seq_len=seq_lens.append, - ) - batch = { - "input_ids": torch.ones(2, 6, dtype=torch.long), - "pixel_values": torch.zeros(2, 3, 4, 4), - "labels": labels, - } - seen_cp_batch = {} - - def _shard(sharder, cp_batch): - """Capture the global model-input mapping before CP transport. - - Args: - sharder: Sharder configured by the VLM recipe. - cp_batch: Mutable model-input mapping whose tensor values have - global batch and sequence extents. - - Returns: - The null context factory and the same input mapping. - """ - del sharder - seen_cp_batch.update(cp_batch) - return nullcontext, cp_batch - - monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _shard) - monkeypatch.setattr(vlm_finetune, "stage_vlm_media_for_pp", lambda *args, **kwargs: nullcontext()) - monkeypatch.setattr(FinetuneRecipeForVLM, "_maybe_set_pp_first_stage_embed_input_meta", lambda self, mi: None) - - FinetuneRecipeForVLM._forward_backward_step( - recipe, 0, batch, loss_buffer=[], num_label_tokens=labels.numel(), num_batches=1 - ) - return seen_cp_batch, seq_lens, recipe.pp.step_batches, schedule.calls - - -def test_forward_backward_step_pp_cp_sunk_model_nonfirst_stage_invokes_hook_keeps_input_ids_full(monkeypatch): - """Regression: a sunk model must invoke its sharder-only hook on NON-first PP - stages under cp>1, so input_ids stays full-length and update_seq_len (which - drives the CP-aware stage metas) sees the FULL seq_len — not the local length - the generic sharder would produce, which would ÷cp a second time and truncate - the inter-stage hidden (the text-decoder RoPE size mismatch).""" - model = _SunkSpyVLM() - seen_cp_batch, seq_lens, step_batches, schedule_calls = _run_nonfirst_stage_fbstep(monkeypatch, model) - - # Hook invoked on the non-first stage (this is the fix). - assert len(model.calls) == 1 - # Sharder-only hook consumes nothing: input_ids stays full-length (seq=6). - assert "input_ids" in seen_cp_batch - assert tuple(seen_cp_batch["input_ids"].shape) == (2, 6) - # All pp ranks feed the FULL seq_len to update_seq_len. - assert seq_lens == [6] - assert step_batches == [{}] - assert schedule_calls[0]["model_input"] is None - - class _FakePPModel: def __init__(self, stage0): self.parts = [stage0] self.pp_batch_size = 4 self.pp_microbatch_size = 2 - self.info = SimpleNamespace(has_last_stage=False, stages=[], schedule=None) + self.scale_grads_in_schedule = False + self.info = SimpleNamespace(has_first_stage=True, has_last_stage=False, stages=[], schedule=None) class _StageWithCPPreembedInForward: @@ -427,6 +168,7 @@ class _StageWithoutCPPrepare: def _patch_pp_setup_minimals(monkeypatch, *, cp_size, stage0, dataloader_calls): monkeypatch.setattr(vlm_finetune, "AutoPipeline", _FakePPModel) + monkeypatch.setattr("nemo_automodel.engine.AutoPipeline", _FakePPModel) monkeypatch.setattr( vlm_finetune, "initialize_distributed", @@ -437,7 +179,7 @@ def _patch_pp_setup_minimals(monkeypatch, *, cp_size, stage0, dataloader_calls): monkeypatch.setattr(vlm_finetune, "StatefulRNG", lambda *args, **kwargs: "rng") monkeypatch.setattr( "nemo_automodel.recipes._typed_config.RecipeConfig.loss_fn", - property(lambda self: SimpleNamespace(build=lambda: "loss_fn")), + property(lambda self: SimpleNamespace(build=lambda: MaskedCrossEntropy(reduction="sum"))), ) monkeypatch.setattr(vlm_finetune, "_supports_logits_to_keep", lambda model: True) monkeypatch.setattr( @@ -452,7 +194,7 @@ def _patch_pp_setup_minimals(monkeypatch, *, cp_size, stage0, dataloader_calls): pp_size=2, ), strategy_config=SimpleNamespace(), - pipeline_config=SimpleNamespace(), + pipeline_config=SimpleNamespace(scale_grads_in_schedule=False), moe_parallel_config=None, activation_checkpointing=False, ), @@ -574,6 +316,8 @@ def test_setup_always_stages_pp_media_under_pp( assert dataloader_calls[0]["pp_n_microbatches"] == expected_pp_n_microbatches assert dataloader_calls[0]["cp_size"] == cp_size + assert trainer.engine.pipeline is trainer.pp + assert trainer.engine.microbatch_size == 1 # ----------------------------------------------------------------------------- diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 154625fd64..b583713003 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -28,10 +28,9 @@ chunk_vlm_media, prepare_vlm_media_for_pp, stage_vlm_media_for_pp, + wrap_vlm_collate_for_pp, ) -from nemo_automodel.components.distributed.cp_vision_frame_shard import CpVisionFrameShardingConfig -from nemo_automodel.components.loggers.metric_logger import MetricsSample -from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler +from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.optim.optimizer import LRSchedulerConfig, build_optimizer_config from nemo_automodel.components.training.step_scheduler import StepSchedulerConfig from nemo_automodel.recipes._typed_config import ( @@ -384,295 +383,177 @@ def forward(self, **batch): return torch.zeros((), requires_grad=True) -@pytest.mark.cuda(False) -def test_run_train_step_supports_tensor_outputs(monkeypatch): +def _build_engine_recipe_for_optim_step(*, pp_enabled: bool = False): + """Build the smallest recipe state needed to exercise the Engine boundary.""" recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device="cpu") + recipe.dist_env = SimpleNamespace(device="cpu", rank=0, is_main=True) recipe.device_mesh = None recipe.moe_mesh = None recipe.loss_fn = object() - model = _TensorModel() - recipe.model_parts = [model] # Now uses model_parts instead of model - recipe.pp_enabled = False # Pipeline parallelism disabled - recipe.optimizer = [_DummyOptimizer()] # Now a list - # ``is_remote_logging_step`` is read by ``_forward_backward_step`` when the - # composite (gemma4 joint drafter) attaches drafter logits; default False - # so non-drafter test paths skip the log line. + recipe.model_parts = [_TensorModel()] + recipe.pp_enabled = pp_enabled + if pp_enabled: + recipe.pp = SimpleNamespace(info=SimpleNamespace(has_first_stage=True)) + recipe.optimizer = [_DummyOptimizer()] recipe.step_scheduler = SimpleNamespace(step=0, epoch=0, is_remote_logging_step=False) recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=lambda: None) recipe.cfg = _Cfg(fp8=None) recipe.lr_scheduler = None recipe.timestamp = 0.0 recipe.distributed_config = None - recipe._dp_allreduce = lambda tensor, include_cp=False: tensor recipe._get_dp_group_size = lambda include_cp=True: 1 recipe._get_cp_group_size = lambda: 1 + recipe.engine = MagicMock() + recipe.engine.forward_backward.return_value = (torch.tensor(0.25), []) + return recipe + +@pytest.mark.cuda(False) +def test_run_train_step_passes_flat_prebatched_datums_to_engine(monkeypatch): + recipe = _build_engine_recipe_for_optim_step(pp_enabled=True) batches = [ { - "labels": torch.tensor([[1, -100]]), - "input_ids": torch.tensor([[1, 2]]), - } + "labels": torch.tensor([[1, -100, 2, -100]]), + "input_ids": torch.tensor([[1, 2, 3, 4]]), + }, + { + "labels": torch.tensor([[-100, 3, -100, 4]]), + "input_ids": torch.tensor([[5, 6, 7, 8]]), + }, ] - - logits_seen = {} - - def fake_calculate_loss(*args, **kwargs): - logits_seen["value"] = kwargs["logits"] - return torch.tensor(1.0, requires_grad=True) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.get_sync_ctx", - lambda model, is_last, defer_fsdp_grad_sync=True: nullcontext(), - ) - - calculate_mock = MagicMock(side_effect=fake_calculate_loss) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_loss", calculate_mock) - - grad_clip_mock = MagicMock(return_value=2.5) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", - grad_clip_mock, - ) - - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_grad_accumulation", - lambda model_parts, pp_enabled: None, - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_final_backward", - lambda model_parts, pp_enabled: None, - ) + finalizer = MagicMock(return_value=2.5) + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", finalizer) metrics = recipe._run_train_optim_step(batches, max_grad_norm=1.0) - assert isinstance(metrics, MetricsSample) - assert logits_seen["value"].requires_grad - grad_clip_mock.assert_called_once() - assert calculate_mock.call_args.kwargs["num_label_tokens"] == 1 - assert metrics.metrics["grad_norm"] == 2.5 - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(1.0) + datums, loss_fn = recipe.engine.forward_backward.call_args.args + assert len(datums) == 2 + assert all(datum.model_inputs.keys() == {"input_ids"} for datum in datums) + assert [datum.loss_fn_inputs["weights"].tolist() for datum in datums] == [ + [[True, False, True, False]], + [[False, True, False, True]], + ] + assert callable(loss_fn) + assert finalizer.call_args.kwargs["num_label_tokens"] is None + assert metrics.metrics["loss"] == pytest.approx(0.25) assert recipe.optimizer[0].step_called assert recipe.optimizer[0].zero_grad_called @pytest.mark.cuda(False) -def test_forward_backward_step_routes_thd_batch_through_te(monkeypatch): - recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device="cpu") - recipe.device_mesh = None - recipe.mesh_context = SimpleNamespace(cp_size=2) - recipe.processor = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=7)) - recipe.model_parts = [_TensorModel()] - recipe.pp_enabled = False - recipe.magi = SimpleNamespace(enabled=False) - recipe.distributed_config = None - recipe.loss_fn = object() - recipe.step_scheduler = SimpleNamespace(is_remote_logging_step=False) - recipe._get_dp_group_size = lambda include_cp=True: 1 - captured = {} - - def make_thd_batch(model, device_mesh, batch, **kwargs): - captured.update(kwargs) - captured["qkv_format"] = batch.get("qkv_format") - return SimpleNamespace(shard=lambda actual: (nullcontext, actual)) - - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ContextParallelSharder", make_thd_batch) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.get_sync_ctx", lambda *args, **kwargs: nullcontext()) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.calculate_loss", - lambda *args, **kwargs: torch.tensor(1.0, requires_grad=True), - ) - - recipe._forward_backward_step( - idx=0, - batch={ - "input_ids": torch.tensor([[1, 2]]), - "labels": torch.tensor([[2, -100]]), - "qkv_format": "thd", - }, - loss_buffer=[], - num_label_tokens=1, - num_batches=1, - ) - - assert captured["qkv_format"] == "thd" - assert "use_te" not in captured - assert "magi" not in captured - assert captured["padding_token_id"] == 7 - - -@pytest.mark.cuda(False) -def test_forward_backward_step_rejects_mrope_thd_with_context_parallelism(): - recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device="cpu") - recipe.device_mesh = None - recipe.mesh_context = SimpleNamespace(cp_size=2) - recipe.model_parts = [_TensorModel()] - recipe.pp_enabled = False - recipe.magi = SimpleNamespace(enabled=False) - - with pytest.raises(NotImplementedError, match="multi-axis mRoPE"): - recipe._forward_backward_step( - idx=0, - batch={ - "input_ids": torch.tensor([[1, 2]]), - "labels": torch.tensor([[2, -100]]), - "position_ids": torch.zeros((3, 1, 2), dtype=torch.long), - "qkv_format": "thd", - }, - loss_buffer=[], - num_label_tokens=1, - num_batches=1, - ) - - -def _build_pp_recipe_for_optim_step(num_label_tokens_in_batch: int): - """Shared setup for _run_train_optim_step tests with pp_enabled=True.""" - recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.dist_env = SimpleNamespace(device="cpu", rank=0, is_main=False) - # No "pp" in dim_names -> src_rank = mesh.reshape(-1)[-1].item(). With rank != src_rank - # and is_main=False, neither distributed send nor recv branch fires. - recipe.device_mesh = SimpleNamespace(mesh=torch.tensor([1]), mesh_dim_names=("dp",)) - recipe.moe_mesh = None - recipe.loss_fn = object() - recipe.model_parts = [_TensorModel()] - recipe.pp_enabled = True - recipe.pp = SimpleNamespace(pp_batch_size=2, pp_microbatch_size=1) - recipe.optimizer = [_DummyOptimizer()] - recipe.step_scheduler = SimpleNamespace(step=0, epoch=0, is_remote_logging_step=False) - recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=lambda: None) - recipe.cfg = _Cfg(fp8=None) - recipe.lr_scheduler = None - recipe.timestamp = 0.0 - recipe.distributed_config = None - recipe._dp_allreduce = lambda tensor, include_cp=False: tensor - recipe._get_dp_group_size = lambda include_cp=True: 1 - recipe._get_cp_group_size = lambda: 1 - - # Build a batch whose (labels != -100).sum() == num_label_tokens_in_batch. - seq = [1] * num_label_tokens_in_batch + [-100] * (4 - num_label_tokens_in_batch) +def test_train_step_logs_joint_drafter_only_on_first_engine_loss_call(monkeypatch): + recipe = _build_engine_recipe_for_optim_step() + recipe.step_scheduler.is_remote_logging_step = True batches = [ - { - "labels": torch.tensor([seq]), - "input_ids": torch.tensor([[1, 2, 3, 4]]), - } + {"labels": torch.tensor([[1, -100, 2]]), "input_ids": torch.tensor([[1, 2, 3]])}, + {"labels": torch.tensor([[-100, 3, 4]]), "input_ids": torch.tensor([[4, 5, 6]])}, ] - return recipe, batches + recipe._compute_vlm_loss = MagicMock(side_effect=[torch.tensor(1.0), torch.tensor(2.0)]) + def forward_backward(datums, loss_fn): + for datum in datums: + loss_fn(object(), datum.loss_fn_inputs) + return torch.tensor(0.25), [] -def _patch_pp_optim_step_dependencies(monkeypatch): + recipe.engine.forward_backward.side_effect = forward_backward monkeypatch.setattr( "nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", - lambda **kwargs: 0.0, - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_grad_accumulation", - lambda model_parts, pp_enabled: None, - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_final_backward", - lambda model_parts, pp_enabled: None, - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_after_first_microbatch", - lambda: None, + MagicMock(return_value=0.0), ) + recipe._run_train_optim_step(batches) -@pytest.mark.cuda(False) -def test_run_train_step_clears_first_microbatch_after_first_batch(monkeypatch): - recipe, _ = _build_pp_recipe_for_optim_step(num_label_tokens_in_batch=2) - batches = [ - { - "labels": torch.tensor([[1, -100, 2, -100]]), - "input_ids": torch.tensor([[1, 2, 3, 4]]), - }, - { - "labels": torch.tensor([[-100, 3, -100, 4]]), - "input_ids": torch.tensor([[5, 6, 7, 8]]), - }, - ] - events = [] + assert recipe._compute_vlm_loss.call_count == 2 + first_call, second_call = recipe._compute_vlm_loss.call_args_list + assert first_call.kwargs["log_drafter"] is True + assert first_call.kwargs["log_denominator"] == 4 + assert second_call.kwargs["log_drafter"] is False + assert second_call.kwargs["log_denominator"] == 4 - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_grad_accumulation", - lambda model_parts, pp_enabled: events.append("prepare"), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_for_final_backward", - lambda model_parts, pp_enabled: events.append("final"), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.prepare_after_first_microbatch", - lambda: events.append("after_first"), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", - lambda **kwargs: 0.0, - ) - def fake_forward_backward_step(idx, batch, loss_buffer, num_label_tokens, num_batches): - events.append(f"forward_{idx}") - loss_buffer.append(torch.tensor(1.0)) - - recipe._forward_backward_step = fake_forward_backward_step +@pytest.mark.cuda(False) +def test_run_train_step_uses_engine_for_empty_supervision(monkeypatch): + recipe = _build_engine_recipe_for_optim_step() + recipe.engine.forward_backward.return_value = (torch.tensor(0.0), []) + finalizer = MagicMock(return_value=0.0) + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", finalizer) - recipe._run_train_optim_step(batches, max_grad_norm=1.0) + batch = {"labels": torch.full((1, 4), -100), "input_ids": torch.arange(4).reshape(1, 4)} + metrics = recipe._run_train_optim_step([batch]) - assert events == ["prepare", "forward_0", "after_first", "final", "forward_1"] - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(1.0) + recipe.engine.forward_backward.assert_called_once() + assert metrics.metrics["loss"] == 0.0 + assert metrics.metrics["num_label_tokens"] == 0 + assert recipe.optimizer[0].step_called -@pytest.mark.cuda(False) -def test_run_train_step_pp_zero_label_tokens_no_nan(monkeypatch): - """Regression for PR #1985: PP reporting loss must be 0.0 (not NaN) when num_label_tokens=0. +def test_make_engine_datum_filters_raw_media_off_first_pipeline_stage(): + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.pp_enabled = True + recipe.pp = SimpleNamespace(info=SimpleNamespace(has_first_stage=False)) + recipe.loss_fn = object() + media_chunks = {"pixel_values": [torch.ones(1, 2)]} + batch = { + "input_ids": torch.tensor([[1, 2]]), + "labels": torch.tensor([[2, -100]]), + "pixel_values": torch.ones(1, 3, 4, 4), + "image_grid_thw": torch.ones(1, 3, dtype=torch.long), + VLM_PP_MEDIA_KEY: media_chunks, + } - With pipeline parallelism enabled, _run_train_optim_step divides reporting_loss by - num_label_tokens. If every label in the batch is the ignore_index (-100), the divisor - is zero and the reported metric would be NaN without the guard at finetune.py:1136. - """ - recipe, batches = _build_pp_recipe_for_optim_step(num_label_tokens_in_batch=0) + datum = recipe._make_engine_datum(batch) - def fake_forward_backward_step(idx, batch, loss_buffer, num_label_tokens, num_batches): - # Mirror the PP path: append a finite per-microbatch sum loss. With the guard, - # this must still yield reporting_loss == 0.0. - loss_buffer.append(torch.tensor(5.0)) + assert datum.model_inputs["input_ids"] is batch["input_ids"] + assert "pixel_values" not in datum.model_inputs + assert "image_grid_thw" not in datum.model_inputs + assert VLM_PP_MEDIA_KEY not in datum.model_inputs - recipe._forward_backward_step = fake_forward_backward_step - _patch_pp_optim_step_dependencies(monkeypatch) + recipe.pp.info.has_first_stage = True + first_stage_datum = recipe._make_engine_datum(batch) + assert first_stage_datum.model_inputs[VLM_PP_MEDIA_KEY] is media_chunks - metrics = recipe._run_train_optim_step(batches, max_grad_norm=1.0) - assert isinstance(metrics, MetricsSample) - assert metrics.metrics["num_label_tokens"] == 0 - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.5) - loss = metrics.metrics["loss"] - assert loss == loss, f"reporting loss must not be NaN, got {loss}" - assert loss == 0.0, f"reporting loss must be 0.0 when num_label_tokens=0, got {loss}" +def test_engine_context_stages_pipeline_media_and_cleans_up(monkeypatch): + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.pp_enabled = True + model = _TensorModel() + recipe.model_parts = [model] + recipe.pp = SimpleNamespace( + info=SimpleNamespace( + has_first_stage=True, + schedule=None, + stages=[SimpleNamespace(is_first=True)], + ) + ) + recipe._cp_vision_frame_sharding_context = nullcontext + input_ids = torch.tensor([[1, 2]]) + model_inputs = { + "input_ids": input_ids, + VLM_PP_MEDIA_KEY: {"pixel_values": [torch.ones(1, 2)]}, + } + with recipe._engine_context(model_inputs): + assert model._vlm_pixel_values_chunks is not None + assert VLM_PP_MEDIA_KEY not in model_inputs -@pytest.mark.cuda(False) -def test_run_train_step_pp_nonzero_label_tokens_divides(monkeypatch): - """PP reporting loss is the summed microbatch loss divided by num_label_tokens.""" - recipe, batches = _build_pp_recipe_for_optim_step(num_label_tokens_in_batch=4) + assert model._vlm_pixel_values_chunks is None + assert model._vlm_chunk_idx is None - def fake_forward_backward_step(idx, batch, loss_buffer, num_label_tokens, num_batches): - loss_buffer.append(torch.tensor(8.0)) - recipe._forward_backward_step = fake_forward_backward_step - _patch_pp_optim_step_dependencies(monkeypatch) +def test_engine_pipeline_loss_reuses_configured_loss_and_thd_metadata(): + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.pp_enabled = True + recipe.pipeline_loss_fn = MagicMock(return_value=torch.tensor(3.0)) + output = torch.randn(1, 2, 4) + labels = torch.tensor([[1, 2]]) + cu_seqlens = torch.tensor([0, 2], dtype=torch.int32) - metrics = recipe._run_train_optim_step(batches, max_grad_norm=1.0) + loss = recipe._engine_loss_fn(output, {"labels": labels, "cu_seqlens": cu_seqlens}) - assert metrics.metrics["num_label_tokens"] == 4 - assert metrics.metrics["loss"] == pytest.approx(8.0 / 4) - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(2.0) + assert loss.item() == 3.0 + assert recipe.pipeline_loss_fn.cu_seqlens is cu_seqlens + recipe.pipeline_loss_fn.assert_called_once_with(output, labels) # ----------------------------------------------------------------------------- @@ -1515,655 +1396,6 @@ def __init__(self): ) -# ----------------------------------------------------------------------------- -# PP Logic tests for _forward_backward_step -# ----------------------------------------------------------------------------- - - -class _MockPPInfo: - """Mock PP info structure.""" - - def __init__(self, has_first_stage=True, has_last_stage=True, n_microbatches=2, add_losses=True): - self.has_first_stage = has_first_stage - self.has_last_stage = has_last_stage - self._n_microbatches = n_microbatches - self._add_losses = add_losses - - # Create a schedule mock that adds losses when called - self.schedule = MagicMock() - self.schedule._n_microbatches = n_microbatches - - def step_side_effect(*args, **kwargs): - if self._add_losses and kwargs.get("losses") is not None: - # Add mock losses for each microbatch - for _ in range(n_microbatches): - kwargs["losses"].append(torch.tensor(0.5)) - - self.schedule.step = MagicMock(side_effect=step_side_effect) - - -class _MockAutoPipeline: - """Mock AutoPipeline for PP testing.""" - - def __init__(self, has_first_stage=True, has_last_stage=True, n_microbatches=2, add_losses=True): - self._info = _MockPPInfo(has_first_stage, has_last_stage, n_microbatches, add_losses) - self.info = self._info - self.step_batches = [] - - def update_seq_len(self, seq_len: int) -> None: - # Dynamic seq-len hook is a no-op in tests; AutoPipeline exposes this for - # variable-length VLM batches. - return None - - def step(self, model_input, *, target=None, losses=None, **kwargs): - """Record and forward an AutoPipeline step. - - Args: - model_input: Tensor of shape [batch, ...] containing the first - pipeline stage's input. - target: Optional tensor of shape [batch, sequence] containing loss - targets. - losses: Optional mutable list populated with scalar loss tensors. - **kwargs: Keyword schedule inputs. Tensor values have arbitrary - model-defined layouts. - - Returns: - The value returned by the schedule mock. - """ - self.step_batches.append(dict(kwargs)) - schedule_args = (model_input,) if self.info.has_first_stage else () - return self.info.schedule.step(*schedule_args, target=target, losses=losses, **kwargs) - - -def _create_pp_recipe(model=None): - """Helper to create a PP recipe bypassing BaseRecipe tracking.""" - if model is None: - model = _TensorModel() - recipe = object.__new__(FinetuneRecipeForVLM) - # Initialize __dict__ directly to bypass BaseRecipe.__setattr__ tracking - recipe.__dict__["__state_tracked"] = set() - recipe.__dict__["dist_env"] = SimpleNamespace(device="cpu") - recipe.__dict__["device_mesh"] = None - recipe.__dict__["moe_mesh"] = None - recipe.__dict__["pp_enabled"] = True - recipe.__dict__["loss_fn"] = MagicMock() - recipe.__dict__["distributed_config"] = None - recipe.__dict__["cp_vision_frame_sharding"] = CpVisionFrameShardingConfig(enabled=True) - recipe.__dict__["model_parts"] = [model] - recipe.__dict__["_get_dp_group_size"] = lambda include_cp=True: 1 - return recipe - - -def _prepare_pp_vlm_batch(batch, n_microbatches=2): - return prepare_vlm_media_for_pp( - batch, - batch_size=batch["input_ids"].shape[0], - n_microbatches=n_microbatches, - ) - - -class TestForwardBackwardStepPP: - """Tests for _forward_backward_step with pipeline parallelism enabled.""" - - @pytest.fixture - def pp_recipe(self): - """Create a recipe configured for PP testing.""" - return _create_pp_recipe() - - def test_pp_skips_validation_forward(self, pp_recipe, monkeypatch): - """Test that PP mode skips forward pass during validation.""" - pp_recipe.pp = _MockAutoPipeline() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch = { - "labels": torch.tensor([[1, 2]]), - "input_ids": torch.tensor([[1, 2]]), - } - loss_buffer = [] - - # Should return early without error - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=2, - num_batches=1, - is_train=False, # Validation mode - ) - - # Loss buffer should be empty (no forward pass) - assert len(loss_buffer) == 0 - - def test_pp_vlm_chunking_equal_images_and_batch(self, pp_recipe, monkeypatch): - """Test VLM pixel_values chunking when n_images == batch_size.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - # image_grid_hws: 4 images, each with different patch counts - image_grid_hws = torch.tensor([[2, 2], [3, 3], [2, 3], [4, 4]]) # patch counts: 4, 9, 6, 16 - total_patches = 4 + 9 + 6 + 16 # = 35 - pixel_values = torch.randn(total_patches, 3, 14, 14) - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_grid_hws": image_grid_hws, - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - captured_chunks = {} - - def step_side_effect(*args, **kwargs): - model = pp_recipe.model_parts[0] - captured_chunks["pixel_values"] = [chunk.clone() for chunk in model._vlm_pixel_values_chunks] - captured_chunks["image_grid"] = [chunk.clone() for chunk in model._vlm_image_grid_hws_chunks] - captured_chunks["chunk_idx"] = model._vlm_chunk_idx - for _ in range(2): - kwargs["losses"].append(torch.tensor(0.5)) - - pp_recipe.pp.info.schedule.step = MagicMock(side_effect=step_side_effect) - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - # Verify chunking happened correctly - model = pp_recipe.model_parts[0] - assert captured_chunks["chunk_idx"] == 0 - assert torch.equal(captured_chunks["pixel_values"][0], pixel_values[:13]) - assert torch.equal(captured_chunks["pixel_values"][1], pixel_values[13:]) - assert torch.equal(captured_chunks["image_grid"][0], image_grid_hws[:2]) - assert torch.equal(captured_chunks["image_grid"][1], image_grid_hws[2:]) - assert model._vlm_pixel_values_chunks is None # Cleared after step - assert model._vlm_image_grid_hws_chunks is None - assert model._vlm_chunk_idx is None - - # Verify schedule.step was called - pp_recipe.pp.info.schedule.step.assert_called_once() - assert pp_recipe.pp.step_batches == [{}] - - # Verify loss was computed - assert len(loss_buffer) == 1 - - def test_pp_step_receives_remaining_kwargs(self, pp_recipe, monkeypatch): - """The recipe passes remaining model kwargs through AutoPipeline.step.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - position_ids = torch.zeros(3, 2, 8, dtype=torch.long) - batch = { - "labels": torch.ones(2, 8, dtype=torch.long), - "input_ids": torch.ones(2, 8, dtype=torch.long), - "position_ids": position_ids, - } - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=[], - num_label_tokens=16, - num_batches=1, - is_train=True, - ) - - assert len(pp_recipe.pp.step_batches) == 1 - assert pp_recipe.pp.step_batches[0].keys() == {"position_ids"} - assert torch.equal(pp_recipe.pp.step_batches[0]["position_ids"], position_ids) - - def test_pp_vlm_chunking_videos_uses_video_grid_and_counts(self, pp_recipe, monkeypatch): - """Video tensors are chunked by per-sample video counts before schedule.step.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - video_grid_thw = torch.tensor([[1, 2, 2], [1, 3, 3], [1, 2, 3], [1, 4, 4]]) - pixel_values_videos = torch.randn(int(video_grid_thw.prod(dim=1).sum().item()), 64) - n_videos_per_sample = torch.tensor([1, 0, 2, 1]) - - def step_side_effect(*args, **kwargs): - model = pp_recipe.model_parts[0] - assert "pixel_values_videos" not in kwargs - assert "video_grid_thw" not in kwargs - assert len(model._vlm_pixel_values_videos_chunks) == 2 - assert len(model._vlm_video_grid_thw_chunks) == 2 - assert model._vlm_video_grid_thw_chunks[0].shape[0] == 1 - assert model._vlm_video_grid_thw_chunks[1].shape[0] == 3 - assert model._vlm_pixel_values_videos_chunks[0].shape[0] == 4 - assert model._vlm_pixel_values_videos_chunks[1].shape[0] == 9 + 6 + 16 - for _ in range(2): - kwargs["losses"].append(torch.tensor(0.5)) - - pp_recipe.pp.info.schedule.step.side_effect = step_side_effect - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values_videos": pixel_values_videos, - "video_grid_thw": video_grid_thw, - "n_videos_per_sample": n_videos_per_sample, - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - model = pp_recipe.model_parts[0] - assert model._vlm_pixel_values_videos_chunks is None - assert model._vlm_video_grid_thw_chunks is None - assert model._vlm_chunk_idx is None - assert len(loss_buffer) == 1 - - def test_pp_vlm_chunking_image_and_video_mixed(self, pp_recipe, monkeypatch): - """When a batch carries both images and videos, both streams chunk independently - but share a single _vlm_chunk_idx initialized once at 0; both clean up to None.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - - # n_images_per_sample=[2,0,1,0]: mb0 (samples 0..1) covers images 0..1; mb1 covers image 2. - image_grid_thw = torch.tensor([[1, 2, 2], [1, 3, 3], [1, 2, 3]]) # patch counts: 4, 9, 6 - pixel_values = torch.randn(int(image_grid_thw.prod(dim=1).sum().item()), 32) - n_images_per_sample = torch.tensor([2, 0, 1, 0]) - - # n_videos_per_sample=[1,0,2,1]: mb0 covers video 0; mb1 covers videos 1..3. - video_grid_thw = torch.tensor([[1, 2, 2], [1, 3, 3], [1, 2, 3], [1, 4, 4]]) # patch counts: 4, 9, 6, 16 - pixel_values_videos = torch.randn(int(video_grid_thw.prod(dim=1).sum().item()), 64) - n_videos_per_sample = torch.tensor([1, 0, 2, 1]) - - def step_side_effect(*args, **kwargs): - model = pp_recipe.model_parts[0] - - # Both modalities are popped before schedule.step so the schedule never - # tries to chunk the misaligned multimodal tensors along dim 0. - assert "pixel_values" not in kwargs - assert "image_grid_hws" not in kwargs - assert "image_grid_thw" not in kwargs - assert "pixel_values_videos" not in kwargs - assert "video_grid_thw" not in kwargs - - assert len(model._vlm_pixel_values_chunks) == 2 - assert len(model._vlm_image_grid_hws_chunks) == 2 - assert model._vlm_image_grid_hws_chunks[0].shape[0] == 2 - assert model._vlm_image_grid_hws_chunks[1].shape[0] == 1 - assert model._vlm_pixel_values_chunks[0].shape[0] == 4 + 9 - assert model._vlm_pixel_values_chunks[1].shape[0] == 6 - - assert len(model._vlm_pixel_values_videos_chunks) == 2 - assert len(model._vlm_video_grid_thw_chunks) == 2 - assert model._vlm_video_grid_thw_chunks[0].shape[0] == 1 - assert model._vlm_video_grid_thw_chunks[1].shape[0] == 3 - assert model._vlm_pixel_values_videos_chunks[0].shape[0] == 4 - assert model._vlm_pixel_values_videos_chunks[1].shape[0] == 9 + 6 + 16 - - # Single shared cursor: image-branch sets it to 0 first, video branch resets to 0 again. - assert model._vlm_chunk_idx == 0 - - for _ in range(2): - kwargs["losses"].append(torch.tensor(0.5)) - - pp_recipe.pp.info.schedule.step.side_effect = step_side_effect - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_grid_thw": image_grid_thw, - "n_images_per_sample": n_images_per_sample, - "pixel_values_videos": pixel_values_videos, - "video_grid_thw": video_grid_thw, - "n_videos_per_sample": n_videos_per_sample, - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - model = pp_recipe.model_parts[0] - assert model._vlm_pixel_values_chunks is None - assert model._vlm_image_grid_hws_chunks is None - assert model._vlm_pixel_values_videos_chunks is None - assert model._vlm_video_grid_thw_chunks is None - assert model._vlm_chunk_idx is None - assert len(loss_buffer) == 1 - - def test_pp_vlm_chunking_with_image_grid_thw(self, pp_recipe, monkeypatch): - """Test VLM pixel_values chunking with image_grid_thw (3D grid) instead of image_grid_hws.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - # image_grid_thw: 4 images with T, H, W dimensions (uses .prod(dim=1) for patch counts) - image_grid_thw = torch.tensor([[1, 2, 2], [1, 3, 3], [1, 2, 3], [1, 4, 4]]) # patch counts: 4, 9, 6, 16 - total_patches = 4 + 9 + 6 + 16 # = 35 - pixel_values = torch.randn(total_patches, 3, 14, 14) - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_grid_thw": image_grid_thw, # Using thw instead of hws - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - # Verify chunking happened correctly - model = pp_recipe.model_parts[0] - assert model._vlm_pixel_values_chunks is None # Cleared after step - assert model._vlm_image_grid_hws_chunks is None - assert model._vlm_chunk_idx is None - - # Verify schedule.step was called - pp_recipe.pp.info.schedule.step.assert_called_once() - - # Verify loss was computed - assert len(loss_buffer) == 1 - - def test_pp_vlm_chunking_qwen35_ep4_pp2_local_batch_images(self, pp_recipe, monkeypatch): - """Qwen3.5 35B EP4/PP2-style local batch keeps proper image chunks during schedule.step.""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - image_grid_thw = torch.tensor([[1, 2, 2], [1, 3, 3]]) - patch_counts = image_grid_thw.prod(dim=1) - pixel_values = torch.arange(int(patch_counts.sum()) * 4, dtype=torch.float32).reshape(-1, 4) - batch = { - "labels": torch.randint(0, 100, (2, 10)), - "input_ids": torch.randint(0, 100, (2, 10)), - "pixel_values": pixel_values, - "image_grid_thw": image_grid_thw, - "n_images_per_sample": torch.tensor([1, 1]), - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - captured_chunks = {} - - def step_side_effect(*args, **kwargs): - model = pp_recipe.model_parts[0] - captured_chunks["pixel_values"] = [chunk.clone() for chunk in model._vlm_pixel_values_chunks] - captured_chunks["image_grid"] = [chunk.clone() for chunk in model._vlm_image_grid_hws_chunks] - for _ in range(2): - kwargs["losses"].append(torch.tensor(0.5)) - - pp_recipe.pp.info.schedule.step = MagicMock(side_effect=step_side_effect) - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=20, - num_batches=1, - is_train=True, - ) - - split_at = int(patch_counts[0].item()) - assert torch.equal(captured_chunks["pixel_values"][0], pixel_values[:split_at]) - assert torch.equal(captured_chunks["pixel_values"][1], pixel_values[split_at:]) - assert torch.equal(captured_chunks["image_grid"][0], image_grid_thw[:1]) - assert torch.equal(captured_chunks["image_grid"][1], image_grid_thw[1:]) - assert pp_recipe.model_parts[0]._vlm_pixel_values_chunks is None - assert len(loss_buffer) == 1 - - def test_pp_vlm_chunking_mismatched_images_raises(self): - """When media cannot be aligned to samples, VLM PP data prep raises.""" - batch_size = 4 - image_grid_hws = torch.tensor([[2, 2], [3, 3]]) - total_patches = 4 + 9 - pixel_values = torch.randn(total_patches, 3, 14, 14) - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_grid_hws": image_grid_hws, - } - - with pytest.raises(ValueError, match="VLM PP chunking cannot align"): - _prepare_pp_vlm_batch(batch) - - def test_pp_vlm_chunking_with_image_sizes(self, pp_recipe, monkeypatch): - """Test VLM pixel_values chunking with image_sizes fallback (e.g., Mistral4-style).""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - # image_sizes: [N_images, 2] — no image_grid_hws or image_grid_thw - image_sizes = torch.tensor([[224, 224], [224, 224], [224, 224], [224, 224]]) - # 4D pixel_values: [N_images, C, H, W] - pixel_values = torch.randn(batch_size, 3, 224, 224) - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_sizes": image_sizes, - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - model = pp_recipe.model_parts[0] - assert model._vlm_pixel_values_chunks is None # Cleared after step - assert model._vlm_image_grid_hws_chunks is None - assert model._vlm_chunk_idx is None - pp_recipe.pp.info.schedule.step.assert_called_once() - assert len(loss_buffer) == 1 - - def test_pp_vlm_chunking_4d_pixel_values(self, pp_recipe, monkeypatch): - """Test VLM pixel_values chunking when pixel_values is 4D [N, C, H, W].""" - pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2) - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch_size = 4 - image_grid_hws = torch.tensor([[224, 224], [224, 224], [224, 224], [224, 224]]) - # 4D pixel_values — triggers the new dim==4 chunking path - pixel_values = torch.randn(batch_size, 3, 224, 224) - - batch = { - "labels": torch.randint(0, 100, (batch_size, 10)), - "input_ids": torch.randint(0, 100, (batch_size, 10)), - "pixel_values": pixel_values, - "image_grid_hws": image_grid_hws, - } - _prepare_pp_vlm_batch(batch) - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=40, - num_batches=1, - is_train=True, - ) - - model = pp_recipe.model_parts[0] - assert model._vlm_pixel_values_chunks is None # Cleared after step - assert model._vlm_image_grid_hws_chunks is None - assert model._vlm_chunk_idx is None - pp_recipe.pp.info.schedule.step.assert_called_once() - assert len(loss_buffer) == 1 - - def test_pp_last_stage_computes_loss(self, pp_recipe, monkeypatch): - """Test that last stage computes and buffers loss.""" - - def mock_schedule_step(*args, **kwargs): - # Simulate loss computation on last stage - if kwargs.get("losses") is not None: - kwargs["losses"].append(torch.tensor(0.5)) - kwargs["losses"].append(torch.tensor(0.3)) - - pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2, add_losses=False) - pp.info.schedule.step = MagicMock(side_effect=mock_schedule_step) - pp_recipe.pp = pp - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch = { - "labels": torch.tensor([[1, 2]]), - "input_ids": torch.tensor([[1, 2]]), - } - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=2, - num_batches=1, - is_train=True, - ) - - # Loss should be sum of microbatch losses - assert len(loss_buffer) == 1 - assert torch.isclose(loss_buffer[0], torch.tensor(0.8)) - - def test_pp_non_last_stage_returns_zero_loss(self, pp_recipe, monkeypatch): - """Test that non-last stage returns zero loss.""" - pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=False, n_microbatches=2) - pp_recipe.pp = pp - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch = { - "labels": torch.tensor([[1, 2]]), - "input_ids": torch.tensor([[1, 2]]), - } - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=2, - num_batches=1, - is_train=True, - ) - - assert len(loss_buffer) == 1 - assert loss_buffer[0].item() == 0.0 - - def test_pp_non_first_stage_skips_input_ids(self, pp_recipe, monkeypatch): - """Test that non-first stage doesn't pass input_ids to schedule.""" - step_calls = [] - - def mock_schedule_step(*args, **kwargs): - step_calls.append((args, kwargs)) - # Add losses so torch.stack doesn't fail - if kwargs.get("losses") is not None: - kwargs["losses"].append(torch.tensor(0.5)) - - pp = _MockAutoPipeline(has_first_stage=False, has_last_stage=True, n_microbatches=2, add_losses=False) - pp.info.schedule.step = MagicMock(side_effect=mock_schedule_step) - pp_recipe.pp = pp - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch = { - "labels": torch.tensor([[1, 2]]), - "input_ids": torch.tensor([[1, 2]]), - } - loss_buffer = [] - - pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=2, - num_batches=1, - is_train=True, - ) - - # Should be called without positional args (no input_ids) - assert len(step_calls) == 1 - args, kwargs = step_calls[0] - assert len(args) == 0 # No positional args - assert "target" in kwargs - - # ----------------------------------------------------------------------------- # FinetuneRecipeForVLM.setup() tests # ----------------------------------------------------------------------------- @@ -2220,389 +1452,6 @@ def get(self, key, default=None): assert max_grad_norm == 0.5 -# ----------------------------------------------------------------------------- -# _forward_backward_step non-PP tests (FusedLinearCE path) -# ----------------------------------------------------------------------------- - - -class _ModelOutput: - """Model output that supports both attribute access and 'in' operator.""" - - def __init__(self, logits, hidden_states=None): - self.logits = logits - self.hidden_states = hidden_states - - def __contains__(self, key): - return hasattr(self, key) and getattr(self, key) is not None - - -class _ModelWithHiddenStates(torch.nn.Module): - """Model that outputs hidden states for FusedLinearCE testing.""" - - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(10, 10) - self.lm_head = torch.nn.Linear(10, 50) - - def forward(self, logits_to_keep=None, **kwargs): - hidden = torch.randn(2, 5, 10) - return _ModelOutput( - logits=self.lm_head(hidden), - hidden_states=[hidden], - ) - - def get_output_embeddings(self): - return self.lm_head - - -def _create_non_pp_recipe(model, device="cpu"): - """Helper to create a non-PP recipe bypassing BaseRecipe tracking.""" - recipe = object.__new__(FinetuneRecipeForVLM) - # Initialize __dict__ directly to bypass BaseRecipe.__setattr__ tracking - recipe.__dict__["__state_tracked"] = set() - recipe.__dict__["dist_env"] = SimpleNamespace(device=device) - recipe.__dict__["device_mesh"] = None - recipe.__dict__["moe_mesh"] = None - recipe.__dict__["pp_enabled"] = False - recipe.__dict__["distributed_config"] = None - recipe.__dict__["cp_vision_frame_sharding"] = CpVisionFrameShardingConfig(enabled=True) - recipe.__dict__["model_parts"] = [model] - recipe.__dict__["_get_dp_group_size"] = lambda include_cp=True: 1 - # ``is_remote_logging_step`` is read by ``_forward_backward_step`` to - # gate the joint-drafter loss-log line; default False so non-drafter - # test paths don't trip on the new attribute. - recipe.__dict__["step_scheduler"] = SimpleNamespace(is_remote_logging_step=False) - return recipe - - -class _DummyCPSubMesh: - def __init__(self, size: int): - self._size = size - self._group = object() - - def size(self) -> int: - return self._size - - def get_group(self): - return self._group - - -class _DummyCPDeviceMesh(dict): - def __init__(self, cp_size: int): - super().__init__() - self["cp"] = _DummyCPSubMesh(cp_size) - self.mesh_dim_names = ["cp"] - - -class _CPPreEmbedModel(torch.nn.Module): - """Sunk model: sharder-only CP hook (consumes nothing; the forward embeds + - shards per microbatch). Records that the hook was invoked.""" - - def __init__(self): - super().__init__() - self.scale = torch.nn.Parameter(torch.tensor(1.0)) - self.hook_calls = [] - - def prepare_model_inputs_for_cp(self, batch, *, num_chunks=1): - self.hook_calls.append(set(batch)) - return {} - - def forward(self, **kwargs): - raise AssertionError("forward should not run: ContextParallelSharder.shard raises first") - - -class _CPPreEmbedStop(RuntimeError): - pass - - -class TestForwardBackwardStepNonPP: - """Tests for _forward_backward_step without pipeline parallelism.""" - - def test_non_pp_cp_invokes_sharder_only_hook_and_keeps_inputs(self, monkeypatch): - # Sunk contract: the non-PP CP path invokes the sharder-only hook, which - # consumes nothing, so input_ids / pixel_values / mm_token_type_ids all - # reach ContextParallelSharder.shard intact (the model embeds + shards them in - # its own forward, not here). - model = _CPPreEmbedModel() - non_pp_recipe = _create_non_pp_recipe(model) - non_pp_recipe.__dict__["device_mesh"] = _DummyCPDeviceMesh(cp_size=2) - - mm_token_type_ids = torch.tensor([[1, 1, 0, 0]]) - - def _capture_cp_batch(sharder, batch): - """Validate the global model-input mapping before CP transport. - - Args: - sharder: Sharder configured by the VLM recipe. - batch: Mutable model-input mapping whose tensor values have - global batch and sequence extents. - """ - del sharder - assert "input_ids" in batch - assert "pixel_values" in batch - assert "mm_token_type_ids" in batch - torch.testing.assert_close(batch["mm_token_type_ids"], mm_token_type_ids) - assert "inputs_embeds" not in batch - raise _CPPreEmbedStop - - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.ContextParallelSharder.shard", - _capture_cp_batch, - ) - - batch = { - "labels": torch.randint(0, 50, (1, 4)), - "input_ids": torch.randint(0, 100, (1, 4)), - "pixel_values": torch.randn(1, 3, 8, 8), - "image_position_ids": torch.zeros(1, 1, 2, dtype=torch.long), - "mm_token_type_ids": mm_token_type_ids, - } - - with pytest.raises(_CPPreEmbedStop): - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=[], - num_label_tokens=4, - num_batches=1, - is_train=False, - ) - assert len(model.hook_calls) == 1 # the CP hook ran before sharding - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="FusedLinearCE requires CUDA") - def test_non_pp_with_fused_linear_ce(self, monkeypatch): - """Test non-PP path with FusedLinearCrossEntropy.""" - from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy - - # Model output class that supports both attribute access and 'in' operator - class ModelOutput: - def __init__(self, logits, hidden_states): - self.logits = logits - self.hidden_states = hidden_states - - def __contains__(self, key): - return hasattr(self, key) - - # Create CUDA model for FusedLinearCE - must use bf16/fp16 for backward - class CudaModelWithHiddenStates(torch.nn.Module): - def __init__(self): - super().__init__() - # Keep lm_head in bfloat16 to match hidden states - self.lm_head = torch.nn.Linear(10, 50) - - def forward(self, logits_to_keep=None, **kwargs): - # FusedLinearCE requires bf16/fp16 hidden states - hidden = torch.randn(2, 5, 10, device="cuda", dtype=torch.bfloat16, requires_grad=True) - # lm_head is already bfloat16, so no conversion needed - return ModelOutput( - logits=self.lm_head(hidden), - hidden_states=[hidden], - ) - - def get_output_embeddings(self): - return self.lm_head - - # Create model and convert entirely to bfloat16 - model = CudaModelWithHiddenStates().cuda().bfloat16() - non_pp_recipe = _create_non_pp_recipe(model, device="cuda") - non_pp_recipe.__dict__["loss_fn"] = FusedLinearCrossEntropy() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.get_sync_ctx", - lambda model, is_last, defer_fsdp_grad_sync=True: nullcontext(), - ) - - batch = { - "labels": torch.randint(0, 50, (2, 5)), - "input_ids": torch.randint(0, 100, (2, 5)), - } - loss_buffer = [] - - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=10, - num_batches=1, - is_train=True, - ) - - assert len(loss_buffer) == 1 - assert isinstance(loss_buffer[0], torch.Tensor) - - def test_non_pp_fused_ce_requires_hidden_states(self, monkeypatch): - """Test that FusedLinearCE raises error when hidden_states not in output.""" - from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy - - # Model output class that supports 'in' operator but has no hidden_states - class ModelOutputNoHiddenStates: - def __init__(self, logits): - self.logits = logits - - def __contains__(self, key): - return hasattr(self, key) - - # Model that doesn't output hidden_states - class BadModel(torch.nn.Module): - def forward(self, logits_to_keep=None, **kwargs): - return ModelOutputNoHiddenStates(logits=torch.randn(2, 5, 50)) - - non_pp_recipe = _create_non_pp_recipe(BadModel()) - non_pp_recipe.__dict__["loss_fn"] = FusedLinearCrossEntropy() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.get_sync_ctx", - lambda model, is_last, defer_fsdp_grad_sync=True: nullcontext(), - ) - - batch = { - "labels": torch.randint(0, 50, (2, 5)), - "input_ids": torch.randint(0, 100, (2, 5)), - } - loss_buffer = [] - - with pytest.raises(ValueError, match="FusedLinearCrossEntropy requires the model to output hidden states"): - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=10, - num_batches=1, - is_train=True, - ) - - def test_non_pp_with_masked_ce(self, monkeypatch): - """Test non-PP path with MaskedCrossEntropy.""" - from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy - - class SimpleModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(10, 50) - - def forward(self, **kwargs): - # Create logits through a layer so gradients can flow - x = torch.randn(2, 5, 10, requires_grad=True) - logits = self.linear(x) - return _ModelOutput(logits=logits, hidden_states=None) - - non_pp_recipe = _create_non_pp_recipe(SimpleModel()) - non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.get_sync_ctx", - lambda model, is_last, defer_fsdp_grad_sync=True: nullcontext(), - ) - - batch = { - "labels": torch.randint(0, 50, (2, 5)), - "input_ids": torch.randint(0, 100, (2, 5)), - } - loss_buffer = [] - - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=10, - num_batches=1, - is_train=True, - ) - - assert len(loss_buffer) == 1 - assert isinstance(loss_buffer[0], torch.Tensor) - - def test_non_pp_validation_mode_no_backward(self, monkeypatch): - """Test that validation mode doesn't call backward.""" - from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy - - # Simple model for this test - class SimpleModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(10, 50) - - def forward(self, **kwargs): - return _ModelOutput(logits=torch.randn(2, 5, 50), hidden_states=None) - - non_pp_recipe = _create_non_pp_recipe(SimpleModel()) - non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - batch = { - "labels": torch.randint(0, 50, (2, 5)), - "input_ids": torch.randint(0, 100, (2, 5)), - } - loss_buffer = [] - - # Should complete without error and not call backward - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=10, - num_batches=1, - is_train=False, # Validation mode - ) - - assert len(loss_buffer) == 1 - - def test_non_pp_handles_dict_batch_values(self, monkeypatch): - """Test that nested dict values in batch are moved to device.""" - from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy - - class SimpleModel(torch.nn.Module): - def forward(self, **kwargs): - return _ModelOutput(logits=torch.randn(2, 5, 50), hidden_states=None) - - non_pp_recipe = _create_non_pp_recipe(SimpleModel()) - non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), - ) - - # Batch with nested dict (like attention_mask dict) - batch = { - "labels": torch.randint(0, 50, (2, 5)), - "input_ids": torch.randint(0, 100, (2, 5)), - "nested": { - "inner_tensor": torch.ones(2, 5), - "none_value": None, - }, - } - loss_buffer = [] - - # Should handle nested dict without error - non_pp_recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=10, - num_batches=1, - is_train=False, - ) - - assert len(loss_buffer) == 1 - - # ----------------------------------------------------------------------------- # build_optimizer returns correct type (diff coverage) # ----------------------------------------------------------------------------- @@ -2796,6 +1645,8 @@ def test_is_recipe_target_accepts_nemo_auto_and_rejects_others(self): def _patch_vlm_setup_minimals(monkeypatch, cp_size): """Patch heavy dependencies so FinetuneRecipeForVLM.setup() runs lightly.""" + from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy + monkeypatch.setattr( "nemo_automodel.recipes.vlm.finetune.initialize_distributed", lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0), @@ -2805,7 +1656,7 @@ def _patch_vlm_setup_minimals(monkeypatch, cp_size): monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.StatefulRNG", lambda *a, **k: "rng") monkeypatch.setattr( "nemo_automodel.recipes._typed_config.RecipeConfig.loss_fn", - property(lambda self: SimpleNamespace(build=lambda: "loss_fn")), + property(lambda self: SimpleNamespace(build=lambda: MaskedCrossEntropy(reduction="sum"))), ) def _stub_build_checkpoint_config(*a, **k): @@ -2922,6 +1773,101 @@ def _minimal_vlm_cfg( return ConfigNode(cfg) +def _patch_vlm_distributed_setup( + monkeypatch, + *, + pp_enabled: bool, + calculate_per_token_loss: bool = False, + scale_grads_in_schedule: bool = False, +): + mesh_context = SimpleNamespace( + pp_enabled=pp_enabled, + device_mesh=None, + moe_mesh=None, + cp_size=1, + pp_size=2 if pp_enabled else 1, + ) + pipeline_config = ( + SimpleNamespace( + scale_grads_in_schedule=scale_grads_in_schedule, + pp_batch_size=1, + pp_microbatch_size=1, + patch_stage_backward_maybe_with_nosync=False, + loss_fn=None, + ) + if pp_enabled + else None + ) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.create_distributed_setup_from_config", + lambda cfg, world_size: SimpleNamespace( + mesh_context=mesh_context, + strategy_config=SimpleNamespace(calculate_per_token_loss=calculate_per_token_loss), + pipeline_config=pipeline_config, + moe_parallel_config=None, + activation_checkpointing=False, + ), + ) + + +def test_vlm_setup_rejects_calculate_per_token_loss(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + _patch_vlm_distributed_setup(monkeypatch, pp_enabled=False, calculate_per_token_loss=True) + + trainer = FinetuneRecipeForVLM(cfg) + with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): + trainer.setup() + + +def test_vlm_setup_rejects_pipeline_schedule_gradient_scaling(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + _patch_vlm_distributed_setup(monkeypatch, pp_enabled=True, scale_grads_in_schedule=True) + + trainer = FinetuneRecipeForVLM(cfg) + with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): + trainer.setup() + + +@pytest.mark.parametrize("local_batch_size", [1, 2]) +def test_vlm_setup_supports_magi_pipeline_only_with_unit_local_batch(monkeypatch, local_batch_size): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False) + cfg.step_scheduler.local_batch_size = local_batch_size + cfg.distributed.pipeline = ConfigNode({"pp_microbatch_size": 1}) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + _patch_vlm_distributed_setup(monkeypatch, pp_enabled=True) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.setup_magi", + lambda *args, **kwargs: SimpleNamespace(enabled=True), + ) + + if local_batch_size == 2: + with pytest.raises(ValueError, match="Magi pipeline training requires"): + FinetuneRecipeForVLM(cfg).setup() + return + + model = DummyModel() + pipeline = object.__new__(AutoPipeline) + pipeline._info = SimpleNamespace( + model_parts=[model], + has_first_stage=True, + has_last_stage=False, + stages=[SimpleNamespace(is_first=True, is_last=False)], + schedule=MagicMock(), + ) + pipeline.scale_grads_in_schedule = False + pipeline.pp_batch_size = 1 + pipeline.pp_microbatch_size = 1 + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.build_model", lambda *args, **kwargs: pipeline) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert trainer.pp is pipeline + assert trainer.engine.pipeline is pipeline + + def test_vlm_setup_applies_prewarm_config(monkeypatch): """VLM setup should apply the typed prewarm config to its parallelized model parts.""" cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False, prewarm={"comm_groups": True}) @@ -2986,7 +1932,7 @@ def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch): assert trainer.engine is not None -def test_vlm_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): +def test_vlm_setup_rejects_cp_with_mtp(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=2, rope_fusion=True) _patch_vlm_setup_minimals(monkeypatch, cp_size=2) model = DummyModel() @@ -2994,12 +1940,11 @@ def test_vlm_setup_keeps_engine_disabled_for_cp_with_mtp(monkeypatch): monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.build_model", lambda *args, **kwargs: model) trainer = FinetuneRecipeForVLM(cfg) - trainer.setup() + with pytest.raises(NotImplementedError, match="MTP with context parallelism"): + trainer.setup() - assert trainer.engine is None - -def test_vlm_setup_keeps_engine_disabled_for_magi(monkeypatch): +def test_vlm_setup_builds_engine_for_magi(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) _patch_vlm_setup_minimals(monkeypatch, cp_size=1) monkeypatch.setattr( @@ -3010,7 +1955,7 @@ def test_vlm_setup_keeps_engine_disabled_for_magi(monkeypatch): trainer = FinetuneRecipeForVLM(cfg) trainer.setup() - assert trainer.engine is None + assert trainer.engine is not None def test_vlm_compute_loss_uses_final_thd_sequence_boundaries(monkeypatch): @@ -3062,16 +2007,18 @@ def test_vlm_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): assert cfg.model.backend.rope_fusion is True -def test_vlm_setup_keeps_engine_disabled_for_loss_without_sum_contract(monkeypatch): +def test_vlm_setup_rejects_loss_without_sum_contract(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + monkeypatch.setattr( + "nemo_automodel.recipes._typed_config.RecipeConfig.loss_fn", + property(lambda self: SimpleNamespace(build=lambda: SimpleNamespace(reduction="mean"))), + ) monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep", lambda _model: True) trainer = FinetuneRecipeForVLM(cfg) - trainer.setup() - - assert trainer.loss_fn == "loss_fn" - assert trainer.engine is None + with pytest.raises(ValueError, match="reduction='sum'"): + trainer.setup() def test_vlm_setup_builds_engine_for_eager_sum_loss(monkeypatch): @@ -3085,6 +2032,7 @@ def test_vlm_setup_builds_engine_for_eager_sum_loss(monkeypatch): assert isinstance(trainer.loss_fn, MaskedCrossEntropy) assert trainer.engine is not None + assert trainer.engine.microbatch_size == 1 def test_vlm_setup_does_not_change_storage_dtype_for_non_kd_recipe(monkeypatch): @@ -3227,6 +2175,95 @@ def test_n_videos_per_sample_packed(self): assert pv_chunks[0].shape[0] == 4 assert pv_chunks[1].shape[0] == 9 + 6 + 16 + def test_variable_resolution_list_chunks_by_sample_counts(self): + pixel_values = [torch.full((3, 8, 9 + index), index, dtype=torch.bfloat16) for index in range(5)] + image_grid = torch.tensor([[1, 2, 2], [1, 2, 3], [1, 3, 3], [1, 4, 2], [1, 2, 4]]) + n_images_per_sample = torch.tensor([2, 0, 1, 2]) + + pv_chunks, ig_chunks = chunk_vlm_media( + pixel_values, + image_grid, + batch_size=4, + n_microbatches=2, + n_images_per_sample=n_images_per_sample, + ) + + assert [[id(value) for value in chunk] for chunk in pv_chunks] == [ + [id(value) for value in pixel_values[:2]], + [id(value) for value in pixel_values[2:]], + ] + assert ig_chunks is not None + assert torch.equal(ig_chunks[0], image_grid[:2]) + assert torch.equal(ig_chunks[1], image_grid[2:]) + + @pytest.mark.parametrize( + ("counts", "grid", "match"), + [ + (torch.tensor([1, 1]), torch.ones(3, 3, dtype=torch.long), "length batch_size"), + (torch.tensor([1, -1, 3]), torch.ones(3, 3, dtype=torch.long), "non-negative"), + (torch.tensor([1, 1, 0]), torch.ones(3, 3, dtype=torch.long), "sum\\(n_images_per_sample\\)"), + (torch.tensor([1, 1, 1]), torch.ones(2, 3, dtype=torch.long), "image_grid.shape\\[0\\]"), + ], + ) + def test_variable_resolution_list_strictly_validates_alignment(self, counts, grid, match): + pixel_values = [torch.ones(3, 8, 8) for _ in range(3)] + + with pytest.raises(ValueError, match=match): + chunk_vlm_media( + pixel_values, + grid, + batch_size=3, + n_microbatches=2, + n_images_per_sample=counts, + ) + + def test_variable_resolution_list_rejects_non_tensor_values(self): + with pytest.raises(TypeError, match="list of tensors"): + chunk_vlm_media( + [torch.ones(3, 8, 8), "not-a-tensor"], + None, + batch_size=2, + n_microbatches=2, + ) + + def test_wrapper_chunks_variable_resolution_image_and_video_lists_without_grids(self): + images = [torch.ones(3, 8, 9 + index) for index in range(4)] + videos = [torch.ones(6, 10, 11 + index) for index in range(3)] + + def collate_fn(_examples): + return { + "input_ids": torch.ones(4, 8, dtype=torch.long), + "pixel_values": images, + "n_images_per_sample": torch.tensor([2, 0, 1, 1]), + "pixel_values_videos": videos, + "n_videos_per_sample": torch.tensor([0, 1, 2, 0]), + } + + prepared = wrap_vlm_collate_for_pp(collate_fn, n_microbatches=2)([{}] * 4) + media = prepared[VLM_PP_MEDIA_KEY] + + assert [[id(value) for value in chunk] for chunk in media["pixel_values"]] == [ + [id(value) for value in images[:2]], + [id(value) for value in images[2:]], + ] + assert [[id(value) for value in chunk] for chunk in media["pixel_values_videos"]] == [ + [id(value) for value in videos[:1]], + [id(value) for value in videos[1:]], + ] + assert "image_grid_hws" not in media + assert "video_grid_thw" not in media + + def test_wrapper_rejects_variable_resolution_list_count_mismatch(self): + def collate_fn(_examples): + return { + "input_ids": torch.ones(3, 8, dtype=torch.long), + "pixel_values": [torch.ones(3, 8, 8) for _ in range(3)], + "n_images_per_sample": torch.tensor([1, 0, 1]), + } + + with pytest.raises(ValueError, match="sum\\(n_images_per_sample\\)=2"): + wrap_vlm_collate_for_pp(collate_fn, n_microbatches=2)([{}] * 3) + def test_uneven_batch_size_general_branch_covers_all_samples(self): """batch_size not divisible by n_microbatches must not drop trailing samples. diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index c806dc4942..a2ac06a13f 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1080,16 +1080,14 @@ def patch_fn(model, name=None, add_backward_hooks=True): assert patch_calls == [] -def test_setup_keeps_engine_disabled_for_loss_without_sum_contract(monkeypatch): +def test_setup_rejects_loss_without_sum_contract(monkeypatch): cfg = _minimal_cfg_with_nvtx(nvtx_value=False) _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._supports_logits_to_keep", lambda _model: True) trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) - trainer.setup() - - assert trainer.loss_fn == "loss_fn" - assert trainer.engine is None + with pytest.raises(ValueError, match="reduction='sum'"): + trainer.setup() def test_setup_builds_engine_for_eager_sum_loss(monkeypatch): @@ -1103,6 +1101,7 @@ def test_setup_builds_engine_for_eager_sum_loss(monkeypatch): assert isinstance(trainer.loss_fn, MaskedCrossEntropy) assert trainer.engine is not None + assert trainer.engine.microbatch_size == 1 def test_setup_builds_engine_for_eager_fused_loss(monkeypatch): @@ -1123,48 +1122,59 @@ def test_setup_builds_engine_for_eager_fused_loss(monkeypatch): assert trainer.loss_fn is fused_loss assert trainer.engine is not None + assert trainer.engine.microbatch_size == 1 @pytest.mark.parametrize( ( "local_has_mtp", - "peer_has_mtp", "cp_size", "packed_sequence_size", "dataloader_emits_thd", + "pipeline_thd_kind", "fused_loss", "scale_grads_in_schedule", - "expect_engine", + "magi_enabled", + "local_batch_size", + "error_match", ), [ - (False, False, 1, 0, False, False, False, True), - (True, False, 1, 0, False, False, False, False), - (False, True, 1, 0, False, False, False, False), - (False, False, 2, 0, False, False, False, False), - (False, False, 1, 8, False, False, False, False), - (False, False, 1, 0, True, False, False, False), - (False, False, 1, 0, False, True, False, False), - (False, False, 1, 0, False, False, True, False), + (False, 1, 0, False, None, False, False, False, 2, None), + (True, 1, 0, False, None, False, False, False, 2, None), + (False, 2, 0, False, None, False, False, False, 2, None), + (False, 1, 8, False, None, False, False, False, 2, None), + (False, 1, 0, True, "native", False, False, False, 2, None), + (False, 1, 0, True, "stock_hf", False, False, False, 2, "do not consume packed document boundaries"), + (False, 1, 0, False, None, True, False, False, 2, None), + (False, 1, 0, False, None, False, True, False, 2, "scale_grads_in_schedule=False"), + (False, 1, 0, False, None, False, False, True, 2, "Magi pipeline training requires"), + (False, 1, 0, False, None, False, False, True, 1, None), ], ) -def test_setup_engine_gate_for_pipeline( +def test_setup_pipeline_engine_matrix( monkeypatch, local_has_mtp, - peer_has_mtp, cp_size, packed_sequence_size, dataloader_emits_thd, + pipeline_thd_kind, fused_loss, scale_grads_in_schedule, - expect_engine, + magi_enabled, + local_batch_size, + error_match, ): from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy cfg = _minimal_cfg_with_nvtx(nvtx_value=False) - cfg.step_scheduler.local_batch_size = 2 - cfg.step_scheduler.global_batch_size = 2 + cfg.step_scheduler.local_batch_size = local_batch_size + cfg.step_scheduler.global_batch_size = local_batch_size cfg.packed_sequence = ConfigNode({"packed_sequence_size": packed_sequence_size}) _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.setup_magi", + lambda *args, **kwargs: SimpleNamespace(enabled=magi_enabled, hf_dispatch=False), + ) fused_loss_fn = FusedLinearCrossEntropy() if fused_loss else None if fused_loss_fn is not None: monkeypatch.setattr( @@ -1173,12 +1183,23 @@ def test_setup_engine_gate_for_pipeline( property(lambda self: SimpleNamespace(build=lambda: fused_loss_fn)), ) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._supports_logits_to_keep", lambda _model: True) + pp_collate_wrapper = object() + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft._build_pp_collate_wrapper", + lambda *_args, **_kwargs: pp_collate_wrapper, + ) + dataloader_build_kwargs = [] + + def build_dataloader(**kwargs): + dataloader_build_kwargs.append(kwargs) + return "dl" + monkeypatch.setattr( RecipeConfig, "dataloader", property( lambda self: SimpleNamespace( - build=lambda **kwargs: "dl", + build=build_dataloader, dataset_builds_on_all_ranks=False, emits_thd=dataloader_emits_thd, seed=42, @@ -1193,11 +1214,16 @@ class DummyAutoPipeline(SimpleNamespace): monkeypatch.setattr("nemo_automodel.engine.AutoPipeline", DummyAutoPipeline) parts = [DummyModel()] parts[0]._pp_return_hidden_states_supported = True + if pipeline_thd_kind == "native": + parts[0].supports_thd = True + elif pipeline_thd_kind == "stock_hf": + parts[0]._te_attention_injected = True + parts[0].forward = MagicMock() if local_has_mtp: parts[0].mtp = nn.Identity() pipeline = DummyAutoPipeline( parts=parts, - pp_batch_size=2, + pp_batch_size=local_batch_size, pp_microbatch_size=1, scale_grads_in_schedule=scale_grads_in_schedule, info=SimpleNamespace( @@ -1219,36 +1245,30 @@ class DummyAutoPipeline(SimpleNamespace): pp_size=2, ), strategy_config=None, - pipeline_config=SimpleNamespace(pp_seq_len=None), + pipeline_config=SimpleNamespace( + pp_seq_len=None, + scale_grads_in_schedule=scale_grads_in_schedule, + ), moe_parallel_config=None, activation_checkpointing=False, ), ) - pp_group = object() - monkeypatch.setattr(TrainFinetuneRecipeForNextTokenPrediction, "_get_pp_group", lambda self: pp_group) - reduced_mtp_flags = [] - - def all_reduce_mtp_flag(flag, *, op, group): - assert op == torch.distributed.ReduceOp.MAX - assert group is pp_group - reduced_mtp_flags.append(bool(flag.item())) - if peer_has_mtp: - flag.fill_(1) - - monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce_mtp_flag) - trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + if error_match is not None: + with pytest.raises(ValueError, match=error_match): + trainer.setup() + return trainer.setup() - assert reduced_mtp_flags == [local_has_mtp] if fused_loss_fn is not None: assert trainer.loss_fn is fused_loss_fn - assert (trainer.engine is not None) is expect_engine - if expect_engine: - assert trainer.engine.pipeline is pipeline + assert trainer.engine is not None + assert trainer.engine.pipeline is pipeline + assert trainer.engine.microbatch_size == 1 + assert dataloader_build_kwargs[0]["collate_wrapper"] is (None if dataloader_emits_thd else pp_collate_wrapper) -def test_setup_keeps_engine_disabled_for_per_token_megatron_fsdp(monkeypatch): +def test_setup_rejects_per_token_megatron_fsdp(monkeypatch): cfg = _minimal_cfg_with_nvtx(nvtx_value=False) _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) monkeypatch.setattr( @@ -1269,9 +1289,30 @@ def test_setup_keeps_engine_disabled_for_per_token_megatron_fsdp(monkeypatch): ) trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) - trainer.setup() + with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): + trainer.setup() + - assert trainer.engine is None +def test_engine_pipeline_loss_reuses_configured_loss_and_thd_metadata(): + recipe = object.__new__(TrainFinetuneRecipeForNextTokenPrediction) + recipe.pp_enabled = True + recipe.pipeline_loss_fn = MagicMock(return_value=torch.tensor(3.0)) + output = object() + labels = torch.tensor([[1, 2, -100]]) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + + loss = recipe._engine_loss_fn( + output, + { + "labels": labels, + "weights": labels.ne(-100), + "cu_seqlens": cu_seqlens, + }, + ) + + assert loss.item() == pytest.approx(3.0) + assert recipe.pipeline_loss_fn.cu_seqlens is cu_seqlens + recipe.pipeline_loss_fn.assert_called_once_with(output, labels) def test_setup_does_not_change_storage_dtype_for_non_kd_recipe(monkeypatch): @@ -1513,7 +1554,7 @@ def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info): # Create the recipe without calling setup recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg) - # Mock out attributes needed for _forward_backward_step + # Mock out attributes needed for the recipe-owned validation forward. # Use object.__setattr__ to bypass the state tracking object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) object.__setattr__(recipe, "device_mesh", None) @@ -1530,12 +1571,12 @@ def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info): ) object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) object.__setattr__(recipe, "te_fp8", None) + object.__setattr__(recipe, "pipeline_loss_fn", None) return recipe -def test_forward_backward_step_pp_uses_eval_for_validation(monkeypatch): - """Test that _forward_backward_step uses schedule.eval() when is_train=False with PP.""" +def test_forward_validation_step_pp_uses_schedule_eval(monkeypatch): from contextlib import nullcontext pp_info = MockPPInfo(has_first_stage=True, has_last_stage=True) @@ -1553,56 +1594,15 @@ def test_forward_backward_step_pp_uses_eval_for_validation(monkeypatch): "labels": torch.tensor([[1, 2, 3]]), } - loss_buffer = [] - recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=None, - num_batches=1, - is_train=False, # Validation mode - ) + loss = recipe._forward_validation_step(batch) # Should use eval, not step assert len(pp_info.schedule.eval_calls) == 1, "schedule.eval() should be called once for validation" assert len(pp_info.schedule.step_calls) == 0, "schedule.step() should not be called for validation" + assert loss.item() == pytest.approx(0.5) -def test_forward_backward_step_pp_uses_step_for_training(monkeypatch): - """Test that _forward_backward_step uses schedule.step() when is_train=True with PP.""" - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=True, has_last_stage=True) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Mock _make_cp_batch_and_ctx to return a no-op context manager - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) - - # Create a minimal batch - batch = { - "input_ids": torch.tensor([[1, 2, 3]]), - "labels": torch.tensor([[1, 2, 3]]), - } - - loss_buffer = [] - recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=None, - num_batches=1, - is_train=True, # Training mode - ) - - # Should use step, not eval - assert len(pp_info.schedule.step_calls) == 1, "schedule.step() should be called once for training" - assert len(pp_info.schedule.eval_calls) == 0, "schedule.eval() should not be called for training" - - -def test_forward_backward_step_pp_non_first_stage_uses_eval_for_validation(monkeypatch): +def test_forward_validation_step_pp_non_first_stage_uses_eval_without_input(monkeypatch): """Test schedule.eval() without input_ids when not on first stage.""" from contextlib import nullcontext @@ -1621,15 +1621,7 @@ def test_forward_backward_step_pp_non_first_stage_uses_eval_for_validation(monke "labels": torch.tensor([[1, 2, 3]]), } - loss_buffer = [] - recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=None, - num_batches=1, - is_train=False, # Validation mode - ) + recipe._forward_validation_step(batch) # Should use eval without input_ids as first positional arg assert len(pp_info.schedule.eval_calls) == 1 @@ -1638,42 +1630,6 @@ def test_forward_backward_step_pp_non_first_stage_uses_eval_for_validation(monke assert "target" in kwargs -def test_forward_backward_step_pp_non_first_stage_uses_step_for_training(monkeypatch): - """Test schedule.step() without input_ids when not on first stage.""" - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=False, has_last_stage=True) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Mock _make_cp_batch_and_ctx to return a no-op context manager - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) - - # Create a minimal batch - batch = { - "input_ids": torch.tensor([[1, 2, 3]]), - "labels": torch.tensor([[1, 2, 3]]), - } - - loss_buffer = [] - recipe._forward_backward_step( - idx=0, - batch=batch, - loss_buffer=loss_buffer, - num_label_tokens=None, - num_batches=1, - is_train=True, # Training mode - ) - - # Should use step without input_ids as first positional arg - assert len(pp_info.schedule.step_calls) == 1 - args, kwargs = pp_info.schedule.step_calls[0] - assert len(args) == 0, "Non-first stage should not pass input_ids as positional arg" - assert "target" in kwargs - - def test_run_validation_epoch_pp_sends_loss_from_last_stage_to_main(monkeypatch): """Test that _run_validation_epoch broadcasts val_loss from last stage to main rank for PP.""" from contextlib import nullcontext @@ -1693,11 +1649,7 @@ def test_run_validation_epoch_pp_sends_loss_from_last_stage_to_main(monkeypatch) # Set dist_env.rank to 0 (last stage and main rank are the same in this test) object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - # Mock the forward_backward_step to populate loss_buffer - def mock_forward_backward_step(idx, batch, *, loss_buffer, num_label_tokens, num_batches, is_train): - loss_buffer.append(torch.tensor(0.5)) - - monkeypatch.setattr(recipe, "_forward_backward_step", mock_forward_backward_step) + monkeypatch.setattr(recipe, "_forward_validation_step", lambda batch: torch.tensor(0.5)) # Mock _dp_allreduce to return the tensor/value def mock_dp_allreduce(val, include_cp=False): @@ -1756,10 +1708,7 @@ def mock_broadcast(tensor): # Main rank (0) is different from last stage (3) object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - def mock_forward_backward_step(idx, batch, *, loss_buffer, num_label_tokens, num_batches, is_train): - loss_buffer.append(torch.tensor(0.0)) # Non-last stage has 0 loss - - monkeypatch.setattr(recipe, "_forward_backward_step", mock_forward_backward_step) + monkeypatch.setattr(recipe, "_forward_validation_step", lambda batch: torch.tensor(0.0)) def mock_dp_allreduce(val, include_cp=False): if isinstance(val, torch.Tensor): @@ -2209,18 +2158,8 @@ def test_log_moe_metrics_detailed_mode_non_detailed_step(): assert "moe/cv_mean" in metrics -class TestRunTrainOptimStepSetsMoEScale: - """Tests that _run_train_optim_step sets MoEAuxLossAutoScaler.main_loss_backward_scale.""" - - def setup_method(self): - from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler - - MoEAuxLossAutoScaler.main_loss_backward_scale = None - - def teardown_method(self): - from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler - - MoEAuxLossAutoScaler.main_loss_backward_scale = None +class TestRunTrainOptimStepUsesEngine: + """Training delegates every supported batch layout to Engine.""" def _make_recipe( self, @@ -2228,7 +2167,6 @@ def _make_recipe( pp_enabled, dp_group_size=4, cp_group_size=1, - pp_microbatches=1, ): from nemo_automodel.components.config.loader import ConfigNode @@ -2272,14 +2210,14 @@ def _make_recipe( pp_info = SimpleNamespace( has_first_stage=True, has_last_stage=True, - schedule=SimpleNamespace(_n_microbatches=pp_microbatches), + schedule=SimpleNamespace(_n_microbatches=1), ) object.__setattr__( recipe, "pp", SimpleNamespace( info=pp_info, - pp_batch_size=pp_microbatches, + pp_batch_size=1, pp_microbatch_size=1, update_seq_len=lambda seq_len: None, ), @@ -2297,45 +2235,20 @@ def _make_recipe( monkeypatch.setattr(recipe, "_get_dp_group_size", lambda include_cp=False: dp_group_size) monkeypatch.setattr(recipe, "_get_cp_group_size", lambda: cp_group_size) - def mock_forward_backward_step(idx, batch, *, loss_buffer, num_label_tokens, num_batches, is_train=True): - loss_buffer.append(torch.tensor(0.5)) - - monkeypatch.setattr(recipe, "_forward_backward_step", mock_forward_backward_step) monkeypatch.setattr( "nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", lambda *a, **k: torch.tensor(1.0), ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.prepare_for_grad_accumulation", lambda *a, **k: None) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.prepare_for_final_backward", lambda *a, **k: None) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.prepare_after_first_microbatch", lambda *a, **k: None) object.__setattr__(recipe, "checkpointer", SimpleNamespace(maybe_wait_for_staging=lambda: None)) object.__setattr__(recipe, "lr_scheduler", None) + object.__setattr__(recipe, "loss_fn", object()) + engine = MagicMock() + engine.forward_backward.return_value = (torch.tensor(0.5), []) + object.__setattr__(recipe, "engine", engine) object.__setattr__(recipe, "timestamp", 0.0) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) return recipe - def test_pp_scale_includes_pipeline_microbatches_and_token_normalization(self, monkeypatch): - from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler - - recipe = self._make_recipe( - monkeypatch, - pp_enabled=True, - dp_group_size=8, - cp_group_size=2, - pp_microbatches=4, - ) - - batches = [ - {"input_ids": torch.tensor([[1, 2, 3, 4]]), "labels": torch.tensor([[1, 2, 3, -100]])}, - {"input_ids": torch.tensor([[5, 6, 7, 8]]), "labels": torch.tensor([[5, 6, 7, -100]])}, - ] - - recipe._run_train_optim_step(batches) - - assert MoEAuxLossAutoScaler.main_loss_backward_scale is not None - # 2 outer batches * 4 PP microbatches = 8 model microbatches. - # Base CP-aware average: 2 / 8. PP post-normalization compensation: 6 / 8. - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.1875) - def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypatch): recipe = self._make_recipe(monkeypatch, pp_enabled=True) batches = [ @@ -2345,7 +2258,6 @@ def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypat datums = [object(), object()] make_datum = MagicMock(side_effect=datums) monkeypatch.setattr(recipe, "_make_engine_datum", make_datum) - monkeypatch.setattr(recipe, "_forward_backward_step", MagicMock(side_effect=AssertionError("legacy path"))) monkeypatch.setattr( recipe, "_broadcast_from_last_pp_stage", @@ -2369,11 +2281,9 @@ def finalize_grads(*args, **kwargs): param_groups=[{"lr": 0.01}], ) object.__setattr__(recipe, "optimizer", [optimizer]) - monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) - metrics = recipe._run_train_optim_step(batches) - engine.forward_backward.assert_called_once_with([[datums[0]], [datums[1]]], recipe._engine_loss_fn) + engine.forward_backward.assert_called_once_with(datums, recipe._engine_loss_fn) assert make_datum.call_count == 2 assert len(finalizer_calls) == 1 assert finalizer_calls[0][1]["num_label_tokens"] is None @@ -2383,7 +2293,7 @@ def finalize_grads(*args, **kwargs): optimizer.zero_grad.assert_called_once_with() assert metrics.metrics["loss"] == pytest.approx(0.25) - def test_pp_thd_batch_uses_legacy_forward_backward(self, monkeypatch): + def test_pp_thd_batch_uses_engine(self, monkeypatch): recipe = self._make_recipe(monkeypatch, pp_enabled=True) batch = { "input_ids": torch.tensor([[1, 2, 3]]), @@ -2391,54 +2301,19 @@ def test_pp_thd_batch_uses_legacy_forward_backward(self, monkeypatch): "qkv_format": "thd", } engine = MagicMock() + engine.forward_backward.return_value = (torch.tensor(0.5), []) object.__setattr__(recipe, "engine", engine) - - def legacy_step(_idx, _batch, *, loss_buffer, **_kwargs): - loss_buffer.append(torch.tensor(0.5)) - - legacy_step = MagicMock(side_effect=legacy_step) - monkeypatch.setattr(recipe, "_forward_backward_step", legacy_step) finalizer = MagicMock(return_value=torch.tensor(1.0)) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", finalizer) - monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) recipe._run_train_optim_step([batch]) - engine.forward_backward.assert_not_called() - legacy_step.assert_called_once() - assert finalizer.call_args.kwargs["num_label_tokens"] == 2 - - @pytest.mark.parametrize("dp_size", [1, 8]) - def test_non_pp_scale_is_independent_of_dp_size(self, monkeypatch, dp_size): - from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler - - recipe = self._make_recipe(monkeypatch, pp_enabled=False, dp_group_size=dp_size) - - batches = [ - {"input_ids": torch.tensor([[1, 2, 3, 4]]), "labels": torch.tensor([[1, 2, 3, -100]])} for _ in range(4) - ] - - recipe._run_train_optim_step(batches) - - assert MoEAuxLossAutoScaler.main_loss_backward_scale is not None - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) - - def test_non_pp_scale_restores_cp_sum(self, monkeypatch): - from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler - - recipe = self._make_recipe( - monkeypatch, - pp_enabled=False, - dp_group_size=8, - cp_group_size=2, - ) - batches = [ - {"input_ids": torch.tensor([[1, 2, 3, 4]]), "labels": torch.tensor([[1, 2, 3, -100]])} for _ in range(4) - ] - - recipe._run_train_optim_step(batches) - - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.5) + engine.forward_backward.assert_called_once() + datums, loss_fn = engine.forward_backward.call_args.args + assert len(datums) == 1 + assert datums[0].model_inputs["qkv_format"] == "thd" + assert loss_fn == recipe._engine_loss_fn + assert finalizer.call_args.kwargs["num_label_tokens"] is None # ----------------------------------------------------------------------------- @@ -2497,8 +2372,8 @@ def test_rope_fusion_disabled_when_cp_gt_1(monkeypatch): assert trainer.engine is not None -@pytest.mark.parametrize(("cp_size", "expect_engine"), [(1, True), (2, False)]) -def test_setup_engine_gate_for_mtp_with_context_parallelism(monkeypatch, cp_size, expect_engine): +def test_setup_builds_engine_for_mtp_without_context_parallelism(monkeypatch): + cp_size = 1 cfg = _minimal_cfg_with_rope_fusion(cp_size=cp_size, rope_fusion=True) _patch_setup_minimals_with_cp(monkeypatch, cp_size=cp_size) model = DummyModel() @@ -2508,10 +2383,23 @@ def test_setup_engine_gate_for_mtp_with_context_parallelism(monkeypatch, cp_size trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) trainer.setup() - assert (trainer.engine is not None) is expect_engine + assert trainer.engine is not None + + +def test_setup_rejects_mtp_with_context_parallelism(monkeypatch): + cp_size = 2 + cfg = _minimal_cfg_with_rope_fusion(cp_size=cp_size, rope_fusion=True) + _patch_setup_minimals_with_cp(monkeypatch, cp_size=cp_size) + model = DummyModel() + model.mtp = nn.Identity() + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.build_model", lambda *args, **kwargs: model) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + with pytest.raises(NotImplementedError, match="MTP with context parallelism"): + trainer.setup() -def test_setup_keeps_engine_disabled_for_magi(monkeypatch): +def test_setup_builds_engine_for_magi(monkeypatch): cfg = _minimal_cfg_with_rope_fusion(cp_size=1, rope_fusion=True) _patch_setup_minimals_with_cp(monkeypatch, cp_size=1) monkeypatch.setattr( @@ -2522,7 +2410,7 @@ def test_setup_keeps_engine_disabled_for_magi(monkeypatch): trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) trainer.setup() - assert trainer.engine is None + assert trainer.engine is not None def test_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): @@ -2764,8 +2652,8 @@ def test_evaluate_failure_is_tolerated(self, monkeypatch): (1, True, True), ], ) -def test_forward_backward_step_model_cp_hook(monkeypatch, cp_size, uses_thd, supports_thd): - """Non-PP training invokes model-owned batch preparation for CP or native THD.""" +def test_forward_validation_step_model_cp_hook(monkeypatch, cp_size, uses_thd, supports_thd): + """Non-PP validation keeps the recipe-owned CP preparation path.""" from contextlib import nullcontext cfg = ConfigNode( @@ -2862,22 +2750,15 @@ def _fake_calc_loss( ) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calc_loss) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_final_hidden_states", lambda out: None) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_sync_ctx", lambda *a, **k: nullcontext()) monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.filter_forward_kwargs", lambda model, batch: batch) batch = {"input_ids": torch.randn(1, 4, 4), "labels": torch.zeros(1, 4, dtype=torch.long)} if uses_thd: batch["qkv_format"] = "thd" - loss_buffer = [] - recipe._forward_backward_step( - idx=0, batch=batch, loss_buffer=loss_buffer, num_label_tokens=None, num_batches=1, is_train=True - ) + loss = recipe._forward_validation_step(batch) assert model.prepared is True assert model.num_chunks == 1 assert captured["logits_is_tensor"] - assert len(loss_buffer) == 1 - assert torch.isfinite(loss_buffer[0]).all() - # backward through the local loss populated grads - assert model.lin.weight.grad is not None - assert torch.isfinite(model.lin.weight.grad).all() + assert torch.isfinite(loss).all() + assert model.lin.weight.grad is None diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 62b605ddec..c4b7a786a7 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -14,6 +14,7 @@ from __future__ import annotations +import sys from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -75,17 +76,34 @@ def __init__(self, size, rank): class _FakeAutoPipeline(AutoPipeline): - def __init__(self, model, *, parts=None, num_microbatches=2, scale_grads=False, events=None): + def __init__( + self, + model, + *, + parts=None, + num_microbatches=2, + scale_grads=False, + events=None, + callback_order=None, + has_last_stage=True, + pp_microbatch_size=1, + ): self.compute_model = model self._parts = parts or [model] self._num_microbatches = num_microbatches + self.pp_microbatch_size = pp_microbatch_size self.scale_grads_in_schedule = scale_grads self.pp_mesh = _SubMesh(2) + self._info = SimpleNamespace(has_last_stage=has_last_stage) self.events = events + self.callback_order = callback_order or list(range(num_microbatches)) self.step_calls = 0 self.backward_calls = 0 self.updated_seq_lens = [] + self.updated_microbatch_sizes = [] + self.updated_input_shapes = [] self.callback_losses = [] + self.prepared_inputs = [] @property def parts(self): @@ -95,28 +113,29 @@ def parts(self): def num_microbatches(self): return self._num_microbatches - def update_seq_len(self, seq_len): + @property + def info(self): + return self._info + + def update_seq_len(self, seq_len, *, microbatch_size=None, input_tensor=None): self.updated_seq_lens.append(seq_len) + self.updated_microbatch_sizes.append(microbatch_size) + self.updated_input_shapes.append(tuple(input_tensor.shape) if input_tensor is not None else None) - def step(self, model_input, *, loss_inputs, loss_fn, return_outputs, **kwargs): + def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): assert return_outputs is False + assert len(model_inputs) == self.num_microbatches + self.prepared_inputs.append(model_inputs) self.step_calls += 1 if self.events is not None: self.events.append("step") - def chunks(value): - if isinstance(value, torch.Tensor) and value.ndim > 0: - return value.chunk(self.num_microbatches, dim=0) - return (value,) * self.num_microbatches - - model_chunks = chunks(model_input) - kwargs_chunks = {name: chunks(value) for name, value in kwargs.items()} - loss_chunks = {name: chunks(value) for name, value in loss_inputs.items()} - for index in range(self.num_microbatches): - model_kwargs = {name: values[index] for name, values in kwargs_chunks.items()} - loss_inputs_mb = {name: values[index] for name, values in loss_chunks.items()} - output = self.compute_model(model_chunks[index], **model_kwargs) - scaled_loss = loss_fn(output, loss_inputs_mb, (model_chunks[index],), model_kwargs) + for index in self.callback_order: + inputs = dict(model_inputs[index]) + primary_name = "inputs_embeds" if "inputs_embeds" in inputs else "input_ids" + primary = inputs.pop(primary_name) + output = self.compute_model(primary, **inputs) + scaled_loss = loss_fn(output, index) self.callback_losses.append(scaled_loss.detach()) scaled_loss.backward() self.backward_calls += 1 @@ -148,7 +167,7 @@ def _datum(values, weights=None) -> Datum: return Datum(model_inputs={"input_ids": values}, loss_fn_inputs={"weights": weights}) -def _identity_loss(output, _loss_inputs, _datums, _model_inputs): +def _identity_loss(output, _loss_inputs): return output @@ -163,7 +182,7 @@ def test_forward_backward_uses_one_denominator_for_the_window(): engine = Engine(model, device="cpu") loss, outputs = engine.forward_backward( - [[_datum([1, 2])], [_datum([3])]], + [_datum([1, 2]), _datum([3])], _identity_loss, ) @@ -174,26 +193,46 @@ def test_forward_backward_uses_one_denominator_for_the_window(): assert model.forward_calls == 2 +def test_forward_backward_groups_flat_datums_by_microbatch_size(): + group_sizes = [] + + def recording_collate(datums): + group_sizes.append(len(datums)) + return collate_datums(datums) + + model = ScaleModel() + loss, _ = Engine( + model, + device="cpu", + microbatch_size=2, + collate_fn=recording_collate, + ).forward_backward([_datum([value]) for value in range(1, 6)], _identity_loss) + + assert group_sizes == [2, 2, 1] + assert model.forward_calls == 3 + assert loss.item() == pytest.approx(3.0) + + def test_raw_thd_packed_collater_is_prepared_by_context_parallel_sharder(): model = ScaleModel() model.backend = SimpleNamespace(attn="te") seen = {} - def loss_fn(output, inputs, _datums, model_inputs): - seen.update(model_inputs) + def loss_fn(output, inputs): + seen.update(inputs) assert inputs["weights"].shape == output.shape == (3,) return output loss, _ = Engine( model, device="cpu", + microbatch_size=2, collate_fn=partial(collate_datums, packed=True), - ).forward_backward([[_datum([1, 2]), _datum([3])]], loss_fn) + ).forward_backward([_datum([1, 2]), _datum([3])], loss_fn) assert loss.item() == pytest.approx(2.0) assert "seq_lens" not in seen assert "seq_lens_padded" not in seen - assert seen["qkv_format"] == "thd" assert seen["cu_seqlens"].tolist() == [0, 2, 3] @@ -204,8 +243,9 @@ def test_raw_thd_requires_a_thd_capable_context_parallel_sharder(): Engine( model, device="cpu", + microbatch_size=2, collate_fn=partial(collate_datums, packed=True), - ).forward_backward([[_datum([1, 2]), _datum([3])]], _identity_loss) + ).forward_backward([_datum([1, 2]), _datum([3])], _identity_loss) assert model.forward_calls == 0 @@ -236,9 +276,10 @@ def prepare_model_inputs_for_cp(self, batch, *, num_chunks): ) } - def forward(self, *args, **kwargs): + def forward(self, input_ids, *args, **kwargs): assert cp_context_active - return super().forward(*args, **kwargs) + assert input_ids.tolist() == [[5, 6, 9, 9]] + return super().forward(input_ids, *args, **kwargs) model = CPModel() mesh = _CPMesh(size=2, rank=1) @@ -266,15 +307,14 @@ def forward(self, *args, **kwargs): lambda grad: grad if cp_context_active else pytest.fail("CP context ended before backward") ) - def loss_fn(output, inputs, _datums, model_inputs): + def loss_fn(output, inputs): assert cp_context_active - assert model_inputs["input_ids"].tolist() == [[5, 6, 9, 9]] assert inputs["target_tokens"].tolist() == [[15, 16, 0, 0]] assert inputs["weights"].tolist() == [[1.0, 1.0, 0.0, 0.0]] torch.testing.assert_close(inputs["advantages"], torch.tensor([[0.5, 0.6, 0.0, 0.0]])) return output - loss, _ = engine.forward_backward([[datum]], loss_fn) + loss, _ = engine.forward_backward([datum], loss_fn) assert loss.item() == pytest.approx(11 / 6) assert model.weight.grad.item() == pytest.approx(11 / 6) @@ -305,7 +345,9 @@ def test_packed_rl_callback_keeps_per_datum_sequence_boundaries(): first.loss_fn_inputs["sequence_scale"] = torch.tensor(2.0) second.loss_fn_inputs["sequence_scale"] = torch.tensor(0.5) - def sequence_loss(output, _loss_inputs, datums, _model_inputs): + datums = [first, second] + + def sequence_loss(output, _loss_inputs): chunks = output.squeeze(0).split([datum.seq_len for datum in datums]) losses = torch.cat([chunk * datum.loss_fn_inputs["sequence_scale"] for chunk, datum in zip(chunks, datums)]) outputs = [ @@ -316,8 +358,9 @@ def sequence_loss(output, _loss_inputs, datums, _model_inputs): loss, outputs = Engine( model, device="cpu", + microbatch_size=2, collate_fn=_model_ready_packed_collate, - ).forward_backward([[first, second]], sequence_loss) + ).forward_backward(datums, sequence_loss) assert loss.item() == pytest.approx(2.5) assert model.weight.grad.item() == pytest.approx(2.5) @@ -328,7 +371,7 @@ def sequence_loss(output, _loss_inputs, datums, _model_inputs): def test_weights_mask_loss_and_denominator(): model = ScaleModel() loss, _ = Engine(model, device="cpu").forward_backward( - [[_datum([1, 100], [1.0, 0.0])], [_datum([3, 5], [0.5, 1.0])]], + [_datum([1, 100], [1.0, 0.0]), _datum([3, 5], [0.5, 1.0])], _identity_loss, ) @@ -336,17 +379,25 @@ def test_weights_mask_loss_and_denominator(): assert model.weight.grad.item() == pytest.approx(3.0) +def test_fractional_weight_sum_below_one_is_not_clamped(): + model = ScaleModel() + loss, _ = Engine(model, device="cpu").forward_backward( + [_datum([2, 4], [0.2, 0.3])], + _identity_loss, + ) + + assert loss.item() == pytest.approx(3.2) + assert model.weight.grad.item() == pytest.approx(3.2) + + def test_loss_fn_outputs_follow_datum_order_and_are_detached(): model = ScaleModel() - def loss_with_outputs(output, _loss_inputs, datums, _model_inputs): - return output, [ - {"first_token": datum.input_ids[0], "model_value": output[index].sum()} - for index, datum in enumerate(datums) - ] + def loss_with_outputs(output, _loss_inputs): + return output, [{"first_token": row.flatten()[0], "model_value": row.sum()} for row in output] - _, outputs = Engine(model, device="cpu").forward_backward( - [[_datum([1, 2]), _datum([3])], [_datum([4])]], + _, outputs = Engine(model, device="cpu", microbatch_size=2).forward_backward( + [_datum([1, 2]), _datum([3]), _datum([4])], loss_with_outputs, ) @@ -357,22 +408,26 @@ def loss_with_outputs(output, _loss_inputs, datums, _model_inputs): def test_loss_fn_outputs_must_align_with_datums(): model = ScaleModel() with pytest.raises(ValueError, match="one mapping per Datum"): - Engine(model, device="cpu").forward_backward( - [[_datum([1]), _datum([2])]], - lambda output, _inputs, _datums, _model_inputs: (output, [{"only": "one"}]), + Engine(model, device="cpu", microbatch_size=2).forward_backward( + [_datum([1]), _datum([2])], + lambda output, _inputs: (output, [{"only": "one"}]), ) assert model.weight.grad is None def test_loss_fn_outputs_must_be_consistent_across_the_window(): - def inconsistent_outputs(output, _inputs, datums, _model_inputs): - if datums[0].input_ids[0].item() == 1: + calls = 0 + + def inconsistent_outputs(output, _inputs): + nonlocal calls + calls += 1 + if calls == 1: return output, [{"value": output.sum()}] return output with pytest.raises(ValueError, match="every microbatch or none"): Engine(ScaleModel(), device="cpu").forward_backward( - [[_datum([1])], [_datum([2])]], + [_datum([1]), _datum([2])], inconsistent_outputs, ) @@ -399,8 +454,7 @@ def test_raw_output_and_loss_inputs_support_an_rl_loss_callback(): ) model = TinyLM() - def policy_loss(logits, inputs, datums, _model_inputs): - assert len(datums) == 1 + def policy_loss(logits, inputs): new_logprobs = -F.cross_entropy( logits.flatten(0, 1), inputs["target_tokens"].flatten(), @@ -410,7 +464,7 @@ def policy_loss(logits, inputs, datums, _model_inputs): losses = -(ratio * inputs["advantages"]) return losses, [{"policy_sum": (losses * inputs["weights"]).sum()}] - loss, outputs = Engine(model, device="cpu").forward_backward([[datum]], policy_loss) + loss, outputs = Engine(model, device="cpu").forward_backward([datum], policy_loss) assert torch.isfinite(loss) assert torch.isfinite(outputs[0]["policy_sum"]) @@ -436,7 +490,7 @@ def sync_context(_model, is_last, _defer): monkeypatch.setattr(engine_module, "get_sync_ctx", sync_context) Engine(ScaleModel(), device="cpu").forward_backward( - [[_datum([1])], [_datum([2])]], + [_datum([1]), _datum([2])], _identity_loss, ) @@ -449,7 +503,7 @@ def test_pipeline_window_uses_schedule_microbatches_and_global_normalization(): backward_calls = 0 @contextmanager - def forward_context(): + def forward_context(_model_inputs): nonlocal active assert not active active = True @@ -476,10 +530,9 @@ def check_backward_context(grad): model.weight.register_hook(check_backward_context) - def loss_fn(output, inputs, datums, model_inputs): + def loss_fn(output, inputs): assert active - assert len(datums) == 1 - assert output.shape == inputs["weights"].shape == model_inputs["input_ids"].shape == (1, 2) + assert output.shape == inputs["weights"].shape == (1, 2) return output loss, outputs = Engine( @@ -490,8 +543,8 @@ def loss_fn(output, inputs, datums, model_inputs): context_fn=forward_context, ).forward_backward( [ - [_datum([[1, 2], [3, 4]])], - [_datum([[5, 6], [7, 8]])], + _datum([[1, 2], [3, 4]]), + _datum([[5, 6], [7, 8]]), ], loss_fn, ) @@ -505,6 +558,12 @@ def loss_fn(output, inputs, datums, model_inputs): assert pipeline.backward_calls == backward_calls == 4 assert model.forward_calls == 4 assert pipeline.updated_seq_lens == [2, 2] + assert [item["input_ids"].shape for call in pipeline.prepared_inputs for item in call] == [ + (1, 2), + (1, 2), + (1, 2), + (1, 2), + ] scaled_losses = torch.stack(pipeline.callback_losses) torch.testing.assert_close( scaled_losses, @@ -549,8 +608,8 @@ def prepare_final(parts, *, pp_enabled): collate_fn=collate_prebatched, ).forward_backward( [ - [_datum([[1], [2]])], - [_datum([[3], [4]])], + _datum([[1], [2]]), + _datum([[3], [4]]), ], _identity_loss, ) @@ -561,49 +620,282 @@ def prepare_final(parts, *, pp_enabled): assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) -def test_pipeline_rejects_per_datum_outputs_before_schedule_backward(): +def test_pipeline_outputs_follow_logical_microbatch_order(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + + _, outputs = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=2, + ).forward_backward( + [_datum([1]), _datum([2])], + lambda output, _inputs: (output, [{"metric": output.sum()}]), + ) + + assert pipeline.step_calls == 1 + assert pipeline.backward_calls == 2 + assert [item["metric"].item() for item in outputs] == [1.0, 2.0] + + +def test_pipeline_default_packed_collater_splits_flat_datums_inside_engine(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = [_datum([1, 2]), _datum([3, 4]), _datum([5, 6]), _datum([7, 8])] + + def loss_with_outputs(output, _loss_inputs): + return output, [ + {"pair_sum": output[..., :2].sum()}, + {"pair_sum": output[..., 2:].sum()}, + ] + + loss, outputs = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=partial(collate_datums, packed=True), + ).forward_backward(datums, loss_with_outputs) + + assert loss.item() == pytest.approx(4.5) + assert model.weight.grad.item() == pytest.approx(4.5) + assert [item["pair_sum"].item() for item in outputs] == [3.0, 7.0, 11.0, 15.0] + assert [item["input_ids"].shape for item in pipeline.prepared_inputs[0]] == [(1, 4), (1, 4)] + + +def test_pipeline_default_packed_collater_rejects_ambiguous_per_datum_loss_fields(): model = ScaleModel() + model.backend = SimpleNamespace(attn="te") pipeline = _FakeAutoPipeline(model, num_microbatches=2) + datums = [ + Datum( + model_inputs={"input_ids": torch.tensor([index, index + 1])}, + loss_fn_inputs={"weights": torch.ones(2), "reward": torch.tensor(float(index))}, + ) + for index in (1, 3, 5, 7) + ] - with pytest.raises(ValueError, match="per-Datum"): + with pytest.raises(NotImplementedError, match="per-Datum loss fields.*reward"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=partial(collate_datums, packed=True), + ).forward_backward(datums, _identity_loss) + + +def test_pipeline_outputs_allow_uneven_datum_counts_at_final_thd_boundaries(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = [_datum([1]), _datum([2]), _datum([3]), _datum([4, 5, 6])] + + def final_thd_collate(items): + lengths = [datum.seq_len for datum in items] + tokens = torch.cat([datum.input_ids for datum in items]) + return ( + { + "input_ids": tokens, + "position_ids": torch.cat([torch.arange(length) for length in lengths]), + "cu_seqlens": torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32), + "max_seqlen": torch.tensor(max(lengths), dtype=torch.int32), + "qkv_format": "thd", + }, + {"weights": torch.cat([datum.loss_fn_inputs["weights"] for datum in items])}, + ) + + def loss_with_outputs(output, _loss_inputs): + first_token = int(output.reshape(-1)[0].item()) + ids = [1, 2, 3] if first_token == 1 else [4] + return output, [{"datum_id": torch.tensor(datum_id)} for datum_id in ids] + + _, outputs = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=final_thd_collate, + ).forward_backward(datums, loss_with_outputs) + + assert [item["datum_id"].item() for item in outputs] == [1, 2, 3, 4] + + +def test_pipeline_output_sync_finds_last_stage_on_physical_rank_zero(monkeypatch): + pipeline = _FakeAutoPipeline(ScaleModel(), has_last_stage=False) + engine = Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()) + pp_group = object() + engine._pp_group_and_size = lambda: (pp_group, 2) + expected = [{"metric": torch.tensor(7.0)}] + + def fake_all_gather_into_tensor(gathered, local, *, group): + torch.testing.assert_close(local, torch.tensor([0, 0], dtype=torch.int64)) + assert group is pp_group + gathered.copy_(torch.tensor([1, 1, 0, 0], dtype=torch.int64)) + + def fake_broadcast_object_list(objects, *, src, group, device): + assert objects == [None] + assert src == 11 + assert group is pp_group + assert device == torch.device("cpu") + objects[0] = expected + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) + monkeypatch.setattr(dist, "get_rank", lambda *, group: 1) + monkeypatch.setattr(dist, "get_global_rank", lambda group, group_rank: 11 if group_rank == 0 else 12) + monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast_object_list) + + result = engine._broadcast_pipeline_outputs([]) + + assert result == expected + + +def test_pipeline_output_sync_skips_object_broadcast_without_outputs(monkeypatch): + pipeline = _FakeAutoPipeline(ScaleModel(), has_last_stage=False) + engine = Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()) + pp_group = object() + engine._pp_group_and_size = lambda: (pp_group, 2) + + def fake_all_gather_into_tensor(gathered, local, *, group): + torch.testing.assert_close(local, torch.tensor([0, 0], dtype=torch.int64)) + assert group is pp_group + gathered.copy_(torch.tensor([1, 0, 0, 0], dtype=torch.int64)) + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) + monkeypatch.setattr( + dist, + "broadcast_object_list", + lambda *_args, **_kwargs: pytest.fail("empty pipeline outputs must not use an object collective"), + ) + + assert engine._broadcast_pipeline_outputs([]) == [] + + +def test_pipeline_prebatched_outputs_require_one_inner_microbatch(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + datum = _datum([[1], [2]]) + + with pytest.raises(ValueError, match="prebatched Datum may return outputs only"): Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), collate_fn=collate_prebatched, ).forward_backward( - [[_datum([[1], [2]])]], - lambda output, _inputs, _datums, _model_inputs: (output, [{"metric": output.sum()}]), + [datum], + lambda output, _inputs: (output, [{"metric": output.sum()}]), ) - assert pipeline.step_calls == 1 - assert pipeline.backward_calls == 0 - assert model.forward_calls == 1 - assert model.weight.grad is None +def test_pipeline_final_thd_embeddings_use_the_token_axis_for_sequence_length(): + pipeline = _FakeAutoPipeline(ScaleModel(), num_microbatches=1) + embeddings = torch.arange(12, dtype=torch.float32).reshape(4, 3) + datum = Datum( + model_inputs={ + "inputs_embeds": embeddings, + "position_ids": torch.arange(4), + "cu_seqlens": torch.tensor([0, 4], dtype=torch.int32), + "max_seqlen": torch.tensor(4, dtype=torch.int32), + "qkv_format": "thd", + }, + loss_fn_inputs={"weights": torch.ones(4)}, + ) + + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward_backward([datum], lambda output, _inputs: output.sum()) -def test_pipeline_rejects_schedule_gradient_scaling_before_forward(): + assert pipeline.updated_seq_lens == [4] + + +def test_pipeline_requires_materialized_padded_microbatch_size_to_match_config(): model = ScaleModel() - pipeline = _FakeAutoPipeline(model, scale_grads=True) + pipeline = _FakeAutoPipeline(model, num_microbatches=2, pp_microbatch_size=2) - with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): - Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()).forward_backward( - [[_datum([1])]], _identity_loss - ) + with pytest.raises(ValueError, match="materialized pipeline microbatch has batch size 1"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward_backward([_datum([[1], [2]])], _identity_loss) assert pipeline.step_calls == 0 assert model.forward_calls == 0 - assert model.weight.grad is None -def test_pipeline_rejects_multiple_datums_in_one_outer_batch_before_forward(): +def test_pipeline_te_thd_keeps_arbitrary_loss_fields_aligned_through_cp(monkeypatch): + class MockTex: + @staticmethod + def thd_get_partitioned_indices(_cu_seqlens, total_tokens, _cp_size, _cp_rank): + assert total_tokens == 4 + return torch.tensor([0, 3]) + + monkeypatch.setitem(sys.modules, "transformer_engine_torch", MockTex) + monkeypatch.setattr(dist, "get_rank", lambda group=None: 0) + model = ScaleModel() - pipeline = _FakeAutoPipeline(model) + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + mesh = _CPMesh(size=2, rank=0) + mesh_context = SimpleNamespace(pp_size=2, cp_size=2, device_mesh=mesh, process_group=None) + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]), + "position_ids": torch.arange(4).expand(2, -1), + "seq_lens": torch.tensor([[4], [4]]), + "seq_lens_padded": torch.tensor([[4], [4]]), + "qkv_format": "thd", + }, + loss_fn_inputs={ + "weights": torch.ones(2, 4), + "advantages": torch.arange(1, 9, dtype=torch.float32).view(2, 4) * 10, + "old_logprobs": -torch.arange(1, 9, dtype=torch.float32).view(2, 4), + }, + ) + engine = Engine( + pipeline, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ) + # This CPU test exercises the exact THD layout only; real CP collectives are + # covered by the distributed tests below. + engine._dp_group_and_size = lambda: (None, 1) + engine._gradient_group_and_size = lambda _group, _size: (None, 1) + seen = [] + + def loss_fn(output, inputs): + seen.append((output.detach().clone(), inputs["advantages"].clone(), inputs["old_logprobs"].clone())) + assert output.shape == inputs["weights"].shape == (1, 2) + return output + + loss, _ = engine.forward_backward([datum], loss_fn) + + assert loss.item() == pytest.approx(18 / 8) + assert len(pipeline.prepared_inputs) == 1 + assert [item["input_ids"].shape for item in pipeline.prepared_inputs[0]] == [(1, 2), (1, 2)] + torch.testing.assert_close(seen[0][0], torch.tensor([[1.0, 4.0]])) + torch.testing.assert_close(seen[0][1], torch.tensor([[10.0, 40.0]])) + torch.testing.assert_close(seen[0][2], torch.tensor([[-1.0, -4.0]])) + torch.testing.assert_close(seen[1][0], torch.tensor([[5.0, 8.0]])) + torch.testing.assert_close(seen[1][1], torch.tensor([[50.0, 80.0]])) + torch.testing.assert_close(seen[1][2], torch.tensor([[-5.0, -8.0]])) - with pytest.raises(ValueError, match="exactly one prebatched Datum"): + +def test_pipeline_rejects_schedule_gradient_scaling_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, scale_grads=True) + + with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): Engine(pipeline, device="cpu", mesh_context=_pipeline_mesh_context()).forward_backward( - [[_datum([1]), _datum([2])]], - _identity_loss, + [_datum([1])], _identity_loss ) assert pipeline.step_calls == 0 @@ -611,28 +903,28 @@ def test_pipeline_rejects_multiple_datums_in_one_outer_batch_before_forward(): assert model.weight.grad is None -def test_pipeline_requires_mesh_context_before_forward(): +def test_pipeline_groups_multiple_flat_datums_into_one_outer_batch(): model = ScaleModel() pipeline = _FakeAutoPipeline(model) - with pytest.raises(ValueError, match="requires mesh_context"): - Engine(pipeline, device="cpu").forward_backward([[_datum([1])]], _identity_loss) + loss, _ = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=2, + ).forward_backward([_datum([1]), _datum([2])], _identity_loss) - assert pipeline.step_calls == 0 - assert model.forward_calls == 0 - assert model.weight.grad is None + assert loss.item() == pytest.approx(1.5) + assert pipeline.step_calls == 1 + assert model.forward_calls == 2 -def test_pipeline_rejects_context_parallelism_before_forward(): +def test_pipeline_requires_mesh_context_before_forward(): model = ScaleModel() pipeline = _FakeAutoPipeline(model) - mesh_context = SimpleNamespace(pp_size=2, cp_size=2, device_mesh=_CPMesh(size=2, rank=0)) - with pytest.raises(NotImplementedError, match="does not yet support context parallelism"): - Engine(pipeline, device="cpu", mesh_context=mesh_context).forward_backward( - [[_datum([1])]], - _identity_loss, - ) + with pytest.raises(ValueError, match="requires mesh_context"): + Engine(pipeline, device="cpu").forward_backward([_datum([1])], _identity_loss) assert pipeline.step_calls == 0 assert model.forward_calls == 0 @@ -643,7 +935,7 @@ def test_forward_context_covers_forward_loss_and_backward(): active = False @contextmanager - def forward_context(): + def forward_context(_model_inputs): nonlocal active active = True try: @@ -659,11 +951,11 @@ def forward(self, input_ids, **kwargs): model = ContextModel() model.weight.register_hook(lambda grad: grad if active else pytest.fail("context ended before backward")) - def loss_fn(output, _inputs, _datums, _model_inputs): + def loss_fn(output, _inputs): assert active return output - Engine(model, device="cpu", context_fn=forward_context).forward_backward([[_datum([1, 2])]], loss_fn) + Engine(model, device="cpu", context_fn=forward_context).forward_backward([_datum([1, 2])], loss_fn) assert not active @@ -671,7 +963,7 @@ def test_window_sets_the_same_moe_aux_scale_as_the_recipes(monkeypatch): monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) Engine(ScaleModel(), device="cpu").forward_backward( - [[_datum([1])], [_datum([2])]], + [_datum([1]), _datum([2])], _identity_loss, ) @@ -711,9 +1003,9 @@ def test_model_specific_collater_keeps_multimodal_inputs_and_gradients(): ), ] model = TinyVLM() - Engine(model, device="cpu", collate_fn=_vlm_collate).forward_backward( - [datums], - lambda output, _inputs, _datums, _model_inputs: output, + Engine(model, device="cpu", microbatch_size=2, collate_fn=_vlm_collate).forward_backward( + datums, + lambda output, _inputs: output, ) assert model.text.weight.grad is not None @@ -729,7 +1021,7 @@ def test_prebatched_datum_keeps_existing_recipe_batch_layout(): loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0], [1.0, 0.0]])}, ) - loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([[datum]], _identity_loss) + loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) assert loss.item() == pytest.approx(2.0) assert model.weight.grad.item() == pytest.approx(2.0) @@ -745,7 +1037,7 @@ def test_prebatched_datum_keeps_vlm_media_layout(): loss_fn_inputs={"weights": torch.ones(2)}, ) - loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([[datum]], _identity_loss) + loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) assert torch.isfinite(loss) assert model.text.weight.grad is not None @@ -756,23 +1048,23 @@ def test_scalar_loss_is_a_local_weighted_sum_numerator(): model = ScaleModel() loss, _ = Engine(model, device="cpu").forward_backward( - [[_datum([1, 100], [1.0, 0.0])], [_datum([3, 5], [0.5, 1.0])]], - lambda output, inputs, _datums, _model_inputs: (output * inputs["weights"]).sum(), + [_datum([1, 100], [1.0, 0.0]), _datum([3, 5], [0.5, 1.0])], + lambda output, inputs: (output * inputs["weights"]).sum(), ) assert loss.item() == pytest.approx(3.0) assert model.weight.grad.item() == pytest.approx(3.0) -def test_zero_weights_fail_before_forward(): +def test_zero_weights_run_graph_connected_zero_backward(): model = ScaleModel() - with pytest.raises(ValueError, match="positive global weight sum"): - Engine(model, device="cpu").forward_backward( - [[_datum([1, 2], [0.0, 0.0])]], - _identity_loss, - ) - assert model.forward_calls == 0 - assert model.weight.grad is None + loss, _ = Engine(model, device="cpu").forward_backward( + [_datum([1, 2], [0.0, 0.0])], + _identity_loss, + ) + assert loss.item() == 0 + assert model.forward_calls == 1 + assert model.weight.grad.item() == 0 def test_pipeline_parallelism_fails_before_forward(): @@ -781,7 +1073,7 @@ def test_pipeline_parallelism_fails_before_forward(): with pytest.raises(NotImplementedError, match="pipeline"): Engine(model, device="cpu", mesh_context=mesh_context).forward_backward( - [[_datum([1])]], + [_datum([1])], _identity_loss, ) @@ -793,7 +1085,7 @@ def test_megatron_fsdp_per_token_loss_mode_fails_before_forward(): model.calculate_per_token_loss = True with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): - Engine(model, device="cpu").forward_backward([[_datum([1])]], _identity_loss) + Engine(model, device="cpu").forward_backward([_datum([1])], _identity_loss) assert model.forward_calls == 0 assert model.weight.grad is None @@ -803,8 +1095,8 @@ def test_loss_shape_must_exactly_match_weights(): model = ScaleModel() with pytest.raises(ValueError, match="exactly the same shape"): Engine(model, device="cpu").forward_backward( - [[_datum([1, 2])]], - lambda output, _inputs, _datums, _model_inputs: output[:, :1], + [_datum([1, 2])], + lambda output, _inputs: output[:, :1], ) assert model.weight.grad is None @@ -819,7 +1111,7 @@ def bad_collate(datums): with pytest.raises(ValueError, match="collate_fn changed"): Engine(model, device="cpu", collate_fn=bad_collate).forward_backward( - [[_datum([1, 2])]], + [_datum([1, 2])], _identity_loss, ) assert model.forward_calls == 0 @@ -829,12 +1121,12 @@ def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) try: model = nn.parallel.DistributedDataParallel(ScaleModel()) - bad_window = [[_datum([1])]] if rank == 0 else [[_datum([1])], [_datum([2])]] + bad_window = [_datum([1])] if rank == 0 else [_datum([1]), _datum([2])] with pytest.raises(ValueError, match="same number of microbatches"): Engine(model, device="cpu").forward_backward(bad_window, _identity_loss) assert model.module.forward_calls == 0 - window = [[_datum([1, 2])], [_datum([3])]] if rank == 0 else [[_datum([4])], [_datum([5, 6])]] + window = [_datum([1, 2]), _datum([3])] if rank == 0 else [_datum([4]), _datum([5, 6])] loss, outputs = Engine(model, device="cpu").forward_backward(window, _identity_loss) assert loss.item() == pytest.approx(3.5) assert outputs == [] @@ -855,29 +1147,23 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size model = _DDPWithCP(_DistributedCPModel()) if dp_size == 1: window = [ - [ - Datum( - model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, - loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0, 0.0, 0.0]])}, - ) - ], - [ - Datum( - model_inputs={"input_ids": torch.tensor([[5, 6, 7, 8]])}, - loss_fn_inputs={"weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]])}, - ) - ], + Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0, 0.0, 0.0]])}, + ), + Datum( + model_inputs={"input_ids": torch.tensor([[5, 6, 7, 8]])}, + loss_fn_inputs={"weights": torch.tensor([[0.0, 0.0, 1.0, 1.0]])}, + ), ] else: dp_rank = get_flat_mesh(mesh_context.device_mesh, "dp").get_local_rank() first = dp_rank * 4 + 1 window = [ - [ - Datum( - model_inputs={"input_ids": torch.arange(first, first + 4).unsqueeze(0)}, - loss_fn_inputs={"weights": torch.ones(1, 4)}, - ) - ] + Datum( + model_inputs={"input_ids": torch.arange(first, first + 4).unsqueeze(0)}, + loss_fn_inputs={"weights": torch.ones(1, 4)}, + ) ] loss, _ = Engine( @@ -913,7 +1199,7 @@ def _mismatched_context_parallel_weights_worker(rank: int, world_size: int, init device="cpu", mesh_context=mesh_context, collate_fn=collate_prebatched, - ).forward_backward([[datum]], _identity_loss) + ).forward_backward([datum], _identity_loss) assert model.module.forward_calls == 0 finally: diff --git a/tests/unit_tests/test_engine_recipe_integration.py b/tests/unit_tests/test_engine_recipe_integration.py index 533ced3574..f3c53c27df 100644 --- a/tests/unit_tests/test_engine_recipe_integration.py +++ b/tests/unit_tests/test_engine_recipe_integration.py @@ -23,6 +23,7 @@ from torch import nn from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy +from nemo_automodel.components.models.nemotron_parse.nemotron_parse_loss import NemotronParseLoss from nemo_automodel.engine import Engine, collate_prebatched from nemo_automodel.recipes.llm.train_ft import TrainFinetuneRecipeForNextTokenPrediction from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM @@ -66,26 +67,31 @@ def zero_grad(self, *args, **kwargs): @pytest.mark.parametrize( - ("recipe_cls", "vlm"), + ("recipe_cls", "vlm", "loss_kind"), [ - (TrainFinetuneRecipeForNextTokenPrediction, False), - (FinetuneRecipeForVLM, True), + (TrainFinetuneRecipeForNextTokenPrediction, False, "masked"), + (FinetuneRecipeForVLM, True, "masked"), + (FinetuneRecipeForVLM, True, "nemotron_parse"), ], ) -def test_recipes_run_one_datum_engine_window_then_one_optimizer_step(recipe_cls, vlm): +def test_recipes_run_one_datum_engine_window_then_one_optimizer_step(recipe_cls, vlm, loss_kind): model = _TinyLM(vlm=vlm) reference = _TinyLM(vlm=vlm) reference.load_state_dict(model.state_dict()) recipe = object.__new__(recipe_cls) recipe.cfg = _Config() - recipe.loss_fn = MaskedCrossEntropy() + recipe.loss_fn = ( + NemotronParseLoss(class_token_start_idx=100, reduction="sum") + if loss_kind == "nemotron_parse" + else MaskedCrossEntropy() + ) recipe.model_parts = [model] recipe.device_mesh = None recipe.moe_mesh = None recipe.pp_enabled = False recipe.dist_env = SimpleNamespace(device=torch.device("cpu"), world_size=1, is_main=True) recipe.distributed_config = SimpleNamespace(defer_fsdp_grad_sync=True) - recipe.engine = Engine(model, device="cpu", collate_fn=collate_prebatched) + recipe.engine = Engine(model, device="cpu", microbatch_size=1, collate_fn=collate_prebatched) optimizer = _CountingSGD(model.parameters()) recipe.optimizer = [optimizer] recipe.lr_scheduler = None From c1ba20fb8dfe7829773a033f68561e5be0645d95 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 17 Aug 2026 02:29:08 -0700 Subject: [PATCH 07/34] refactor(recipes): defer context-parallel MTP support Signed-off-by: HuiyingLi --- .../qwen3_5_4b_cp2_vision_frame_shard.yaml | 3 -- .../qwen3_5_122b_128k_ep8cp32.yaml | 2 -- .../qwen3_6_35b_medpix_ep8cp2_4k.yaml | 2 -- nemo_automodel/recipes/base_recipe.py | 36 ------------------- nemo_automodel/recipes/llm/train_ft.py | 2 -- nemo_automodel/recipes/vlm/finetune.py | 1 - tests/unit_tests/recipes/test_base_recipe.py | 23 ------------ .../recipes/test_finetune_vlm_helpers.py | 12 ------- tests/unit_tests/recipes/test_train_ft.py | 13 ------- 9 files changed, 94 deletions(-) diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml index 47944ccdaa..ecaa88f2ee 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_cp2_vision_frame_shard.yaml @@ -59,9 +59,6 @@ model: rope_fusion: false attn_implementation: sdpa torch_dtype: bfloat16 - # MTP's future-token shift is not yet context-parallel aware. - num_nextn_predict_layers: 0 - processor: _target_: transformers.AutoProcessor.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-4B diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml index 3af3c425c3..d47dff48b9 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_122b_128k_ep8cp32.yaml @@ -70,8 +70,6 @@ model: text_config: # FusedLinearCrossEntropy consumes the decoder hidden states directly. output_hidden_states: true - # MTP's future-token shift is not yet context-parallel aware. - num_nextn_predict_layers: 0 # Qwen3.5 stores MTP expert projections as per-expert tensors in HF checkpoints. mtp_expert_hf_layout: split backend: diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml index b363b07f82..1e841b9d73 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b_medpix_ep8cp2_4k.yaml @@ -31,8 +31,6 @@ model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.6-35B-A3B trust_remote_code: false - # MTP's future-token shift is not yet context-parallel aware. - num_nextn_predict_layers: 0 backend: _target_: nemo_automodel.components.models.common.BackendConfig attn: te diff --git a/nemo_automodel/recipes/base_recipe.py b/nemo_automodel/recipes/base_recipe.py index c948b5cbfe..69520f9e2e 100644 --- a/nemo_automodel/recipes/base_recipe.py +++ b/nemo_automodel/recipes/base_recipe.py @@ -756,42 +756,6 @@ def _get_cp_group_size(self): return 1 return device_mesh["cp"].size() - def _validate_mtp_context_parallelism(self, model_parts: list[nn.Module]) -> None: - """Reject MTP until future-token shifts are CP-aware. - - MTP currently rolls inputs and labels after context-parallel sharding. - A local roll cannot recover the next token across CP rank boundaries, - and round-robin layouts introduce additional false adjacencies. The - model-wide reduction makes the decision identical on every PP stage, - including stages that do not locally own the MTP module. - """ - mesh_context = getattr(self, "mesh_context", None) - if mesh_context is None or mesh_context.cp_size <= 1: - return - - modules = ( - module - for part in model_parts - for module in (part.modules() if callable(getattr(part, "modules", None)) else (part,)) - ) - local_has_mtp = any( - getattr(module, "mtp", None) is not None - or bool(getattr(getattr(module, "mtp_config", None), "enabled", False)) - for module in modules - ) - enabled = torch.tensor( - int(local_has_mtp), - dtype=torch.int32, - device=getattr(getattr(self, "dist_env", None), "device", torch.device("cpu")), - ) - if dist.is_initialized(): - dist.all_reduce(enabled, op=dist.ReduceOp.MAX, group=getattr(mesh_context, "process_group", None)) - if bool(enabled.item()): - raise NotImplementedError( - "MTP with context parallelism is not supported because future-token shifts are not CP-aware; " - "set the model's MTP layer count to 0 or use cp_size=1" - ) - def _set_moe_aux_loss_backward_scale(self, *, num_batches: int, num_label_tokens: int) -> None: """Set the per-microbatch MoE auxiliary-loss scale for one optimizer step. diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 865b68c2b7..0e951fe546 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -693,8 +693,6 @@ def setup(self): self.model_parts = [model] self.pp = None - self._validate_mtp_context_parallelism(self.model_parts) - # Loss-function capability check self.loss_fn = _maybe_downgrade_loss_fn(self.loss_fn, self.model_parts[0], self.pp is not None) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 8beb181d6e..c5c31a4cdc 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -574,7 +574,6 @@ def setup(self): else: self.model_parts = [model] self.pp = None - self._validate_mtp_context_parallelism(self.model_parts) self.pipeline_loss_fn = None if self.pp_enabled: self._configure_pipeline_loss_fn() diff --git a/tests/unit_tests/recipes/test_base_recipe.py b/tests/unit_tests/recipes/test_base_recipe.py index 93cd4df5f1..38351971cb 100644 --- a/tests/unit_tests/recipes/test_base_recipe.py +++ b/tests/unit_tests/recipes/test_base_recipe.py @@ -328,29 +328,6 @@ def fake_all_reduce(tensor, op=None, group=None): assert calls[0][1] is None -def test_mtp_cp_validation_reaches_consensus_across_pipeline_stages(monkeypatch): - """A stage without the MTP module must still reject when a peer owns it.""" - group = object() - recipe = SimpleNamespace( - mesh_context=SimpleNamespace(cp_size=2, process_group=group), - dist_env=SimpleNamespace(device=torch.device("cpu")), - ) - part = nn.Linear(2, 2) - calls = [] - - def fake_all_reduce(flag, op=None, group=None): - calls.append((op, group)) - flag.fill_(1) - - monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) - monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) - - with pytest.raises(NotImplementedError, match="MTP with context parallelism"): - BaseRecipe._validate_mtp_context_parallelism(recipe, [part]) - - assert calls == [(torch.distributed.ReduceOp.MAX, group)] - - def test_optimizer_checkpoint_part_ids_use_global_pipeline_stage_indices(tmp_path): recipe_inst = _ToyRecipe(tmp_path) recipe_inst.pp = SimpleNamespace( diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index b583713003..170d3f68b3 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -1932,18 +1932,6 @@ def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch): assert trainer.engine is not None -def test_vlm_setup_rejects_cp_with_mtp(monkeypatch): - cfg = _minimal_vlm_cfg(cp_size=2, rope_fusion=True) - _patch_vlm_setup_minimals(monkeypatch, cp_size=2) - model = DummyModel() - model.mtp = nn.Identity() - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.build_model", lambda *args, **kwargs: model) - - trainer = FinetuneRecipeForVLM(cfg) - with pytest.raises(NotImplementedError, match="MTP with context parallelism"): - trainer.setup() - - def test_vlm_setup_builds_engine_for_magi(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) _patch_vlm_setup_minimals(monkeypatch, cp_size=1) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index a2ac06a13f..dd9adab1b6 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -2386,19 +2386,6 @@ def test_setup_builds_engine_for_mtp_without_context_parallelism(monkeypatch): assert trainer.engine is not None -def test_setup_rejects_mtp_with_context_parallelism(monkeypatch): - cp_size = 2 - cfg = _minimal_cfg_with_rope_fusion(cp_size=cp_size, rope_fusion=True) - _patch_setup_minimals_with_cp(monkeypatch, cp_size=cp_size) - model = DummyModel() - model.mtp = nn.Identity() - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.build_model", lambda *args, **kwargs: model) - - trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) - with pytest.raises(NotImplementedError, match="MTP with context parallelism"): - trainer.setup() - - def test_setup_builds_engine_for_magi(monkeypatch): cfg = _minimal_cfg_with_rope_fusion(cp_size=1, rope_fusion=True) _patch_setup_minimals_with_cp(monkeypatch, cp_size=1) From 9c1dc33e93924d9260623081c27a31a6041f96a8 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Mon, 17 Aug 2026 03:41:01 -0700 Subject: [PATCH 08/34] fix(engine): update pipeline metadata before batch contexts Signed-off-by: HuiyingLi --- nemo_automodel/engine/__init__.py | 127 +++++++++++++++++++----------- tests/unit_tests/test_engine.py | 61 ++++++++++++++ 2 files changed, 142 insertions(+), 46 deletions(-) diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 359e606b32..43e99d9aba 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -418,60 +418,95 @@ def _pipeline_step( outputs_by_microbatch: list[list[dict[str, Any]] | None] = [None] * self.pipeline.num_microbatches returns_outputs: bool | None = None - with self.context_fn(model_inputs), cp_context(): - model_microbatches, loss_microbatches = self._materialize_pipeline_microbatches(model_inputs, loss_inputs) - primary = model_microbatches[0].get("inputs_embeds", model_microbatches[0].get("input_ids")) + with cp_context(): + primary_name = _primary_name(model_inputs) + primary = model_inputs[primary_name] if not isinstance(primary, torch.Tensor) or primary.ndim == 0: raise ValueError("pipeline Engine requires a tensor input_ids or inputs_embeds") - if model_microbatches[0].get("qkv_format") == "thd" and self.pipeline.num_microbatches == 1: - seq_len = primary.shape[0] + + num_microbatches = self.pipeline.num_microbatches + is_thd = model_inputs.get("qkv_format") == "thd" + if num_microbatches == 1: + primary_microbatch = primary + elif is_thd: + if primary.shape[0] != num_microbatches: + raise ValueError( + f"THD sharder produced {primary.shape[0]} chunks, " + f"expected {num_microbatches} pipeline microbatches" + ) + primary_microbatch = primary.narrow(0, 0, 1) + else: + batch_size = primary.shape[0] + if batch_size % num_microbatches != 0: + raise ValueError( + f"pipeline outer batch size {batch_size} must be divisible by {num_microbatches} microbatches" + ) + materialized_batch_size = batch_size // num_microbatches + if materialized_batch_size != self.pipeline.pp_microbatch_size: + raise ValueError( + f"materialized pipeline microbatch has batch size {materialized_batch_size}, " + f"but AutoPipeline is configured for pp_microbatch_size={self.pipeline.pp_microbatch_size}" + ) + primary_microbatch = primary.narrow(0, 0, materialized_batch_size) + + if is_thd and num_microbatches == 1: + seq_len = primary_microbatch.shape[0] else: - seq_len = primary.shape[1] if primary.ndim >= 2 else primary.shape[0] - effective_microbatch_size = 1 if model_microbatches[0].get("qkv_format") == "thd" else primary.shape[0] + seq_len = primary_microbatch.shape[1] if primary_microbatch.ndim >= 2 else primary_microbatch.shape[0] + effective_microbatch_size = 1 if is_thd else primary_microbatch.shape[0] self.pipeline.update_seq_len( seq_len, microbatch_size=effective_microbatch_size, - input_tensor=primary, + input_tensor=primary_microbatch, ) - def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: - nonlocal returns_outputs - loss_inputs_mb = loss_microbatches[microbatch_index] - result = loss_fn(output, loss_inputs_mb) - has_outputs = isinstance(result, tuple) - if returns_outputs is None: - returns_outputs = has_outputs - elif returns_outputs != has_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - if isinstance(result, tuple): - losses, batch_outputs = result - if ( - not isinstance(batch_outputs, Sequence) - or isinstance(batch_outputs, (str, bytes)) - or not all(isinstance(item, Mapping) for item in batch_outputs) - ): - raise ValueError("loss_fn outputs must be a sequence of mappings") - if len(datums) == 1 and self.pipeline.num_microbatches > 1: - raise ValueError( - "a prebatched Datum may return outputs only when num_microbatches=1 because " - "its inner sample boundaries are not part of the Datum contract" - ) - outputs_by_microbatch[microbatch_index] = [_detach(dict(item)) for item in batch_outputs] - else: - losses = result - numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) - if zero_denominator: - numerator = numerator * 0 - local_loss_sum.add_(numerator.detach().to(torch.float64)) - return numerator * (grad_group_size / denominator) - - losses = [] if self.pipeline.info.has_last_stage else None - self.pipeline.step_microbatches( - model_microbatches, - loss_fn=pipeline_loss, - losses=losses, - return_outputs=False, - ) + # Batch contexts may inspect the stage metadata to decide whether a + # real forward will be consumed for dynamic shape inference. Update + # that metadata first so VLM media cursors are not reset after the + # first actual pipeline microbatch. + with self.context_fn(model_inputs): + model_microbatches, loss_microbatches = self._materialize_pipeline_microbatches( + model_inputs, loss_inputs + ) + + def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: + nonlocal returns_outputs + loss_inputs_mb = loss_microbatches[microbatch_index] + result = loss_fn(output, loss_inputs_mb) + has_outputs = isinstance(result, tuple) + if returns_outputs is None: + returns_outputs = has_outputs + elif returns_outputs != has_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + if isinstance(result, tuple): + losses, batch_outputs = result + if ( + not isinstance(batch_outputs, Sequence) + or isinstance(batch_outputs, (str, bytes)) + or not all(isinstance(item, Mapping) for item in batch_outputs) + ): + raise ValueError("loss_fn outputs must be a sequence of mappings") + if len(datums) == 1 and self.pipeline.num_microbatches > 1: + raise ValueError( + "a prebatched Datum may return outputs only when num_microbatches=1 because " + "its inner sample boundaries are not part of the Datum contract" + ) + outputs_by_microbatch[microbatch_index] = [_detach(dict(item)) for item in batch_outputs] + else: + losses = result + numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) + if zero_denominator: + numerator = numerator * 0 + local_loss_sum.add_(numerator.detach().to(torch.float64)) + return numerator * (grad_group_size / denominator) + + losses = [] if self.pipeline.info.has_last_stage else None + self.pipeline.step_microbatches( + model_microbatches, + loss_fn=pipeline_loss, + losses=losses, + return_outputs=False, + ) outputs: list[dict[str, Any]] = [] if self.pipeline.info.has_last_stage and returns_outputs: diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index c4b7a786a7..32ad50fa00 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -30,6 +30,7 @@ from nemo_automodel import Datum as PublicDatum from nemo_automodel import Engine as PublicEngine from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.datasets.vlm.pp_media import VLM_PP_MEDIA_KEY, stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel.sharder import ( ContextParallelSharder, @@ -573,6 +574,66 @@ def loss_fn(output, inputs): assert not active +def test_pipeline_stage_metadata_prevents_real_forward_from_resetting_media_cursor(): + class MediaModel(ScaleModel): + def __init__(self): + super().__init__() + self.consumed_media = [] + + def forward(self, input_ids, **kwargs): + chunk = self._vlm_pixel_values_chunks[self._vlm_chunk_idx] + self._vlm_chunk_idx += 1 + self.consumed_media.append(int(chunk.item())) + assert int(input_ids.item()) == int(chunk.item()) + return super().forward(input_ids, **kwargs) + + model = MediaModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + stage = SimpleNamespace(is_first=True, _user_meta=None) + schedule = SimpleNamespace(_stage_forward_initialized=False) + pipeline._info = SimpleNamespace( + has_first_stage=True, + has_last_stage=True, + stages=[stage], + schedule=schedule, + ) + original_update_seq_len = pipeline.update_seq_len + + def update_seq_len(seq_len, *, microbatch_size=None, input_tensor=None): + original_update_seq_len(seq_len, microbatch_size=microbatch_size, input_tensor=input_tensor) + stage._user_meta = SimpleNamespace(inputs=(object(),), outputs=(object(),)) + + pipeline.update_seq_len = update_seq_len + + @contextmanager + def batch_context(model_inputs): + with stage_vlm_media_for_pp(pipeline, [model], model_inputs): + yield + + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1], [2]]), + VLM_PP_MEDIA_KEY: { + "pixel_values": [torch.tensor([1.0]), torch.tensor([2.0])], + "image_grid_hws": [torch.ones(1), torch.ones(1)], + }, + }, + loss_fn_inputs={"weights": torch.ones(2, 1)}, + ) + + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + context_fn=batch_context, + ).forward_backward([datum], _identity_loss) + + assert model.consumed_media == [1, 2] + assert pipeline.updated_seq_lens == [1] + assert pipeline.step_calls == 1 + + def test_pipeline_lifecycle_and_moe_scale_cover_outer_and_inner_microbatches(monkeypatch): events = [] model = ScaleModel() From 9c591d0dbac9cefd357feb96a70c73b24725be1d Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Thu, 20 Aug 2026 14:44:09 -0700 Subject: [PATCH 09/34] feat(engine): add forward-only Datum execution Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 75 +- nemo_automodel/engine/__init__.py | 237 +++++- nemo_automodel/recipes/llm/train_ft.py | 192 +---- nemo_automodel/recipes/vlm/finetune.py | 126 +--- .../context_parallel/run_packed_pp.py | 125 +++- .../pipelining/test_autopipeline.py | 150 +++- .../recipes/test_finetune_vlm_cp_wiring.py | 170 +---- .../recipes/test_finetune_vlm_helpers.py | 247 ++----- tests/unit_tests/recipes/test_train_ft.py | 672 +++--------------- tests/unit_tests/test_engine.py | 182 ++++- 10 files changed, 973 insertions(+), 1203 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index c3ebc62aa2..4ff29105a0 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -389,9 +389,68 @@ def step_microbatches( Returns: The value returned by the underlying PyTorch pipeline schedule. """ + return self._run_prepared_microbatches( + model_inputs, + loss_fn=loss_fn, + losses=losses, + return_outputs=return_outputs, + schedule_method="step", + ) + + def eval_microbatches( + self, + model_inputs: list[dict[str, Any]], + *, + loss_fn: Callable[[Any, int], Any], + losses: list[torch.Tensor] | None = None, + return_outputs: bool = True, + ) -> Any: + """Run forward-only evaluation over already prepared model microbatches. + + This is the forward-only counterpart to :meth:`step_microbatches`. + The caller owns microbatch preparation, while the pipeline schedule owns + pipeline communication and execution order. The underlying + ``schedule.eval`` path does not run backward. + + Args: + model_inputs: Exactly :attr:`num_microbatches` complete model-input + mappings. Each mapping contains exactly one of ``input_ids`` or + ``inputs_embeds``; all remaining items are model keyword inputs. + loss_fn: Callback invoked as ``loss_fn(output, microbatch_index)``. + The index identifies the corresponding item in ``model_inputs`` + regardless of the schedule's execution order. + losses: Mutable list populated by the schedule on the last stage. + return_outputs: Whether the last stage returns merged model outputs + when supported by the installed PyTorch version. + + Returns: + The value returned by the underlying PyTorch pipeline schedule. + + Raises: + NotImplementedError: If the installed PyTorch pipeline schedule + does not provide forward-only ``eval`` execution. + """ + return self._run_prepared_microbatches( + model_inputs, + loss_fn=loss_fn, + losses=losses, + return_outputs=return_outputs, + schedule_method="eval", + ) + + def _run_prepared_microbatches( + self, + model_inputs: list[dict[str, Any]], + *, + loss_fn: Callable[[Any, int], Any], + losses: list[torch.Tensor] | None, + return_outputs: bool, + schedule_method: Literal["step", "eval"], + ) -> Any: + """Run one schedule method with an exact prepared-microbatch split.""" schedule = self._info.schedule if schedule is None: - raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") + raise RuntimeError("AutoPipeline.build() must be called before running a prepared PP schedule") if len(model_inputs) != self.num_microbatches: raise ValueError(f"Expected {self.num_microbatches} model input microbatches, got {len(model_inputs)}") @@ -413,6 +472,18 @@ def step_microbatches( def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: return loss_fn(output, int(microbatch_id.item())) + run_schedule = getattr(schedule, schedule_method, None) + if not callable(run_schedule): + if schedule_method == "eval": + raise NotImplementedError( + "forward-only pipeline execution requires a PyTorch pipeline schedule with eval(); " + "upgrade PyTorch or use non-pipeline Engine.forward" + ) + raise RuntimeError("PyTorch pipeline schedule has no callable step method") + # ``schedule.eval`` forwards arbitrary kwargs to ``schedule.step`` and + # does not list ``return_outputs`` explicitly in supported PyTorch + # releases. Inspect step for both modes so eval can still suppress the + # otherwise-unused merged output without breaking older step APIs. schedule_options = ( {"return_outputs": return_outputs} if "return_outputs" in inspect.signature(schedule.step).parameters @@ -423,7 +494,7 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) schedule._loss_fn = indexed_loss try: - return schedule.step( + return run_schedule( target=torch.arange(self.num_microbatches, device=self.device), losses=losses, **schedule_options, diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index a65b5d5edd..766ae2c776 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -18,6 +18,7 @@ from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass from typing import Any import torch @@ -50,7 +51,7 @@ _LOSS_FIELD_PREFIX = "__engine_loss__" _LOSS_METADATA = ("cu_seqlens", "cu_seqlens_padded", "max_seqlen", "padding_mask") -__all__ = ["Engine", "collate_prebatched"] +__all__ = ["Engine", "ForwardResult", "collate_prebatched"] def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextManager[Any]: @@ -78,15 +79,41 @@ def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], dict[str, t return dict(datum.model_inputs), dict(datum.loss_fn_inputs) +@dataclass(frozen=True) +class ForwardResult: + """Forward-only loss statistics and per-Datum outputs. + + ``loss_sum`` and ``weight_sum`` are complete across model-parallel CP and + PP ranks, but remain local to one data-parallel replica. The Engine adds no + per-call DP loss-statistic collective, so callers reduce the two sums once + at the end of an evaluation epoch. Distributed model wrappers may still + communicate during forward and therefore retain their own call-alignment + requirements. ``loss_fn_outputs`` remains local to the replica's input + Datums; PP stages receive identical detached copies, while any CP-local + tensor layout inside those mappings remains caller defined. + + Attributes: + loss_sum: Detached weighted numerator for this Datum window. + weight_sum: Detached full-sequence weight denominator for this window. + loss_fn_outputs: Detached per-Datum mappings in input order. + """ + + loss_sum: torch.Tensor + weight_sum: torch.Tensor + loss_fn_outputs: list[dict[str, Any]] + + class Engine: - """Run model forward/backward over one optimizer accumulation window. + """Run model forward or forward/backward over Datum windows. The model and distributed topology are already constructed when they are - passed here. The Engine owns batching, global weight normalization, - gradient-accumulation synchronization, and backward. It deliberately does - not zero, clip, finalize expert gradients, or step them; callers choose the - optimizer boundary and retain the repository's existing distributed - gradient-finalization path. + passed here. The Engine owns batching and model-parallel execution. + :meth:`forward` performs evaluation without gradients; + :meth:`forward_backward` additionally owns global weight normalization, + gradient-accumulation synchronization, and backward. The Engine + deliberately does not zero, clip, finalize expert gradients, or step them; + callers choose the optimizer boundary and retain the repository's existing + distributed gradient-finalization path. Args: model: An already configured and distributed model, or a built @@ -157,6 +184,114 @@ def __init__( self.context_fn = context_fn self.defer_fsdp_grad_sync = defer_fsdp_grad_sync + @torch.no_grad() + def forward( + self, + datums: Sequence[Datum], + loss_fn: LossFn, + ) -> ForwardResult: + """Run a forward-only Datum window. + + The Engine groups and collates the flat Datum sequence exactly as in + :meth:`forward_backward`, then applies the same device movement, + context-parallel layout, packed metadata, batch contexts, and pipeline + microbatch materialization. Model parts run in evaluation mode and no + autograd graph, gradient synchronization, or backward lifecycle is + created. + + Loss statistics are reduced only across model-parallel CP and PP + ranks. They deliberately remain local to one DP replica, avoiding an + Engine-introduced per-call DP loss collective; callers perform one + dataset-level DP reduction afterwards. Distributed wrappers such as + DDP/FSDP may still require aligned forward calls for their own model + communication. + + Args: + datums: Flat sequence of Datum items for this forward-only window. + loss_fn: Computes a per-element loss tensor or scalar local + weighted numerator from the raw model output and CP-local loss + inputs. It may additionally return one output mapping per + Datum. + + Returns: + Detached model-parallel-complete loss statistics and per-Datum + outputs. The sums are local to one data-parallel replica. + """ + microbatches = self._group_datums(datums) + self._validate_execution_parallelism() + cp_group, cp_size = self._cp_group_and_size() + self._validate_window_size_across_group(len(microbatches), cp_group, cp_size) + weight_sum = self._local_weight_sum(microbatches) + zero_weight_sum = bool(weight_sum == 0) + self._validate_pipeline_window(len(microbatches), weight_sum) + + for part in self.model_parts: + part.eval() + inner_microbatches = self.pipeline.num_microbatches if self.pipeline is not None else 1 + local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) + loss_fn_outputs: list[dict[str, Any]] = [] + returns_outputs: bool | None = None + + for batch_datums in microbatches: + cp_context, model_inputs, loss_inputs = self._prepare_batch(batch_datums, inner_microbatches) + if self.pipeline is not None: + batch_returns_outputs, batch_outputs = self._pipeline_execute( + model_inputs, + loss_inputs, + batch_datums, + loss_fn, + local_loss_sum, + cp_context, + backward_scale=None, + zero_weight_sum=zero_weight_sum, + ) + if batch_returns_outputs is not None: + if returns_outputs is None: + returns_outputs = batch_returns_outputs + elif returns_outputs != batch_returns_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + loss_fn_outputs.extend(batch_outputs) + continue + + loss_inputs = _with_loss_metadata(model_inputs, loss_inputs) + with self.context_fn(model_inputs), cp_context(): + forward_inputs = filter_forward_kwargs(self.model, model_inputs) + output = self.model(**forward_inputs) + result = loss_fn(output, loss_inputs) + has_outputs = isinstance(result, tuple) + if returns_outputs is None: + returns_outputs = has_outputs + elif returns_outputs != has_outputs: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + if isinstance(result, tuple): + losses, outputs = result + if ( + not isinstance(outputs, Sequence) + or isinstance(outputs, (str, bytes)) + or len(outputs) != len(batch_datums) + or not all(isinstance(item, Mapping) for item in outputs) + ): + raise ValueError("loss_fn outputs must contain one mapping per Datum") + loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) + else: + losses = result + numerator = _weighted_numerator(losses, loss_inputs["weights"]) + if zero_weight_sum: + numerator = numerator * 0 + local_loss_sum.add_(numerator.detach().to(torch.float64)) + + if cp_size > 1: + dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=cp_group) + pp_group, pp_size = self._pp_group_and_size() + if pp_size > 1: + dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=pp_group) + + return ForwardResult( + loss_sum=local_loss_sum.detach(), + weight_sum=weight_sum.detach(), + loss_fn_outputs=loss_fn_outputs, + ) + def forward_backward( self, datums: Sequence[Datum], @@ -237,16 +372,20 @@ def forward_backward( cp_context, model_inputs, loss_inputs = self._prepare_batch(datums, inner_microbatches) if self.pipeline is not None: - batch_returns_outputs, batch_outputs = self._pipeline_step( + backward_scale = ( + safe_denominator.new_zeros(()) + if zero_denominator + else safe_denominator.new_tensor(grad_group_size) / safe_denominator + ) + batch_returns_outputs, batch_outputs = self._pipeline_execute( model_inputs, loss_inputs, datums, loss_fn, - safe_denominator, - zero_denominator, - grad_group_size, local_loss_sum, cp_context, + backward_scale=backward_scale, + zero_weight_sum=zero_denominator, ) if batch_returns_outputs is not None: if returns_outputs is None: @@ -301,9 +440,9 @@ def forward_backward( def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: if not isinstance(datums, Sequence) or isinstance(datums, (str, bytes)) or not datums: - raise ValueError("forward_backward requires a non-empty flat sequence of Datum") + raise ValueError("Engine requires a non-empty flat sequence of Datum") if not all(isinstance(datum, Datum) for datum in datums): - raise TypeError("forward_backward received a value that is not a Datum") + raise TypeError("Engine received a value that is not a Datum") return [ list(datums[start : start + self.microbatch_size]) for start in range(0, len(datums), self.microbatch_size) ] @@ -496,18 +635,36 @@ def _attach_mtp_cp_inputs( ) return result - def _pipeline_step( + def _pipeline_execute( self, model_inputs: dict[str, Any], loss_inputs: LossInputs, datums: Sequence[Datum], loss_fn: LossFn, - denominator: torch.Tensor, - zero_denominator: bool, - grad_group_size: int, local_loss_sum: torch.Tensor, cp_context: Callable[[], AbstractContextManager[Any]], + *, + backward_scale: torch.Tensor | None, + zero_weight_sum: bool, ) -> tuple[bool | None, list[dict[str, Any]]]: + """Run prepared pipeline microbatches in training or forward-only mode. + + Args: + model_inputs: CP-prepared outer-batch model inputs. + loss_inputs: CP-prepared outer-batch loss inputs. + datums: Datum items represented by the outer batch. + loss_fn: Model-output loss callback. + local_loss_sum: Accumulator updated with detached numerators. + cp_context: Context covering the complete pipeline schedule. + backward_scale: Multiplier returned to the training schedule for + backward, or ``None`` to run the forward-only schedule. + zero_weight_sum: Whether reporting numerators must be forced to + graph-connected zero. + + Returns: + Whether the callback returned outputs, and its detached outputs in + logical Datum order. + """ outputs_by_microbatch: list[list[dict[str, Any]] | None] = [None] * self.pipeline.num_microbatches returns_outputs: bool | None = None @@ -588,13 +745,16 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: else: losses = result numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) - if zero_denominator: + if zero_weight_sum: numerator = numerator * 0 local_loss_sum.add_(numerator.detach().to(torch.float64)) - return numerator * (grad_group_size / denominator) + return numerator if backward_scale is None else numerator * backward_scale losses = [] if self.pipeline.info.has_last_stage else None - self.pipeline.step_microbatches( + run_microbatches = ( + self.pipeline.eval_microbatches if backward_scale is None else self.pipeline.step_microbatches + ) + run_microbatches( model_microbatches, loss_fn=pipeline_loss, losses=losses, @@ -715,6 +875,8 @@ def _materialize_pipeline_microbatches( return model_microbatches, loss_microbatches def _validate_parallelism(self) -> None: + """Validate topology plus backward-specific distributed contracts.""" + self._validate_execution_parallelism() if any( bool(getattr(module, "calculate_per_token_loss", False)) for part in self.model_parts @@ -724,6 +886,11 @@ def _validate_parallelism(self) -> None: "Engine.forward_backward requires averaged distributed gradients; " "MegatronFSDP calculate_per_token_loss=True uses summed gradients" ) + if self.pipeline is not None and self.pipeline.scale_grads_in_schedule: + raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") + + def _validate_execution_parallelism(self) -> None: + """Validate model-parallel topology shared by forward and backward.""" if ( self.pipeline is not None and self._cp_size() > 1 @@ -736,8 +903,6 @@ def _validate_parallelism(self) -> None: raise NotImplementedError( "MTP with context and pipeline parallelism is not supported; use PP size 1 or CP size 1" ) - if self.pipeline is not None and self.pipeline.scale_grads_in_schedule: - raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") if self.pipeline is not None and self.mesh_context is None: raise ValueError("pipeline Engine requires mesh_context") if self.mesh_context is None: @@ -750,6 +915,19 @@ def _validate_parallelism(self) -> None: def _cp_size(self) -> int: return self.mesh_context.cp_size if self.mesh_context is not None else 1 + def _cp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: + if ( + self.mesh_context is None + or self.mesh_context.device_mesh is None + or self.mesh_context.cp_size <= 1 + or not dist.is_available() + or not dist.is_initialized() + ): + return None, 1 + cp_mesh = self.mesh_context.device_mesh["cp"] + size = int(cp_mesh.size()) + return (cp_mesh.get_group() if size > 1 else None), size + def _dp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: if self.mesh_context is not None and self.mesh_context.device_mesh is not None: dp_mesh = get_flat_mesh(self.mesh_context.device_mesh, "dp") @@ -789,9 +967,7 @@ def _validate_pipeline_window(self, size: int, denominator: torch.Tensor) -> Non if not bool((gathered[:, 0] == gathered[0, 0]).all()): raise ValueError(f"pipeline stages must use the same outer window size; got {gathered[:, 0].tolist()}") if not torch.allclose(gathered[:, 1], gathered[0, 1].expand(pp_size), rtol=1e-8, atol=1e-12): - raise ValueError( - f"pipeline stages must use the same DP-reduced weight denominator; got {gathered[:, 1].tolist()}" - ) + raise ValueError(f"pipeline stages must use the same weight denominator; got {gathered[:, 1].tolist()}") def _global_weight_sum( self, @@ -799,6 +975,13 @@ def _global_weight_sum( dp_group: dist.ProcessGroup | None, dp_size: int, ) -> torch.Tensor: + denominator = self._local_weight_sum(microbatches) + if dp_size > 1: + dist.all_reduce(denominator, op=dist.ReduceOp.SUM, group=dp_group) + return denominator + + def _local_weight_sum(self, microbatches: list[list[Datum]]) -> torch.Tensor: + """Return the full-sequence denominator for one DP replica.""" local_sum = 0.0 for datum in (datum for microbatch in microbatches for datum in microbatch): weights = datum.loss_fn_inputs.get("weights") @@ -810,8 +993,6 @@ def _global_weight_sum( denominator = torch.tensor(local_sum, dtype=torch.float64, device=self.device) self._validate_weight_sum_across_cp(denominator) - if dp_size > 1: - dist.all_reduce(denominator, op=dist.ReduceOp.SUM, group=dp_group) return denominator def _validate_window_size_across_group( @@ -826,7 +1007,7 @@ def _validate_window_size_across_group( sizes = torch.empty(group_size, dtype=torch.int64, device=self.device) dist.all_gather_into_tensor(sizes, local_size, group=group) if not bool((sizes == sizes[0]).all()): - raise ValueError(f"every gradient rank must use the same number of microbatches; got {sizes.tolist()}") + raise ValueError(f"every participating rank must use the same number of microbatches; got {sizes.tolist()}") def _validate_weight_sum_across_cp(self, local_weight_sum: torch.Tensor) -> None: """Verify that CP replicas started from the same full-sequence weights.""" diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index e64d5cb978..4de5860d4a 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -59,7 +59,6 @@ from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.loader import DataloaderConfig from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig -from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi from nemo_automodel.components.distributed.init_utils import initialize_distributed from nemo_automodel.components.distributed.mesh import MeshContext @@ -91,7 +90,6 @@ from nemo_automodel.components.utils.model_utils import ( _supports_logits_to_keep, _supports_seq_lens, - filter_forward_kwargs, resolve_trust_remote_code, ) from nemo_automodel.engine import Engine, collate_prebatched @@ -951,8 +949,8 @@ def _configure_pipeline_loss_fn(self): last_stage_model, grad_reduce_group=self._get_dp_group(include_cp=True), ) - # Validation still executes the schedule directly. Training supplies the - # same loss through Engine's per-microbatch callback. + # Engine supplies this loss through its training and forward-only + # per-microbatch callbacks. self.pp.info.schedule._loss_fn = self.pipeline_loss_fn def _setup_qat(self, cfg, model_parts: list[nn.Module]): @@ -1071,71 +1069,6 @@ def run_train_validation_loop(self): self._partial_cuda_graph_capture_pending = False # ------------------ helpers ------------------ - def _prepare_validation_batch(self, batch: dict[str, Any]): - """Move and CP-prepare one validation batch. - - Args: - batch: Worker-collated inputs. ``input_ids`` and ``labels`` have - shape [batch, sequence], while optional ``position_ids`` has - shape [batch, sequence] or [axes, batch, sequence]. Packed THD - token tensors have shape [tokens]. - - Returns: - The CP context factory; a CP-local model-input mapping; CP-local - labels; and optional per-depth MTP targets. The model mapping's - token IDs and validity masks have shape [batch, local_sequence], - while multi-axis positions have shape [axes, batch, - local_sequence]. Labels and each target tensor have shape [batch, - local_sequence]. Packed outputs use the corresponding CP-local - [tokens] layout. - """ - batch = { - k: ( - {dk: dv.to(self.dist_env.device, non_blocking=True) for dk, dv in v.items() if dv is not None} - if isinstance(v, dict) - else (v.to(self.dist_env.device, non_blocking=True) if isinstance(v, torch.Tensor) else v) - ) - for k, v in batch.items() - } - model = self.model_parts[0] if hasattr(self, "model_parts") else None - supports = getattr(model, "supports", None) - mtp_cp_enabled = ( - not self.pp_enabled and self._get_cp_group_size() > 1 and bool(getattr(supports, "mtp_enabled", False)) - ) - if mtp_cp_enabled and not bool(getattr(supports, "supports_mtp_cp", False)): - raise NotImplementedError( - f"{type(model).__name__} declares supports_mtp_cp=False; " - "MTP target preparation for context parallelism is unavailable" - ) - cp_sharder = ContextParallelSharder( - model, - self.device_mesh, - batch, - padding_token_id=self.tokenizer.pad_token_id if self.tokenizer else 0, - num_chunks=self.pp.pp_batch_size // self.pp.pp_microbatch_size if self.pp_enabled else 1, - ) - mtp_cp_inputs = ( - model.prepare_mtp_inputs_for_cp(batch, ignore_index=self.cfg.mtp.ignore_index) if mtp_cp_enabled else None - ) - train_ctx, batch = cp_sharder.shard(batch) - mtp_per_depth_targets = None - if mtp_cp_inputs is not None: - batch["mtp_per_depth_input_ids"] = tuple( - cp_sharder.shard_token_tensor(ids, seq_dim=1, fill=0) for ids in mtp_cp_inputs.input_ids - ) - batch["mtp_per_depth_position_ids"] = tuple( - cp_sharder.shard_token_tensor(ids, seq_dim=mtp_cp_inputs.position_ids_seq_dim, fill=0) - for ids in mtp_cp_inputs.position_ids - ) - batch["mtp_per_depth_valid_masks"] = tuple( - cp_sharder.shard_token_tensor(mask, seq_dim=1, fill=False) for mask in mtp_cp_inputs.valid_masks - ) - mtp_per_depth_targets = tuple( - cp_sharder.shard_token_tensor(targets, seq_dim=1, fill=self.cfg.mtp.ignore_index) - for targets in mtp_cp_inputs.targets - ) - return train_ctx, batch, batch.pop("labels"), mtp_per_depth_targets - def _compute_causal_lm_loss(self, output, labels, model_inputs, *, num_label_tokens, is_train): """Compute the recipe's causal-LM and optional MTP loss. @@ -1262,80 +1195,24 @@ def _engine_loss_fn( is_train=True, ) - def _forward_validation_step(self, batch: dict[str, Any]) -> torch.Tensor: - """Run one recipe-owned forward-only validation step. - - Args: - batch: Worker-collated inputs and labels. Padded token tensors have - shape [batch, sequence]; packed THD token tensors have shape [tokens]. - - Returns: - Detached scalar local loss-sum tensor. - """ - train_ctx, batch, labels, mtp_per_depth_targets = self._prepare_validation_batch(batch) - fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext() - + def _engine_validation_loss_fn( + self, + output: Any, + loss_inputs: dict[str, torch.Tensor | tuple[torch.Tensor, ...]], + ) -> torch.Tensor: + """Compute the validation loss numerator without training reductions.""" if self.pp_enabled: - with train_ctx(), fp8_ctx: - losses = [] if self.pp.info.has_last_stage else None - if self.pp.info.has_last_stage: - masked_labels = labels.clone() - targets = masked_labels - else: - targets = None - - input_ids = batch.pop("input_ids") - - # Update PP stage shapes for the current batch's seq_len. - # This is a no-op when the length hasn't changed. - self.pp.update_seq_len(input_ids.shape[1]) - - # Filter out None values and empty dicts from batch to avoid PP chunking errors - batch_filtered = { - k: v for k, v in batch.items() if v is not None and not (isinstance(v, dict) and len(v) == 0) - } - # Hand the THD ``cu_seqlens`` to the PP loss to mask cross-sequence boundaries — - # the fallback when the model emits no per-microbatch seq_idx tail (which the loss - # prefers). One cu_seqlens encodes a single shared layout, so it is only correct at - # one pack/microbatch per step; the seq_idx tail handles differing per-microbatch boundaries. - cu_seqlens = batch_filtered.get("cu_seqlens") - if isinstance(cu_seqlens, torch.Tensor) and cu_seqlens.dim() == 2: - cu_seqlens = cu_seqlens.squeeze(0) # [1, T] -> [T] - if self.pipeline_loss_fn is not None: - self.pipeline_loss_fn.cu_seqlens = cu_seqlens - if self.pp.info.has_first_stage: - self.pp.info.schedule.eval(input_ids, target=targets, losses=losses, **batch_filtered) - else: - self.pp.info.schedule.eval(target=targets, losses=losses, **batch_filtered) - - if self.pp.info.has_last_stage: - return torch.sum(torch.stack(losses)).detach() - return torch.zeros((), device=self.dist_env.device) - - model = self.model_parts[0] - with train_ctx(), fp8_ctx: - loss_inputs = dict(batch) - if mtp_per_depth_targets is not None: - loss_inputs["mtp_per_depth_targets"] = mtp_per_depth_targets - batch = filter_forward_kwargs(model, batch) - if isinstance(self.loss_fn, FusedLinearCrossEntropy): - out = model(logits_to_keep=1, **batch) - else: - out = model(**batch) - return self._compute_causal_lm_loss( - out, - labels, - loss_inputs, - num_label_tokens=None, - is_train=False, - ).detach() - - def _broadcast_from_last_pp_stage(self, tensor: torch.Tensor) -> torch.Tensor: - """Broadcast a PP last-stage scalar to the other ranks in its pipeline group.""" - pp_group = self.device_mesh["pp"].get_group() - pp_src_rank = torch.distributed.get_global_rank(pp_group, torch.distributed.get_world_size(pp_group) - 1) - torch.distributed.broadcast(tensor, src=pp_src_rank, group=pp_group) - return tensor + if self.pipeline_loss_fn is None: + raise RuntimeError("The last pipeline stage has no configured causal-LM loss") + self.pipeline_loss_fn.cu_seqlens = loss_inputs.get("cu_seqlens") + return self.pipeline_loss_fn(output, loss_inputs["labels"]) + return self._compute_causal_lm_loss( + output, + loss_inputs["labels"], + loss_inputs, + num_label_tokens=None, + is_train=False, + ) def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: float | None = None) -> MetricsSample: """Execute a single training step. @@ -1464,37 +1341,26 @@ def _run_validation_epoch(self, val_dataloader): """Run one pass over a single validation dataloader. Args: - val_name: Name of the validation dataset. val_dataloader: DataLoader for the validation dataset. """ with ScopedRNG(seed=1, ranked=True): for mp in self.model_parts: mp.eval() - total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device) - total_num_label_tokens = 0 + total_loss = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) + total_num_label_tokens = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) for batch in val_dataloader: - num_label_tokens = (batch["labels"] != -100).sum().item() - total_loss += self._forward_validation_step(batch).item() - total_num_label_tokens += num_label_tokens - - total_loss = self._dp_allreduce(total_loss, include_cp=True) - total_num_label_tokens = self._dp_allreduce( - torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) - ).item() + result = self.engine.forward([self._make_engine_datum(batch)], self._engine_validation_loss_fn) + total_loss += result.loss_sum + total_num_label_tokens += result.weight_sum + + # Engine.forward has already reconstructed CP shards and synchronized + # PP stages. Only independent DP validation shards remain to combine. + total_loss = self._dp_allreduce(total_loss) + total_num_label_tokens = int(self._dp_allreduce(total_num_label_tokens).item()) val_loss = total_loss / max(total_num_label_tokens, 1e-8) - # For PP, send val_loss and num_label_tokens from last stage to main rank - if self.pp_enabled: - val_loss = val_loss.to(self.dist_env.device) - # On non-last ranks total_num_label_tokens is 0; this tensor is just a recv buffer. - pp_num_tokens = torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) - val_loss = self._broadcast_from_last_pp_stage(val_loss) - pp_num_tokens = self._broadcast_from_last_pp_stage(pp_num_tokens) - if self.dist_env.is_main: - total_num_label_tokens = pp_num_tokens.item() - val_loss = val_loss.item() if isinstance(val_loss, torch.Tensor) else val_loss metrics = { diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index ca17fcd3ad..e799b615fe 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -50,7 +50,6 @@ from nemo_automodel.components.datasets.datum import Datum from nemo_automodel.components.datasets.vlm.pp_media import VLM_PP_MEDIA_KEY, stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig -from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi from nemo_automodel.components.distributed.cp_vision_frame_shard import ( CpVisionFrameShardingConfig, @@ -80,7 +79,7 @@ scale_grads_and_clip_grad_norm, ) from nemo_automodel.components.utils.compile_utils import build_compile_config -from nemo_automodel.components.utils.model_utils import VLM_INPUT_KEYS, _supports_logits_to_keep, filter_forward_kwargs +from nemo_automodel.components.utils.model_utils import VLM_INPUT_KEYS, _supports_logits_to_keep from nemo_automodel.engine import Engine, collate_prebatched from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config, shard_optimizers_for_megatron_fsdp from nemo_automodel.recipes._typed_config import RecipeConfig @@ -978,6 +977,31 @@ def _engine_loss_fn( log_denominator=log_denominator, ) + def _engine_validation_loss_fn( + self, + out: Any, + loss_inputs: Mapping[str, torch.Tensor | tuple[torch.Tensor, ...]], + ) -> torch.Tensor: + """Compute the validation loss numerator without training reductions.""" + labels = cast(torch.Tensor, loss_inputs["labels"]) + cu_seqlens = cast(torch.Tensor | None, loss_inputs.get("cu_seqlens")) + if self.pp_enabled: + if self.pipeline_loss_fn is None: + raise RuntimeError("The last pipeline stage has no configured causal-LM loss") + self.pipeline_loss_fn.cu_seqlens = cu_seqlens + return self.pipeline_loss_fn(out, labels) + return self._compute_vlm_loss( + out=out, + labels=labels, + num_label_tokens=None, + is_train=False, + cu_seqlens=cu_seqlens, + mtp_per_depth_targets=cast( + tuple[torch.Tensor, ...] | None, + loss_inputs.get("mtp_per_depth_targets"), + ), + ) + @contextmanager def _cp_vision_frame_sharding_context(self): """Publish the CP-only group while a VLM forward may run its vision tower.""" @@ -1168,94 +1192,18 @@ def _run_validation_epoch(self, val_dataloader): for mp in self.model_parts: mp.eval() - total_loss = 0.0 - total_tokens = 0 - total_num_label_tokens = 0 + total_loss = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) + total_num_label_tokens = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) for batch in val_dataloader: - batch = { - k: (v.to(self.dist_env.device, non_blocking=True) if isinstance(v, torch.Tensor) else v) - for k, v in batch.items() - } - num_label_tokens = (batch["labels"] != -100).sum().item() - - model = self.model_parts[0] - cp_sharder = ContextParallelSharder( - model, - self.device_mesh, - batch, - invoke_pre_embed=not self.pp_enabled, - ) - supports = getattr(model, "supports", None) - mtp_cp_inputs = None - if ( - not self.pp_enabled - and self._get_cp_group_size() > 1 - and bool(getattr(supports, "mtp_enabled", False)) - ): - if not bool(getattr(supports, "supports_mtp_cp", False)): - raise NotImplementedError( - f"{type(model).__name__} declares supports_mtp_cp=False; " - "MTP target preparation for context parallelism is unavailable" - ) - mtp_cp_inputs = model.prepare_mtp_inputs_for_cp( - batch, - ignore_index=self.cfg.mtp.ignore_index, - ) - train_ctx, batch = cp_sharder.shard(batch) - mtp_per_depth_targets = None - if mtp_cp_inputs is not None: - batch["mtp_per_depth_input_ids"] = tuple( - cp_sharder.shard_token_tensor(ids, seq_dim=1, fill=0) for ids in mtp_cp_inputs.input_ids - ) - batch["mtp_per_depth_position_ids"] = tuple( - cp_sharder.shard_token_tensor( - ids, - seq_dim=mtp_cp_inputs.position_ids_seq_dim, - fill=0, - ) - for ids in mtp_cp_inputs.position_ids - ) - batch["mtp_per_depth_valid_masks"] = tuple( - cp_sharder.shard_token_tensor(mask, seq_dim=1, fill=False) for mask in mtp_cp_inputs.valid_masks - ) - mtp_per_depth_targets = tuple( - cp_sharder.shard_token_tensor( - targets, - seq_dim=1, - fill=self.cfg.mtp.ignore_index, - ) - for targets in mtp_cp_inputs.targets - ) - labels = batch.pop("labels") - with self._cp_vision_frame_sharding_context(), train_ctx(): - cu_seqlens = None if mtp_per_depth_targets is not None else batch.get("cu_seqlens") - batch = filter_forward_kwargs(model, batch) - if isinstance(self.loss_fn, FusedLinearCrossEntropy): - out = model(logits_to_keep=1, **batch) - else: - out = model(**batch) - local_loss = self._compute_vlm_loss( - out=out, - labels=labels, - num_label_tokens=num_label_tokens, - is_train=False, - cu_seqlens=cu_seqlens, - mtp_per_depth_targets=mtp_per_depth_targets, - ) - total_num_label_tokens += num_label_tokens - - total_loss += local_loss.item() * num_label_tokens - total_tokens += num_label_tokens - - # Aggregate across ranks if distributed is initialized - total_loss = self._dp_allreduce(torch.FloatTensor([total_loss]), include_cp=True).item() - # `num_label_tokens` is measured before CP sharding, so each CP rank - # contributes the full sequence token count while `total_loss` is - # reconstructed from CP-sharded loss sums. Do not sum tokens over CP. - total_tokens = self._dp_allreduce(torch.LongTensor([total_tokens])).item() - total_num_label_tokens = self._dp_allreduce(torch.LongTensor([total_num_label_tokens])).item() - - val_loss = total_loss / max(total_tokens, 1e-8) + result = self.engine.forward([self._make_engine_datum(batch)], self._engine_validation_loss_fn) + total_loss += result.loss_sum + total_num_label_tokens += result.weight_sum + + # Engine.forward has already reconstructed CP shards. Only independent + # DP validation shards remain to combine (VLM PP validation stays disabled). + total_loss = self._dp_allreduce(total_loss).item() + total_num_label_tokens = int(self._dp_allreduce(total_num_label_tokens).item()) + val_loss = total_loss / max(total_num_label_tokens, 1e-8) return MetricsSample( step=self.step_scheduler.step, diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index 8a3e6fe6ee..c615d3ff76 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -16,9 +16,12 @@ The test runs the same two-microbatch, four-document update through PP=2 from both raw THD metadata (``seq_lens``) and final THD metadata (``cu_seqlens``). -Each path must match a native eager Llama in loss and every local stage -gradient. A final padded two-Datum update verifies that Engine broadcasts the -callback's per-Datum mappings to both pipeline ranks in logical input order. +Each path runs training, forward-only evaluation, then training again on the +same pipeline. Evaluation must match a native eager Llama in summed loss and +weight statistics without creating gradients; both surrounding training calls +must match eager loss plus every local-stage gradient. A final padded two-Datum +update verifies that Engine broadcasts the callback's per-Datum mappings to +both pipeline ranks in logical input order. Run with:: @@ -119,7 +122,26 @@ def _token_losses(output, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: def _eager_reference( device: torch.device, -) -> tuple[torch.Tensor, dict[str, torch.Tensor], dict[str, object], torch.Tensor, torch.Tensor]: +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + dict[str, torch.Tensor], + dict[str, object], + torch.Tensor, + torch.Tensor, +]: + """Build eager eval and training references for the shared packed batch. + + Args: + device: CUDA device on which the reference model and batch run. + + Returns: + The scalar eval loss sum, scalar eval weight sum, scalar normalized + training loss, per-parameter gradient mapping, raw model-input mapping, + labels of shape [batch, sequence], and weights of shape [batch, sequence]. + Raw token and position tensors have shape [batch, sequence]. + """ raw_inputs, labels, weights = _raw_batch(device) prepared = make_cp_batch_for_te( None, @@ -128,13 +150,26 @@ def _eager_reference( prepared_labels = prepared.pop("labels") model = _build_model(device) + model.eval() + with torch.no_grad(): + eval_output = model(**prepared) + eval_token_losses = _token_losses( + eval_output, + {"labels": prepared_labels, "weights": weights.reshape(-1)}, + ) + eval_loss_sum = (eval_token_losses * weights.reshape(-1)).sum() + eval_weight_sum = weights.sum() + if any(parameter.grad is not None for parameter in model.parameters()): + raise AssertionError("eager forward-only reference unexpectedly created parameter gradients") + + model.train() output = model(**prepared) token_losses = _token_losses(output, {"labels": prepared_labels, "weights": weights.reshape(-1)}) loss = (token_losses * weights.reshape(-1)).sum() / weights.sum() loss.backward() grads = {name: parameter.grad.detach().clone() for name, parameter in model.named_parameters()} del model - return loss.detach(), grads, raw_inputs, labels, weights + return eval_loss_sum.detach(), eval_weight_sum.detach(), loss.detach(), grads, raw_inputs, labels, weights def _build_pipeline( @@ -184,12 +219,35 @@ def _run_thd_layout( layout: str, device: torch.device, mesh_context: MeshContext, + reference_eval_loss_sum: torch.Tensor, + reference_eval_weight_sum: torch.Tensor, reference_loss: torch.Tensor, reference_grads: dict[str, torch.Tensor], raw_inputs: dict[str, object], labels: torch.Tensor, weights: torch.Tensor, ) -> AutoPipeline: + """Run training, forward-only, then training parity for one THD layout. + + Args: + layout: ``raw`` for batch-major pre-THD tensors or ``final`` for the + flattened model-ready THD stream. + device: CUDA device on which this physical pipeline rank runs. + mesh_context: Runtime mesh whose PP axis has size two. + reference_eval_loss_sum: Scalar eager forward-only weighted numerator. + reference_eval_weight_sum: Scalar eager full-sequence weight sum. + reference_loss: Scalar eager normalized training loss. + reference_grads: Mapping from parameter names to eager gradient tensors + with each parameter's native shape. + raw_inputs: Mapping whose token and position tensors have shape [batch, + sequence] before THD flattening. + labels: Target token IDs of shape [batch, sequence]. + weights: Token weights of shape [batch, sequence]. + + Returns: + The two-stage pipeline after the second training pass. Every local + parameter has its native gradient shape. + """ pipeline = _build_pipeline(device, mesh_context) if layout == "raw": model_inputs = _clone_mapping(raw_inputs) @@ -209,20 +267,55 @@ def _run_thd_layout( model_inputs=model_inputs, loss_fn_inputs={"labels": loss_labels, "weights": loss_weights}, ) - loss, outputs = Engine( + engine = Engine( pipeline, device=device, mesh_context=mesh_context, collate_fn=collate_prebatched, - ).forward_backward([datum], _token_losses) + ) + + pre_eval_loss, pre_eval_outputs = engine.forward_backward([datum], _token_losses) + torch.testing.assert_close(pre_eval_loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) + assert pre_eval_outputs == [] + pre_eval_grad_diff = _assert_local_grad_parity(pipeline, reference_grads) + for part in pipeline.parts: + part.zero_grad(set_to_none=True) + + forward_result = engine.forward([datum], _token_losses) + + torch.testing.assert_close( + forward_result.loss_sum.float(), + reference_eval_loss_sum.float(), + atol=4e-2, + rtol=2e-3, + ) + torch.testing.assert_close( + forward_result.weight_sum.float(), + reference_eval_weight_sum.float(), + atol=0, + rtol=0, + ) + assert forward_result.loss_fn_outputs == [] + if any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): + raise AssertionError(f"PP2 {layout} THD forward-only evaluation unexpectedly created parameter gradients") + if any(part.training for part in pipeline.parts): + raise AssertionError(f"PP2 {layout} THD forward-only evaluation did not keep every model part in eval mode") + + # Reuse the exact pipeline immediately. Together with the training call + # above, this proves both train->eval and eval->train schedule transitions + # restore temporary split/loss callbacks and backward state. + loss, outputs = engine.forward_backward([datum], _token_losses) torch.testing.assert_close(loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) assert outputs == [] grad_diff = _assert_local_grad_parity(pipeline, reference_grads) if dist.get_rank() == 0: print( - f"PP2 {layout} THD parity passed " - f"(loss diff={(loss.float() - reference_loss.float()).abs().item():.6f}, grad max={grad_diff:.6f})" + f"PP2 {layout} THD forward+backward parity passed " + f"(eval loss-sum diff=" + f"{(forward_result.loss_sum.float() - reference_eval_loss_sum.float()).abs().item():.6f}, " + f"train loss diff={(loss.float() - reference_loss.float()).abs().item():.6f}, " + f"pre/post-eval grad max={pre_eval_grad_diff:.6f}/{grad_diff:.6f})" ) return pipeline @@ -285,11 +378,21 @@ def main() -> None: world_size=dist.get_world_size(), ) try: - reference_loss, reference_grads, raw_inputs, labels, weights = _eager_reference(device) + ( + reference_eval_loss_sum, + reference_eval_weight_sum, + reference_loss, + reference_grads, + raw_inputs, + labels, + weights, + ) = _eager_reference(device) _run_thd_layout( "raw", device, mesh_context, + reference_eval_loss_sum, + reference_eval_weight_sum, reference_loss, reference_grads, raw_inputs, @@ -301,6 +404,8 @@ def main() -> None: "final", device, mesh_context, + reference_eval_loss_sum, + reference_eval_weight_sum, reference_loss, reference_grads, raw_inputs, diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 5e4f168edb..5ae884b2c5 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -253,6 +253,9 @@ def __init__(self, *, fail_on_step: bool = False, invoke_loss: bool = False): self.losses_during_step = None self.return_outputs_during_step = None self.loss_results = [] + self.step_calls = 0 + self.eval_calls = 0 + self.split_inputs_calls = 0 def _split_inputs(self, args, kwargs=None): return split_args_kwargs_into_chunks( @@ -262,7 +265,7 @@ def _split_inputs(self, args, kwargs=None): kwargs_chunk_spec=self._kwargs_chunk_spec, ) - def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): + def _run_schedule(self, *args, target=None, losses=None, return_outputs=True, **kwargs): """Split schedule inputs using the chunk spec active during the call. Args: @@ -287,6 +290,7 @@ def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): self.return_outputs_during_step = return_outputs if self.fail_on_step: raise RuntimeError("schedule failed") + self.split_inputs_calls += 1 self.args_split, self.kwargs_split = self._split_inputs(args, kwargs) if self.invoke_loss: assert target is not None @@ -295,6 +299,26 @@ def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): self.loss_results.append(self._loss_fn(torch.tensor(float(index)), target_chunks[index])) return "schedule-result" + def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): + self.step_calls += 1 + return self._run_schedule( + *args, + target=target, + losses=losses, + return_outputs=return_outputs, + **kwargs, + ) + + def eval(self, *args, target=None, losses=None, return_outputs=True, **kwargs): + self.eval_calls += 1 + return self._run_schedule( + *args, + target=target, + losses=losses, + return_outputs=return_outputs, + **kwargs, + ) + class _LegacyStepSchedule(_KwargsChunkSchedule): """Schedule with the PyTorch 2.6-2.9 step signature.""" @@ -308,6 +332,33 @@ def step(self, *args, target=None, losses=None, **kwargs): return super().step(*args, target=target, losses=losses, **kwargs) +class _LegacyEvalSchedule(_LegacyStepSchedule): + """Schedule with an eval signature that predates return_outputs.""" + + def __init__(self): + super().__init__() + self.received_return_outputs = False + + def eval(self, *args, target=None, losses=None, **kwargs): + self.received_return_outputs = "return_outputs" in kwargs + self.eval_calls += 1 + return self._run_schedule(*args, target=target, losses=losses, **kwargs) + + +class _ForwardingEvalSchedule(_KwargsChunkSchedule): + """Current PyTorch shape: eval forwards kwargs to a newer step API.""" + + def eval(self, *args, target=None, losses=None, **kwargs): + self.eval_calls += 1 + return self._run_schedule(*args, target=target, losses=losses, **kwargs) + + +class _NoEvalSchedule(_LegacyStepSchedule): + """PyTorch pipeline schedule shape before forward-only eval existed.""" + + eval = None + + class TestAutoPipelineKwargsChunkSpec: def _pipeline_with_parts(self, *parts: nn.Module, schedule=None, has_first_stage: bool = True): ap = AutoPipeline( @@ -435,6 +486,84 @@ def test_step_microbatches_does_not_forward_return_outputs_to_older_pytorch(self assert schedule.received_return_outputs is False + def test_eval_microbatches_uses_forward_only_schedule_with_exact_prepared_split(self): + schedule = _KwargsChunkSchedule(invoke_loss=True) + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + input_ids = [torch.full((1, 8), index, dtype=torch.long) for index in range(2)] + metadata = [object(), object()] + model_inputs = [ + { + "input_ids": input_ids[index], + "metadata": metadata[index], + } + for index in range(2) + ] + original_split_inputs = schedule._split_inputs + original_loss_fn = schedule._loss_fn + seen = [] + losses = [] + + def loss_fn(output, index): + seen.append((output, index)) + return output + + result = ap.eval_microbatches(model_inputs, loss_fn=loss_fn, losses=losses, return_outputs=False) + + assert result == "schedule-result" + assert schedule.eval_calls == 1 + assert schedule.step_calls == 0 + assert schedule.split_inputs_calls == 1 + assert schedule.args_split[0][0] is input_ids[0] + assert schedule.args_split[1][0] is input_ids[1] + assert schedule.kwargs_split[0]["metadata"] is metadata[0] + assert schedule.kwargs_split[1]["metadata"] is metadata[1] + assert schedule.target_during_step.tolist() == [0, 1] + assert schedule.losses_during_step is losses + assert schedule.return_outputs_during_step is False + assert [(output.item(), index) for output, index in seen] == [(1.0, 1), (0.0, 0)] + assert schedule._split_inputs == original_split_inputs + assert schedule._loss_fn is original_loss_fn + + def test_eval_microbatches_does_not_forward_return_outputs_to_older_pytorch(self): + schedule = _LegacyEvalSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + ap.eval_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + return_outputs=False, + ) + + assert schedule.eval_calls == 1 + assert schedule.step_calls == 0 + assert schedule.received_return_outputs is False + + def test_eval_microbatches_uses_step_capability_when_eval_forwards_kwargs(self): + schedule = _ForwardingEvalSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + ap.eval_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + return_outputs=False, + ) + + assert schedule.eval_calls == 1 + assert schedule.step_calls == 0 + assert schedule.return_outputs_during_step is False + + def test_eval_microbatches_fails_clearly_when_pytorch_has_no_eval_schedule(self): + schedule = _NoEvalSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + with pytest.raises(NotImplementedError, match=r"schedule with eval\(\)"): + ap.eval_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + ) + + assert schedule.step_calls == 0 + @pytest.mark.parametrize( "model_inputs", [ @@ -466,6 +595,25 @@ def test_step_microbatches_restores_schedule_state_after_failure(self): assert schedule._split_inputs == original_split_inputs assert schedule._loss_fn is original_loss_fn + def test_eval_microbatches_restores_schedule_state_after_failure(self): + schedule = _KwargsChunkSchedule(fail_on_step=True) + original_split_inputs = schedule._split_inputs + original_loss_fn = schedule._loss_fn + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + + with pytest.raises(RuntimeError, match="schedule failed"): + ap.eval_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + ) + + assert schedule.eval_calls == 1 + assert schedule.step_calls == 0 + assert schedule.split_inputs_during_step is not original_split_inputs + assert schedule.loss_fn_during_step is not original_loss_fn + assert schedule._split_inputs == original_split_inputs + assert schedule._loss_fn is original_loss_fn + def test_only_canonical_model_part_supplies_chunk_policy(self): ap = self._pipeline_with_parts( _KwargsChunkHookPart({"position_ids": 1}), diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py index a38a4d1602..3d97911f16 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -17,7 +17,7 @@ Training forward/backward and CP sharding are owned by :class:`Engine`. These tests cover the VLM recipe responsibilities that remain around that core: pipeline media staging setup, vision-frame context publication, and the -recipe-owned validation forward path. +validation handoff plus epoch-level DP aggregation. """ from __future__ import annotations @@ -35,21 +35,6 @@ from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM -def _identity_cp_shard(sharder, batch): - """Bypass CP transport while preserving constructor-side strategy resolution. - - Args: - sharder: Sharder whose resolved strategy is not exercised by this test. - batch: Mutable model-input mapping whose tensor values retain their - existing shapes. - - Returns: - The null context factory and the same input mapping. - """ - del sharder - return nullcontext, batch - - class _UnsupportedVisionModel: supports_cp_vision_frame_sharding = False @@ -321,88 +306,36 @@ def test_setup_always_stages_pp_media_under_pp( # ----------------------------------------------------------------------------- -# val-side wiring (the bug-fix territory) +# validation Engine boundary # ----------------------------------------------------------------------------- -class _ShardLabelsOnEnter: - def __init__(self, labels, local_labels): - self.labels = labels - self.local_labels = local_labels - - def __enter__(self): - self.labels.resize_(self.local_labels.shape) - self.labels.copy_(self.local_labels) - - def __exit__(self, exc_type, exc, tb): - return False - - -def test_val_counts_label_tokens_inside_cp_context_after_labels_are_sharded(): - """Validation must count label tokens after CP has exposed the local shard.""" - labels = torch.tensor([[1, 2, -100, 4]]) - batch = {"labels": labels} - local_labels = torch.tensor([[1, -100]]) - - def train_ctx(): - return _ShardLabelsOnEnter(labels, local_labels) - - labels = batch.pop("labels") - pre_context_count = (labels != -100).sum().item() - with train_ctx(): - local_num_label_tokens = (labels != -100).sum().item() - - assert pre_context_count == 3 - assert local_num_label_tokens == 1 - - -def test_val_pos_ids_uses_dist_env_device_not_model_device(): - """Reproduce the bug fix at finetune.py:1281 — val must use - ``self.dist_env.device``, not ``self.model_parts[0].device`` which - AttributeErrors on FSDP-wrapped models.""" - - class _FSDPWrapped: - # Intentionally has NO ``.device`` attribute (mirrors real FSDP wrapper). - def __getattr__(self, name): - if name == "device": - raise AttributeError("'FSDPWrapped' object has no attribute 'device'") - raise AttributeError(name) - - model = _FSDPWrapped() - dist_env = SimpleNamespace(device=torch.device("cpu")) - - # The fixed line: - pos = torch.arange(0, 4).unsqueeze(0).to(dist_env.device) - assert pos.device.type == "cpu" - - # The buggy line would have raised: - with pytest.raises(AttributeError, match="no attribute 'device'"): - _ = torch.arange(0, 4).unsqueeze(0).to(model.device) - - def test_run_validation_epoch_does_not_sum_tokens_over_cp(monkeypatch): - """``total_loss`` is all-reduced with include_cp=True, but ``total_tokens`` - (measured pre-CP-shard) must NOT include CP — otherwise val_loss is scaled - down by cp_size. Guards the fix at finetune.py:_run_validation_epoch.""" - from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM - - # No-op replacements for the heavy collaborators. + """Engine returns CP-complete sums, so the epoch reduces only over DP.""" monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext()) - monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _identity_cp_shard) - monkeypatch.setattr(vlm_finetune, "filter_forward_kwargs", lambda model, batch: batch) - monkeypatch.setattr(vlm_finetune, "calculate_loss", lambda *a, **k: torch.tensor(2.0)) class _Model(torch.nn.Module): - def eval(self): # noqa: D401 - return self + def prepare_model_inputs_for_cp(self, *args, **kwargs): + raise AssertionError("the recipe must delegate CP preparation to Engine.forward") + + def forward(self, *args, **kwargs): + raise AssertionError("the recipe must delegate model execution to Engine.forward") + + engine_calls = [] - def forward(self, **batch): - return SimpleNamespace(logits=torch.zeros(1, 4, 8), hidden_states=None) + class _Engine: + def forward(self, datums, loss_fn): + engine_calls.append((datums, loss_fn)) + return SimpleNamespace( + loss_sum=torch.tensor(6.0, dtype=torch.float64), + weight_sum=torch.tensor(3.0, dtype=torch.float64), + loss_fn_outputs=[], + ) recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) recipe.model_parts = [_Model()] recipe.loss_fn = object() # not a FusedLinearCrossEntropy - recipe.device_mesh = None # CP inactive -> skip pre-embed branch + recipe.engine = _Engine() recipe.pp_enabled = False recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) recipe.step_scheduler = SimpleNamespace(step=3, epoch=1) @@ -425,61 +358,12 @@ def _fake_allreduce(tensor, include_cp=False): result = recipe._run_validation_epoch([batch]) - # total_loss all-reduced WITH cp; total_tokens and num_label_tokens WITHOUT. - loss_call = allreduce_calls[0] - tokens_call = allreduce_calls[1] - assert loss_call[1] is True, "total_loss must include CP ranks" - assert tokens_call[1] is False, "total_tokens must NOT be summed over CP ranks" - # val_loss = (2.0 * 3 tokens) / 3 tokens == 2.0 - assert result.metrics["val_loss"] == pytest.approx(2.0) - - -def test_run_validation_epoch_cp_active_runs_pre_embed(monkeypatch): - """With CP active and a model exposing prepare_model_inputs_for_cp, the - validation loop must invoke the model's sharder-only CP hook before sharding. - Guards finetune.py:_run_validation_epoch CP pre-embed branch.""" - from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM - - monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext()) - monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _identity_cp_shard) - monkeypatch.setattr(vlm_finetune, "filter_forward_kwargs", lambda model, batch: batch) - monkeypatch.setattr(vlm_finetune, "calculate_loss", lambda *a, **k: torch.tensor(2.0)) - - pre_embed_calls = [] - - class _Model(torch.nn.Module): - def eval(self): - return self - - def prepare_model_inputs_for_cp(self, batch, *, num_chunks=1): # sharder-only hook - pre_embed_calls.append(set(batch)) - return {} - - def forward(self, **batch): - return SimpleNamespace(logits=torch.zeros(1, 4, 8), hidden_states=None) - - class _DM(dict): - mesh_dim_names = ["cp"] - - recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.model_parts = [_Model()] - recipe.loss_fn = object() - recipe.device_mesh = _DM(cp=SimpleNamespace(size=lambda: 2, get_group=lambda: "cp-group")) - recipe.cp_vision_frame_sharding = CpVisionFrameShardingConfig(enabled=True) - recipe.pp_enabled = False - recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) - recipe.step_scheduler = SimpleNamespace(step=3, epoch=1) - recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.001}])] - recipe._maybe_add_drafter_loss = lambda *, base_loss, **kwargs: base_loss - recipe._dp_allreduce = lambda tensor, include_cp=False: tensor - - batch = { - "input_ids": torch.tensor([[1, 2, 3, 4]]), - "pixel_values": torch.randn(1, 3, 8, 8), - "labels": torch.tensor([[1, 2, -100, 4]]), - } - - result = recipe._run_validation_epoch([batch]) - - assert pre_embed_calls, "the CP hook (prepare_model_inputs_for_cp) must run when CP is active" + assert len(engine_calls) == 1 + datums, loss_fn = engine_calls[0] + assert len(datums) == 1 + assert datums[0].model_inputs["input_ids"] is batch["input_ids"] + assert datums[0].loss_fn_inputs["labels"] is batch["labels"] + assert loss_fn == recipe._engine_validation_loss_fn + assert len(allreduce_calls) == 2 + assert all(include_cp is False for _, include_cp in allreduce_calls) assert result.metrics["val_loss"] == pytest.approx(2.0) diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 8071f192ce..c411e7fb5a 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -2032,184 +2032,95 @@ def test_vlm_engine_loss_uses_explicit_mtp_targets_instead_of_thd_boundaries(mon ) -def test_vlm_validation_shards_global_mtp_inputs_and_targets(monkeypatch): - """Validation prepares MTP futures globally and reuses the main CP layout.""" - from nemo_automodel.components.models.common.mtp import prepare_mtp_context_parallel_inputs - - events = [] - captured = {} - local_indices = torch.tensor([0, 1, 4, 5]) +def test_vlm_engine_validation_loss_uses_eval_path(): + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.pp_enabled = False + expected = torch.tensor(5.0) + compute_loss = MagicMock(return_value=expected) + recipe._compute_vlm_loss = compute_loss + + output = object() + labels = torch.tensor([[1, 2, -100]]) + targets = (torch.tensor([[2, -100, -100]]),) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + loss_inputs = { + "labels": labels, + "weights": labels.ne(-100), + "cu_seqlens": cu_seqlens, + "mtp_per_depth_targets": targets, + } - class _MTPValidationModel(nn.Module): - def __init__(self): - super().__init__() - self.scale = nn.Parameter(torch.tensor(1.0)) - self.supports = SimpleNamespace(mtp_enabled=True, supports_mtp_cp=False) - self.prepared_before_shard = False - - def prepare_mtp_inputs_for_cp(self, batch, *, ignore_index=-100): - """Prepare global MTP tensors after model-owned position metadata. - - Args: - batch: Mutable global mapping whose token tensors have shape - [batch, sequence]. - ignore_index: Target fill value at invalid future positions. - - Returns: - Global per-depth MTP tensors with shape [batch, sequence]. - """ - events.append("prepare_mtp") - self.prepared_before_shard = True - torch.testing.assert_close(batch["position_ids"], torch.tensor([[0, 1, 2, 0, 1, 2]])) - return prepare_mtp_context_parallel_inputs( - batch, - num_depths=1, - ignore_index=ignore_index, - ) + loss = recipe._engine_validation_loss_fn(output, loss_inputs) - def forward( - self, - input_ids, - position_ids, - pixel_values, - mtp_per_depth_input_ids, - mtp_per_depth_position_ids, - mtp_per_depth_valid_masks, - ): - """Capture CP-local VLM and MTP model inputs. - - Args: - input_ids: Tensor of shape [batch, local_sequence]. - position_ids: Tensor of shape [batch, local_sequence]. - pixel_values: Tensor of shape [batch, channels, height, width]. - mtp_per_depth_input_ids: Per-depth tensors of shape - [batch, local_sequence]. - mtp_per_depth_position_ids: Per-depth tensors of shape - [batch, local_sequence]. - mtp_per_depth_valid_masks: Per-depth boolean tensors of shape - [batch, local_sequence]. - - Returns: - Model output with logits and per-depth MTP logits of shape - [batch, local_sequence, vocab]. - """ - captured["model_input_ids"] = input_ids.detach().clone() - captured["position_ids"] = position_ids.detach().clone() - captured["pixel_values"] = pixel_values.detach().clone() - captured["mtp_input_ids"] = tuple(value.detach().clone() for value in mtp_per_depth_input_ids) - captured["mtp_position_ids"] = tuple(value.detach().clone() for value in mtp_per_depth_position_ids) - captured["mtp_valid_masks"] = tuple(value.detach().clone() for value in mtp_per_depth_valid_masks) - logits = self.scale * input_ids.float().unsqueeze(-1) - return SimpleNamespace( - logits=logits, - mtp_per_depth_h=None, - mtp_per_depth_logits=[logits], - mtp_loss_scaling_factor=1.0, - ) + assert loss is expected + compute_loss.assert_called_once_with( + out=output, + labels=labels, + num_label_tokens=None, + is_train=False, + cu_seqlens=cu_seqlens, + mtp_per_depth_targets=targets, + ) - model = _MTPValidationModel() - - class _FakeContextParallelSharder: - def __init__(self, resolved_model, device_mesh, batch, *, invoke_pre_embed): - """Materialize model-owned global position metadata. - - Args: - resolved_model: Model that owns the CP preparation hook. - device_mesh: Runtime device mesh; unused by this CPU fake. - batch: Mutable global mapping whose token tensors have shape - [batch, sequence]. - invoke_pre_embed: Whether model-owned preparation is enabled. - """ - del device_mesh - assert resolved_model is model - assert invoke_pre_embed is True - events.append("sharder_init") - batch["position_ids"] = torch.tensor([[0, 1, 2, 0, 1, 2]]) - - def shard(self, batch): - """Select one deterministic CP-local token layout. - - Args: - batch: Global mapping whose token tensors have shape - [batch, sequence]. - - Returns: - Context factory and mapping whose token tensors have shape - [batch, local_sequence]. - """ - events.append("shard") - assert model.prepared_before_shard - local_batch = dict(batch) - for key in ("input_ids", "labels", "position_ids"): - local_batch[key] = local_batch[key].index_select(1, local_indices) - local_batch.pop("seq_lens") - local_batch.pop("seq_lens_padded") - return nullcontext, local_batch - - def shard_token_tensor(self, tensor, seq_dim=1, fill=None): - """Apply the captured local token indices to an auxiliary tensor. - - Args: - tensor: Global tensor of shape [batch, sequence]. - seq_dim: Sequence axis in ``tensor``. - fill: Padding value; unused because this fake does not pad. - - Returns: - Tensor of shape [batch, local_sequence]. - """ - del fill - return tensor.index_select(seq_dim, local_indices) +def test_vlm_validation_uses_engine_forward_and_aggregates_uneven_batches(monkeypatch): recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.cfg = SimpleNamespace(mtp=SimpleNamespace(ignore_index=-100, scaling_factor=1.0)) + recipe.model_parts = [MagicMock()] recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) - recipe.device_mesh = None - recipe.pp_enabled = False - recipe.model_parts = [model] - recipe.loss_fn = object() recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.01}])] recipe.step_scheduler = SimpleNamespace(step=3, epoch=1) - recipe._get_cp_group_size = lambda: 2 - recipe._get_dp_group = lambda include_cp=False: None - recipe._dp_allreduce = lambda tensor, include_cp=False: tensor - - base_loss = MagicMock(return_value=torch.tensor(1.0)) - mtp_loss = MagicMock(return_value=torch.tensor(0.5)) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ContextParallelSharder", _FakeContextParallelSharder) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_loss", base_loss) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.calculate_mtp_loss", mtp_loss) + recipe.pp_enabled = False + recipe.loss_fn = object() - batch = { - "input_ids": torch.tensor([[10, 11, 12, 20, 21, 22]]), - "labels": torch.tensor([[11, 12, -100, 21, 22, -100]]), - "pixel_values": torch.ones(1, 3, 2, 2), - "seq_lens": torch.tensor([[3, 3, -1000]]), - "seq_lens_padded": torch.tensor([[3, 3, -1000]]), - "cu_seqlens": torch.tensor([0, 3, 6], dtype=torch.int32), - } + engine = MagicMock() + engine.forward.side_effect = [ + SimpleNamespace( + loss_sum=torch.tensor(4.0, dtype=torch.float64), + weight_sum=torch.tensor(2.0, dtype=torch.float64), + loss_fn_outputs=[], + ), + SimpleNamespace( + loss_sum=torch.tensor(9.0, dtype=torch.float64), + weight_sum=torch.tensor(3.0, dtype=torch.float64), + loss_fn_outputs=[], + ), + ] + recipe.engine = engine + allreduce = MagicMock(side_effect=lambda tensor, **kwargs: tensor) + recipe._dp_allreduce = allreduce + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.ScopedRNG", + lambda **kwargs: nullcontext(), + ) - with pytest.raises(NotImplementedError, match="supports_mtp_cp=False"): - recipe._run_validation_epoch([batch]) - assert events == ["sharder_init"] - assert model.prepared_before_shard is False - - events.clear() - model.supports.supports_mtp_cp = True - metrics = recipe._run_validation_epoch([batch]) - - assert events == ["sharder_init", "prepare_mtp", "shard"] - assert captured["model_input_ids"].tolist() == [[10, 11, 21, 22]] - assert captured["position_ids"].tolist() == [[0, 1, 1, 2]] - assert captured["pixel_values"].shape == (1, 3, 2, 2) - assert captured["mtp_input_ids"][0].tolist() == [[11, 12, 22, 0]] - assert captured["mtp_position_ids"][0].tolist() == [[1, 2, 2, 0]] - assert captured["mtp_valid_masks"][0].tolist() == [[True, True, True, False]] - assert mtp_loss.call_args.kwargs["mtp_per_depth_targets"][0].tolist() == [[12, -100, -100, -100]] - assert mtp_loss.call_args.kwargs["cu_seqlens"] is None - assert base_loss.call_args.kwargs["num_label_tokens"] == 4 - assert metrics.metrics["val_loss"] == pytest.approx(1.5) - assert metrics.metrics["num_label_tokens"] == 4 - assert model.scale.grad is None + batches = [ + { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[1, 2, -100]]), + "pixel_values": torch.ones(1, 3, 2, 2), + }, + { + "input_ids": torch.tensor([[4, 5, 6, 7]]), + "labels": torch.tensor([[3, 4, 5, -100]]), + "pixel_values": torch.zeros(1, 3, 2, 2), + }, + ] + metrics = recipe._run_validation_epoch(batches) + + assert engine.forward.call_count == 2 + for call, batch in zip(engine.forward.call_args_list, batches): + datums, loss_fn = call.args + assert len(datums) == 1 + assert datums[0].model_inputs["input_ids"] is batch["input_ids"] + assert datums[0].model_inputs["pixel_values"] is batch["pixel_values"] + assert datums[0].loss_fn_inputs["labels"] is batch["labels"] + torch.testing.assert_close(datums[0].loss_fn_inputs["weights"], batch["labels"].ne(-100)) + assert loss_fn == recipe._engine_validation_loss_fn + assert allreduce.call_count == 2 + assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) + assert metrics.metrics["val_loss"] == pytest.approx(13.0 / 5.0) + assert metrics.metrics["num_label_tokens"] == pytest.approx(5.0) def test_vlm_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index b9051fb166..f77fdca67a 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1548,244 +1548,79 @@ def test_compute_trust_remote_code_falls_back_to_resolve(): # ----------------- -class MockSchedule: - """Mock PP schedule that tracks step/eval calls.""" - - def __init__(self): - self.step_calls = [] - self.eval_calls = [] - - def step(self, *args, **kwargs): - self.step_calls.append((args, kwargs)) - # Populate losses list if provided - if "losses" in kwargs and kwargs["losses"] is not None: - kwargs["losses"].append(torch.tensor(0.5)) - - def eval(self, *args, **kwargs): - self.eval_calls.append((args, kwargs)) - # Populate losses list if provided - if "losses" in kwargs and kwargs["losses"] is not None: - kwargs["losses"].append(torch.tensor(0.5)) - - -class MockPPInfo: - """Mock PP info with configurable first/last stage flags.""" - - def __init__(self, has_first_stage=True, has_last_stage=True): - self.has_first_stage = has_first_stage - self.has_last_stage = has_last_stage - self.schedule = MockSchedule() - +def test_engine_validation_pipeline_loss_reuses_configured_loss_and_thd_metadata(): + recipe = object.__new__(TrainFinetuneRecipeForNextTokenPrediction) + recipe.pp_enabled = True + recipe.pipeline_loss_fn = MagicMock(return_value=torch.tensor(3.0)) + output = object() + labels = torch.tensor([[1, 2, -100]]) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) -def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info): - """Create a minimal TrainFinetuneRecipeForNextTokenPrediction for PP testing.""" - cfg = ConfigNode( + loss = recipe._engine_validation_loss_fn( + output, { - "nvtx": False, - "model": {}, - "dataloader": {"collate_fn": "nemo_automodel.components.datasets.utils.default_collater"}, - "dataset": {}, - "validation_dataloader": {}, - "step_scheduler": {"local_batch_size": 1, "global_batch_size": 1}, - "optimizer": {}, - "loss_fn": {}, - "checkpoint": {"best_metric_key": "default"}, - "distributed": {"cp_size": 1}, - "autopipeline": {"pp_microbatch_size": 1}, - } - ) - - # Minimal stubs so we can create the recipe - monkeypatch.setattr( - "nemo_automodel.recipes.llm.train_ft.initialize_distributed", - lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0), - ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None) - - # Create the recipe without calling setup - recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg) - - # Mock out attributes needed for the recipe-owned validation forward. - # Use object.__setattr__ to bypass the state tracking - object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - object.__setattr__(recipe, "device_mesh", None) - object.__setattr__(recipe, "pp_enabled", True) - object.__setattr__( - recipe, - "pp", - SimpleNamespace( - info=pp_info, - pp_batch_size=1, - pp_microbatch_size=1, - update_seq_len=lambda seq_len: None, - ), - ) - object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) - object.__setattr__(recipe, "te_fp8", None) - object.__setattr__(recipe, "pipeline_loss_fn", None) - - return recipe - - -def test_forward_validation_step_pp_uses_schedule_eval(monkeypatch): - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=True, has_last_stage=True) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Mock _make_cp_batch_and_ctx to return a no-op context manager - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) - - # Create a minimal batch - batch = { - "input_ids": torch.tensor([[1, 2, 3]]), - "labels": torch.tensor([[1, 2, 3]]), - } - - loss = recipe._forward_validation_step(batch) - - # Should use eval, not step - assert len(pp_info.schedule.eval_calls) == 1, "schedule.eval() should be called once for validation" - assert len(pp_info.schedule.step_calls) == 0, "schedule.step() should not be called for validation" - assert loss.item() == pytest.approx(0.5) - - -def test_forward_validation_step_pp_non_first_stage_uses_eval_without_input(monkeypatch): - """Test schedule.eval() without input_ids when not on first stage.""" - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=False, has_last_stage=True) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Mock _make_cp_batch_and_ctx to return a no-op context manager - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) - - # Create a minimal batch - batch = { - "input_ids": torch.tensor([[1, 2, 3]]), - "labels": torch.tensor([[1, 2, 3]]), - } - - recipe._forward_validation_step(batch) - - # Should use eval without input_ids as first positional arg - assert len(pp_info.schedule.eval_calls) == 1 - args, kwargs = pp_info.schedule.eval_calls[0] - assert len(args) == 0, "Non-first stage should not pass input_ids as positional arg" - assert "target" in kwargs - - -def test_run_validation_epoch_pp_sends_loss_from_last_stage_to_main(monkeypatch): - """Test that _run_validation_epoch broadcasts val_loss from last stage to main rank for PP.""" - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=True, has_last_stage=True) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Set up recipe attributes for validation - use object.__setattr__ to bypass state tracking - object.__setattr__(recipe, "model_parts", [DummyModel()]) - object.__setattr__(recipe, "step_scheduler", SimpleNamespace(step=1, epoch=0)) - object.__setattr__(recipe, "optimizer", [SimpleNamespace(param_groups=[{"lr": 0.01}])]) - - # Stub the PP last-stage broadcast helper (post-d96f1b20 the recipe broadcasts - # within the PP group instead of doing send/recv to global rank 0). - monkeypatch.setattr(recipe, "_broadcast_from_last_pp_stage", lambda t: t) - - # Set dist_env.rank to 0 (last stage and main rank are the same in this test) - object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - - monkeypatch.setattr(recipe, "_forward_validation_step", lambda batch: torch.tensor(0.5)) - - # Mock _dp_allreduce to return the tensor/value - def mock_dp_allreduce(val, include_cp=False): - if isinstance(val, torch.Tensor): - return val - return torch.tensor(val) - - monkeypatch.setattr(recipe, "_dp_allreduce", mock_dp_allreduce) - - # Mock _make_cp_batch_and_ctx - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) - - # Mock ScopedRNG - monkeypatch.setattr( - "nemo_automodel.recipes.llm.train_ft.ScopedRNG", - lambda **kwargs: MagicMock(__enter__=lambda s: s, __exit__=lambda s, *a: None), + "labels": labels, + "weights": labels.ne(-100), + "cu_seqlens": cu_seqlens, + }, ) - # Create a simple dataloader that yields one batch - val_dataloader = [{"input_ids": torch.tensor([[1, 2, 3]]), "labels": torch.tensor([[1, 2, 3]])}] - - result = recipe._run_validation_epoch(val_dataloader) - - # Verify result is a MetricsSample with val_loss - assert "val_loss" in result.metrics - # val_loss should be a float, not a tensor - assert isinstance(result.metrics["val_loss"], float) - - -def test_run_validation_epoch_pp_main_rank_receives_from_last_stage(monkeypatch): - """Test that main rank receives val_loss from last stage via the PP broadcast helper.""" - from contextlib import nullcontext - - pp_info = MockPPInfo(has_first_stage=True, has_last_stage=False) - recipe = _create_minimal_recipe_for_pp_test(monkeypatch, pp_info) - - # Set up recipe attributes - use object.__setattr__ to bypass state tracking - object.__setattr__(recipe, "model_parts", [DummyModel()]) - object.__setattr__(recipe, "step_scheduler", SimpleNamespace(step=1, epoch=0)) - object.__setattr__(recipe, "optimizer", [SimpleNamespace(param_groups=[{"lr": 0.01}])]) - - # Track calls to the PP last-stage broadcast helper and simulate the last - # stage's value of 0.5 propagating into the non-last-stage tensor. - broadcast_calls = [] - - def mock_broadcast(tensor): - broadcast_calls.append(tensor) - tensor.fill_(0.5) - return tensor - - monkeypatch.setattr(recipe, "_broadcast_from_last_pp_stage", mock_broadcast) - - # Main rank (0) is different from last stage (3) - object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - - monkeypatch.setattr(recipe, "_forward_validation_step", lambda batch: torch.tensor(0.0)) - - def mock_dp_allreduce(val, include_cp=False): - if isinstance(val, torch.Tensor): - return val - return torch.tensor(val) + assert loss.item() == pytest.approx(3.0) + assert recipe.pipeline_loss_fn.cu_seqlens is cu_seqlens + recipe.pipeline_loss_fn.assert_called_once_with(output, labels) - monkeypatch.setattr(recipe, "_dp_allreduce", mock_dp_allreduce) - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None), - ) +def test_run_validation_epoch_uses_engine_forward_for_pp_complete_results(monkeypatch): + recipe = TrainFinetuneRecipeForNextTokenPrediction.__new__(TrainFinetuneRecipeForNextTokenPrediction) + recipe.model_parts = [MagicMock()] + recipe.dist_env = SimpleNamespace(device=torch.device("cpu"), is_main=True) + recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.01}])] + recipe.step_scheduler = SimpleNamespace(step=1, epoch=0) + recipe.pp_enabled = True + recipe.pipeline_loss_fn = MagicMock() + recipe.loss_fn = object() + recipe.tool_call_evaluator = None + engine = MagicMock() + engine.forward.side_effect = [ + SimpleNamespace( + loss_sum=torch.tensor(4.0, dtype=torch.float64), + weight_sum=torch.tensor(2.0, dtype=torch.float64), + loss_fn_outputs=[], + ), + SimpleNamespace( + loss_sum=torch.tensor(9.0, dtype=torch.float64), + weight_sum=torch.tensor(3.0, dtype=torch.float64), + loss_fn_outputs=[], + ), + ] + recipe.engine = engine + allreduce = MagicMock(side_effect=lambda tensor, **kwargs: tensor) + recipe._dp_allreduce = allreduce + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) monkeypatch.setattr( "nemo_automodel.recipes.llm.train_ft.ScopedRNG", - lambda **kwargs: MagicMock(__enter__=lambda s: s, __exit__=lambda s, *a: None), + lambda **kwargs: nullcontext(), ) - val_dataloader = [{"input_ids": torch.tensor([[1, 2, 3]]), "labels": torch.tensor([[1, 2, 3]])}] - - result = recipe._run_validation_epoch(val_dataloader) + batches = [ + {"input_ids": torch.tensor([[1, 2, 3]]), "labels": torch.tensor([[1, 2, -100]])}, + {"input_ids": torch.tensor([[4, 5, 6, 7]]), "labels": torch.tensor([[4, 5, 6, -100]])}, + ] + metrics = recipe._run_validation_epoch(batches) - # Main rank should have invoked the PP broadcast helper to pull val_loss - # and pp_num_tokens from the last PP stage (two calls total). - assert len(broadcast_calls) >= 1, "Main rank should broadcast val_loss from the last PP stage" - assert isinstance(result.metrics["val_loss"], float) + assert engine.forward.call_count == 2 + for call, batch in zip(engine.forward.call_args_list, batches): + datums, loss_fn = call.args + assert len(datums) == 1 + assert datums[0].model_inputs["input_ids"] is batch["input_ids"] + assert datums[0].loss_fn_inputs["labels"] is batch["labels"] + torch.testing.assert_close(datums[0].loss_fn_inputs["weights"], batch["labels"].ne(-100)) + assert loss_fn == recipe._engine_validation_loss_fn + assert allreduce.call_count == 2 + assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) + assert metrics.metrics["val_loss"] == pytest.approx(13.0 / 5.0) + assert metrics.metrics["num_label_tokens"] == pytest.approx(5.0) # ----------------- @@ -2273,9 +2108,6 @@ def _make_recipe( update_seq_len=lambda seq_len: None, ), ) - # Stub the PP last-stage broadcast helper (post-d96f1b20 the recipe - # broadcasts inside the PP group instead of using send/recv). - monkeypatch.setattr(recipe, "_broadcast_from_last_pp_stage", lambda t: t) object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) monkeypatch.setattr( @@ -2309,12 +2141,6 @@ def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypat datums = [object(), object()] make_datum = MagicMock(side_effect=datums) monkeypatch.setattr(recipe, "_make_engine_datum", make_datum) - monkeypatch.setattr( - recipe, - "_broadcast_from_last_pp_stage", - MagicMock(side_effect=AssertionError("Engine loss must not be broadcast again")), - ) - engine = MagicMock() engine.forward_backward.return_value = (torch.tensor(0.25), []) object.__setattr__(recipe, "engine", engine) @@ -2683,362 +2509,28 @@ def test_evaluate_failure_is_tolerated(self, monkeypatch): assert out.metrics["tool_call/has_call"] == 0.0 -@pytest.mark.parametrize( - ("cp_size", "uses_thd", "supports_thd"), - [ - (2, False, False), - (1, True, True), - ], -) -def test_forward_validation_step_model_cp_hook(monkeypatch, cp_size, uses_thd, supports_thd): - """Non-PP validation keeps the recipe-owned CP preparation path.""" - from contextlib import nullcontext - - cfg = ConfigNode( - { - "nvtx": False, - "model": {}, - "dataloader": {"collate_fn": "nemo_automodel.components.datasets.utils.default_collater"}, - "dataset": {}, - "validation_dataloader": {}, - "step_scheduler": {"local_batch_size": 1, "global_batch_size": 1}, - "optimizer": {}, - "loss_fn": {}, - "checkpoint": {"best_metric_key": "default"}, - "distributed": {"cp_size": cp_size}, - "autopipeline": {"pp_microbatch_size": 1}, - } - ) - monkeypatch.setattr( - "nemo_automodel.recipes.llm.train_ft.initialize_distributed", - lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0), - ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None) - recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg) - - class _CPModel(nn.Module): - def __init__(self): - super().__init__() - self.lin = nn.Linear(4, 8) - self.prepared = False - self.supports = SimpleNamespace(mtp_enabled=False, supports_mtp_cp=False) - - def prepare_model_inputs_for_cp(self, batch, **kwargs): - self.prepared = True - self.num_chunks = kwargs.get("num_chunks") - from nemo_automodel.components.distributed.context_parallel.sharder import ( - ContextParallelSharder, - contiguous_local_indices, - ) - - return { - "cp_sharder": ContextParallelSharder( - shard_batch=lambda cp_mesh, tp_mesh, batch, **k: (nullcontext, batch, None), - local_token_global_indices=contiguous_local_indices, - ) - } - - def forward(self, **batch): - logits = self.lin(batch["input_ids"].float()) - return SimpleNamespace(logits=logits) - - model = _CPModel() - model.supports_thd = supports_thd - - # The hook gate reads the CP size from the mesh (the runtime truth), not - # from the config: fake a mesh whose "cp" dim matches the parametrization. - class _FakeSubMesh: - def __init__(self, size): - self._size = size - - def size(self): - return self._size - - def get_group(self): - return None - - def get_local_rank(self): - return 0 - - fake_mesh = {"cp": _FakeSubMesh(cp_size)} - fake_mesh = type("_FakeDeviceMesh", (dict,), {"mesh_dim_names": ("cp",)})(fake_mesh) - object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) - object.__setattr__(recipe, "device_mesh", fake_mesh) - object.__setattr__(recipe, "pp_enabled", False) - object.__setattr__(recipe, "magi", SimpleNamespace(enabled=False)) - object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) - object.__setattr__(recipe, "te_fp8", None) - object.__setattr__(recipe, "model_parts", [model]) - object.__setattr__(recipe, "distributed_config", SimpleNamespace(defer_fsdp_grad_sync=True)) - object.__setattr__(recipe, "loss_fn", object()) # not FusedLinearCrossEntropy - object.__setattr__(recipe, "_get_dp_group", lambda include_cp=False: None) - object.__setattr__(recipe, "_get_dp_group_size", lambda include_cp=False: 1) - - captured = {} - - def _fake_calc_loss( - loss_fn, *, logits, labels, model, hidden_states, lm_weight, num_label_tokens, grad_reduce_group - ): - captured["logits_is_tensor"] = isinstance(logits, torch.Tensor) - assert lm_weight is None - return logits.mean() - - monkeypatch.setattr( - "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", - lambda device_mesh, batch, *a, **k: (nullcontext, batch, None), - ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calc_loss) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_final_hidden_states", lambda out: None) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.filter_forward_kwargs", lambda model, batch: batch) - - batch = {"input_ids": torch.randn(1, 4, 4), "labels": torch.zeros(1, 4, dtype=torch.long)} - if uses_thd: - batch["qkv_format"] = "thd" - loss = recipe._forward_validation_step(batch) +def test_engine_validation_loss_uses_eval_path(): + recipe = TrainFinetuneRecipeForNextTokenPrediction.__new__(TrainFinetuneRecipeForNextTokenPrediction) + recipe.pp_enabled = False + expected = torch.tensor(5.0) + compute_loss = MagicMock(return_value=expected) + recipe._compute_causal_lm_loss = compute_loss - assert model.prepared is True - assert model.num_chunks == 1 - assert captured["logits_is_tensor"] - assert torch.isfinite(loss).all() - assert model.lin.weight.grad is None + output = object() + labels = torch.tensor([[1, 2, -100]]) + loss_inputs = { + "labels": labels, + "weights": labels.ne(-100), + "cu_seqlens": torch.tensor([0, 3], dtype=torch.int32), + } + loss = recipe._engine_validation_loss_fn(output, loss_inputs) -def test_forward_validation_step_shards_global_mtp_inputs_and_targets(monkeypatch): - """Validation shifts MTP inputs globally and reuses the model-input CP layout.""" - from nemo_automodel.components.models.common.mtp import ( - MTPConfig, - prepare_mtp_context_parallel_inputs, + assert loss is expected + compute_loss.assert_called_once_with( + output, + labels, + loss_inputs, + num_label_tokens=None, + is_train=False, ) - - captured = {} - local_indices = torch.tensor([0, 1, 4, 5]) - - class _MTPModel(nn.Module): - def __init__(self): - super().__init__() - self.scale = nn.Parameter(torch.tensor(1.0)) - self.mtp_config = MTPConfig(num_layers=1, layer_pattern="*") - self.cp_prepared = False - self.mtp_prepared = False - self.supports = SimpleNamespace(mtp_enabled=True, supports_mtp_cp=True) - - def prepare_model_inputs_for_cp(self, batch, *, num_chunks=1): - """Materialize global multi-axis positions before MTP preparation. - - Args: - batch: Mutable mapping whose token tensors have shape [batch, - sequence]. - num_chunks: Number of downstream sharding chunks. - - Returns: - Empty update mapping; ``batch`` receives ``position_ids`` of - shape [axes, batch, sequence] in place. - """ - assert num_chunks == 1 - assert not self.mtp_prepared - self.cp_prepared = True - base_positions = torch.tensor([[0, 1, 2, 0, 1, 2]]) - batch["position_ids"] = torch.stack( - (base_positions, base_positions + 10, base_positions + 20), - dim=0, - ) - return {} - - def prepare_mtp_inputs_for_cp(self, batch, *, ignore_index=-100): - """Build global MTP tensors after multi-axis positions exist. - - Args: - batch: Mutable mapping whose token tensors have shape [batch, - sequence] and positions have shape [axes, batch, sequence]. - ignore_index: Fill value for invalid target positions. - - Returns: - Global per-depth MTP tensors matching the batch token axes. - """ - assert self.cp_prepared - assert batch["position_ids"].shape == (3, 1, 6) - self.mtp_prepared = True - return prepare_mtp_context_parallel_inputs( - batch, - num_depths=self.mtp_config.num_layers, - ignore_index=ignore_index, - ) - - def forward( - self, - input_ids, - *, - mtp_per_depth_input_ids, - mtp_per_depth_position_ids, - mtp_per_depth_valid_masks, - **kwargs, - ): - """Emit MTP outputs while recording the CP-local auxiliary inputs. - - Args: - input_ids: Token IDs of shape [batch, local_sequence]. - mtp_per_depth_input_ids: Per-depth token IDs, each of shape - [batch, local_sequence]. - mtp_per_depth_position_ids: Per-depth positions, each of shape - [axes, batch, local_sequence]. - mtp_per_depth_valid_masks: Per-depth masks, each of shape - [batch, local_sequence]. - **kwargs: Remaining CP-local model tensors. - - Returns: - Model output whose logits and MTP hidden states have shape - [batch, local_sequence, hidden]. - """ - captured["model_input_ids"] = input_ids.detach().clone() - captured["mtp_input_ids"] = tuple(t.detach().clone() for t in mtp_per_depth_input_ids) - captured["mtp_position_ids"] = tuple(t.detach().clone() for t in mtp_per_depth_position_ids) - captured["mtp_valid_masks"] = tuple(t.detach().clone() for t in mtp_per_depth_valid_masks) - hidden = self.scale * input_ids.float().unsqueeze(-1) - return SimpleNamespace( - logits=hidden, - mtp_per_depth_h=[hidden], - mtp_per_depth_logits=None, - mtp_loss_scaling_factor=1.0, - ) - - model = _MTPModel() - - class _FakeContextParallelSharder: - def __init__(self, resolved_model, device_mesh, batch, **kwargs): - """Run the model-owned CP preparation phase without sharding. - - Args: - resolved_model: Model that owns the CP preparation hook. - device_mesh: Unused fake device mesh. - batch: Mutable global batch whose token tensors have shape - [batch, sequence]. - **kwargs: Sharder options including the number of chunks. - """ - del device_mesh - assert resolved_model is model - assert not model.mtp_prepared - assert batch["input_ids"].shape == (1, 6) - assert batch["seq_lens_padded"].tolist() == [[3, 3, -1000]] - updates = model.prepare_model_inputs_for_cp(batch, num_chunks=kwargs["num_chunks"]) - batch.update(updates) - - def shard(self, batch): - """Select this fake rank's token positions from every main input. - - Args: - batch: Global model-input mapping with token tensors of shape - [batch, sequence] and positions of shape [axes, batch, - sequence]. - - Returns: - A null context factory and mapping in the CP-local token layout. - """ - assert model.mtp_prepared - local_batch = dict(batch) - for key in ("input_ids", "labels"): - local_batch[key] = local_batch[key].index_select(1, local_indices) - local_batch["position_ids"] = local_batch["position_ids"].index_select(2, local_indices) - local_batch.pop("seq_lens") - local_batch.pop("seq_lens_padded") - return nullcontext, local_batch - - def shard_token_tensor(self, tensor, seq_dim=1, fill=None): - """Apply the captured CP index selection to an auxiliary tensor. - - Args: - tensor: Tensor of shape [batch, sequence] or [axes, batch, - sequence]. - seq_dim: Sequence axis selected by ``local_indices``. - fill: Padding value required by the production sharder. - - Returns: - Tensor with the same axis order and a CP-local sequence axis. - """ - if tensor.ndim == 3: - captured["mtp_position_seq_dim"] = seq_dim - if tensor.dtype == torch.bool: - captured["mtp_valid_mask_fill"] = fill - return tensor.index_select(seq_dim, local_indices) - - recipe = TrainFinetuneRecipeForNextTokenPrediction.__new__(TrainFinetuneRecipeForNextTokenPrediction) - object.__setattr__( - recipe, - "cfg", - SimpleNamespace(mtp=SimpleNamespace(ignore_index=-100, scaling_factor=1.0)), - ) - object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"))) - object.__setattr__(recipe, "device_mesh", object()) - object.__setattr__(recipe, "pp_enabled", False) - object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) - object.__setattr__(recipe, "te_fp8", None) - object.__setattr__(recipe, "model_parts", [model]) - object.__setattr__(recipe, "distributed_config", SimpleNamespace(defer_fsdp_grad_sync=True)) - object.__setattr__(recipe, "loss_fn", object()) - object.__setattr__(recipe, "_get_cp_group_size", lambda: 2) - object.__setattr__(recipe, "_get_dp_group_size", lambda include_cp=False: 1) - - def _fake_calculate_loss(loss_fn, *, logits, **kwargs): - """Return a differentiable scalar from local logits. - - Args: - loss_fn: Unused configured loss object. - logits: Tensor of shape [batch, local_sequence, hidden]. - **kwargs: Remaining loss inputs. - - Returns: - Scalar loss tensor. - """ - del loss_fn, kwargs - return logits.sum() * 0.0 - - def _fake_calculate_mtp_loss(loss_fn, *, mtp_per_depth_h, mtp_per_depth_targets, cu_seqlens, **kwargs): - """Capture precomputed targets and return a scalar auxiliary loss. - - Args: - loss_fn: Unused configured loss object. - mtp_per_depth_h: Per-depth hidden tensors of shape [batch, - local_sequence, hidden]. - mtp_per_depth_targets: Per-depth targets of shape [batch, - local_sequence]. - cu_seqlens: Optional packed-boundary tensor; must be absent when - precomputed targets are supplied. - **kwargs: Remaining MTP loss inputs. - - Returns: - Scalar MTP loss tensor. - """ - del loss_fn, kwargs - captured["mtp_targets"] = tuple(t.detach().clone() for t in mtp_per_depth_targets) - captured["cu_seqlens"] = cu_seqlens - return sum(hidden.sum() for hidden in mtp_per_depth_h) * 0.01 - - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.ContextParallelSharder", _FakeContextParallelSharder) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calculate_loss) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_mtp_loss", _fake_calculate_mtp_loss) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_final_hidden_states", lambda out: None) - - batch = { - "input_ids": torch.tensor([[10, 11, 12, 20, 21, 22]]), - "labels": torch.tensor([[11, 12, -100, 21, 22, -100]]), - "seq_lens": torch.tensor([[3, 3, -1000]]), - "seq_lens_padded": torch.tensor([[3, 3, -1000]]), - } - model.supports.supports_mtp_cp = False - with pytest.raises(NotImplementedError, match="supports_mtp_cp=False"): - recipe._forward_validation_step(batch) - assert not model.cp_prepared - assert not model.mtp_prepared - model.supports.supports_mtp_cp = True - - loss = recipe._forward_validation_step(batch) - - assert captured["model_input_ids"].tolist() == [[10, 11, 21, 22]] - assert captured["mtp_input_ids"][0].tolist() == [[11, 12, 22, 0]] - assert captured["mtp_position_ids"][0][0].tolist() == [[1, 2, 2, 0]] - assert captured["mtp_position_seq_dim"] == 2 - assert captured["mtp_valid_masks"][0].tolist() == [[True, True, True, False]] - assert captured["mtp_valid_mask_fill"] is False - assert captured["mtp_targets"][0].tolist() == [[12, -100, -100, -100]] - assert captured["cu_seqlens"] is None - assert torch.isfinite(loss) - assert model.scale.grad is None diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 6eff482398..2b2d07b2e1 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -42,7 +42,7 @@ from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.models.common.mtp import prepare_mtp_context_parallel_inputs from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler -from nemo_automodel.engine import Engine, collate_prebatched +from nemo_automodel.engine import Engine, ForwardResult, collate_prebatched class ScaleModel(nn.Module): @@ -100,6 +100,7 @@ def __init__( self.events = events self.callback_order = callback_order or list(range(num_microbatches)) self.step_calls = 0 + self.eval_calls = 0 self.backward_calls = 0 self.updated_seq_lens = [] self.updated_microbatch_sizes = [] @@ -142,6 +143,22 @@ def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): scaled_loss.backward() self.backward_calls += 1 + def eval_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): + assert return_outputs is False + assert len(model_inputs) == self.num_microbatches + self.prepared_inputs.append(model_inputs) + self.eval_calls += 1 + if self.events is not None: + self.events.append("eval") + + for index in self.callback_order: + inputs = dict(model_inputs[index]) + primary_name = "inputs_embeds" if "inputs_embeds" in inputs else "input_ids" + primary = inputs.pop(primary_name) + output = self.compute_model(primary, **inputs) + loss = loss_fn(output, index) + self.callback_losses.append(loss.detach()) + def _pipeline_mesh_context(): return SimpleNamespace(pp_size=2, cp_size=1, device_mesh=None, process_group=None) @@ -178,6 +195,69 @@ def test_engine_and_datum_are_lazy_top_level_exports(): assert PublicDatum is Datum +def test_forward_runs_eval_without_grad_lifecycle_and_returns_local_statistics(monkeypatch): + class EvalModel(ScaleModel): + def forward(self, input_ids: torch.Tensor, **kwargs) -> torch.Tensor: + assert not self.training + assert not torch.is_grad_enabled() + return super().forward(input_ids, **kwargs) + + model = EvalModel() + model.weight.grad = torch.tensor(7.0) + monkeypatch.setattr( + engine_module, + "prepare_for_grad_accumulation", + lambda *_args, **_kwargs: pytest.fail("forward-only execution must not prepare gradient accumulation"), + ) + monkeypatch.setattr( + engine_module, + "prepare_for_final_backward", + lambda *_args, **_kwargs: pytest.fail("forward-only execution must not prepare backward"), + ) + monkeypatch.setattr( + engine_module, + "get_sync_ctx", + lambda *_args, **_kwargs: pytest.fail("forward-only execution must not enter a gradient sync context"), + ) + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", torch.tensor(9.0)) + + def loss_with_outputs(output, _loss_inputs): + assert not torch.is_grad_enabled() + return output, [{"value": output.sum()}] + + engine = Engine(model, device="cpu") + engine._dp_group_and_size = lambda: pytest.fail("forward must not synchronize data-parallel replicas") + result = engine.forward( + [_datum([1, 100], [1.0, 0.0]), _datum([3], [0.5])], + loss_with_outputs, + ) + + assert isinstance(result, ForwardResult) + assert result.loss_sum.item() == pytest.approx(2.5) + assert result.weight_sum.item() == pytest.approx(1.5) + assert [item["value"].item() for item in result.loss_fn_outputs] == [101.0, 3.0] + assert all(not item["value"].requires_grad for item in result.loss_fn_outputs) + assert model.weight.grad.item() == 7.0 + assert model.forward_calls == 2 + assert not model.training + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == 9.0 + + +def test_forward_zero_weight_window_still_executes_without_gradients(): + model = ScaleModel() + + result = Engine(model, device="cpu").forward( + [_datum([1, 2], [0.0, 0.0])], + lambda output, _inputs: output.sum(), + ) + + assert result.loss_sum.item() == 0 + assert result.weight_sum.item() == 0 + assert result.loss_fn_outputs == [] + assert model.forward_calls == 1 + assert model.weight.grad is None + + def test_forward_backward_uses_one_denominator_for_the_window(): model = ScaleModel() initial_weight = model.weight.detach().clone() @@ -324,7 +404,8 @@ def loss_fn(output, inputs): assert not cp_context_active -def test_context_parallel_prepares_mtp_futures_before_sharding(): +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +def test_context_parallel_prepares_mtp_futures_before_sharding(execution): class MTPModel(ScaleModel): def __init__(self): super().__init__() @@ -393,7 +474,7 @@ def loss_fn(output, loss_inputs): captured_loss_inputs.update(loss_inputs) return output - loss, _ = engine.forward_backward([datum], loss_fn) + result = getattr(engine, execution)([datum], loss_fn) input_ids, position_ids, valid_masks = model.mtp_forward_inputs @@ -414,8 +495,14 @@ def loss_fn(output, loss_inputs): [[21, 22, 23, -7]], [[22, 23, -7, -7]], ] - assert loss.item() == pytest.approx(46 / 8) - assert model.weight.grad.item() == pytest.approx(46 / 8) + if execution == "forward": + assert result.loss_sum.item() == pytest.approx(46) + assert result.weight_sum.item() == pytest.approx(8) + assert model.weight.grad is None + else: + loss, _ = result + assert loss.item() == pytest.approx(46 / 8) + assert model.weight.grad.item() == pytest.approx(46 / 8) def test_context_parallel_rejects_mtp_without_model_capability(): @@ -709,7 +796,60 @@ def loss_fn(output, inputs): assert not active -def test_pipeline_stage_metadata_prevents_real_forward_from_resetting_media_cursor(): +def test_pipeline_forward_uses_eval_microbatches_without_backward_and_orders_outputs(): + model = ScaleModel() + other_part = ScaleModel() + pipeline = _FakeAutoPipeline( + model, + parts=[model, other_part], + num_microbatches=2, + callback_order=[1, 0], + ) + + result = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=2, + ).forward( + [_datum([1, 2]), _datum([3, 4])], + lambda output, _inputs: (output, [{"first_token": output.flatten()[0]}]), + ) + + assert result.loss_sum.item() == pytest.approx(10.0) + assert result.weight_sum.item() == pytest.approx(4.0) + assert [item["first_token"].item() for item in result.loss_fn_outputs] == [1.0, 3.0] + assert pipeline.eval_calls == 1 + assert pipeline.step_calls == 0 + assert pipeline.backward_calls == 0 + assert model.weight.grad is None + assert model.forward_calls == 2 + assert not model.training + assert not other_part.training + torch.testing.assert_close(torch.stack(pipeline.callback_losses), torch.tensor([7.0, 3.0])) + + +def test_forward_does_not_apply_backward_only_parallelism_restrictions(): + model = ScaleModel() + model.calculate_per_token_loss = True + eager_result = Engine(model, device="cpu").forward([_datum([1])], _identity_loss) + + pipeline_model = ScaleModel() + pipeline = _FakeAutoPipeline(pipeline_model, num_microbatches=1, scale_grads=True) + pipeline_result = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + ).forward([_datum([2])], _identity_loss) + + assert eager_result.loss_sum.item() == 1.0 + assert pipeline_result.loss_sum.item() == 2.0 + assert pipeline.eval_calls == 1 + assert pipeline.step_calls == 0 + + +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +def test_pipeline_stage_metadata_prevents_real_forward_from_resetting_media_cursor(execution): class MediaModel(ScaleModel): def __init__(self): super().__init__() @@ -756,17 +896,19 @@ def batch_context(model_inputs): loss_fn_inputs={"weights": torch.ones(2, 1)}, ) - Engine( + engine = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), collate_fn=collate_prebatched, context_fn=batch_context, - ).forward_backward([datum], _identity_loss) + ) + getattr(engine, execution)([datum], _identity_loss) assert model.consumed_media == [1, 2] assert pipeline.updated_seq_lens == [1] - assert pipeline.step_calls == 1 + assert pipeline.step_calls == int(execution == "forward_backward") + assert pipeline.eval_calls == int(execution == "forward") def test_pipeline_lifecycle_and_moe_scale_cover_outer_and_inner_microbatches(monkeypatch): @@ -1316,6 +1458,14 @@ def bad_collate(datums): def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) try: + # With an unwrapped model, forward-only execution adds no DP + # collectives and keeps each rank's statistics local. DDP/FSDP wrappers + # may impose their own aligned-call requirement. + forward_window = [_datum([1])] if rank == 0 else [_datum([2]), _datum([3])] + forward_result = Engine(ScaleModel(), device="cpu").forward(forward_window, _identity_loss) + assert forward_result.loss_sum.item() == pytest.approx(1.0 if rank == 0 else 5.0) + assert forward_result.weight_sum.item() == pytest.approx(1.0 if rank == 0 else 2.0) + model = nn.parallel.DistributedDataParallel(ScaleModel()) bad_window = [_datum([1])] if rank == 0 else [_datum([1]), _datum([2])] with pytest.raises(ValueError, match="same number of microbatches"): @@ -1371,6 +1521,20 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size assert loss.item() == pytest.approx(4.5) assert model.module.weight.grad.item() == pytest.approx(4.5) + + model.module.weight.grad = None + forward_result = Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ).forward(window, _identity_loss) + expected_sum = ( + 18.0 if dp_size == 1 else 10.0 + 16.0 * get_flat_mesh(mesh_context.device_mesh, "dp").get_local_rank() + ) + assert forward_result.loss_sum.item() == pytest.approx(expected_sum) + assert forward_result.weight_sum.item() == pytest.approx(4.0) + assert model.module.weight.grad is None finally: dist.destroy_process_group() From d563fc99a76f2034f6377b991898484c7ba6fea2 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Thu, 20 Aug 2026 17:32:23 -0700 Subject: [PATCH 10/34] feat(engine): route explicit loss input layouts Signed-off-by: HuiyingLi --- nemo_automodel/__init__.py | 2 + nemo_automodel/components/datasets/datum.py | 154 ++++- nemo_automodel/engine/__init__.py | 573 +++++++++++++++--- nemo_automodel/recipes/llm/train_ft.py | 6 +- nemo_automodel/recipes/vlm/finetune.py | 6 +- .../context_parallel/run_packed_pp.py | 244 +++++++- tests/unit_tests/datasets/test_datum.py | 181 ++++++ .../recipes/test_finetune_vlm_helpers.py | 9 + tests/unit_tests/recipes/test_train_ft.py | 5 + tests/unit_tests/test_engine.py | 345 +++++++++-- 10 files changed, 1381 insertions(+), 144 deletions(-) diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 143aae8a5d..3a0a69d04a 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -40,8 +40,10 @@ _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"), "NeMoAutoModelForCausalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForCausalLM"), "NeMoAutoModelForImageTextToText": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForImageTextToText"), "NeMoAutoModelForMultimodalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForMultimodalLM"), diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index 96bcf18075..d701253a6f 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -16,7 +16,10 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType from typing import Any import torch @@ -30,7 +33,70 @@ CROSS_ENTROPY_IGNORE_IDX = -100 -__all__ = ["Datum", "collate_datums"] +__all__ = ["CollatedLossInputs", "Datum", "LossInputLayout", "collate_datums"] + + +class LossInputLayout(str, Enum): + """How one loss input relates to the Datums being collated. + + ``PER_TOKEN`` values follow token padding, packing, CP sharding, and PP + microbatching. ``PER_DATUM`` values contain one scalar for each outer + :class:`Datum`; CP ranks receive the same scalars, so a loss callback uses + them with its CP-local token contribution rather than returning a repeated + full-sequence scalar. ``REPLICATED`` values are batch-level metadata copied + unchanged to every CP/PP microbatch. + """ + + PER_TOKEN = "per_token" + PER_DATUM = "per_datum" + REPLICATED = "replicated" + + +class CollatedLossInputs(dict[str, torch.Tensor]): + """Collated loss tensors with layout metadata outside the tensor mapping. + + This remains a normal ``dict`` for source compatibility. ``layouts`` is a + complete mapping over the dictionary's initial keys. ``item_to_datum`` + maps each padded row or valid THD sequence back to the input Datum index; + the current Engine accepts the identity mapping (one item per Datum, in + input order). It is ``None`` when a collater cannot expose its inner + boundaries, as with an already-prebatched batch. ``copy()`` preserves the + side channel; converting this object to a plain ``dict`` intentionally + drops it and opts back into the Engine's conservative legacy inference. + """ + + def __init__( + self, + values: Mapping[str, torch.Tensor], + *, + layouts: Mapping[str, LossInputLayout], + item_to_datum: tuple[int, ...] | None, + ) -> None: + super().__init__(values) + if set(layouts) != set(self): + raise ValueError("layouts must contain exactly the CollatedLossInputs keys") + if not all(isinstance(layout, LossInputLayout) for layout in layouts.values()): + raise TypeError("every loss input layout must be a LossInputLayout") + resolved_item_to_datum = None if item_to_datum is None else tuple(item_to_datum) + if resolved_item_to_datum is not None and not all(isinstance(index, int) for index in resolved_item_to_datum): + raise TypeError("item_to_datum must contain integer Datum indices") + + # Store a normal dict so the public collate result remains pickleable + # across DataLoader worker boundaries. Expose only a read-only view. + self._layouts = dict(layouts) + self.item_to_datum = resolved_item_to_datum + + @property + def layouts(self) -> Mapping[str, LossInputLayout]: + """Complete, read-only field-layout mapping.""" + return MappingProxyType(self._layouts) + + def copy(self) -> CollatedLossInputs: + """Return a shallow copy that retains the layout side channel.""" + return type(self)(self, layouts=self.layouts, item_to_datum=self.item_to_datum) + + def __copy__(self) -> CollatedLossInputs: + return self.copy() @dataclass(init=False) @@ -49,12 +115,16 @@ class Datum: model-specific because LLM and VLM processors emit different fields. loss_fn_inputs: Tensor values consumed by the loss function. + loss_fn_input_layouts: Optional semantic layouts for loss fields. The + canonical collater infers omitted fields using its legacy + token-aligned-versus-scalar rules. input_ids: Deprecated convenience spelling for the old text-only API. It cannot be combined with ``model_inputs``. """ model_inputs: dict[str, Any] loss_fn_inputs: dict[str, torch.Tensor] = field(default_factory=dict) + loss_fn_input_layouts: dict[str, LossInputLayout] = field(default_factory=dict) def __init__( self, @@ -62,6 +132,7 @@ def __init__( loss_fn_inputs: dict[str, torch.Tensor] | None = None, *, input_ids: torch.Tensor | list[int] | None = None, + loss_fn_input_layouts: Mapping[str, LossInputLayout] | None = None, ) -> None: # Preserve the old positional ``Datum(input_ids, loss_fn_inputs)`` form # while downstream users move to the model-ready mapping. @@ -79,6 +150,7 @@ def __init__( self.model_inputs = dict(model_inputs) self.loss_fn_inputs = dict(loss_fn_inputs or {}) + self.loss_fn_input_layouts = dict(loss_fn_input_layouts or {}) self.__post_init__() def __post_init__(self) -> None: @@ -97,6 +169,12 @@ def __post_init__(self) -> None: if not isinstance(value, torch.Tensor): self.loss_fn_inputs[key] = torch.as_tensor(value) + unknown_layouts = set(self.loss_fn_input_layouts) - set(self.loss_fn_inputs) + if unknown_layouts: + raise ValueError(f"loss_fn_input_layouts contains unknown loss inputs: {sorted(unknown_layouts)}") + if not all(isinstance(layout, LossInputLayout) for layout in self.loss_fn_input_layouts.values()): + raise TypeError("every loss_fn_input_layouts value must be a LossInputLayout") + @property def input_ids(self) -> torch.Tensor: """The text token sequence, retained for source compatibility.""" @@ -160,7 +238,7 @@ def collate_datums( packed: bool = False, pad_seq_len_divisible: int | None = None, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX, -) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: +) -> tuple[dict[str, Any], CollatedLossInputs]: """Collate text Datums into separate model and loss inputs. This is the default text collater. Callers with model-specific VLM @@ -174,9 +252,12 @@ def collate_datums( ignore_index: Label fill value used internally by the THD collater. Returns: - ``(model_inputs, loss_fn_inputs)``. Per-token loss inputs have shape - ``[B, T]`` in padded mode and ``[1, total_tokens]`` in packed mode; - scalar loss inputs have shape ``[B]``. + ``(model_inputs, loss_fn_inputs)``. The second item remains a ``dict`` + and also exposes complete ``layouts`` and ``item_to_datum`` metadata. + Per-token loss inputs have shape ``[B, T]`` in padded mode and + ``[1, total_tokens]`` in packed mode; per-Datum scalar loss inputs have + shape ``[B]``. Replicated inputs retain one copy of their original + shape. """ if not datums: raise ValueError("collate_datums requires at least one Datum") @@ -213,21 +294,66 @@ def collate_datums( width = int(model_inputs["input_ids"].shape[-1]) loss_inputs: dict[str, torch.Tensor] = {} + loss_layouts: dict[str, LossInputLayout] = {} for key in sorted(loss_keys): values = [datum.loss_fn_inputs[key] for datum in datums] - per_token = all(value.ndim == 1 and value.shape[0] == datum.seq_len for value, datum in zip(values, datums)) - if per_token: + declared_layouts = [datum.loss_fn_input_layouts[key] for datum in datums if key in datum.loss_fn_input_layouts] + if declared_layouts and len(declared_layouts) != len(datums): + raise ValueError(f"every Datum must declare the layout for loss input {key!r}, or none may declare it") + explicit_layouts = set(declared_layouts) + if len(explicit_layouts) > 1: + raise ValueError(f"every Datum must use the same explicit layout for loss input {key!r}") + explicit_layout = next(iter(explicit_layouts), None) + + token_aligned = [value.ndim == 1 and value.shape[0] == datum.seq_len for value, datum in zip(values, datums)] + if explicit_layout is LossInputLayout.PER_TOKEN and not all(token_aligned): + shapes = [tuple(value.shape) for value in values] + raise ValueError( + f"PER_TOKEN loss input {key!r} must be 1-D and match each Datum's token length; got {shapes}" + ) + + scalar_per_datum = [value.numel() == 1 for value in values] + if explicit_layout is LossInputLayout.PER_DATUM and not all(scalar_per_datum): + shapes = [tuple(value.shape) for value in values] + raise ValueError(f"PER_DATUM loss input {key!r} must contain one value per Datum; got {shapes}") + + layout = explicit_layout + if layout is None: + if all(token_aligned): + layout = LossInputLayout.PER_TOKEN + elif all(scalar_per_datum): + layout = LossInputLayout.PER_DATUM + else: + shapes = [tuple(value.shape) for value in values] + raise ValueError( + f"the default collater only supports scalar or 1-D token-aligned loss inputs; {key!r} has {shapes}" + ) + + loss_layouts[key] = layout + if layout is LossInputLayout.PER_TOKEN: if packed: loss_inputs[key] = torch.cat(values).unsqueeze(0) else: loss_inputs[key] = torch.stack([F.pad(value, (0, width - value.shape[0])) for value in values]) continue - if not all(value.numel() == 1 for value in values): - shapes = [tuple(value.shape) for value in values] - raise ValueError( - f"the default collater only supports scalar or 1-D token-aligned loss inputs; {key!r} has {shapes}" - ) - loss_inputs[key] = torch.stack([value.reshape(()) for value in values]) + if layout is LossInputLayout.PER_DATUM: + loss_inputs[key] = torch.stack([value.reshape(()) for value in values]) + continue - return model_inputs, loss_inputs + first = values[0] + if not all( + value.shape == first.shape + and value.dtype == first.dtype + and value.device == first.device + and torch.equal(value, first) + for value in values[1:] + ): + raise ValueError(f"REPLICATED loss input {key!r} must have the same value in every Datum") + loss_inputs[key] = first + + return model_inputs, CollatedLossInputs( + loss_inputs, + layouts=loss_layouts, + item_to_datum=tuple(range(len(datums))), + ) diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 766ae2c776..3a16ae1ffc 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -25,7 +25,12 @@ import torch.distributed as dist from torch import nn -from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel.components.datasets.datum import ( + CollatedLossInputs, + Datum, + LossInputLayout, + collate_datums, +) from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh @@ -40,7 +45,10 @@ ) from nemo_automodel.components.utils.model_utils import filter_forward_kwargs -CollateFn = Callable[[list[Datum]], tuple[dict[str, Any], dict[str, torch.Tensor]]] +CollateFn = Callable[ + [list[Datum]], + tuple[dict[str, Any], dict[str, torch.Tensor] | CollatedLossInputs], +] LossInputValue = torch.Tensor | tuple[torch.Tensor, ...] LossInputs = dict[str, LossInputValue] LossFn = Callable[ @@ -58,7 +66,7 @@ def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextMana return nullcontext() -def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: +def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLossInputs | dict[str, torch.Tensor]]: """Return one already-collated Datum without changing its layout. The Datum represents the whole prebatched item. Consequently, one @@ -76,7 +84,27 @@ def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], dict[str, t if len(datums) != 1: raise ValueError("collate_prebatched expects exactly one Datum per microbatch") datum = datums[0] - return dict(datum.model_inputs), dict(datum.loss_fn_inputs) + if set(datum.loss_fn_input_layouts) == set(datum.loss_fn_inputs): + loss_inputs: CollatedLossInputs | dict[str, torch.Tensor] = CollatedLossInputs( + datum.loss_fn_inputs, + layouts=datum.loss_fn_input_layouts, + item_to_datum=None, + ) + else: + # Source-compatible prebatched callers without complete metadata keep + # the legacy inference path. It remains fail-closed for ambiguous + # packed per-Datum fields. + loss_inputs = dict(datum.loss_fn_inputs) + return dict(datum.model_inputs), loss_inputs + + +@dataclass(frozen=True) +class _LossBatchLayout: + """Field semantics and logical-item routing captured at collate time.""" + + fields: Mapping[str, LossInputLayout] + item_to_datum: tuple[int, ...] | None + unresolved_fields: frozenset[str] = frozenset() @dataclass(frozen=True) @@ -150,9 +178,14 @@ class Engine: uses one inner pipeline microbatch; recipes enforce that configuration. Packed pipeline batches with multiple inner microbatches must split at sequence boundaries into equal-width token chunks. Token-aligned loss - fields follow those chunks. Per-Datum scalar loss fields do not yet - carry enough boundary metadata to be split in that layout and are - rejected instead of being replicated silently. + fields follow those chunks. Layout-aware collaters may additionally + route per-Datum scalar fields from THD sequence boundaries and preserve + replicated fields unchanged. A ``PER_DATUM`` field is copied to every + CP rank; a loss callback combines it with that rank's CP-local token + contribution so scalar numerators remain additive across CP. Custom or + prebatched collaters that hide the sequence-to-Datum relationship + remain fail-closed for per-Datum fields instead of guessing inner + sample boundaries. """ def __init__( @@ -233,7 +266,9 @@ def forward( returns_outputs: bool | None = None for batch_datums in microbatches: - cp_context, model_inputs, loss_inputs = self._prepare_batch(batch_datums, inner_microbatches) + cp_context, model_inputs, loss_inputs, loss_batch_layout = self._prepare_batch( + batch_datums, inner_microbatches + ) if self.pipeline is not None: batch_returns_outputs, batch_outputs = self._pipeline_execute( model_inputs, @@ -242,6 +277,7 @@ def forward( loss_fn, local_loss_sum, cp_context, + loss_batch_layout, backward_scale=None, zero_weight_sum=zero_weight_sum, ) @@ -369,7 +405,7 @@ def forward_backward( if is_last: prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) - cp_context, model_inputs, loss_inputs = self._prepare_batch(datums, inner_microbatches) + cp_context, model_inputs, loss_inputs, loss_batch_layout = self._prepare_batch(datums, inner_microbatches) if self.pipeline is not None: backward_scale = ( @@ -384,6 +420,7 @@ def forward_backward( loss_fn, local_loss_sum, cp_context, + loss_batch_layout, backward_scale=backward_scale, zero_weight_sum=zero_denominator, ) @@ -447,11 +484,92 @@ def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: list(datums[start : start + self.microbatch_size]) for start in range(0, len(datums), self.microbatch_size) ] + @staticmethod + def _resolve_loss_batch_layout( + datums: list[Datum], + model_inputs: Mapping[str, Any], + loss_inputs: Mapping[str, LossInputValue], + ) -> _LossBatchLayout: + """Resolve collater metadata without exposing it to the loss callback.""" + if isinstance(loss_inputs, CollatedLossInputs): + if set(loss_inputs.layouts) != set(loss_inputs): + raise ValueError("CollatedLossInputs.layouts must describe every loss field exactly once") + item_to_datum = loss_inputs.item_to_datum + if item_to_datum is not None and item_to_datum != tuple(range(len(datums))): + raise ValueError( + "Engine currently requires collater item_to_datum to preserve outer Datum order; " + f"got {list(item_to_datum)} for {len(datums)} Datums" + ) + return _LossBatchLayout( + fields=dict(loss_inputs.layouts), + item_to_datum=item_to_datum, + ) + + weights = loss_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("collate_fn must return a Tensor loss input named 'weights'") + + fields: dict[str, LossInputLayout] = {} + unresolved: set[str] = set() + for name, value in loss_inputs.items(): + declared = {datum.loss_fn_input_layouts[name] for datum in datums if name in datum.loss_fn_input_layouts} + if len(declared) > 1: + raise ValueError(f"Datum items disagree on the loss layout for field {name!r}") + if declared: + if not all(name in datum.loss_fn_input_layouts for datum in datums): + raise ValueError(f"every Datum must declare the loss layout for field {name!r}") + fields[name] = next(iter(declared)) + continue + + if isinstance(value, torch.Tensor) and _loss_sequence_dim(dict(model_inputs), value) is not None: + fields[name] = LossInputLayout.PER_TOKEN + continue + + # Preserve source compatibility for older prepared/custom batches + # whose output weights are not token-shaped. This is deliberately + # unresolved rather than a new public PER_DATUM-weight contract: + # padded PP keeps its historical shape slicing, while packed PP + # fails closed without explicit token weights. + if name == "weights": + fields[name] = LossInputLayout.REPLICATED + unresolved.add(name) + continue + + datum_values = [datum.loss_fn_inputs.get(name) for datum in datums] + if ( + isinstance(value, torch.Tensor) + and value.ndim > 0 + and value.shape[0] == len(datums) + and all(isinstance(item, torch.Tensor) and item.numel() == 1 for item in datum_values) + ): + fields[name] = LossInputLayout.PER_DATUM + else: + fields[name] = LossInputLayout.REPLICATED + unresolved.add(name) + + # Legacy padded collaters historically preserve one input row per + # Datum. Keep that path source compatible; packed collaters must opt in + # explicitly because a THD sequence is not necessarily an outer Datum. + item_to_datum: tuple[int, ...] | None = None + primary = model_inputs.get("inputs_embeds", model_inputs.get("input_ids")) + if ( + model_inputs.get("qkv_format") != "thd" + and isinstance(primary, torch.Tensor) + and primary.ndim > 0 + and primary.shape[0] == len(datums) + ): + item_to_datum = tuple(range(len(datums))) + return _LossBatchLayout( + fields=fields, + item_to_datum=item_to_datum, + unresolved_fields=frozenset(unresolved), + ) + def _prepare_batch( self, datums: list[Datum], num_pipeline_microbatches: int, - ) -> tuple[Callable[[], AbstractContextManager[Any]], dict[str, Any], LossInputs]: + ) -> tuple[Callable[[], AbstractContextManager[Any]], dict[str, Any], LossInputs, _LossBatchLayout]: """Collate, move, and CP-shard one outer batch. Args: @@ -461,21 +579,40 @@ def _prepare_batch( the prepared outer batch must materialize. Returns: - The CP context factory, CP-local model inputs, and CP-local loss - inputs. Token-aligned model and loss tensors use the same padded, - packed THD, Magi, or model-owned local sequence layout. + The CP context factory, CP-local model inputs, CP-local loss + inputs, and collated field-layout metadata. Token-aligned model + and loss tensors use the same padded, packed THD, Magi, or + model-owned local sequence layout. """ - model_inputs, loss_inputs = self.collate_fn(datums) + model_inputs, collated_loss_inputs = self.collate_fn(datums) + loss_batch_layout = self._resolve_loss_batch_layout(datums, model_inputs, collated_loss_inputs) + loss_inputs = dict(collated_loss_inputs) + weight_layout = loss_batch_layout.fields.get("weights") + if weight_layout is not LossInputLayout.PER_TOKEN and not ( + weight_layout is LossInputLayout.REPLICATED and "weights" in loss_batch_layout.unresolved_fields + ): + raise ValueError("loss input 'weights' must use the PER_TOKEN layout") + if "labels" in loss_inputs and loss_batch_layout.fields["labels"] is not LossInputLayout.PER_TOKEN: + raise ValueError("loss input 'labels' must use the PER_TOKEN layout") self._validate_collated_weights(datums, loss_inputs) + token_reference_name, loss_seq_dim = self._validate_loss_batch_layout( + datums, + model_inputs, + loss_inputs, + loss_batch_layout.fields, + ) model_inputs = _to_device(model_inputs, self.device) loss_inputs = _to_device(loss_inputs, self.device) - full_weights = loss_inputs["weights"] - loss_seq_dim = _loss_sequence_dim(model_inputs, full_weights) + token_reference = ( + loss_inputs[token_reference_name] + if token_reference_name is not None + else _model_token_template(model_inputs) + ) cp_batch = dict(model_inputs) labels = loss_inputs.get("labels") cp_batch["labels"] = ( - labels.clone() if isinstance(labels, torch.Tensor) else torch.zeros_like(full_weights, dtype=torch.long) + labels.clone() if isinstance(labels, torch.Tensor) else torch.zeros_like(token_reference, dtype=torch.long) ) is_thd = cp_batch.get("qkv_format") == "thd" position_ids = cp_batch.get("position_ids") @@ -491,24 +628,31 @@ def _prepare_batch( thd_loss_fields: list[str] = [] if is_thd: - ambiguous_per_datum_fields = [ + if "weights" in loss_batch_layout.unresolved_fields: + raise ValueError("packed THD execution requires token-aligned loss weights") + unresolved_non_token_fields = [ name - for name, value in loss_inputs.items() - if name != "labels" - and isinstance(value, torch.Tensor) - and value.ndim > 0 - and value.shape[0] == len(datums) - and len(datums) > 1 - and not _is_token_aligned(value, full_weights) + for name in loss_batch_layout.unresolved_fields + if loss_batch_layout.fields[name] is not LossInputLayout.PER_TOKEN + ] + if num_pipeline_microbatches > 1 and unresolved_non_token_fields: + raise NotImplementedError( + "packed pipeline microbatching requires an explicit PER_DATUM or REPLICATED layout for " + f"non-token loss fields {unresolved_non_token_fields}" + ) + per_datum_fields = [ + name for name, layout in loss_batch_layout.fields.items() if layout is LossInputLayout.PER_DATUM ] - if num_pipeline_microbatches > 1 and ambiguous_per_datum_fields: + if num_pipeline_microbatches > 1 and per_datum_fields and loss_batch_layout.item_to_datum is None: raise NotImplementedError( - "packed pipeline microbatching cannot yet split per-Datum loss fields " - f"{ambiguous_per_datum_fields}; use token-aligned fields or a prepared collater" + "packed pipeline microbatching requires collater item_to_datum metadata for per-Datum " + f"loss fields {per_datum_fields}; use collate_datums or an explicitly layout-aware collater" ) for name, value in loss_inputs.items(): - if name == "labels" or not _is_token_aligned(value, full_weights): + if name == "labels" or loss_batch_layout.fields[name] is not LossInputLayout.PER_TOKEN: continue + if not _is_token_aligned(value, token_reference): + raise ValueError(f"per-token loss field {name!r} does not match the collated token layout") key = f"{_LOSS_FIELD_PREFIX}{name}" if key in cp_batch: raise ValueError(f"model inputs contain reserved Engine key {key!r}") @@ -547,12 +691,27 @@ def _prepare_batch( ) loss_inputs = local_loss_inputs else: - loss_inputs = self._shard_loss_inputs(sharder, loss_inputs, loss_seq_dim) + loss_inputs = self._shard_loss_inputs( + sharder, + loss_inputs, + loss_seq_dim, + loss_batch_layout.fields, + token_reference, + loss_batch_layout.unresolved_fields, + ) if labels is not None: loss_inputs["labels"] = local_labels if mtp_cp_inputs is not None: loss_inputs = self._attach_mtp_cp_inputs(sharder, mtp_cp_inputs, model_inputs, loss_inputs) - return cp_context, model_inputs, loss_inputs + loss_batch_layout = _LossBatchLayout( + fields={ + **loss_batch_layout.fields, + "mtp_per_depth_targets": LossInputLayout.PER_TOKEN, + }, + item_to_datum=loss_batch_layout.item_to_datum, + unresolved_fields=loss_batch_layout.unresolved_fields, + ) + return cp_context, model_inputs, loss_inputs, loss_batch_layout def _prepare_mtp_cp_inputs(self, batch: dict[str, Any]) -> MTPContextParallelInputs | None: """Prepare global MTP future-token tensors before CP shards the batch. @@ -643,6 +802,7 @@ def _pipeline_execute( loss_fn: LossFn, local_loss_sum: torch.Tensor, cp_context: Callable[[], AbstractContextManager[Any]], + loss_batch_layout: _LossBatchLayout, *, backward_scale: torch.Tensor | None, zero_weight_sum: bool, @@ -656,6 +816,8 @@ def _pipeline_execute( loss_fn: Model-output loss callback. local_loss_sum: Accumulator updated with detached numerators. cp_context: Context covering the complete pipeline schedule. + loss_batch_layout: Semantic layout of every loss field plus the + collater's logical item-to-Datum routing, when available. backward_scale: Multiplier returned to the training schedule for backward, or ``None`` to run the forward-only schedule. zero_weight_sum: Whether reporting numerators must be forced to @@ -666,6 +828,7 @@ def _pipeline_execute( logical Datum order. """ outputs_by_microbatch: list[list[dict[str, Any]] | None] = [None] * self.pipeline.num_microbatches + outputs_by_datum: list[dict[str, Any] | None] = [None] * len(datums) returns_outputs: bool | None = None with cp_context(): @@ -715,8 +878,13 @@ def _pipeline_execute( # that metadata first so VLM media cursors are not reset after the # first actual pipeline microbatch. with self.context_fn(model_inputs): - model_microbatches, loss_microbatches = self._materialize_pipeline_microbatches( - model_inputs, loss_inputs + model_microbatches, loss_microbatches, datum_indices_by_microbatch = ( + self._materialize_pipeline_microbatches( + model_inputs, + loss_inputs, + loss_batch_layout, + num_datums=len(datums), + ) ) def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: @@ -736,12 +904,31 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: or not all(isinstance(item, Mapping) for item in batch_outputs) ): raise ValueError("loss_fn outputs must be a sequence of mappings") - if len(datums) == 1 and self.pipeline.num_microbatches > 1: + datum_indices = ( + None + if datum_indices_by_microbatch is None + else datum_indices_by_microbatch[microbatch_index] + ) + if datum_indices is None and len(datums) == 1 and self.pipeline.num_microbatches > 1: raise ValueError( "a prebatched Datum may return outputs only when num_microbatches=1 because " "its inner sample boundaries are not part of the Datum contract" ) - outputs_by_microbatch[microbatch_index] = [_detach(dict(item)) for item in batch_outputs] + detached_outputs = [_detach(dict(item)) for item in batch_outputs] + if datum_indices is not None: + if len(detached_outputs) != len(datum_indices): + raise ValueError( + f"pipeline loss_fn returned {len(detached_outputs)} outputs for microbatch " + f"{microbatch_index}, expected {len(datum_indices)} from its Datum mapping" + ) + for datum_index, item in zip(datum_indices, detached_outputs): + if outputs_by_datum[datum_index] is not None: + raise RuntimeError( + f"pipeline returned more than one output for Datum {datum_index}" + ) + outputs_by_datum[datum_index] = item + else: + outputs_by_microbatch[microbatch_index] = detached_outputs else: losses = result numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) @@ -763,14 +950,19 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: outputs: list[dict[str, Any]] = [] if self.pipeline.info.has_last_stage and returns_outputs: - if any(items is None for items in outputs_by_microbatch): - raise RuntimeError("pipeline schedule did not evaluate loss_fn for every logical microbatch") - outputs = [item for items in outputs_by_microbatch if items is not None for item in items] - if len(outputs) != len(datums): - raise ValueError( - f"pipeline loss_fn returned {len(outputs)} outputs across the outer batch, " - f"expected one for each of its {len(datums)} Datums" - ) + if datum_indices_by_microbatch is not None: + if any(item is None for item in outputs_by_datum): + raise RuntimeError("pipeline schedule did not return exactly one output for every Datum") + outputs = [item for item in outputs_by_datum if item is not None] + else: + if any(items is None for items in outputs_by_microbatch): + raise RuntimeError("pipeline schedule did not evaluate loss_fn for every logical microbatch") + outputs = [item for items in outputs_by_microbatch if items is not None for item in items] + if len(outputs) != len(datums): + raise ValueError( + f"pipeline loss_fn returned {len(outputs)} outputs across the outer batch, " + f"expected one for each of its {len(datums)} Datums" + ) outputs = self._broadcast_pipeline_outputs(outputs) return bool(outputs), outputs @@ -808,7 +1000,10 @@ def _materialize_pipeline_microbatches( self, model_inputs: dict[str, Any], loss_inputs: LossInputs, - ) -> tuple[list[dict[str, Any]], list[LossInputs]]: + loss_batch_layout: _LossBatchLayout, + *, + num_datums: int, + ) -> tuple[list[dict[str, Any]], list[LossInputs], list[tuple[int, ...]] | None]: """Split one CP-prepared outer batch into exact pipeline inputs. Args: @@ -817,22 +1012,40 @@ def _materialize_pipeline_microbatches( [microbatches, tokens, ...]. loss_inputs: CP-local loss tensors. Token-aligned fields have the same leading token axes as the primary model tensor. + loss_batch_layout: Semantic field layouts and optional collater + item-to-Datum mapping. + num_datums: Number of outer Datum items represented by this batch. Returns: - Parallel lists of complete model and loss mappings, each with - exactly ``pipeline.num_microbatches`` items. Tensor slicing returns - views that retain a size-one pipeline microbatch axis. + Parallel lists of complete model and loss mappings, plus optional + logical Datum indices for each microbatch. Every list has exactly + ``pipeline.num_microbatches`` items. Token slicing retains a + size-one pipeline microbatch axis. """ num_microbatches = self.pipeline.num_microbatches if num_microbatches == 1: - return [dict(model_inputs)], [_with_loss_metadata(model_inputs, loss_inputs)] + datum_indices = [tuple(range(num_datums))] + if loss_batch_layout.item_to_datum is not None: + datum_indices = _pipeline_datum_indices( + [model_inputs], + loss_batch_layout.item_to_datum, + num_datums=num_datums, + is_thd=model_inputs.get("qkv_format") == "thd", + ) + assert datum_indices is not None + return ( + [dict(model_inputs)], + [_with_loss_metadata(model_inputs, loss_inputs)], + datum_indices, + ) primary_name = _primary_name(model_inputs) primary = model_inputs[primary_name] if not isinstance(primary, torch.Tensor) or primary.ndim == 0: raise ValueError(f"pipeline Engine requires tensor {primary_name}") - if model_inputs.get("qkv_format") == "thd": + is_thd = model_inputs.get("qkv_format") == "thd" + if is_thd: if primary.shape[0] != num_microbatches: raise ValueError( f"THD sharder produced {primary.shape[0]} chunks, expected {num_microbatches} pipeline microbatches" @@ -840,9 +1053,6 @@ def _materialize_pipeline_microbatches( model_microbatches = [ _select_chunk(model_inputs, index, num_microbatches) for index in range(num_microbatches) ] - loss_microbatches = [ - _select_chunk(loss_inputs, index, num_microbatches) for index in range(num_microbatches) - ] else: batch_size = primary.shape[0] if batch_size % num_microbatches != 0: @@ -863,16 +1073,33 @@ def _materialize_pipeline_microbatches( _slice_batch_mapping(model_inputs, index, num_microbatches, batch_size, custom_dims) for index in range(num_microbatches) ] - loss_microbatches = [ - _slice_batch_mapping(loss_inputs, index, num_microbatches, batch_size) - for index in range(num_microbatches) - ] + + datum_indices_by_microbatch = _pipeline_datum_indices( + model_microbatches, + loss_batch_layout.item_to_datum, + num_datums=num_datums, + is_thd=is_thd, + ) + loss_microbatches = [ + _materialize_loss_mapping( + loss_inputs, + loss_batch_layout.fields, + loss_batch_layout.unresolved_fields, + index=index, + num_chunks=num_microbatches, + batch_size=None if is_thd else batch_size, + datum_indices=(None if datum_indices_by_microbatch is None else datum_indices_by_microbatch[index]), + num_datums=num_datums, + is_thd=is_thd, + ) + for index in range(num_microbatches) + ] loss_microbatches = [ _with_loss_metadata(model_microbatch, loss_microbatch) for model_microbatch, loss_microbatch in zip(model_microbatches, loss_microbatches) ] - return model_microbatches, loss_microbatches + return model_microbatches, loss_microbatches, datum_indices_by_microbatch def _validate_parallelism(self) -> None: """Validate topology plus backward-specific distributed contracts.""" @@ -1022,9 +1249,12 @@ def _validate_weight_sum_across_cp(self, local_weight_sum: torch.Tensor) -> None def _shard_loss_inputs( self, sharder: ContextParallelSharder, - loss_inputs: dict[str, torch.Tensor], + loss_inputs: LossInputs, seq_dim: int | None, - ) -> dict[str, torch.Tensor]: + layouts: Mapping[str, LossInputLayout], + token_reference: torch.Tensor, + unresolved_fields: frozenset[str], + ) -> LossInputs: """Apply the model batch's CP token layout to loss-only tensors. Args: @@ -1034,40 +1264,93 @@ def _shard_loss_inputs( leading token axes is sharded identically. seq_dim: Sequence axis in the pre-CP loss tensors, or ``None`` when the weights do not follow the model's token axes. + layouts: Explicit semantic layout for every loss field. Only + ``PER_TOKEN`` fields follow the context-parallel token shard. + token_reference: Tensor carrying the full collated token axes. + unresolved_fields: Legacy fields whose semantics were not declared + by their collater. Non-token legacy weights cannot cross a CP + layout change safely. Returns: Loss tensors in the model output's CP-local token layout. Non-token tensors are returned unchanged; the input mapping is not mutated. """ - weights = loss_inputs["weights"] - layout = sharder.shard_layout - if self._cp_size() == 1 and layout is None: - return {name: value for name, value in loss_inputs.items() if name != "labels"} + shard_layout = sharder.shard_layout layout_changed = self._cp_size() > 1 or ( - layout is not None + shard_layout is not None and ( - layout.input_row_shape is not None - or layout.input_token_stream_positions is not None - or layout.original_seq_len != layout.padded_seq_len + shard_layout.input_row_shape is not None + or shard_layout.input_token_stream_positions is not None + or shard_layout.original_seq_len != shard_layout.padded_seq_len ) ) + if "weights" in unresolved_fields and layout_changed: + raise ValueError("context-parallel loss weights must match the model's token axes") + if self._cp_size() == 1 and shard_layout is None: + return {name: value for name, value in loss_inputs.items() if name != "labels"} if seq_dim is None: - if layout_changed: - raise ValueError("context-parallel loss weights must match the model's token axes") - return dict(loss_inputs) + if any(field_layout is LossInputLayout.PER_TOKEN for field_layout in layouts.values()): + raise ValueError("context-parallel per-token loss inputs must match the model's token axes") + return {name: value for name, value in loss_inputs.items() if name != "labels"} - local: dict[str, torch.Tensor] = {} + local: LossInputs = {} for name, value in loss_inputs.items(): if name == "labels": continue - token_aligned = value.ndim >= weights.ndim and tuple(value.shape[: weights.ndim]) == tuple(weights.shape) - local[name] = sharder.shard_token_tensor(value, seq_dim=seq_dim, fill=0) if token_aligned else value + if layouts[name] is not LossInputLayout.PER_TOKEN: + local[name] = value + continue + if not isinstance(value, torch.Tensor): + raise TypeError(f"per-token loss field {name!r} must be a Tensor before CP sharding") + token_aligned = _is_token_aligned(value, token_reference) + if not token_aligned: + raise ValueError(f"per-token loss field {name!r} does not match the collated token layout") + local[name] = sharder.shard_token_tensor(value, seq_dim=seq_dim, fill=0) return local + @staticmethod + def _validate_loss_batch_layout( + datums: Sequence[Datum], + model_inputs: Mapping[str, Any], + loss_inputs: Mapping[str, LossInputValue], + layouts: Mapping[str, LossInputLayout], + ) -> tuple[str | None, int | None]: + """Validate semantic loss layouts before any CP/PP transformation.""" + if set(layouts) != set(loss_inputs): + raise ValueError("loss input layouts must describe every collated loss field exactly once") + weights = loss_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("collate_fn must return a Tensor loss input named 'weights'") + token_fields = [name for name, layout in layouts.items() if layout is LossInputLayout.PER_TOKEN] + token_reference_name = "weights" if layouts.get("weights") is LossInputLayout.PER_TOKEN else None + if token_reference_name is None and "labels" in token_fields: + token_reference_name = "labels" + if token_reference_name is None and token_fields: + token_reference_name = token_fields[0] + token_reference = loss_inputs.get(token_reference_name) if token_reference_name is not None else None + if token_reference_name is not None and not isinstance(token_reference, torch.Tensor): + raise TypeError(f"per-token loss field {token_reference_name!r} must be a Tensor") + loss_seq_dim = ( + _loss_sequence_dim(model_inputs, token_reference) if isinstance(token_reference, torch.Tensor) else None + ) + if token_reference_name is not None and loss_seq_dim is None: + raise ValueError(f"per-token loss field {token_reference_name!r} must match the primary model token axes") + + for name, value in loss_inputs.items(): + layout = layouts[name] + if layout is LossInputLayout.PER_TOKEN: + if not isinstance(value, torch.Tensor) or not _is_token_aligned(value, token_reference): + raise ValueError(f"per-token loss field {name!r} does not match the collated token layout") + elif layout is LossInputLayout.PER_DATUM: + if not isinstance(value, torch.Tensor) or value.ndim == 0 or value.shape[0] != len(datums): + shape = tuple(value.shape) if isinstance(value, torch.Tensor) else type(value).__name__ + raise ValueError(f"per-Datum loss field {name!r} must have leading size {len(datums)}, got {shape}") + return token_reference_name, loss_seq_dim + @staticmethod def _validate_collated_weights( datums: list[Datum], - loss_inputs: dict[str, torch.Tensor], + loss_inputs: Mapping[str, LossInputValue], ) -> None: weights = loss_inputs.get("weights") if not isinstance(weights, torch.Tensor): @@ -1094,26 +1377,44 @@ def _to_device(value: Any, device: torch.device) -> Any: return value -def _loss_sequence_dim(model_inputs: dict[str, Any], weights: torch.Tensor) -> int | None: +def _model_token_template(model_inputs: Mapping[str, Any]) -> torch.Tensor: + """Return a tensor with exactly the primary model input's token axes.""" + primary_name = _primary_name(model_inputs) + primary = model_inputs[primary_name] + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError("model primary input must be a non-scalar Tensor") + if primary_name == "inputs_embeds": + if primary.ndim < 2 or primary.shape[-1] == 0: + raise ValueError("inputs_embeds must contain token and hidden dimensions") + return primary.select(-1, 0) + return primary + + +def _loss_sequence_dim(model_inputs: dict[str, Any], value: torch.Tensor) -> int | None: """Find the sequence axis shared by primary model tokens and loss weights. Args: model_inputs: Pre-CP model mapping whose ``input_ids`` has shape ``[batch, sequence]`` or ``[tokens]``, or whose ``inputs_embeds`` has shape ``[batch, sequence, hidden]``. - weights: Loss weights of shape ``[batch, sequence]`` or ``[tokens]``. + value: Candidate token-aligned tensor. It may have trailing feature + dimensions after the primary input's token axes. Returns: - The sequence axis in ``weights``, or ``None`` when the layouts do not + The sequence axis in ``value``, or ``None`` when the layouts do not describe the same token stream. """ - primary = model_inputs.get("inputs_embeds", model_inputs.get("input_ids")) - if not isinstance(primary, torch.Tensor): + try: + token_template = _model_token_template(model_inputs) + except ValueError: return None - if primary.ndim >= 2 and weights.ndim >= 2 and tuple(weights.shape[:2]) == tuple(primary.shape[:2]): - return 1 - if primary.ndim == 1 and weights.ndim == 1 and weights.shape == primary.shape: - return 0 + token_dims = token_template.ndim + if ( + token_dims in {1, 2} + and value.ndim >= token_dims + and tuple(value.shape[:token_dims]) == tuple(token_template.shape) + ): + return token_dims - 1 return None @@ -1230,6 +1531,112 @@ def slice_value(value: Any, dim: int | None = None) -> Any: } +def _pipeline_datum_indices( + model_microbatches: Sequence[Mapping[str, Any]], + item_to_datum: tuple[int, ...] | None, + *, + num_datums: int, + is_thd: bool, +) -> list[tuple[int, ...]] | None: + """Map each materialized model microbatch back to its outer Datums.""" + if item_to_datum is None: + if is_thd: + return None + row_count = sum(_padded_microbatch_size(microbatch) for microbatch in model_microbatches) + if row_count != num_datums: + return None + item_to_datum = tuple(range(num_datums)) + + if len(item_to_datum) != num_datums or sorted(item_to_datum) != list(range(num_datums)): + raise ValueError( + "collater item_to_datum must contain every outer Datum index exactly once; " + f"got {list(item_to_datum)} for {num_datums} Datums" + ) + + counts = [ + _thd_microbatch_sequence_count(microbatch) if is_thd else _padded_microbatch_size(microbatch) + for microbatch in model_microbatches + ] + if sum(counts) != len(item_to_datum): + raise ValueError( + "collater item_to_datum does not match the materialized model items; " + f"microbatch counts are {counts}, mapping has {len(item_to_datum)} entries" + ) + + result: list[tuple[int, ...]] = [] + start = 0 + for count in counts: + result.append(item_to_datum[start : start + count]) + start += count + return result + + +def _padded_microbatch_size(model_inputs: Mapping[str, Any]) -> int: + primary = model_inputs[_primary_name(model_inputs)] + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError("padded pipeline microbatch requires a batched primary tensor") + return int(primary.shape[0]) + + +def _thd_microbatch_sequence_count(model_inputs: Mapping[str, Any]) -> int: + cu_seqlens = model_inputs.get("cu_seqlens") + if not isinstance(cu_seqlens, torch.Tensor): + raise ValueError("packed per-Datum routing requires tensor cu_seqlens in every pipeline microbatch") + valid = cu_seqlens.reshape(-1) + valid = valid[valid >= 0] + if valid.numel() < 2: + raise ValueError("packed pipeline microbatch must contain at least one sequence") + return int(valid.numel() - 1) + + +def _materialize_loss_mapping( + loss_inputs: Mapping[str, LossInputValue], + layouts: Mapping[str, LossInputLayout], + unresolved_fields: frozenset[str], + *, + index: int, + num_chunks: int, + batch_size: int | None, + datum_indices: tuple[int, ...] | None, + num_datums: int, + is_thd: bool, +) -> LossInputs: + """Materialize token, Datum, and replicated loss fields by semantics.""" + result: LossInputs = {} + for name, value in loss_inputs.items(): + if name in unresolved_fields: + if is_thd: + raise NotImplementedError( + f"packed pipeline microbatching requires an explicit layout for loss field {name!r}" + ) + assert batch_size is not None + result[name] = _slice_batch_mapping({name: value}, index, num_chunks, batch_size)[name] + continue + layout = layouts[name] + if layout is LossInputLayout.PER_TOKEN: + if is_thd: + result[name] = _select_chunk(value, index, num_chunks) + else: + assert batch_size is not None + result[name] = _slice_batch_mapping({name: value}, index, num_chunks, batch_size)[name] + elif layout is LossInputLayout.PER_DATUM: + if datum_indices is None: + raise NotImplementedError( + f"pipeline microbatching cannot route per-Datum loss field {name!r} without " + "collater item_to_datum metadata" + ) + if not isinstance(value, torch.Tensor) or value.ndim == 0 or value.shape[0] != num_datums: + shape = tuple(value.shape) if isinstance(value, torch.Tensor) else type(value).__name__ + raise ValueError(f"per-Datum loss field {name!r} must have leading size {num_datums}, got {shape}") + indices = torch.tensor(datum_indices, dtype=torch.long, device=value.device) + result[name] = value.index_select(0, indices) + elif layout is LossInputLayout.REPLICATED: + result[name] = value + else: # pragma: no cover - normalized by Datum/CollatedLossInputs + raise ValueError(f"unsupported loss layout {layout!r} for field {name!r}") + return result + + def _weighted_numerator(losses: Any, weights: torch.Tensor) -> torch.Tensor: if not isinstance(losses, torch.Tensor): raise TypeError("loss_fn must return a Tensor, optionally followed by per-Datum outputs") diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 4de5860d4a..d0d5eb8190 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -56,7 +56,7 @@ from nemo_automodel._transformers.utils import apply_cache_compatibility_patches from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.components.cuda_graphs import PartialCudaGraphManager -from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.datasets.datum import Datum, LossInputLayout from nemo_automodel.components.datasets.loader import DataloaderConfig from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi @@ -1163,6 +1163,10 @@ def _make_engine_datum(self, batch: dict[str, Any]) -> Datum: return Datum( model_inputs=model_inputs, loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, + loss_fn_input_layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + }, ) def _engine_loss_fn( diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index e799b615fe..27d77385bb 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -47,7 +47,7 @@ ) from nemo_automodel._transformers.utils import apply_cache_compatibility_patches, resolve_get_rope_index from nemo_automodel.components.config._arg_parser import parse_args_and_load_config -from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.datasets.datum import Datum, LossInputLayout from nemo_automodel.components.datasets.vlm.pp_media import VLM_PP_MEDIA_KEY, stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi @@ -932,6 +932,10 @@ def _make_engine_datum(self, batch: dict[str, Any]) -> Datum: return Datum( model_inputs=model_inputs, loss_fn_inputs={"labels": labels, "weights": labels.ne(-100)}, + loss_fn_input_layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + }, ) def _engine_loss_fn( diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index c615d3ff76..1af694da16 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -12,20 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Two-GPU Engine/AutoPipeline parity for packed THD Llama. +"""Engine/AutoPipeline parity for packed THD Llama. The test runs the same two-microbatch, four-document update through PP=2 from both raw THD metadata (``seq_lens``) and final THD metadata (``cu_seqlens``). Each path runs training, forward-only evaluation, then training again on the same pipeline. Evaluation must match a native eager Llama in summed loss and weight statistics without creating gradients; both surrounding training calls -must match eager loss plus every local-stage gradient. A final padded two-Datum -update verifies that Engine broadcasts the callback's per-Datum mappings to -both pipeline ranks in logical input order. +must match eager loss plus every local-stage gradient. Additional flat-Datum +forwards cover explicit per-token, per-Datum, and replicated loss layouts for +raw 2+2 and final-THD 3+1 document splits. A final padded two-Datum update +verifies that Engine broadcasts callback mappings to both pipeline ranks in +logical input order and that the same pipeline can return to training. Run with:: torchrun --standalone --nproc-per-node=2 run_packed_pp.py + +Set ``CP_SIZE=2`` and use four ranks to run the forward-only PP2 x CP2 +explicit-layout checks:: + + CP_SIZE=2 torchrun --standalone --nproc-per-node=4 run_packed_pp.py """ from __future__ import annotations @@ -38,9 +45,17 @@ import torch.nn.functional as F from transformers import LlamaConfig -from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.datasets.datum import ( + CollatedLossInputs, + Datum, + LossInputLayout, + collate_datums, +) from nemo_automodel.components.distributed.config import FSDP2Config -from nemo_automodel.components.distributed.context_parallel.utils import make_cp_batch_for_te +from nemo_automodel.components.distributed.context_parallel.utils import ( + attach_te_context_parallel, + make_cp_batch_for_te, +) from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.models.common import BackendConfig @@ -172,6 +187,103 @@ def _eager_reference( return eval_loss_sum.detach(), eval_weight_sum.detach(), loss.detach(), grads, raw_inputs, labels, weights +def _explicit_layout_datums(device: torch.device, lengths: list[int]) -> list[Datum]: + """Build flat Datums carrying all three explicit loss-input layouts.""" + datums = [] + token_start = 1 + for datum_index, length in enumerate(lengths): + input_ids = torch.arange(token_start, token_start + length, device=device) + datums.append( + Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs={ + "labels": (input_ids + 1) % VOCAB_SIZE, + "weights": torch.ones(length, dtype=torch.float32, device=device), + "advantages": input_ids.to(torch.float32) / 10, + "old_logprobs": -input_ids.to(torch.float32) / 20, + "sample_id": torch.tensor((datum_index + 1) * 11, device=device), + # Leading size equals the number of PP microbatches. It + # must remain a complete replicated vector in both. + "global_coefficients": torch.tensor([0.25, 0.75], device=device), + }, + loss_fn_input_layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + "advantages": LossInputLayout.PER_TOKEN, + "old_logprobs": LossInputLayout.PER_TOKEN, + "sample_id": LossInputLayout.PER_DATUM, + "global_coefficients": LossInputLayout.REPLICATED, + }, + ) + ) + token_start += length + return datums + + +def _raw_explicit_layout_collate(datums: list[Datum]): + return collate_datums(datums, packed=True) + + +def _final_explicit_layout_collate(datums: list[Datum]): + """Produce a model-ready flat THD stream while retaining layout metadata.""" + model_inputs, loss_inputs = collate_datums(datums, packed=True) + lengths = torch.tensor( + [datum.seq_len for datum in datums], dtype=torch.int32, device=model_inputs["input_ids"].device + ) + final_model_inputs = { + "input_ids": model_inputs["input_ids"].reshape(-1), + "position_ids": model_inputs["position_ids"].reshape(-1), + "cu_seqlens": F.pad(lengths.cumsum(0), (1, 0)).to(torch.int32), + "max_seqlen": lengths.max(), + "qkv_format": "thd", + } + final_loss_inputs = { + name: value.reshape(-1) if loss_inputs.layouts[name] is LossInputLayout.PER_TOKEN else value + for name, value in loss_inputs.items() + } + return final_model_inputs, CollatedLossInputs( + final_loss_inputs, + layouts=loss_inputs.layouts, + item_to_datum=loss_inputs.item_to_datum, + ) + + +def _explicit_layout_losses(output, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Use both RL-style token fields so their routing affects the numerator.""" + return _token_losses(output, loss_inputs) + 0.01 * loss_inputs["advantages"] + 0.02 * loss_inputs["old_logprobs"] + + +def _explicit_layout_eager_reference( + device: torch.device, + datums: list[Datum], + collate_fn, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute full-stream numerator and denominator for a layout-aware batch.""" + model_inputs, loss_inputs = collate_fn(datums) + prepared = make_cp_batch_for_te( + None, + {**_clone_mapping(model_inputs), "labels": loss_inputs["labels"].clone()}, + ) + prepared = { + name: value.to(device) if isinstance(value, torch.Tensor) else value for name, value in prepared.items() + } + prepared_labels = prepared.pop("labels") + eager_loss_inputs = { + "labels": prepared_labels, + "weights": loss_inputs["weights"].reshape(-1).to(device), + "advantages": loss_inputs["advantages"].reshape(-1).to(device), + "old_logprobs": loss_inputs["old_logprobs"].reshape(-1).to(device), + } + + model = _build_model(device).eval() + with torch.no_grad(): + losses = _explicit_layout_losses(model(**prepared), eager_loss_inputs) + loss_sum = (losses * eager_loss_inputs["weights"]).sum() + weight_sum = eager_loss_inputs["weights"].sum() + del model + return loss_sum.detach(), weight_sum.detach() + + def _build_pipeline( device: torch.device, mesh_context: MeshContext, @@ -189,6 +301,12 @@ def _build_pipeline( pp_seq_len=SEQ_LEN, ).build(model, loss_fn=_token_losses) del model + if mesh_context.cp_size > 1: + cp_mesh = mesh_context.device_mesh["cp"] + tp_mesh = mesh_context.device_mesh["tp"] + configured = sum(attach_te_context_parallel(part, cp_mesh, tp_mesh) for part in pipeline.parts) + if configured == 0: + raise AssertionError("PP2 x CP2 functional configured no Transformer Engine attention modules") return pipeline @@ -320,6 +438,95 @@ def _run_thd_layout( return pipeline +def _run_explicit_loss_layout_forward( + pipeline: AutoPipeline, + layout: str, + device: torch.device, + mesh_context: MeshContext, +) -> None: + """Validate semantic loss routing for one real two-stage packed pipeline.""" + if layout == "raw": + lengths = [4, 4, 4, 4] if mesh_context.cp_size > 1 else [2, 2, 2, 2] + collate_fn = _raw_explicit_layout_collate + expected_microbatch_ids = {(11, 22), (33, 44)} + elif layout == "final": + # Keep the uneven 3+1 Datum split while making every sequence length + # divisible by TE's 2*CP head/tail partition count under CP2. + lengths = [4, 4, 4, 12] if mesh_context.cp_size > 1 else [1, 1, 2, 4] + collate_fn = _final_explicit_layout_collate + expected_microbatch_ids = {(11, 22, 33), (44,)} + else: + raise ValueError(f"unknown explicit loss layout: {layout}") + + for part in pipeline.parts: + part.zero_grad(set_to_none=True) + + datums = _explicit_layout_datums(device, lengths) + reference_loss_sum, reference_weight_sum = _explicit_layout_eager_reference(device, datums, collate_fn) + expected_tokens_by_ids = {} + for ids in expected_microbatch_ids: + indices = [sample_id // 11 - 1 for sample_id in ids] + expected_tokens_by_ids[ids] = { + "advantages": torch.cat([datums[index].loss_fn_inputs["advantages"] for index in indices]), + "old_logprobs": torch.cat([datums[index].loss_fn_inputs["old_logprobs"] for index in indices]), + } + + def loss_with_outputs(output, loss_inputs): + sample_ids = tuple(int(value) for value in loss_inputs["sample_id"].tolist()) + if sample_ids not in expected_microbatch_ids: + raise AssertionError(f"PP2 {layout} routed unexpected sample IDs {sample_ids}") + expected_tokens = expected_tokens_by_ids[sample_ids] + valid_cu_seqlens = loss_inputs["cu_seqlens"].reshape(-1) + valid_cu_seqlens = valid_cu_seqlens[valid_cu_seqlens >= 0] + assert valid_cu_seqlens.numel() - 1 == len(sample_ids) + if mesh_context.cp_size > 1: + import transformer_engine_torch as tex + + cp_mesh = mesh_context.device_mesh["cp"] + local_indices = tex.thd_get_partitioned_indices( + valid_cu_seqlens.to(torch.int32), + int(valid_cu_seqlens[-1].item()), + cp_mesh.size(), + cp_mesh.get_local_rank(), + ).to(device=device, dtype=torch.long) + expected_tokens = {name: value.index_select(0, local_indices) for name, value in expected_tokens.items()} + torch.testing.assert_close(loss_inputs["advantages"].reshape(-1), expected_tokens["advantages"]) + torch.testing.assert_close(loss_inputs["old_logprobs"].reshape(-1), expected_tokens["old_logprobs"]) + torch.testing.assert_close( + loss_inputs["global_coefficients"], + torch.tensor([0.25, 0.75], device=device), + ) + assert loss_inputs["global_coefficients"].shape == (2,) + losses = _explicit_layout_losses(output, loss_inputs) + return losses, [{"sample_id": value} for value in loss_inputs["sample_id"]] + + result = Engine( + pipeline, + device=device, + mesh_context=mesh_context, + microbatch_size=4, + collate_fn=collate_fn, + ).forward(datums, loss_with_outputs) + + torch.testing.assert_close(result.loss_sum.float(), reference_loss_sum.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(result.weight_sum.float(), reference_weight_sum.float(), atol=0, rtol=0) + output_ids = torch.stack([item["sample_id"] for item in result.loss_fn_outputs]).to(torch.long) + expected_ids = torch.tensor([11, 22, 33, 44], device=device) + torch.testing.assert_close(output_ids, expected_ids) + gathered = [torch.empty_like(output_ids) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, output_ids) + assert all(torch.equal(ids, expected_ids) for ids in gathered) + if any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): + raise AssertionError(f"PP2 {layout} explicit-layout forward unexpectedly created gradients") + if dist.get_rank() == 0: + datum_counts = "2+2" if layout == "raw" else "3+1" + print( + f"PP2 x CP{mesh_context.cp_size} {layout} explicit " + f"PER_TOKEN/PER_DATUM/REPLICATED routing passed ({datum_counts} Datums; " + f"loss_sum={result.loss_sum.item():.6f}, weight_sum={result.weight_sum.item():.1f})" + ) + + def _run_padded_output_broadcast(pipeline: AutoPipeline, device: torch.device, mesh_context: MeshContext) -> None: for part in pipeline.parts: part.zero_grad(set_to_none=True) @@ -369,15 +576,30 @@ def main() -> None: local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) - if dist.get_world_size() != 2: - raise ValueError("packed PP functional requires exactly two ranks") + cp_size = int(os.environ.get("CP_SIZE", "1")) + if cp_size not in {1, 2}: + raise ValueError(f"packed PP functional supports CP_SIZE 1 or 2, got {cp_size}") + expected_world_size = 2 * cp_size + if dist.get_world_size() != expected_world_size: + raise ValueError( + f"packed PP functional with CP_SIZE={cp_size} requires {expected_world_size} ranks, " + f"got {dist.get_world_size()}" + ) mesh_context = MeshContext.build( FSDP2Config(), - ParallelismSizes(dp_size=1, pp_size=2), + ParallelismSizes(dp_size=1, pp_size=2, cp_size=cp_size), world_size=dist.get_world_size(), ) try: + if cp_size > 1: + pipeline = _build_pipeline(device, mesh_context) + _run_explicit_loss_layout_forward(pipeline, "raw", device, mesh_context) + dist.barrier() + _run_explicit_loss_layout_forward(pipeline, "final", device, mesh_context) + dist.barrier() + return + ( reference_eval_loss_sum, reference_eval_weight_sum, @@ -413,6 +635,10 @@ def main() -> None: weights, ) dist.barrier() + _run_explicit_loss_layout_forward(final_pipeline, "raw", device, mesh_context) + dist.barrier() + _run_explicit_loss_layout_forward(final_pipeline, "final", device, mesh_context) + dist.barrier() _run_padded_output_broadcast(final_pipeline, device, mesh_context) dist.barrier() finally: diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index 34eb78c06f..b0a58bc989 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -12,12 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pickle +from copy import copy, deepcopy + import pytest import torch from nemo_automodel.components.datasets.datum import ( CROSS_ENTROPY_IGNORE_IDX, + CollatedLossInputs, Datum, + LossInputLayout, collate_datums, ) @@ -77,6 +82,30 @@ def test_datum_accepts_model_specific_inputs(): assert datum.seq_len == 2 +def test_datum_accepts_optional_loss_input_layouts(): + datum = Datum( + input_ids=torch.tensor([1, 2]), + loss_fn_inputs={"weights": torch.ones(2)}, + loss_fn_input_layouts={"weights": LossInputLayout.PER_TOKEN}, + ) + assert datum.loss_fn_input_layouts == {"weights": LossInputLayout.PER_TOKEN} + + +def test_datum_rejects_invalid_loss_input_layouts(): + with pytest.raises(ValueError, match="unknown loss inputs"): + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"weights": torch.ones(1)}, + loss_fn_input_layouts={"missing": LossInputLayout.PER_TOKEN}, + ) + with pytest.raises(TypeError, match="must be a LossInputLayout"): + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"weights": torch.ones(1)}, + loss_fn_input_layouts={"weights": "per_token"}, # type: ignore[dict-item] + ) + + def test_to_features_applies_masking_convention(): feats = _toy_datums()[0].to_features() assert feats["input_ids"] == [10, 11, 12] @@ -101,6 +130,13 @@ def test_to_features_native_python_ints(): def test_collate_padded_uses_default_collater_schema(): batch, loss_inputs = collate_datums(_toy_datums()) + assert isinstance(loss_inputs, CollatedLossInputs) + assert loss_inputs.layouts == { + "advantages": LossInputLayout.PER_TOKEN, + "target_tokens": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } + assert loss_inputs.item_to_datum == (0, 1) assert batch["input_ids"].shape == (2, 3) assert batch["input_ids"][1].tolist() == [20, 21, 0] # right-pad assert "labels" not in batch @@ -148,6 +184,151 @@ def test_collate_packed_per_sample_side_input_is_one_per_datum(): batch, loss_inputs = collate_datums(datums, packed=True) assert batch["input_ids"].shape == (1, 5) assert loss_inputs["advantages"].tolist() == pytest.approx([0.5, 0.9]) + assert loss_inputs.layouts == {"advantages": LossInputLayout.PER_DATUM} + # Logical THD items are the two valid sequences, not the one physical row. + assert loss_inputs.item_to_datum == (0, 1) + + +def test_collated_loss_inputs_copy_preserves_side_channel_and_dict_compatibility(): + result = collate_datums(_toy_datums()) + loss_inputs = result[1] + + assert isinstance(result, tuple) + for copied in ( + loss_inputs.copy(), + copy(loss_inputs), + deepcopy(loss_inputs), + pickle.loads(pickle.dumps(loss_inputs)), # noqa: S301 - trusted in-process round trip + ): + assert isinstance(copied, CollatedLossInputs) + assert copied.keys() == loss_inputs.keys() + assert all(torch.equal(copied[key], loss_inputs[key]) for key in loss_inputs) + assert copied.layouts == loss_inputs.layouts + assert copied.item_to_datum == loss_inputs.item_to_datum + + +def test_collated_loss_inputs_requires_complete_read_only_layouts(): + with pytest.raises(ValueError, match="exactly"): + CollatedLossInputs( + {"weights": torch.ones(2)}, + layouts={}, + item_to_datum=(0,), + ) + + loss_inputs = CollatedLossInputs( + {"weights": torch.ones(2)}, + layouts={"weights": LossInputLayout.PER_TOKEN}, + item_to_datum=(index for index in [0]), # type: ignore[arg-type] + ) + assert loss_inputs.item_to_datum == (0,) + with pytest.raises(TypeError): + loss_inputs.layouts["weights"] = LossInputLayout.REPLICATED # type: ignore[index] + + +def test_collate_explicit_per_datum_overrides_single_token_shape_inference(): + datum = Datum( + input_ids=torch.tensor([7]), + loss_fn_inputs={"advantage": torch.tensor([0.5])}, + loss_fn_input_layouts={"advantage": LossInputLayout.PER_DATUM}, + ) + + _, loss_inputs = collate_datums([datum]) + + assert loss_inputs["advantage"].shape == (1,) + assert loss_inputs.layouts == {"advantage": LossInputLayout.PER_DATUM} + + +@pytest.mark.parametrize( + ("layout", "value", "message"), + [ + (LossInputLayout.PER_TOKEN, torch.tensor(1.0), "PER_TOKEN"), + (LossInputLayout.PER_DATUM, torch.ones(2), "PER_DATUM"), + ], +) +def test_collate_validates_explicit_loss_input_layout(layout, value, message): + datum = Datum( + input_ids=torch.tensor([1, 2]), + loss_fn_inputs={"field": value}, + loss_fn_input_layouts={"field": layout}, + ) + + with pytest.raises(ValueError, match=message): + collate_datums([datum]) + + +def test_collate_rejects_conflicting_explicit_layouts(): + datums = [ + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"field": torch.tensor([0.5])}, + loss_fn_input_layouts={"field": LossInputLayout.PER_TOKEN}, + ), + Datum( + input_ids=torch.tensor([2]), + loss_fn_inputs={"field": torch.tensor([0.9])}, + loss_fn_input_layouts={"field": LossInputLayout.PER_DATUM}, + ), + ] + + with pytest.raises(ValueError, match="same explicit layout"): + collate_datums(datums) + + +def test_collate_rejects_partially_declared_layouts(): + datums = [ + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"field": torch.tensor([0.5])}, + loss_fn_input_layouts={"field": LossInputLayout.PER_TOKEN}, + ), + Datum( + input_ids=torch.tensor([2]), + loss_fn_inputs={"field": torch.tensor([0.9])}, + ), + ] + + with pytest.raises(ValueError, match="every Datum must declare"): + collate_datums(datums) + + +def test_collate_replicated_input_keeps_one_identical_value(): + shared = torch.tensor([0.1, 0.2]) + datums = [ + Datum( + input_ids=torch.tensor([1, 2]), + loss_fn_inputs={"coefficients": shared}, + loss_fn_input_layouts={"coefficients": LossInputLayout.REPLICATED}, + ), + Datum( + input_ids=torch.tensor([3]), + loss_fn_inputs={"coefficients": shared.clone()}, + loss_fn_input_layouts={"coefficients": LossInputLayout.REPLICATED}, + ), + ] + + _, loss_inputs = collate_datums(datums, packed=True) + + assert loss_inputs["coefficients"] is shared + assert loss_inputs.layouts == {"coefficients": LossInputLayout.REPLICATED} + assert loss_inputs.item_to_datum == (0, 1) + + +def test_collate_replicated_input_requires_equal_values(): + datums = [ + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"coefficient": torch.tensor(0.1)}, + loss_fn_input_layouts={"coefficient": LossInputLayout.REPLICATED}, + ), + Datum( + input_ids=torch.tensor([2]), + loss_fn_inputs={"coefficient": torch.tensor(0.2)}, + loss_fn_input_layouts={"coefficient": LossInputLayout.REPLICATED}, + ), + ] + + with pytest.raises(ValueError, match="same value"): + collate_datums(datums) def test_collate_carries_per_token_float_side_inputs(): diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index c411e7fb5a..4cc70d51d5 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -22,6 +22,7 @@ import torch.nn as nn from nemo_automodel.components.config.loader import ConfigNode +from nemo_automodel.components.datasets.datum import LossInputLayout from nemo_automodel.components.datasets.vlm.pp_media import ( VLM_PP_MEDIA_KEY, chunk_step3_media, @@ -508,6 +509,10 @@ def test_make_engine_datum_filters_raw_media_off_first_pipeline_stage(): assert "pixel_values" not in datum.model_inputs assert "image_grid_thw" not in datum.model_inputs assert VLM_PP_MEDIA_KEY not in datum.model_inputs + assert datum.loss_fn_input_layouts == { + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } recipe.pp.info.has_first_stage = True first_stage_datum = recipe._make_engine_datum(batch) @@ -2116,6 +2121,10 @@ def test_vlm_validation_uses_engine_forward_and_aggregates_uneven_batches(monkey assert datums[0].model_inputs["pixel_values"] is batch["pixel_values"] assert datums[0].loss_fn_inputs["labels"] is batch["labels"] torch.testing.assert_close(datums[0].loss_fn_inputs["weights"], batch["labels"].ne(-100)) + assert datums[0].loss_fn_input_layouts == { + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } assert loss_fn == recipe._engine_validation_loss_fn assert allreduce.call_count == 2 assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index f77fdca67a..eb2cb14dd5 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -25,6 +25,7 @@ import torch.nn as nn from nemo_automodel.components.config.loader import ConfigNode +from nemo_automodel.components.datasets.datum import LossInputLayout # Skip decorator for tests that require CUDA requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -1616,6 +1617,10 @@ def test_run_validation_epoch_uses_engine_forward_for_pp_complete_results(monkey assert datums[0].model_inputs["input_ids"] is batch["input_ids"] assert datums[0].loss_fn_inputs["labels"] is batch["labels"] torch.testing.assert_close(datums[0].loss_fn_inputs["weights"], batch["labels"].ne(-100)) + assert datums[0].loss_fn_input_layouts == { + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } assert loss_fn == recipe._engine_validation_loss_fn assert allreduce.call_count == 2 assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 2b2d07b2e1..af776bfaac 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -27,9 +27,16 @@ from torch import nn import nemo_automodel.engine as engine_module +from nemo_automodel import CollatedLossInputs as PublicCollatedLossInputs from nemo_automodel import Datum as PublicDatum from nemo_automodel import Engine as PublicEngine -from nemo_automodel.components.datasets.datum import Datum, collate_datums +from nemo_automodel import LossInputLayout as PublicLossInputLayout +from nemo_automodel.components.datasets.datum import ( + CollatedLossInputs, + Datum, + LossInputLayout, + collate_datums, +) from nemo_automodel.components.datasets.vlm.pp_media import VLM_PP_MEDIA_KEY, stage_vlm_media_for_pp from nemo_automodel.components.distributed.config import MegatronFSDPConfig from nemo_automodel.components.distributed.context_parallel.sharder import ( @@ -193,6 +200,8 @@ def _identity_loss(output, _loss_inputs): def test_engine_and_datum_are_lazy_top_level_exports(): assert PublicEngine is Engine assert PublicDatum is Datum + assert PublicLossInputLayout is LossInputLayout + assert PublicCollatedLossInputs is CollatedLossInputs def test_forward_runs_eval_without_grad_lifecycle_and_returns_local_statistics(monkeypatch): @@ -404,6 +413,63 @@ def loss_fn(output, inputs): assert not cp_context_active +def test_context_parallel_rejects_legacy_non_token_weights(): + class CPModel(ScaleModel): + def prepare_model_inputs_for_cp(self, batch, *, num_chunks): + assert num_chunks == 1 + return { + "cp_sharder": ContextParallelSharder( + shard_batch=lambda *args, **kwargs: shard_batch_contiguous( + *args, + pad_multiple=1, + **kwargs, + ), + local_token_global_indices=contiguous_local_indices, + ) + } + + model = CPModel() + mesh_context = SimpleNamespace(pp_size=1, cp_size=2, device_mesh=_CPMesh(size=2, rank=0)) + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={ + "labels": torch.tensor([[2, 3, 4, -100]]), + "weights": torch.tensor([1.0]), + }, + ) + + with pytest.raises(ValueError, match="context-parallel loss weights must match"): + Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + ).forward([datum], _identity_loss) + assert model.forward_calls == 0 + + +def test_packed_thd_rejects_legacy_non_token_weights_without_pp_splitting(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "position_ids": torch.tensor([[0, 1, 2, 3]]), + "seq_lens": torch.tensor([[4]], dtype=torch.int32), + "seq_lens_padded": torch.tensor([[4]], dtype=torch.int32), + "qkv_format": "thd", + }, + loss_fn_inputs={ + "labels": torch.tensor([[2, 3, 4, -100]]), + "weights": torch.tensor([1.0]), + }, + ) + + with pytest.raises(ValueError, match="packed THD.*token-aligned loss weights"): + Engine(model, device="cpu", collate_fn=collate_prebatched).forward([datum], _identity_loss) + assert model.forward_calls == 0 + + @pytest.mark.parametrize("execution", ["forward", "forward_backward"]) def test_context_parallel_prepares_mtp_futures_before_sharding(execution): class MTPModel(ScaleModel): @@ -1003,62 +1069,236 @@ def loss_with_outputs(output, _loss_inputs): assert [item["input_ids"].shape for item in pipeline.prepared_inputs[0]] == [(1, 4), (1, 4)] -def test_pipeline_default_packed_collater_rejects_ambiguous_per_datum_loss_fields(): +def _packed_layout_datums(lengths: list[int]) -> list[Datum]: + """Build flat Datums whose three loss layouts are easy to distinguish.""" + datums = [] + token_start = 1 + for datum_index, length in enumerate(lengths): + input_ids = torch.arange(token_start, token_start + length) + datums.append( + Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs={ + "weights": torch.ones(length), + "advantages": input_ids.to(torch.float32) * 10, + "old_logprobs": -input_ids.to(torch.float32), + "sample_id": torch.tensor((datum_index + 1) * 11), + # Its leading extent deliberately collides with PP=2. A + # shape-based splitter would silently turn this into one + # value per microbatch instead of replicating it intact. + "global_coefficients": torch.tensor([701.0, 709.0]), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "advantages": LossInputLayout.PER_TOKEN, + "old_logprobs": LossInputLayout.PER_TOKEN, + "sample_id": LossInputLayout.PER_DATUM, + "global_coefficients": LossInputLayout.REPLICATED, + }, + ) + ) + token_start += length + return datums + + +def _final_thd_layout_collate(datums: list[Datum]): + """Convert canonical packed output to final THD without dropping layout metadata.""" + model_inputs, loss_inputs = collate_datums(datums, packed=True) + lengths = torch.tensor([datum.seq_len for datum in datums], dtype=torch.int32) + final_model_inputs = { + "input_ids": model_inputs["input_ids"].reshape(-1), + "position_ids": model_inputs["position_ids"].reshape(-1), + "cu_seqlens": F.pad(lengths.cumsum(0), (1, 0)).to(torch.int32), + "max_seqlen": lengths.max(), + "qkv_format": "thd", + } + final_loss_inputs = { + name: value.reshape(-1) if loss_inputs.layouts[name] is LossInputLayout.PER_TOKEN else value + for name, value in loss_inputs.items() + } + return final_model_inputs, CollatedLossInputs( + final_loss_inputs, + layouts=loss_inputs.layouts, + item_to_datum=loss_inputs.item_to_datum, + ) + + +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +def test_pipeline_raw_thd_routes_explicit_loss_layouts_and_outputs_in_datum_order(execution): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = _packed_layout_datums([2, 2, 2, 2]) + seen = [] + + def loss_with_outputs(output, loss_inputs): + torch.testing.assert_close(loss_inputs["global_coefficients"], torch.tensor([701.0, 709.0])) + assert loss_inputs["global_coefficients"].shape == (2,) + torch.testing.assert_close(loss_inputs["advantages"], output * 10) + torch.testing.assert_close(loss_inputs["old_logprobs"], -output) + sample_ids = loss_inputs["sample_id"].clone() + seen.append(sample_ids.tolist()) + return output, [{"sample_id": sample_id} for sample_id in sample_ids] + + engine = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=partial(collate_datums, packed=True), + ) + result = getattr(engine, execution)(datums, loss_with_outputs) + + assert seen == [[33, 44], [11, 22]] + if execution == "forward": + assert result.loss_sum.item() == pytest.approx(36.0) + assert result.weight_sum.item() == pytest.approx(8.0) + outputs = result.loss_fn_outputs + assert pipeline.eval_calls == 1 + assert pipeline.step_calls == 0 + assert pipeline.backward_calls == 0 + assert model.weight.grad is None + else: + loss, outputs = result + assert loss.item() == pytest.approx(4.5) + assert pipeline.eval_calls == 0 + assert pipeline.step_calls == 1 + assert pipeline.backward_calls == 2 + assert model.weight.grad.item() == pytest.approx(4.5) + assert [item["sample_id"].item() for item in outputs] == [11, 22, 33, 44] + + +def test_pipeline_packed_legacy_collater_metadata_stripping_fails_closed(): model = ScaleModel() model.backend = SimpleNamespace(attn="te") pipeline = _FakeAutoPipeline(model, num_microbatches=2) - datums = [ - Datum( - model_inputs={"input_ids": torch.tensor([index, index + 1])}, - loss_fn_inputs={"weights": torch.ones(2), "reward": torch.tensor(float(index))}, - ) - for index in (1, 3, 5, 7) - ] + datums = _packed_layout_datums([2, 2, 2, 2]) - with pytest.raises(NotImplementedError, match="per-Datum loss fields.*reward"): + def strip_layout_metadata(items): + model_inputs, loss_inputs = collate_datums(items, packed=True) + return model_inputs, dict(loss_inputs) + + with pytest.raises(NotImplementedError, match="item_to_datum metadata"): Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), microbatch_size=4, - collate_fn=partial(collate_datums, packed=True), - ).forward_backward(datums, _identity_loss) + collate_fn=strip_layout_metadata, + ).forward(datums, _identity_loss) + + assert pipeline.eval_calls == 0 -def test_pipeline_outputs_allow_uneven_datum_counts_at_final_thd_boundaries(): +def test_pipeline_final_thd_routes_three_plus_one_datums_and_restores_output_order(): model = ScaleModel() model.backend = SimpleNamespace(attn="te") pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) - datums = [_datum([1]), _datum([2]), _datum([3]), _datum([4, 5, 6])] + datums = _packed_layout_datums([1, 1, 2, 4]) + seen = [] - def final_thd_collate(items): - lengths = [datum.seq_len for datum in items] - tokens = torch.cat([datum.input_ids for datum in items]) - return ( - { - "input_ids": tokens, - "position_ids": torch.cat([torch.arange(length) for length in lengths]), - "cu_seqlens": torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32), - "max_seqlen": torch.tensor(max(lengths), dtype=torch.int32), - "qkv_format": "thd", - }, - {"weights": torch.cat([datum.loss_fn_inputs["weights"] for datum in items])}, - ) + def loss_with_outputs(output, loss_inputs): + torch.testing.assert_close(loss_inputs["global_coefficients"], torch.tensor([701.0, 709.0])) + assert loss_inputs["global_coefficients"].shape == (2,) + torch.testing.assert_close(loss_inputs["advantages"], output * 10) + torch.testing.assert_close(loss_inputs["old_logprobs"], -output) + sample_ids = loss_inputs["sample_id"].clone() + seen.append(sample_ids.tolist()) + return output, [{"sample_id": sample_id} for sample_id in sample_ids] - def loss_with_outputs(output, _loss_inputs): - first_token = int(output.reshape(-1)[0].item()) - ids = [1, 2, 3] if first_token == 1 else [4] - return output, [{"datum_id": torch.tensor(datum_id)} for datum_id in ids] - - _, outputs = Engine( + result = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), microbatch_size=4, - collate_fn=final_thd_collate, - ).forward_backward(datums, loss_with_outputs) + collate_fn=_final_thd_layout_collate, + ).forward(datums, loss_with_outputs) + + assert result.loss_sum.item() == pytest.approx(36.0) + assert result.weight_sum.item() == pytest.approx(8.0) + assert seen == [[44], [11, 22, 33]] + assert [item["sample_id"].item() for item in result.loss_fn_outputs] == [11, 22, 33, 44] + + +def test_pipeline_final_thd_rejects_wrong_outputs_even_when_window_total_matches(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = _packed_layout_datums([1, 1, 2, 4]) + + with pytest.raises(ValueError, match="returned 2 outputs for microbatch 1, expected 1"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=_final_thd_layout_collate, + ).forward( + datums, + lambda output, _loss_inputs: (output, [{"wrong": 0}, {"wrong": 1}]), + ) + + +def test_engine_rejects_collater_item_reordering_before_execution(): + datums = _packed_layout_datums([2, 2]) + + def reordered_collate(items): + model_inputs, loss_inputs = collate_datums(items) + return model_inputs, CollatedLossInputs( + loss_inputs, + layouts=loss_inputs.layouts, + item_to_datum=(1, 0), + ) + + with pytest.raises(ValueError, match="preserve outer Datum order"): + Engine( + ScaleModel(), + device="cpu", + microbatch_size=2, + collate_fn=reordered_collate, + ).forward(datums, _identity_loss) + + +def test_engine_keeps_weights_as_a_per_token_contract(): + datum = Datum( + model_inputs={"input_ids": torch.tensor([1])}, + loss_fn_inputs={"weights": torch.tensor([1.0])}, + loss_fn_input_layouts={"weights": LossInputLayout.PER_DATUM}, + ) + + with pytest.raises(ValueError, match="weights.*PER_TOKEN"): + Engine(ScaleModel(), device="cpu").forward([datum], _identity_loss) - assert [item["datum_id"].item() for item in outputs] == [1, 2, 3, 4] + +def test_pipeline_single_microbatch_validates_thd_item_count(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=1) + datum = _datum([1, 2]) + + def two_sequence_collate(_items): + return ( + { + "input_ids": torch.tensor([1, 2]), + "position_ids": torch.tensor([0, 0]), + "cu_seqlens": torch.tensor([0, 1, 2], dtype=torch.int32), + "max_seqlen": torch.tensor(1, dtype=torch.int32), + "qkv_format": "thd", + }, + CollatedLossInputs( + {"weights": torch.ones(2)}, + layouts={"weights": LossInputLayout.PER_TOKEN}, + item_to_datum=(0,), + ), + ) + + with pytest.raises(ValueError, match="item_to_datum does not match"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=two_sequence_collate, + ).forward([datum], _identity_loss) def test_pipeline_output_sync_finds_last_stage_on_physical_rank_zero(monkeypatch): @@ -1128,6 +1368,39 @@ def test_pipeline_prebatched_outputs_require_one_inner_microbatch(): ) +def test_pipeline_prebatched_per_datum_field_without_item_mapping_is_rejected(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + datum = Datum( + model_inputs={ + "input_ids": torch.arange(1, 9), + "position_ids": torch.tensor([0, 1, 2, 3, 0, 1, 2, 3]), + "cu_seqlens": torch.tensor([0, 4, 8], dtype=torch.int32), + "max_seqlen": torch.tensor(4, dtype=torch.int32), + "qkv_format": "thd", + }, + loss_fn_inputs={ + "weights": torch.ones(8), + "sample_id": torch.tensor([17]), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "sample_id": LossInputLayout.PER_DATUM, + }, + ) + + with pytest.raises(NotImplementedError, match="item_to_datum metadata"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward([datum], _identity_loss) + + assert pipeline.eval_calls == 0 + + def test_pipeline_final_thd_embeddings_use_the_token_axis_for_sequence_length(): pipeline = _FakeAutoPipeline(ScaleModel(), num_microbatches=1) embeddings = torch.arange(12, dtype=torch.float32).reshape(4, 3) From bfe6c955904165b8ba04c18850a37044c35eeffb Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Thu, 20 Aug 2026 21:26:41 -0700 Subject: [PATCH 11/34] feat(engine): own optimizer updates Signed-off-by: HuiyingLi --- nemo_automodel/engine/__init__.py | 128 ++++++++++- nemo_automodel/recipes/llm/train_ft.py | 82 +++---- nemo_automodel/recipes/vlm/finetune.py | 73 ++---- .../moe/test_experts_ep_tp_grad_parity.py | 51 +++-- .../recipes/test_finetune_vlm_helpers.py | 72 ++++-- tests/unit_tests/recipes/test_train_ft.py | 78 +++++-- .../test_train_ft_partial_cuda_graphs.py | 2 +- tests/unit_tests/test_engine.py | 209 +++++++++++++++++- .../test_engine_recipe_integration.py | 11 +- 9 files changed, 523 insertions(+), 183 deletions(-) diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 3a16ae1ffc..28db6b4338 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -19,7 +19,7 @@ from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass -from typing import Any +from typing import Any, TypeVar import torch import torch.distributed as dist @@ -38,10 +38,13 @@ from nemo_automodel.components.distributed.utils import get_sync_ctx from nemo_automodel.components.models.common.mtp import MTPContextParallelInputs from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler +from nemo_automodel.components.optim.scheduler import OptimizerParamScheduler from nemo_automodel.components.training.utils import ( + get_expert_tp_replication_factor, prepare_after_first_microbatch, prepare_for_final_backward, prepare_for_grad_accumulation, + scale_grads_and_clip_grad_norm, ) from nemo_automodel.components.utils.model_utils import filter_forward_kwargs @@ -58,14 +61,23 @@ _LOSS_FIELD_PREFIX = "__engine_loss__" _LOSS_METADATA = ("cu_seqlens", "cu_seqlens_padded", "max_seqlen", "padding_mask") +_T = TypeVar("_T") -__all__ = ["Engine", "ForwardResult", "collate_prebatched"] +__all__ = ["Engine", "ForwardResult", "OptimStepResult", "collate_prebatched"] def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextManager[Any]: return nullcontext() +def _as_tuple(value: _T | Sequence[_T] | None) -> tuple[_T, ...]: + if value is None: + return () + if isinstance(value, Sequence): + return tuple(value) + return (value,) + + def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLossInputs | dict[str, torch.Tensor]]: """Return one already-collated Datum without changing its layout. @@ -131,6 +143,22 @@ class ForwardResult: loss_fn_outputs: list[dict[str, Any]] +@dataclass(frozen=True) +class OptimStepResult: + """Statistics from one completed optimizer step. + + Attributes: + grad_norm: Gradient norm reported before clipping. This is a scalar + tensor on the gradients' device, or ``0.0`` when clipping is + disabled by ``max_grad_norm=None``. + learning_rates: Learning rates of every optimizer parameter group after + the configured schedulers advance. + """ + + grad_norm: torch.Tensor | float + learning_rates: tuple[float, ...] + + class Engine: """Run model forward or forward/backward over Datum windows. @@ -138,10 +166,13 @@ class Engine: passed here. The Engine owns batching and model-parallel execution. :meth:`forward` performs evaluation without gradients; :meth:`forward_backward` additionally owns global weight normalization, - gradient-accumulation synchronization, and backward. The Engine - deliberately does not zero, clip, finalize expert gradients, or step them; - callers choose the optimizer boundary and retain the repository's existing - distributed gradient-finalization path. + gradient-accumulation synchronization, and backward. When optimizers are + provided, :meth:`optim_step` owns distributed gradient finalization, + clipping, parameter updates, gradient clearing, model post-step hooks, and + LR-scheduler advancement. One :meth:`forward_backward` call represents the + complete optimizer accumulation window whose gradients :meth:`optim_step` + consumes. Dynamic loss scaling and overflow-skipped updates are not part of + this contract. Args: model: An already configured and distributed model, or a built @@ -167,6 +198,13 @@ class Engine: pipeline schedule. Recipes use it for FP8 and model input staging. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the final microbatch. + optimizers: Already-built optimizer or optimizers for these model parts. + The Engine retains the same objects; it does not build or copy them. + lr_schedulers: Already-built optimizer parameter scheduler or schedulers. + They advance once after a completed optimizer update. + max_grad_norm: Maximum gradient norm. ``None`` preserves gradients + without clipping while still running distributed expert-gradient + finalization. Note: Context-parallel input layout and transport are delegated to @@ -200,6 +238,9 @@ def __init__( mtp_ignore_index: int = -100, context_fn: Callable[[dict[str, Any]], AbstractContextManager[Any]] = _nullcontext_for_batch, defer_fsdp_grad_sync: bool = True, + optimizers: torch.optim.Optimizer | Sequence[torch.optim.Optimizer] | None = None, + lr_schedulers: OptimizerParamScheduler | Sequence[OptimizerParamScheduler] | None = None, + max_grad_norm: float | None = 1.0, ) -> None: if isinstance(microbatch_size, bool) or not isinstance(microbatch_size, int) or microbatch_size <= 0: raise ValueError(f"microbatch_size must be a positive integer, got {microbatch_size!r}") @@ -216,6 +257,9 @@ def __init__( self.mtp_ignore_index = mtp_ignore_index self.context_fn = context_fn self.defer_fsdp_grad_sync = defer_fsdp_grad_sync + self.optimizers = _as_tuple(optimizers) + self.lr_schedulers = _as_tuple(lr_schedulers) + self.max_grad_norm = max_grad_norm @torch.no_grad() def forward( @@ -475,6 +519,78 @@ def forward_backward( loss = (local_loss_sum / safe_denominator).detach() return loss, loss_fn_outputs + @torch.no_grad() + def optim_step( + self, + *, + before_optimizer_step: Callable[[], None] | None = None, + ) -> OptimStepResult: + """Finalize accumulated gradients and perform one optimizer update. + + Gradient normalization performed by :meth:`forward_backward` is not + repeated here. This method applies the repository's model-parallel + expert-gradient correction and global clipping once, then invokes an + optional mutation fence before any optimizer changes. Async + checkpointers use that fence to preserve ``finalize/clip -> wait -> + step`` overlap. + + Args: + before_optimizer_step: Optional callback invoked exactly once after + gradient finalization and clipping, but before the first + optimizer step. If it raises, parameters, optimizer state, + model post-step state, and schedulers remain untouched; the + finalized gradients remain available. + + Returns: + Gradient norm and post-scheduler learning rates for the completed + optimizer update. + + Raises: + RuntimeError: If this Engine was constructed without optimizers. + """ + if not self.optimizers: + raise RuntimeError("Engine.optim_step requires at least one optimizer") + if before_optimizer_step is not None and not callable(before_optimizer_step): + raise TypeError("before_optimizer_step must be callable or None") + + device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None + moe_mesh = self.mesh_context.moe_mesh if self.mesh_context is not None else None + dp_group, dp_size = self._dp_group_and_size() + _, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) + pp_enabled = self.pipeline is not None + grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=self.max_grad_norm, + model_parts=self.model_parts, + norm_type=2.0, + pp_enabled=pp_enabled, + device_mesh=device_mesh, + moe_mesh=moe_mesh, + ep_axis_name="ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None, + pp_axis_name="pp" if pp_enabled else None, + foreach=True, + num_label_tokens=None, + dp_group_size=grad_group_size, + expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, device_mesh), + ) + + if before_optimizer_step is not None: + before_optimizer_step() + + for optimizer in self.optimizers: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + for part in self.model_parts: + update_moe_gate_bias = getattr(part, "update_moe_gate_bias", None) + if callable(update_moe_gate_bias): + update_moe_gate_bias() + + for scheduler in self.lr_schedulers: + scheduler.step(1) + + learning_rates = tuple(float(group["lr"]) for optimizer in self.optimizers for group in optimizer.param_groups) + return OptimStepResult(grad_norm=grad_norm, learning_rates=learning_rates) + def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: if not isinstance(datums, Sequence) or isinstance(datums, (str, bytes)) or not datums: raise ValueError("Engine requires a non-empty flat sequence of Datum") diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index d0d5eb8190..6704a2e061 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -78,11 +78,7 @@ from nemo_automodel.components.quantization.fp8 import build_fp8_config from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG -from nemo_automodel.components.training.utils import ( - count_tail_padding, - get_expert_tp_replication_factor, - scale_grads_and_clip_grad_norm, -) +from nemo_automodel.components.training.utils import count_tail_padding from nemo_automodel.components.utils.compile_utils import ( build_compile_config, ) @@ -737,21 +733,6 @@ def setup(self): _, self.tokenizer = _build_tokenizer(self.cfg.model, self.cfg.dataset) if getattr(self.loss_fn, "reduction", None) != "sum": raise ValueError("Engine-backed finetuning requires a loss with reduction='sum'") - self.engine = Engine( - self.pp if self.pp_enabled else self.model_parts[0], - device=self.dist_env.device, - mesh_context=self.mesh_context, - microbatch_size=1, - collate_fn=collate_prebatched, - padding_token_id=(self.tokenizer.pad_token_id if self.tokenizer is not None else 0) or 0, - mtp_ignore_index=self.cfg.mtp.ignore_index, - context_fn=( - (lambda _model_inputs: self.te_fp8.maybe_te_autocast()) - if self.te_fp8 is not None - else (lambda _model_inputs: nullcontext()) - ), - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) attn_implementation = None if ( self.cfg.get("packed_sequence.packed_sequence_size", 0) > 0 @@ -823,6 +804,25 @@ def materialize_loader(config): else None ) + self.engine = Engine( + self.pp if self.pp_enabled else self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + optimizers=self.optimizer, + lr_schedulers=self.lr_scheduler, + max_grad_norm=self.max_grad_norm, + microbatch_size=1, + collate_fn=collate_prebatched, + padding_token_id=(self.tokenizer.pad_token_id if self.tokenizer is not None else 0) or 0, + mtp_ignore_index=self.cfg.mtp.ignore_index, + context_fn=( + (lambda _model_inputs: self.te_fp8.maybe_te_autocast()) + if self.te_fp8 is not None + else (lambda _model_inputs: nullcontext()) + ), + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), + ) + # Log model, parameter counts, norms, optimizer and scheduler self._log_model_and_optimizer_details(self.model_parts, self.optimizer, self.lr_scheduler) @@ -1016,7 +1016,7 @@ def run_train_validation_loop(self): for batches in self.step_scheduler: # If QAT delayed fake-quant is configured, enable after threshold self._enable_qat_if_delayed(self.step_scheduler.step) - train_log_data = self._run_train_optim_step(batches, self.max_grad_norm) + train_log_data = self._run_train_optim_step(batches) # Capture outside the microbatch loop and only after the # eager optimizer step has completed. This leaves no # pending checkpoint recomputation or GA backward work. @@ -1218,14 +1218,12 @@ def _engine_validation_loss_fn( is_train=False, ) - def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: float | None = None) -> MetricsSample: + def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: """Execute a single training step. Args: batches: Worker-collated optimizer window. Padded token tensors use shape [batch, sequence]; packed tensors use their THD token layout. - max_grad_norm: Gradient clipping norm. Optional, if None will not clip gradients. - Returns: Metrics for the completed optimizer step. """ @@ -1246,37 +1244,7 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: fl [self._make_engine_datum(batch) for batch in batches], self._engine_loss_fn, ) - - grad_norm = scale_grads_and_clip_grad_norm( - max_grad_norm, - self.model_parts, - norm_type=2.0, - pp_enabled=self.pp_enabled, - device_mesh=self.device_mesh, - moe_mesh=self.moe_mesh, - ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, - pp_axis_name="pp" if self.pp_enabled else None, - foreach=True, - num_label_tokens=None, - dp_group_size=self._get_dp_group_size(include_cp=True), - expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, self.device_mesh), - ) - - # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api - # self.model_parts[0].finish_grad_sync() - - self.checkpointer.maybe_wait_for_staging() - for opt in self.optimizer: - opt.step() - opt.zero_grad() - - if hasattr(self.model_parts[0], "update_moe_gate_bias"): - for mp in self.model_parts: - mp.update_moe_gate_bias() - - if self.lr_scheduler is not None: - for scheduler in self.lr_scheduler: - scheduler.step(1) + step_result = self.engine.optim_step(before_optimizer_step=self.checkpointer.maybe_wait_for_staging) # Precompute FP8 scales fp8_config = self.cfg.get("fp8", None) @@ -1329,8 +1297,8 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: fl epoch=self.step_scheduler.epoch, metrics={ "loss": reporting_loss, - "grad_norm": grad_norm, - "lr": self.optimizer[0].param_groups[0]["lr"], + "grad_norm": step_result.grad_norm, + "lr": step_result.learning_rates[0], "mem": torch.cuda.max_memory_allocated() / 1024**3, "tps": tps, "tps_per_gpu": tps / self._get_cp_group_size() / max(self._get_dp_group_size(), 1), diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 27d77385bb..72e8f28375 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -73,11 +73,7 @@ from nemo_automodel.components.quantization.fp8 import build_fp8_config from nemo_automodel.components.training.model_output_utils import get_final_hidden_states from nemo_automodel.components.training.rng import ScopedRNG, StatefulRNG -from nemo_automodel.components.training.utils import ( - count_tail_padding, - get_expert_tp_replication_factor, - scale_grads_and_clip_grad_norm, -) +from nemo_automodel.components.training.utils import count_tail_padding from nemo_automodel.components.utils.compile_utils import build_compile_config from nemo_automodel.components.utils.model_utils import VLM_INPUT_KEYS, _supports_logits_to_keep from nemo_automodel.engine import Engine, collate_prebatched @@ -642,17 +638,6 @@ def setup(self): if getattr(self.loss_fn, "reduction", None) != "sum": raise ValueError("Engine-backed VLM finetuning requires a loss with reduction='sum'") padding_token_id = getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0 - self.engine = Engine( - self.pp if self.pp_enabled else self.model_parts[0], - device=self.dist_env.device, - mesh_context=self.mesh_context, - microbatch_size=1, - collate_fn=collate_prebatched, - padding_token_id=padding_token_id, - mtp_ignore_index=self.cfg.mtp.ignore_index, - context_fn=self._engine_context, - defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), - ) # Build validation dataloader if the config provides it self.val_dataloader = None @@ -688,6 +673,21 @@ def setup(self): else None ) + self.engine = Engine( + self.pp if self.pp_enabled else self.model_parts[0], + device=self.dist_env.device, + mesh_context=self.mesh_context, + microbatch_size=1, + collate_fn=collate_prebatched, + padding_token_id=padding_token_id, + mtp_ignore_index=self.cfg.mtp.ignore_index, + context_fn=self._engine_context, + defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True), + optimizers=self.optimizer, + lr_schedulers=self.lr_scheduler, + max_grad_norm=self.max_grad_norm, + ) + # Log model, parameter counts, norms, optimizer and scheduler self._log_model_and_optimizer_details(self.model_parts, self.optimizer, self.lr_scheduler) @@ -723,7 +723,7 @@ def run_train_validation_loop(self): for epoch in self.step_scheduler.epochs: self.step_scheduler.set_epoch(epoch) for batch_idx, batches in enumerate(self.step_scheduler): - log_data = self._run_train_optim_step(batches, self.max_grad_norm) + log_data = self._run_train_optim_step(batches) # log self.log_train_metrics(log_data) self._update_progress_bar(pbar, log_data.metrics) @@ -1067,15 +1067,13 @@ def _configure_pipeline_loss_fn(self): ) self.pp.info.schedule._loss_fn = self.pipeline_loss_fn - def _run_train_optim_step(self, batches: list[dict[str, Any]], max_grad_norm: float | None = None) -> MetricsSample: + def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: """Execute a single training step. Args: batches: Processor-collated optimizer window. Padded token tensors use shape [batch, sequence]; packed tensors use their THD token layout, and media tensors retain model-specific layouts. - max_grad_norm: Gradient clipping norm. Optional, if None will not clip gradients. - Returns: Metrics for the completed optimizer step. """ @@ -1112,37 +1110,10 @@ def engine_loss_fn( engine_loss_fn, ) - grad_norm = scale_grads_and_clip_grad_norm( - max_grad_norm=max_grad_norm, - model_parts=self.model_parts, - norm_type=2.0, - pp_enabled=self.pp_enabled, - device_mesh=self.device_mesh, - moe_mesh=self.moe_mesh, - ep_axis_name="ep" if self.moe_mesh is not None and "ep" in self.moe_mesh.mesh_dim_names else None, - pp_axis_name="pp" if self.pp_enabled else None, - foreach=True, - num_label_tokens=None, - dp_group_size=self._get_dp_group_size(include_cp=True), - expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, self.device_mesh), + step_result = self.engine.optim_step( + before_optimizer_step=self.checkpointer.maybe_wait_for_staging, ) - # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api - # self.model.finish_grad_sync() - - self.checkpointer.maybe_wait_for_staging() - for opt in self.optimizer: - opt.step() - opt.zero_grad(set_to_none=True) - - if hasattr(self.model_parts[0], "update_moe_gate_bias"): - for mp in self.model_parts: - mp.update_moe_gate_bias() - - if self.lr_scheduler is not None: - for scheduler in self.lr_scheduler: - scheduler.step(1) - # Precompute FP8 scales fp8_config = self.cfg.get("fp8", None) if ( @@ -1170,8 +1141,8 @@ def engine_loss_fn( epoch=self.step_scheduler.epoch, metrics={ "loss": reporting_loss, - "grad_norm": grad_norm, - "lr": self.optimizer[0].param_groups[0]["lr"], + "grad_norm": step_result.grad_norm, + "lr": step_result.learning_rates[0], "mem": torch.cuda.max_memory_allocated() / 1024**3, "tps": tps, "tps_per_gpu": tps / self._get_cp_group_size() / max(self._get_dp_group_size(), 1), diff --git a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py index 7dc5a3cb6b..482f87c37f 100644 --- a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py +++ b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py @@ -40,10 +40,6 @@ from nemo_automodel.components.distributed.mesh import MeshContext from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.experts import GroupedExperts -from nemo_automodel.components.training.utils import ( - get_expert_tp_replication_factor, - scale_grads_and_clip_grad_norm, -) from nemo_automodel.engine import Engine, collate_prebatched _TP_SIZE = 2 @@ -209,12 +205,18 @@ def loss_fn(output: torch.Tensor, loss_inputs: dict[str, torch.Tensor]) -> torch observed["output"] = output.detach() return output.square() - engine_loss, _ = Engine( - _ExpertModel(experts), + model = _ExpertModel(experts) + model._nemo_moe_tp_requires_replica_sync = True + optimizer = torch.optim.SGD(model.parameters(), lr=_LEARNING_RATE) + engine = Engine( + model, device="cpu", mesh_context=MeshContext.from_meshes(world_mesh, ep_mesh), collate_fn=collate_prebatched, - ).forward_backward([datum], loss_fn) + optimizers=optimizer, + max_grad_norm=1e6, + ) + engine_loss, _ = engine.forward_backward([datum], loss_fn) torch.testing.assert_close(observed["output"], y_ref, rtol=1e-4, atol=1e-5) torch.testing.assert_close(engine_loss, loss_ref.to(torch.float64), rtol=1e-5, atol=1e-7) @@ -235,27 +237,24 @@ def loss_fn(output: torch.Tensor, loss_inputs: dict[str, torch.Tensor]) -> torch experts.down_projs.grad.to_local(), _TP_SIZE * down_grad_ref_local, rtol=1e-4, atol=1e-5 ) - # The recipe-side scaling must remove exactly that factor. With no FSDP - # gradient averaging in this test, dp_group_size=1 and no ep_shard axis - # make the TP replication factor the only expert divisor. - replication_factor = get_expert_tp_replication_factor([experts], world_mesh) - assert replication_factor == _TP_SIZE - grad_norm = scale_grads_and_clip_grad_norm( - max_grad_norm=1e6, - model_parts=[experts], - moe_mesh=ep_mesh, - ep_axis_name="ep", - dp_group_size=1, - expert_tp_replication_factor=replication_factor, - ) - torch.testing.assert_close( - experts.gate_and_up_projs.grad.to_local(), gate_up_grad_ref_local, rtol=1e-4, atol=1e-5 - ) - torch.testing.assert_close(experts.down_projs.grad.to_local(), down_grad_ref_local, rtol=1e-4, atol=1e-5) - torch.testing.assert_close(grad_norm, reference_grad_norm, rtol=1e-5, atol=1e-7) + # Engine.optim_step must remove exactly that factor before any parameter + # mutation. Capture the corrected gradients at its staging-fence hook; + # the method then performs the SGD update and clears all gradients. + corrected_grads: dict[str, torch.Tensor] = {} + + def capture_corrected_grads() -> None: + corrected_grads["gate_up"] = experts.gate_and_up_projs.grad.to_local().detach().clone() + corrected_grads["down"] = experts.down_projs.grad.to_local().detach().clone() + + step_result = engine.optim_step(before_optimizer_step=capture_corrected_grads) + torch.testing.assert_close(corrected_grads["gate_up"], gate_up_grad_ref_local, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(corrected_grads["down"], down_grad_ref_local, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(step_result.grad_norm, reference_grad_norm, rtol=1e-5, atol=1e-7) + assert step_result.learning_rates == (_LEARNING_RATE,) + assert experts.gate_and_up_projs.grad is None + assert experts.down_projs.grad is None torch.optim.SGD(reference_experts.parameters(), lr=_LEARNING_RATE).step() - torch.optim.SGD(experts.parameters(), lr=_LEARNING_RATE).step() torch.testing.assert_close( experts.gate_and_up_projs.to_local(), reference_experts.gate_and_up_projs[start:end], diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 4cc70d51d5..7c8a9f3247 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -407,11 +407,12 @@ def _build_engine_recipe_for_optim_step(*, pp_enabled: bool = False): recipe._get_cp_group_size = lambda: 1 recipe.engine = MagicMock() recipe.engine.forward_backward.return_value = (torch.tensor(0.25), []) + recipe.engine.optim_step.return_value = SimpleNamespace(grad_norm=2.5, learning_rates=(0.01,)) return recipe @pytest.mark.cuda(False) -def test_run_train_step_passes_flat_prebatched_datums_to_engine(monkeypatch): +def test_run_train_step_passes_flat_prebatched_datums_to_engine(): recipe = _build_engine_recipe_for_optim_step(pp_enabled=True) batches = [ { @@ -423,10 +424,7 @@ def test_run_train_step_passes_flat_prebatched_datums_to_engine(monkeypatch): "input_ids": torch.tensor([[5, 6, 7, 8]]), }, ] - finalizer = MagicMock(return_value=2.5) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", finalizer) - - metrics = recipe._run_train_optim_step(batches, max_grad_norm=1.0) + metrics = recipe._run_train_optim_step(batches) datums, loss_fn = recipe.engine.forward_backward.call_args.args assert len(datums) == 2 @@ -436,14 +434,18 @@ def test_run_train_step_passes_flat_prebatched_datums_to_engine(monkeypatch): [[False, True, False, True]], ] assert callable(loss_fn) - assert finalizer.call_args.kwargs["num_label_tokens"] is None + recipe.engine.optim_step.assert_called_once_with( + before_optimizer_step=recipe.checkpointer.maybe_wait_for_staging, + ) assert metrics.metrics["loss"] == pytest.approx(0.25) - assert recipe.optimizer[0].step_called - assert recipe.optimizer[0].zero_grad_called + assert metrics.metrics["grad_norm"] == pytest.approx(2.5) + assert metrics.metrics["lr"] == pytest.approx(0.01) + assert not recipe.optimizer[0].step_called + assert not recipe.optimizer[0].zero_grad_called @pytest.mark.cuda(False) -def test_train_step_logs_joint_drafter_only_on_first_engine_loss_call(monkeypatch): +def test_train_step_logs_joint_drafter_only_on_first_engine_loss_call(): recipe = _build_engine_recipe_for_optim_step() recipe.step_scheduler.is_remote_logging_step = True batches = [ @@ -458,10 +460,6 @@ def forward_backward(datums, loss_fn): return torch.tensor(0.25), [] recipe.engine.forward_backward.side_effect = forward_backward - monkeypatch.setattr( - "nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", - MagicMock(return_value=0.0), - ) recipe._run_train_optim_step(batches) @@ -474,19 +472,45 @@ def forward_backward(datums, loss_fn): @pytest.mark.cuda(False) -def test_run_train_step_uses_engine_for_empty_supervision(monkeypatch): +def test_run_train_step_uses_engine_for_empty_supervision(): recipe = _build_engine_recipe_for_optim_step() recipe.engine.forward_backward.return_value = (torch.tensor(0.0), []) - finalizer = MagicMock(return_value=0.0) - monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.scale_grads_and_clip_grad_norm", finalizer) - batch = {"labels": torch.full((1, 4), -100), "input_ids": torch.arange(4).reshape(1, 4)} metrics = recipe._run_train_optim_step([batch]) recipe.engine.forward_backward.assert_called_once() + recipe.engine.optim_step.assert_called_once() assert metrics.metrics["loss"] == 0.0 assert metrics.metrics["num_label_tokens"] == 0 - assert recipe.optimizer[0].step_called + assert not recipe.optimizer[0].step_called + + +@pytest.mark.cuda(False) +def test_run_train_step_keeps_fp8_precompute_after_engine_optim_step(monkeypatch): + recipe = _build_engine_recipe_for_optim_step() + recipe.cfg = _Cfg( + fp8={ + "enabled": True, + "precompute_float8_dynamic_scale_for_fsdp": True, + } + ) + recipe.device_mesh = {"dp_shard": SimpleNamespace(size=lambda: 2)} + events = [] + + def optim_step(**kwargs): + events.append("optim_step") + return SimpleNamespace(grad_norm=2.5, learning_rates=(0.01,)) + + recipe.engine.optim_step.side_effect = optim_step + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.precompute_float8_dynamic_scale_for_fsdp", + lambda model: events.append(("fp8_precompute", model)), + ) + + batch = {"labels": torch.tensor([[1, 2]]), "input_ids": torch.tensor([[3, 4]])} + recipe._run_train_optim_step([batch]) + + assert events == ["optim_step", ("fp8_precompute", recipe.model_parts[0])] def test_make_engine_datum_filters_raw_media_off_first_pipeline_stage(): @@ -1937,6 +1961,18 @@ def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch): assert trainer.engine is not None +def test_vlm_setup_binds_optimizer_state_to_engine(monkeypatch): + cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False) + _patch_vlm_setup_minimals(monkeypatch, cp_size=1) + + trainer = FinetuneRecipeForVLM(cfg) + trainer.setup() + + assert trainer.engine.optimizers == tuple(trainer.optimizer) + assert trainer.engine.lr_schedulers == tuple(trainer.lr_scheduler or ()) + assert trainer.engine.max_grad_norm == trainer.max_grad_norm + + def test_vlm_setup_builds_engine_for_magi(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) _patch_vlm_setup_minimals(monkeypatch, cp_size=1) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index eb2cb14dd5..a1ff944189 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1109,6 +1109,9 @@ def test_setup_builds_engine_for_eager_sum_loss(monkeypatch): assert trainer.engine is not None assert trainer.engine.microbatch_size == 1 assert trainer.engine.mtp_ignore_index == -321 + assert trainer.engine.optimizers == tuple(trainer.optimizer) + assert trainer.engine.lr_schedulers == tuple(trainer.lr_scheduler or ()) + assert trainer.engine.max_grad_norm == trainer.max_grad_norm def test_setup_builds_engine_for_eager_fused_loss(monkeypatch): @@ -2123,15 +2126,12 @@ def _make_recipe( monkeypatch.setattr(recipe, "_get_dp_group_size", lambda include_cp=False: dp_group_size) monkeypatch.setattr(recipe, "_get_cp_group_size", lambda: cp_group_size) - monkeypatch.setattr( - "nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", - lambda *a, **k: torch.tensor(1.0), - ) object.__setattr__(recipe, "checkpointer", SimpleNamespace(maybe_wait_for_staging=lambda: None)) object.__setattr__(recipe, "lr_scheduler", None) object.__setattr__(recipe, "loss_fn", object()) engine = MagicMock() engine.forward_backward.return_value = (torch.tensor(0.5), []) + engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) object.__setattr__(recipe, "engine", engine) object.__setattr__(recipe, "timestamp", 0.0) monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) @@ -2148,15 +2148,8 @@ def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypat monkeypatch.setattr(recipe, "_make_engine_datum", make_datum) engine = MagicMock() engine.forward_backward.return_value = (torch.tensor(0.25), []) + engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.02,)) object.__setattr__(recipe, "engine", engine) - - finalizer_calls = [] - - def finalize_grads(*args, **kwargs): - finalizer_calls.append((args, kwargs)) - return torch.tensor(1.0) - - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", finalize_grads) optimizer = SimpleNamespace( step=MagicMock(), zero_grad=MagicMock(), @@ -2166,14 +2159,13 @@ def finalize_grads(*args, **kwargs): metrics = recipe._run_train_optim_step(batches) engine.forward_backward.assert_called_once_with(datums, recipe._engine_loss_fn) + engine.optim_step.assert_called_once_with(before_optimizer_step=recipe.checkpointer.maybe_wait_for_staging) assert make_datum.call_count == 2 - assert len(finalizer_calls) == 1 - assert finalizer_calls[0][1]["num_label_tokens"] is None - assert finalizer_calls[0][1]["pp_enabled"] is True - assert finalizer_calls[0][1]["pp_axis_name"] == "pp" - optimizer.step.assert_called_once_with() - optimizer.zero_grad.assert_called_once_with() + optimizer.step.assert_not_called() + optimizer.zero_grad.assert_not_called() assert metrics.metrics["loss"] == pytest.approx(0.25) + assert metrics.metrics["grad_norm"] == pytest.approx(1.0) + assert metrics.metrics["lr"] == pytest.approx(0.02) def test_pp_thd_batch_uses_engine(self, monkeypatch): recipe = self._make_recipe(monkeypatch, pp_enabled=True) @@ -2184,18 +2176,62 @@ def test_pp_thd_batch_uses_engine(self, monkeypatch): } engine = MagicMock() engine.forward_backward.return_value = (torch.tensor(0.5), []) + engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) object.__setattr__(recipe, "engine", engine) - finalizer = MagicMock(return_value=torch.tensor(1.0)) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", finalizer) recipe._run_train_optim_step([batch]) engine.forward_backward.assert_called_once() + engine.optim_step.assert_called_once_with(before_optimizer_step=recipe.checkpointer.maybe_wait_for_staging) datums, loss_fn = engine.forward_backward.call_args.args assert len(datums) == 1 assert datums[0].model_inputs["qkv_format"] == "thd" assert loss_fn == recipe._engine_loss_fn - assert finalizer.call_args.kwargs["num_label_tokens"] is None + + def test_fp8_scale_precompute_stays_after_engine_optim_step(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=False) + object.__setattr__( + recipe, + "cfg", + ConfigNode( + { + "fp8": { + "enabled": True, + "precompute_float8_dynamic_scale_for_fsdp": True, + } + } + ), + ) + + class _DeviceMesh: + def __getitem__(self, name): + assert name == "dp_shard" + return SimpleNamespace(size=lambda: 2) + + object.__setattr__(recipe, "device_mesh", _DeviceMesh()) + events = [] + recipe.checkpointer.maybe_wait_for_staging = lambda: events.append("checkpoint_wait") + + def optim_step(*, before_optimizer_step): + events.append("optim_step_start") + before_optimizer_step() + events.append("optim_step_done") + return SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) + + recipe.engine.optim_step.side_effect = optim_step + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.precompute_float8_dynamic_scale_for_fsdp", + lambda model: events.append(("fp8_precompute", model)), + ) + + recipe._run_train_optim_step([{"input_ids": torch.tensor([[1, 2]]), "labels": torch.tensor([[2, -100]])}]) + + assert events == [ + "optim_step_start", + "checkpoint_wait", + "optim_step_done", + ("fp8_precompute", recipe.model_parts[0]), + ] # ----------------------------------------------------------------------------- diff --git a/tests/unit_tests/recipes/test_train_ft_partial_cuda_graphs.py b/tests/unit_tests/recipes/test_train_ft_partial_cuda_graphs.py index 5a48fc6454..7d2e60b67a 100644 --- a/tests/unit_tests/recipes/test_train_ft_partial_cuda_graphs.py +++ b/tests/unit_tests/recipes/test_train_ft_partial_cuda_graphs.py @@ -95,7 +95,7 @@ def __iter__(self): recipe.partial_cuda_graph_manager = manager recipe._partial_cuda_graph_capture_pending = True recipe._enable_qat_if_delayed = lambda _step: None - recipe._run_train_optim_step = lambda batches, _norm: ( + recipe._run_train_optim_step = lambda batches: ( events.append(("train-step", tuple(batches))) or SimpleNamespace(metrics={"loss": 1.0}) ) recipe._collect_moe_load_balance = lambda: None diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index af776bfaac..91b044c9d7 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -49,7 +49,7 @@ from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.models.common.mtp import prepare_mtp_context_parallel_inputs from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler -from nemo_automodel.engine import Engine, ForwardResult, collate_prebatched +from nemo_automodel.engine import Engine, ForwardResult, OptimStepResult, collate_prebatched class ScaleModel(nn.Module): @@ -84,6 +84,12 @@ def __init__(self, size, rank): self.mesh_dim_names = ("cp", "tp") +class _NamedMesh(dict): + def __init__(self, names, **axes): + super().__init__(axes) + self.mesh_dim_names = tuple(names) + + class _FakeAutoPipeline(AutoPipeline): def __init__( self, @@ -204,6 +210,207 @@ def test_engine_and_datum_are_lazy_top_level_exports(): assert PublicCollatedLossInputs is CollatedLossInputs +def test_optim_step_clips_updates_and_clears_real_gradients(): + model = ScaleModel() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=0.5) + + engine.forward_backward([_datum([4])], _identity_loss) + torch.testing.assert_close(model.weight.grad, torch.tensor(4.0)) + + result = engine.optim_step() + + assert isinstance(result, OptimStepResult) + torch.testing.assert_close(result.grad_norm, torch.tensor(4.0, dtype=torch.float64)) + torch.testing.assert_close(model.weight, torch.tensor(0.95)) + assert model.weight.grad is None + assert result.learning_rates == (0.1,) + + +def test_optim_step_orders_multi_optimizer_topology_and_post_step_work(monkeypatch): + events = [] + + class GateModel(ScaleModel): + def __init__(self, name): + super().__init__() + self.name = name + + def update_moe_gate_bias(self): + events.append(f"gate-{self.name}") + + class RecordingSGD(torch.optim.SGD): + def __init__(self, name, parameters, lr): + self.name = name + super().__init__(parameters, lr=lr) + + def step(self, closure=None): + events.append(f"step-{self.name}") + return super().step(closure) + + def zero_grad(self, set_to_none=True): + assert set_to_none is True + events.append(f"zero-{self.name}") + return super().zero_grad(set_to_none=set_to_none) + + class RecordingScheduler: + def __init__(self, name, optimizer): + self.name = name + self.optimizer = optimizer + + def step(self, increment): + assert increment == 1 + events.append(f"scheduler-{self.name}") + for group in self.optimizer.param_groups: + group["lr"] += 0.01 + + first = GateModel("first") + second = GateModel("second") + first._nemo_moe_tp_requires_replica_sync = True + first.weight.grad = torch.tensor(2.0) + second.weight.grad = torch.tensor(3.0) + pipeline = _FakeAutoPipeline(first, parts=[first, second]) + first_optimizer = RecordingSGD("first", first.parameters(), lr=0.1) + second_optimizer = RecordingSGD("second", second.parameters(), lr=0.2) + first_scheduler = RecordingScheduler("first", first_optimizer) + second_scheduler = RecordingScheduler("second", second_optimizer) + device_mesh = _NamedMesh(("cp", "tp"), cp=_SubMesh(2), tp=_SubMesh(4)) + moe_mesh = _NamedMesh(("ep_shard", "ep"), ep_shard=_SubMesh(2), ep=_SubMesh(2)) + engine = Engine( + pipeline, + device="cpu", + mesh_context=SimpleNamespace( + device_mesh=device_mesh, + moe_mesh=moe_mesh, + cp_size=2, + pp_size=2, + process_group=None, + ), + optimizers=[first_optimizer, second_optimizer], + lr_schedulers=[first_scheduler, second_scheduler], + max_grad_norm=None, + ) + engine._dp_group_and_size = lambda: (None, 2) + engine._gradient_group_and_size = lambda _group, _size: (None, 8) + observed = {} + + def finalize(**kwargs): + events.append("finalize") + observed.update(kwargs) + return torch.tensor(7.0) + + monkeypatch.setattr(engine_module, "scale_grads_and_clip_grad_norm", finalize) + + def before_optimizer_step(): + events.append("before-step") + + result = engine.optim_step(before_optimizer_step=before_optimizer_step) + + assert engine.optimizers == (first_optimizer, second_optimizer) + assert engine.lr_schedulers == (first_scheduler, second_scheduler) + assert events == [ + "finalize", + "before-step", + "step-first", + "zero-first", + "step-second", + "zero-second", + "gate-first", + "gate-second", + "scheduler-first", + "scheduler-second", + ] + assert observed == { + "max_grad_norm": None, + "model_parts": [first, second], + "norm_type": 2.0, + "pp_enabled": True, + "device_mesh": device_mesh, + "moe_mesh": moe_mesh, + "ep_axis_name": "ep", + "pp_axis_name": "pp", + "foreach": True, + "num_label_tokens": None, + "dp_group_size": 8, + "expert_tp_replication_factor": 4, + } + assert result.grad_norm.item() == pytest.approx(7.0) + assert result.learning_rates == pytest.approx((0.11, 0.21)) + assert first.weight.grad is None + assert second.weight.grad is None + + +def test_optim_step_callback_failure_preserves_optimizer_and_post_step_state(monkeypatch): + events = [] + + class GateModel(ScaleModel): + def update_moe_gate_bias(self): + events.append("gate") + + class RecordingSGD(torch.optim.SGD): + def step(self, closure=None): + events.append("step") + return super().step(closure) + + def zero_grad(self, set_to_none=True): + events.append("zero") + return super().zero_grad(set_to_none=set_to_none) + + class RecordingScheduler: + def step(self, increment): + events.append("scheduler") + + model = GateModel() + model.weight.grad = torch.tensor(2.0) + optimizer = RecordingSGD(model.parameters(), lr=0.1) + engine = Engine( + model, + device="cpu", + optimizers=optimizer, + lr_schedulers=RecordingScheduler(), + ) + + def finalize(**_kwargs): + events.append("finalize") + return torch.tensor(2.0) + + monkeypatch.setattr(engine_module, "scale_grads_and_clip_grad_norm", finalize) + + def fail_before_step(): + events.append("before-step") + raise ValueError("staging failed") + + with pytest.raises(ValueError, match="staging failed"): + engine.optim_step(before_optimizer_step=fail_before_step) + + assert events == ["finalize", "before-step"] + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + torch.testing.assert_close(model.weight.grad, torch.tensor(2.0)) + + +def test_optim_step_requires_an_optimizer(monkeypatch): + monkeypatch.setattr( + engine_module, + "scale_grads_and_clip_grad_norm", + lambda **_kwargs: pytest.fail("gradient finalization must not run without an optimizer"), + ) + + with pytest.raises(RuntimeError, match="requires at least one optimizer"): + Engine(ScaleModel(), device="cpu").optim_step() + + +def test_optim_step_rejects_noncallable_mutation_fence_before_finalization(monkeypatch): + model = ScaleModel() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + monkeypatch.setattr( + engine_module, + "scale_grads_and_clip_grad_norm", + lambda **_kwargs: pytest.fail("invalid callback must fail before gradient finalization"), + ) + + with pytest.raises(TypeError, match="before_optimizer_step must be callable or None"): + Engine(model, device="cpu", optimizers=optimizer).optim_step(before_optimizer_step=object()) + + def test_forward_runs_eval_without_grad_lifecycle_and_returns_local_statistics(monkeypatch): class EvalModel(ScaleModel): def forward(self, input_ids: torch.Tensor, **kwargs) -> torch.Tensor: diff --git a/tests/unit_tests/test_engine_recipe_integration.py b/tests/unit_tests/test_engine_recipe_integration.py index f3c53c27df..29c8f79413 100644 --- a/tests/unit_tests/test_engine_recipe_integration.py +++ b/tests/unit_tests/test_engine_recipe_integration.py @@ -91,10 +91,17 @@ def test_recipes_run_one_datum_engine_window_then_one_optimizer_step(recipe_cls, recipe.pp_enabled = False recipe.dist_env = SimpleNamespace(device=torch.device("cpu"), world_size=1, is_main=True) recipe.distributed_config = SimpleNamespace(defer_fsdp_grad_sync=True) - recipe.engine = Engine(model, device="cpu", microbatch_size=1, collate_fn=collate_prebatched) optimizer = _CountingSGD(model.parameters()) recipe.optimizer = [optimizer] recipe.lr_scheduler = None + recipe.engine = Engine( + model, + device="cpu", + microbatch_size=1, + collate_fn=collate_prebatched, + optimizers=recipe.optimizer, + max_grad_norm=None, + ) recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=lambda: None) recipe.step_scheduler = SimpleNamespace(step=1, epoch=0, is_remote_logging_step=False) recipe.timestamp = time.perf_counter() - 1.0 @@ -135,7 +142,7 @@ def local_reduce(value, include_cp=False): reference_loss.backward() reference_optimizer.step() - metrics = recipe._run_train_optim_step(batches, max_grad_norm=None) + metrics = recipe._run_train_optim_step(batches) assert metrics.metrics["loss"] == pytest.approx(reference_loss.item()) assert model.forward_calls == 2 From be6c0c806c488c076ae9d21784f899f5173ab61d Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Thu, 20 Aug 2026 22:13:34 -0700 Subject: [PATCH 12/34] refactor(engine): simplify execution plumbing Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 11 - nemo_automodel/engine/__init__.py | 202 ++++++++---------- nemo_automodel/recipes/llm/train_ft.py | 25 +-- nemo_automodel/recipes/vlm/finetune.py | 26 +-- .../pipelining/test_autopipeline.py | 68 ++---- .../recipes/test_finetune_vlm_cp_wiring.py | 64 ------ .../recipes/test_finetune_vlm_helpers.py | 5 + tests/unit_tests/recipes/test_train_ft.py | 61 +----- tests/unit_tests/test_engine.py | 72 ------- 9 files changed, 125 insertions(+), 409 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index 4ff29105a0..f913d85143 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -305,7 +305,6 @@ def step( *, target: torch.Tensor | None = None, losses: list[torch.Tensor] | None = None, - return_outputs: bool = True, **kwargs: Any, ) -> Any: """Run one pipeline schedule step with model-owned input chunking. @@ -317,9 +316,6 @@ def step( ranks without the last pipeline stage. losses: Mutable list populated with scalar loss tensors, or ``None`` on ranks without the last pipeline stage. - return_outputs: Whether the last pipeline stage returns merged model - outputs when supported by the installed PyTorch version. Tensor - layouts are defined by the underlying model. **kwargs: Keyword schedule inputs. Tensor values may have arbitrary model-defined layouts; model-owned metadata identifies any nonstandard batch axis. @@ -333,18 +329,12 @@ def step( schedule_args = (model_input,) if self._info.has_first_stage else () kwargs_chunk_spec = self._get_schedule_kwargs_chunk_spec(kwargs) - schedule_options = ( - {"return_outputs": return_outputs} - if "return_outputs" in inspect.signature(schedule.step).parameters - else {} - ) if kwargs_chunk_spec is None: return schedule.step( *schedule_args, target=target, losses=losses, - **schedule_options, **kwargs, ) @@ -355,7 +345,6 @@ def step( *schedule_args, target=target, losses=losses, - **schedule_options, **kwargs, ) finally: diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 28db6b4338..fe3805366f 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -337,25 +337,12 @@ def forward( with self.context_fn(model_inputs), cp_context(): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - result = loss_fn(output, loss_inputs) - has_outputs = isinstance(result, tuple) - if returns_outputs is None: - returns_outputs = has_outputs - elif returns_outputs != has_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - if isinstance(result, tuple): - losses, outputs = result - if ( - not isinstance(outputs, Sequence) - or isinstance(outputs, (str, bytes)) - or len(outputs) != len(batch_datums) - or not all(isinstance(item, Mapping) for item in outputs) - ): + numerator, outputs = _parse_loss_result(loss_fn(output, loss_inputs), loss_inputs["weights"]) + returns_outputs = _update_output_mode(returns_outputs, outputs) + if outputs is not None: + if len(outputs) != len(batch_datums): raise ValueError("loss_fn outputs must contain one mapping per Datum") - loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) - else: - losses = result - numerator = _weighted_numerator(losses, loss_inputs["weights"]) + loss_fn_outputs.extend(outputs) if zero_weight_sum: numerator = numerator * 0 local_loss_sum.add_(numerator.detach().to(torch.float64)) @@ -483,25 +470,12 @@ def forward_backward( ): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - result = loss_fn(output, loss_inputs) - has_outputs = isinstance(result, tuple) - if returns_outputs is None: - returns_outputs = has_outputs - elif returns_outputs != has_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - if isinstance(result, tuple): - losses, outputs = result - if ( - not isinstance(outputs, Sequence) - or isinstance(outputs, (str, bytes)) - or len(outputs) != len(datums) - or not all(isinstance(item, Mapping) for item in outputs) - ): + numerator, outputs = _parse_loss_result(loss_fn(output, loss_inputs), loss_inputs["weights"]) + returns_outputs = _update_output_mode(returns_outputs, outputs) + if outputs is not None: + if len(outputs) != len(datums): raise ValueError("loss_fn outputs must contain one mapping per Datum") - loss_fn_outputs.extend(_detach(dict(item)) for item in outputs) - else: - losses = result - numerator = _weighted_numerator(losses, loss_inputs["weights"]) + loss_fn_outputs.extend(outputs) if zero_denominator: numerator = numerator * 0 (numerator * (grad_group_size / safe_denominator)).backward() @@ -948,44 +922,16 @@ def _pipeline_execute( returns_outputs: bool | None = None with cp_context(): - primary_name = _primary_name(model_inputs) - primary = model_inputs[primary_name] - if not isinstance(primary, torch.Tensor) or primary.ndim == 0: - raise ValueError("pipeline Engine requires a tensor input_ids or inputs_embeds") - + primary_microbatch, is_thd, batch_size = self._plan_pipeline_batch(model_inputs) num_microbatches = self.pipeline.num_microbatches - is_thd = model_inputs.get("qkv_format") == "thd" - if num_microbatches == 1: - primary_microbatch = primary - elif is_thd: - if primary.shape[0] != num_microbatches: - raise ValueError( - f"THD sharder produced {primary.shape[0]} chunks, " - f"expected {num_microbatches} pipeline microbatches" - ) - primary_microbatch = primary.narrow(0, 0, 1) - else: - batch_size = primary.shape[0] - if batch_size % num_microbatches != 0: - raise ValueError( - f"pipeline outer batch size {batch_size} must be divisible by {num_microbatches} microbatches" - ) - materialized_batch_size = batch_size // num_microbatches - if materialized_batch_size != self.pipeline.pp_microbatch_size: - raise ValueError( - f"materialized pipeline microbatch has batch size {materialized_batch_size}, " - f"but AutoPipeline is configured for pp_microbatch_size={self.pipeline.pp_microbatch_size}" - ) - primary_microbatch = primary.narrow(0, 0, materialized_batch_size) - - if is_thd and num_microbatches == 1: - seq_len = primary_microbatch.shape[0] - else: - seq_len = primary_microbatch.shape[1] if primary_microbatch.ndim >= 2 else primary_microbatch.shape[0] - effective_microbatch_size = 1 if is_thd else primary_microbatch.shape[0] + seq_len = ( + primary_microbatch.shape[0] + if is_thd and num_microbatches == 1 + else primary_microbatch.shape[min(primary_microbatch.ndim - 1, 1)] + ) self.pipeline.update_seq_len( seq_len, - microbatch_size=effective_microbatch_size, + microbatch_size=1 if is_thd else primary_microbatch.shape[0], input_tensor=primary_microbatch, ) @@ -1000,26 +946,19 @@ def _pipeline_execute( loss_inputs, loss_batch_layout, num_datums=len(datums), + is_thd=is_thd, + batch_size=batch_size, ) ) def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: nonlocal returns_outputs loss_inputs_mb = loss_microbatches[microbatch_index] - result = loss_fn(output, loss_inputs_mb) - has_outputs = isinstance(result, tuple) - if returns_outputs is None: - returns_outputs = has_outputs - elif returns_outputs != has_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - if isinstance(result, tuple): - losses, batch_outputs = result - if ( - not isinstance(batch_outputs, Sequence) - or isinstance(batch_outputs, (str, bytes)) - or not all(isinstance(item, Mapping) for item in batch_outputs) - ): - raise ValueError("loss_fn outputs must be a sequence of mappings") + numerator, batch_outputs = _parse_loss_result( + loss_fn(output, loss_inputs_mb), loss_inputs_mb["weights"] + ) + returns_outputs = _update_output_mode(returns_outputs, batch_outputs) + if batch_outputs is not None: datum_indices = ( None if datum_indices_by_microbatch is None @@ -1030,24 +969,20 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: "a prebatched Datum may return outputs only when num_microbatches=1 because " "its inner sample boundaries are not part of the Datum contract" ) - detached_outputs = [_detach(dict(item)) for item in batch_outputs] if datum_indices is not None: - if len(detached_outputs) != len(datum_indices): + if len(batch_outputs) != len(datum_indices): raise ValueError( - f"pipeline loss_fn returned {len(detached_outputs)} outputs for microbatch " + f"pipeline loss_fn returned {len(batch_outputs)} outputs for microbatch " f"{microbatch_index}, expected {len(datum_indices)} from its Datum mapping" ) - for datum_index, item in zip(datum_indices, detached_outputs): + for datum_index, item in zip(datum_indices, batch_outputs): if outputs_by_datum[datum_index] is not None: raise RuntimeError( f"pipeline returned more than one output for Datum {datum_index}" ) outputs_by_datum[datum_index] = item else: - outputs_by_microbatch[microbatch_index] = detached_outputs - else: - losses = result - numerator = _weighted_numerator(losses, loss_inputs_mb["weights"]) + outputs_by_microbatch[microbatch_index] = batch_outputs if zero_weight_sum: numerator = numerator * 0 local_loss_sum.add_(numerator.detach().to(torch.float64)) @@ -1112,6 +1047,37 @@ def _broadcast_pipeline_outputs(self, outputs: list[dict[str, Any]]) -> list[dic raise RuntimeError("pipeline output synchronization received invalid per-Datum outputs") return _to_device(received, self.device) + def _plan_pipeline_batch(self, model_inputs: Mapping[str, Any]) -> tuple[torch.Tensor, bool, int | None]: + """Validate the PP outer-batch shape once for metadata and slicing.""" + primary_name = _primary_name(model_inputs) + primary = model_inputs[primary_name] + if not isinstance(primary, torch.Tensor) or primary.ndim == 0: + raise ValueError("pipeline Engine requires a tensor input_ids or inputs_embeds") + + num_microbatches = self.pipeline.num_microbatches + is_thd = model_inputs.get("qkv_format") == "thd" + if num_microbatches == 1: + return primary, is_thd, None + if is_thd: + if primary.shape[0] != num_microbatches: + raise ValueError( + f"THD sharder produced {primary.shape[0]} chunks, expected {num_microbatches} pipeline microbatches" + ) + return primary.narrow(0, 0, 1), True, None + + batch_size = primary.shape[0] + if batch_size % num_microbatches != 0: + raise ValueError( + f"pipeline outer batch size {batch_size} must be divisible by {num_microbatches} microbatches" + ) + materialized_batch_size = batch_size // num_microbatches + if materialized_batch_size != self.pipeline.pp_microbatch_size: + raise ValueError( + f"materialized pipeline microbatch has batch size {materialized_batch_size}, " + f"but AutoPipeline is configured for pp_microbatch_size={self.pipeline.pp_microbatch_size}" + ) + return primary.narrow(0, 0, materialized_batch_size), False, batch_size + def _materialize_pipeline_microbatches( self, model_inputs: dict[str, Any], @@ -1119,6 +1085,8 @@ def _materialize_pipeline_microbatches( loss_batch_layout: _LossBatchLayout, *, num_datums: int, + is_thd: bool, + batch_size: int | None, ) -> tuple[list[dict[str, Any]], list[LossInputs], list[tuple[int, ...]] | None]: """Split one CP-prepared outer batch into exact pipeline inputs. @@ -1146,7 +1114,7 @@ def _materialize_pipeline_microbatches( [model_inputs], loss_batch_layout.item_to_datum, num_datums=num_datums, - is_thd=model_inputs.get("qkv_format") == "thd", + is_thd=is_thd, ) assert datum_indices is not None return ( @@ -1155,32 +1123,12 @@ def _materialize_pipeline_microbatches( datum_indices, ) - primary_name = _primary_name(model_inputs) - primary = model_inputs[primary_name] - if not isinstance(primary, torch.Tensor) or primary.ndim == 0: - raise ValueError(f"pipeline Engine requires tensor {primary_name}") - - is_thd = model_inputs.get("qkv_format") == "thd" if is_thd: - if primary.shape[0] != num_microbatches: - raise ValueError( - f"THD sharder produced {primary.shape[0]} chunks, expected {num_microbatches} pipeline microbatches" - ) model_microbatches = [ _select_chunk(model_inputs, index, num_microbatches) for index in range(num_microbatches) ] else: - batch_size = primary.shape[0] - if batch_size % num_microbatches != 0: - raise ValueError( - f"pipeline outer batch size {batch_size} must be divisible by {num_microbatches} microbatches" - ) - materialized_batch_size = batch_size // num_microbatches - if materialized_batch_size != self.pipeline.pp_microbatch_size: - raise ValueError( - f"materialized pipeline microbatch has batch size {materialized_batch_size}, " - f"but AutoPipeline is configured for pp_microbatch_size={self.pipeline.pp_microbatch_size}" - ) + assert batch_size is not None custom_dims: dict[str, int] = {} chunk_dims = getattr(self.model, "get_pipeline_kwargs_chunk_dims", None) if chunk_dims is not None: @@ -1770,6 +1718,30 @@ def _weighted_numerator(losses: Any, weights: torch.Tensor) -> torch.Tensor: return numerator +def _parse_loss_result( + result: torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], + weights: torch.Tensor, +) -> tuple[torch.Tensor, list[dict[str, Any]] | None]: + """Normalize one loss callback result without applying Datum routing.""" + if not isinstance(result, tuple): + return _weighted_numerator(result, weights), None + losses, outputs = result + if ( + not isinstance(outputs, Sequence) + or isinstance(outputs, (str, bytes)) + or not all(isinstance(item, Mapping) for item in outputs) + ): + raise ValueError("loss_fn outputs must be a sequence of mappings") + return _weighted_numerator(losses, weights), [_detach(dict(item)) for item in outputs] + + +def _update_output_mode(previous: bool | None, outputs: list[dict[str, Any]] | None) -> bool: + current = outputs is not None + if previous is not None and previous != current: + raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") + return current + + def _detach(value: Any) -> Any: """Detach tensor leaves without changing an output record's structure.""" if isinstance(value, torch.Tensor): diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 6704a2e061..6aa46d0b99 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -1088,10 +1088,6 @@ def _compute_causal_lm_loss(self, output, labels, model_inputs, *, num_label_tok Scalar local causal-LM plus MTP loss sum. """ model = self.model_parts[0] - if self.pp_enabled: - model = next( - model_part for model_part, stage in zip(self.model_parts, self.pp.info.stages) if stage.is_last - ) grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None hidden_states = get_final_hidden_states(output) if isinstance(self.loss_fn, FusedLinearCrossEntropy) and hidden_states is None: @@ -1173,6 +1169,8 @@ def _engine_loss_fn( self, output: Any, loss_inputs: dict[str, torch.Tensor | tuple[torch.Tensor, ...]], + *, + is_train: bool = True, ) -> torch.Tensor: """Compute a local summed causal-LM loss for Engine normalization. @@ -1182,6 +1180,7 @@ def _engine_loss_fn( loss_inputs: CP-local labels and weights with matching token axes, optional packed-sequence metadata, and optional per-depth MTP targets whose tensors share the labels' token axes. + is_train: Whether the loss participates in training gradients. Returns: Scalar local loss-sum tensor. @@ -1196,7 +1195,7 @@ def _engine_loss_fn( loss_inputs["labels"], loss_inputs, num_label_tokens=None, - is_train=True, + is_train=is_train, ) def _engine_validation_loss_fn( @@ -1205,18 +1204,7 @@ def _engine_validation_loss_fn( loss_inputs: dict[str, torch.Tensor | tuple[torch.Tensor, ...]], ) -> torch.Tensor: """Compute the validation loss numerator without training reductions.""" - if self.pp_enabled: - if self.pipeline_loss_fn is None: - raise RuntimeError("The last pipeline stage has no configured causal-LM loss") - self.pipeline_loss_fn.cu_seqlens = loss_inputs.get("cu_seqlens") - return self.pipeline_loss_fn(output, loss_inputs["labels"]) - return self._compute_causal_lm_loss( - output, - loss_inputs["labels"], - loss_inputs, - num_label_tokens=None, - is_train=False, - ) + return self._engine_loss_fn(output, loss_inputs, is_train=False) def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: """Execute a single training step. @@ -1316,9 +1304,6 @@ def _run_validation_epoch(self, val_dataloader): val_dataloader: DataLoader for the validation dataset. """ with ScopedRNG(seed=1, ranked=True): - for mp in self.model_parts: - mp.eval() - total_loss = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) total_num_label_tokens = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 72e8f28375..eacd5582fb 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -943,6 +943,7 @@ def _engine_loss_fn( out: Any, loss_inputs: Mapping[str, torch.Tensor | tuple[torch.Tensor, ...]], *, + is_train: bool = True, log_drafter: bool = False, log_denominator: int | float | None = None, ) -> torch.Tensor: @@ -955,6 +956,7 @@ def _engine_loss_fn( have shape [batch, sequence] or [tokens]. Optional ``cu_seqlens`` has shape [num_sequences + 1] or [1, num_sequences + 1], and each ``mtp_per_depth_targets`` tensor matches the labels' local layout. + is_train: Whether the loss participates in training gradients. Returns: Scalar local loss-sum tensor. @@ -974,7 +976,7 @@ def _engine_loss_fn( out=out, labels=labels, num_label_tokens=None, - is_train=True, + is_train=is_train, cu_seqlens=cu_seqlens, mtp_per_depth_targets=mtp_per_depth_targets, log_drafter=log_drafter, @@ -987,24 +989,7 @@ def _engine_validation_loss_fn( loss_inputs: Mapping[str, torch.Tensor | tuple[torch.Tensor, ...]], ) -> torch.Tensor: """Compute the validation loss numerator without training reductions.""" - labels = cast(torch.Tensor, loss_inputs["labels"]) - cu_seqlens = cast(torch.Tensor | None, loss_inputs.get("cu_seqlens")) - if self.pp_enabled: - if self.pipeline_loss_fn is None: - raise RuntimeError("The last pipeline stage has no configured causal-LM loss") - self.pipeline_loss_fn.cu_seqlens = cu_seqlens - return self.pipeline_loss_fn(out, labels) - return self._compute_vlm_loss( - out=out, - labels=labels, - num_label_tokens=None, - is_train=False, - cu_seqlens=cu_seqlens, - mtp_per_depth_targets=cast( - tuple[torch.Tensor, ...] | None, - loss_inputs.get("mtp_per_depth_targets"), - ), - ) + return self._engine_loss_fn(out, loss_inputs, is_train=False) @contextmanager def _cp_vision_frame_sharding_context(self): @@ -1164,9 +1149,6 @@ def _run_validation_epoch(self, val_dataloader): Metrics for the completed validation epoch. """ with ScopedRNG(seed=1, ranked=True): - for mp in self.model_parts: - mp.eval() - total_loss = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) total_num_label_tokens = torch.zeros((), dtype=torch.float64, device=self.dist_env.device) for batch in val_dataloader: diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 5ae884b2c5..9b93d1e47c 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -404,23 +404,12 @@ def test_step_splits_mrope_position_ids_on_model_owned_batch_axis(self): def test_step_without_model_hook_uses_pytorch_default_chunking(self): ap = self._pipeline_with_parts(nn.Module()) - original_loss_fn = ap.info.schedule._loss_fn ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) assert ap.info.schedule.kwargs_chunk_spec_during_step is None assert ap.info.schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) assert ap.info.schedule._kwargs_chunk_spec is None - assert ap.info.schedule._loss_fn is original_loss_fn - assert ap.info.schedule.return_outputs_during_step is True - - def test_step_does_not_forward_return_outputs_to_older_pytorch(self): - schedule = _LegacyStepSchedule() - ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - - ap.step(torch.zeros(2, 8), return_outputs=False) - - assert schedule.received_return_outputs is False def test_step_microbatches_passes_prepared_inputs_without_resplitting(self): schedule = _KwargsChunkSchedule(invoke_loss=True) @@ -474,11 +463,22 @@ def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): assert schedule.args_split == [(), ()] assert all("inputs_embeds" not in kwargs for kwargs in schedule.kwargs_split) - def test_step_microbatches_does_not_forward_return_outputs_to_older_pytorch(self): - schedule = _LegacyStepSchedule() + @pytest.mark.parametrize( + ("method_name", "schedule_cls"), + [ + ("step_microbatches", _LegacyStepSchedule), + ("eval_microbatches", _LegacyEvalSchedule), + ], + ) + def test_prepared_microbatches_do_not_forward_return_outputs_to_older_pytorch( + self, + method_name, + schedule_cls, + ): + schedule = schedule_cls() ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - ap.step_microbatches( + getattr(ap, method_name)( [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], loss_fn=Mock(), return_outputs=False, @@ -524,20 +524,6 @@ def loss_fn(output, index): assert schedule._split_inputs == original_split_inputs assert schedule._loss_fn is original_loss_fn - def test_eval_microbatches_does_not_forward_return_outputs_to_older_pytorch(self): - schedule = _LegacyEvalSchedule() - ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - - ap.eval_microbatches( - [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], - loss_fn=Mock(), - return_outputs=False, - ) - - assert schedule.eval_calls == 1 - assert schedule.step_calls == 0 - assert schedule.received_return_outputs is False - def test_eval_microbatches_uses_step_capability_when_eval_forwards_kwargs(self): schedule = _ForwardingEvalSchedule() ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) @@ -578,37 +564,21 @@ def test_step_microbatches_validates_prepared_inputs(self, model_inputs): with pytest.raises(ValueError, match="Expected 2|exactly one"): ap.step_microbatches(model_inputs, loss_fn=Mock()) - def test_step_microbatches_restores_schedule_state_after_failure(self): + @pytest.mark.parametrize("method_name", ["step_microbatches", "eval_microbatches"]) + def test_prepared_microbatches_restore_schedule_state_after_failure(self, method_name): schedule = _KwargsChunkSchedule(fail_on_step=True) original_split_inputs = schedule._split_inputs original_loss_fn = schedule._loss_fn ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) with pytest.raises(RuntimeError, match="schedule failed"): - ap.step_microbatches( + getattr(ap, method_name)( [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], loss_fn=Mock(), ) - assert schedule.split_inputs_during_step is not original_split_inputs - assert schedule.loss_fn_during_step is not original_loss_fn - assert schedule._split_inputs == original_split_inputs - assert schedule._loss_fn is original_loss_fn - - def test_eval_microbatches_restores_schedule_state_after_failure(self): - schedule = _KwargsChunkSchedule(fail_on_step=True) - original_split_inputs = schedule._split_inputs - original_loss_fn = schedule._loss_fn - ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - - with pytest.raises(RuntimeError, match="schedule failed"): - ap.eval_microbatches( - [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], - loss_fn=Mock(), - ) - - assert schedule.eval_calls == 1 - assert schedule.step_calls == 0 + assert schedule.eval_calls == int(method_name == "eval_microbatches") + assert schedule.step_calls == int(method_name == "step_microbatches") assert schedule.split_inputs_during_step is not original_split_inputs assert schedule.loss_fn_during_step is not original_loss_fn assert schedule._split_inputs == original_split_inputs diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py index 3d97911f16..79541613b0 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -303,67 +303,3 @@ def test_setup_always_stages_pp_media_under_pp( assert dataloader_calls[0]["cp_size"] == cp_size assert trainer.engine.pipeline is trainer.pp assert trainer.engine.microbatch_size == 1 - - -# ----------------------------------------------------------------------------- -# validation Engine boundary -# ----------------------------------------------------------------------------- - - -def test_run_validation_epoch_does_not_sum_tokens_over_cp(monkeypatch): - """Engine returns CP-complete sums, so the epoch reduces only over DP.""" - monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext()) - - class _Model(torch.nn.Module): - def prepare_model_inputs_for_cp(self, *args, **kwargs): - raise AssertionError("the recipe must delegate CP preparation to Engine.forward") - - def forward(self, *args, **kwargs): - raise AssertionError("the recipe must delegate model execution to Engine.forward") - - engine_calls = [] - - class _Engine: - def forward(self, datums, loss_fn): - engine_calls.append((datums, loss_fn)) - return SimpleNamespace( - loss_sum=torch.tensor(6.0, dtype=torch.float64), - weight_sum=torch.tensor(3.0, dtype=torch.float64), - loss_fn_outputs=[], - ) - - recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) - recipe.model_parts = [_Model()] - recipe.loss_fn = object() # not a FusedLinearCrossEntropy - recipe.engine = _Engine() - recipe.pp_enabled = False - recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) - recipe.step_scheduler = SimpleNamespace(step=3, epoch=1) - recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.001}])] - recipe._maybe_add_drafter_loss = lambda *, base_loss, **kwargs: base_loss - - allreduce_calls = [] - - def _fake_allreduce(tensor, include_cp=False): - allreduce_calls.append((tensor.tolist(), include_cp)) - return tensor - - recipe._dp_allreduce = _fake_allreduce - - # One batch, 3 supervised tokens (-100 ignored). - batch = { - "input_ids": torch.tensor([[1, 2, 3, 4]]), - "labels": torch.tensor([[1, 2, -100, 4]]), - } - - result = recipe._run_validation_epoch([batch]) - - assert len(engine_calls) == 1 - datums, loss_fn = engine_calls[0] - assert len(datums) == 1 - assert datums[0].model_inputs["input_ids"] is batch["input_ids"] - assert datums[0].loss_fn_inputs["labels"] is batch["labels"] - assert loss_fn == recipe._engine_validation_loss_fn - assert len(allreduce_calls) == 2 - assert all(include_cp is False for _, include_cp in allreduce_calls) - assert result.metrics["val_loss"] == pytest.approx(2.0) diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 7c8a9f3247..54d42bcf9b 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -465,8 +465,10 @@ def forward_backward(datums, loss_fn): assert recipe._compute_vlm_loss.call_count == 2 first_call, second_call = recipe._compute_vlm_loss.call_args_list + assert first_call.kwargs["is_train"] is True assert first_call.kwargs["log_drafter"] is True assert first_call.kwargs["log_denominator"] == 4 + assert second_call.kwargs["is_train"] is True assert second_call.kwargs["log_drafter"] is False assert second_call.kwargs["log_denominator"] == 4 @@ -2101,6 +2103,8 @@ def test_vlm_engine_validation_loss_uses_eval_path(): is_train=False, cu_seqlens=cu_seqlens, mtp_per_depth_targets=targets, + log_drafter=False, + log_denominator=None, ) @@ -2164,6 +2168,7 @@ def test_vlm_validation_uses_engine_forward_and_aggregates_uneven_batches(monkey assert loss_fn == recipe._engine_validation_loss_fn assert allreduce.call_count == 2 assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) + recipe.model_parts[0].eval.assert_not_called() assert metrics.metrics["val_loss"] == pytest.approx(13.0 / 5.0) assert metrics.metrics["num_label_tokens"] == pytest.approx(5.0) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index a1ff944189..6108d8e11d 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1627,6 +1627,7 @@ def test_run_validation_epoch_uses_engine_forward_for_pp_complete_results(monkey assert loss_fn == recipe._engine_validation_loss_fn assert allreduce.call_count == 2 assert all("include_cp" not in call.kwargs for call in allreduce.call_args_list) + recipe.model_parts[0].eval.assert_not_called() assert metrics.metrics["val_loss"] == pytest.approx(13.0 / 5.0) assert metrics.metrics["num_label_tokens"] == pytest.approx(5.0) @@ -2059,75 +2060,23 @@ def _make_recipe( self, monkeypatch, pp_enabled, - dp_group_size=4, - cp_group_size=1, ): - from nemo_automodel.components.config.loader import ConfigNode - - cfg = ConfigNode( - { - "nvtx": False, - "model": {}, - "dataloader": {"collate_fn": "nemo_automodel.components.datasets.utils.default_collater"}, - "dataset": {}, - "validation_dataloader": {}, - "step_scheduler": {"local_batch_size": 1, "global_batch_size": 1}, - "optimizer": {}, - "loss_fn": {}, - "checkpoint": {"best_metric_key": "default"}, - "distributed": {"cp_size": 1}, - "autopipeline": {"pp_microbatch_size": 1}, - } - ) - monkeypatch.setattr( - "nemo_automodel.recipes.llm.train_ft.initialize_distributed", - lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0), - ) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None) - recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg) - + recipe = object.__new__(TrainFinetuneRecipeForNextTokenPrediction) + object.__setattr__(recipe, "cfg", ConfigNode({"fp8": None})) object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True)) object.__setattr__(recipe, "device_mesh", None) - object.__setattr__(recipe, "moe_mesh", None) object.__setattr__(recipe, "pp_enabled", pp_enabled) - object.__setattr__(recipe, "te_fp8", None) object.__setattr__(recipe, "model_parts", [nn.Linear(4, 4)]) - object.__setattr__( - recipe, - "optimizer", - [SimpleNamespace(step=lambda: None, zero_grad=lambda: None, param_groups=[{"lr": 0.01}])], - ) - object.__setattr__(recipe, "lr_schedulers", []) object.__setattr__(recipe, "step_scheduler", SimpleNamespace(step=1, epoch=0)) - - if pp_enabled: - pp_info = SimpleNamespace( - has_first_stage=True, - has_last_stage=True, - schedule=SimpleNamespace(_n_microbatches=1), - ) - object.__setattr__( - recipe, - "pp", - SimpleNamespace( - info=pp_info, - pp_batch_size=1, - pp_microbatch_size=1, - update_seq_len=lambda seq_len: None, - ), - ) - object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) - monkeypatch.setattr( recipe, "_dp_allreduce", lambda val, include_cp=False: val if isinstance(val, torch.Tensor) else torch.tensor(val), ) - monkeypatch.setattr(recipe, "_get_dp_group_size", lambda include_cp=False: dp_group_size) - monkeypatch.setattr(recipe, "_get_cp_group_size", lambda: cp_group_size) + monkeypatch.setattr(recipe, "_get_dp_group_size", lambda include_cp=False: 4) + monkeypatch.setattr(recipe, "_get_cp_group_size", lambda: 1) object.__setattr__(recipe, "checkpointer", SimpleNamespace(maybe_wait_for_staging=lambda: None)) - object.__setattr__(recipe, "lr_scheduler", None) object.__setattr__(recipe, "loss_fn", object()) engine = MagicMock() engine.forward_backward.return_value = (torch.tensor(0.5), []) diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 91b044c9d7..4333b0fb9f 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -818,52 +818,6 @@ def prepare_mtp_inputs_for_cp(self, batch, *, ignore_index): assert model.forward_calls == 0 -def _model_ready_packed_collate(datums): - model_inputs, loss_inputs = collate_datums(datums, packed=True) - lengths = torch.tensor([datum.seq_len for datum in datums], dtype=torch.int32) - model_inputs = { - "input_ids": model_inputs["input_ids"].flatten(), - "position_ids": model_inputs["position_ids"].flatten(), - "cu_seqlens": F.pad(lengths.cumsum(0), (1, 0)), - "max_seqlen": lengths.max(), - "qkv_format": "thd", - } - loss_inputs = { - key: value.flatten() if value.ndim == 2 and value.shape[0] == 1 else value for key, value in loss_inputs.items() - } - return model_inputs, loss_inputs - - -def test_packed_rl_callback_keeps_per_datum_sequence_boundaries(): - model = ScaleModel() - first = _datum([1, 2]) - second = _datum([3]) - first.loss_fn_inputs["sequence_scale"] = torch.tensor(2.0) - second.loss_fn_inputs["sequence_scale"] = torch.tensor(0.5) - - datums = [first, second] - - def sequence_loss(output, _loss_inputs): - chunks = output.squeeze(0).split([datum.seq_len for datum in datums]) - losses = torch.cat([chunk * datum.loss_fn_inputs["sequence_scale"] for chunk, datum in zip(chunks, datums)]) - outputs = [ - {"sequence_sum": chunk.sum(), "sequence_length": datum.seq_len} for chunk, datum in zip(chunks, datums) - ] - return losses, outputs - - loss, outputs = Engine( - model, - device="cpu", - microbatch_size=2, - collate_fn=_model_ready_packed_collate, - ).forward_backward(datums, sequence_loss) - - assert loss.item() == pytest.approx(2.5) - assert model.weight.grad.item() == pytest.approx(2.5) - assert [item["sequence_length"] for item in outputs] == [2, 1] - assert [item["sequence_sum"].item() for item in outputs] == pytest.approx([3.0, 3.0]) - - def test_weights_mask_loss_and_denominator(): model = ScaleModel() loss, _ = Engine(model, device="cpu").forward_backward( @@ -1250,32 +1204,6 @@ def test_pipeline_outputs_follow_logical_microbatch_order(): assert [item["metric"].item() for item in outputs] == [1.0, 2.0] -def test_pipeline_default_packed_collater_splits_flat_datums_inside_engine(): - model = ScaleModel() - model.backend = SimpleNamespace(attn="te") - pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) - datums = [_datum([1, 2]), _datum([3, 4]), _datum([5, 6]), _datum([7, 8])] - - def loss_with_outputs(output, _loss_inputs): - return output, [ - {"pair_sum": output[..., :2].sum()}, - {"pair_sum": output[..., 2:].sum()}, - ] - - loss, outputs = Engine( - pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - microbatch_size=4, - collate_fn=partial(collate_datums, packed=True), - ).forward_backward(datums, loss_with_outputs) - - assert loss.item() == pytest.approx(4.5) - assert model.weight.grad.item() == pytest.approx(4.5) - assert [item["pair_sum"].item() for item in outputs] == [3.0, 7.0, 11.0, 15.0] - assert [item["input_ids"].shape for item in pipeline.prepared_inputs[0]] == [(1, 4), (1, 4)] - - def _packed_layout_datums(lengths: list[int]) -> list[Datum]: """Build flat Datums whose three loss layouts are easy to distinguish.""" datums = [] From e5af34261b4f4b3368d57e8cfb6269dfac788dee Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 01:08:41 -0700 Subject: [PATCH 13/34] feat(engine): restore context-parallel token outputs Signed-off-by: HuiyingLi --- nemo_automodel/__init__.py | 2 + .../distributed/context_parallel/sharder.py | 44 +- .../distributed/context_parallel/utils.py | 73 +- nemo_automodel/engine/__init__.py | 1001 +++++++++++++++-- nemo_automodel/engine/outputs.py | 149 +++ .../L2_PP_CP_Dense_Packed_Test.sh | 23 + .../context_parallel/run_packed_pp.py | 194 +++- .../context_parallel/test_context_parallel.py | 10 + .../unit_tests/distributed/test_cp_sharder.py | 66 ++ tests/unit_tests/distributed/test_cp_utils.py | 88 ++ tests/unit_tests/test_engine.py | 409 ++++++- tests/unit_tests/test_engine_outputs.py | 109 ++ 12 files changed, 1976 insertions(+), 192 deletions(-) create mode 100644 nemo_automodel/engine/outputs.py create mode 100755 tests/functional_tests/context_parallel/L2_PP_CP_Dense_Packed_Test.sh create mode 100644 tests/unit_tests/test_engine_outputs.py diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 3a0a69d04a..489c6c1624 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -44,6 +44,7 @@ "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"), @@ -60,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"), diff --git a/nemo_automodel/components/distributed/context_parallel/sharder.py b/nemo_automodel/components/distributed/context_parallel/sharder.py index 1c3b09c125..61929f303d 100644 --- a/nemo_automodel/components/distributed/context_parallel/sharder.py +++ b/nemo_automodel/components/distributed/context_parallel/sharder.py @@ -228,6 +228,10 @@ class ShardLayout: input_token_stream_positions: For layouts that reposition tokens (DSV4 packed repad), the per-row map from input position to padded output column (-1 = an input pad slot whose token was dropped). + chunk_layouts: Per-chunk layouts when a backend prepares multiple + independent token streams for pipeline microbatches. Each child + uses its own token coordinate space; callers select one explicitly + with ``gather_token_tensor(..., chunk_index=index)``. """ local_token_global_indices: torch.Tensor | None = None @@ -235,6 +239,7 @@ class ShardLayout: padded_seq_len: int | None = None input_row_shape: tuple[int, ...] | None = None input_token_stream_positions: torch.Tensor | None = None + chunk_layouts: tuple["ShardLayout", ...] | None = None class ContextParallelSharder: @@ -364,8 +369,14 @@ def shard(self, batch: dict[str, Any]) -> tuple[Callable, dict[str, Any]]: ) return ctx, batch - def _indices(self, padded_seq_len: int, device) -> torch.Tensor: - layout = self.shard_layout or _NO_SHARD_LAYOUT + def _indices( + self, + padded_seq_len: int, + device, + *, + layout: "ShardLayout | None" = None, + ) -> torch.Tensor: + layout = layout or self.shard_layout or _NO_SHARD_LAYOUT captured = layout.local_token_global_indices if captured is not None: # Data-dependent layout: use the partition the shard reported, and @@ -447,6 +458,7 @@ def gather_token_tensor( seq_dim: int = 1, trim: bool = False, fill: float | int | None = None, + chunk_index: int | None = None, ) -> torch.Tensor: """Differentiably gather a token-aligned local shard to global order. @@ -456,10 +468,23 @@ def gather_token_tensor( reported position map (``fill`` for input positions whose tokens were dropped, e.g. re-padded pack slots). Raises when no layout is present (nothing to trim to). + + A chunked THD batch contains one independent token coordinate space per + pipeline microbatch. Such layouts require ``chunk_index`` so the gather + uses the exact partition reported for that chunk; a top-level chunked + layout is never flattened or guessed. """ layout = self.shard_layout or _NO_SHARD_LAYOUT + if layout.chunk_layouts is not None: + if chunk_index is None: + raise ValueError("chunk_index is required for a chunked CP token layout") + if chunk_index < 0 or chunk_index >= len(layout.chunk_layouts): + raise IndexError(f"chunk_index {chunk_index} is out of range for {len(layout.chunk_layouts)} CP chunks") + layout = layout.chunk_layouts[chunk_index] + elif chunk_index is not None: + raise ValueError("chunk_index is only valid for a chunked CP token layout") padded_seq_len = tensor.shape[seq_dim] * (self._cp_mesh.size() if self._cp_mesh is not None else 1) - indices = self._indices(padded_seq_len, tensor.device) + indices = self._indices(padded_seq_len, tensor.device, layout=layout) full = gather_token_tensor_by_indices(self._cp_mesh, tensor, indices, seq_dim=seq_dim) if not trim: return full @@ -473,8 +498,17 @@ def gather_token_tensor( if fill is None: raise ValueError("trimming to input coordinates on a repositioned layout requires `fill`") positions = layout.input_token_stream_positions.to(full.device) - out = full.gather(1, positions.clamp(min=0).to(torch.long)) - return out.masked_fill(positions < 0, fill) + gather_indices = positions.clamp(min=0).to(torch.long) + padding_mask = positions < 0 + while gather_indices.ndim < full.ndim: + gather_indices = gather_indices.unsqueeze(-1) + padding_mask = padding_mask.unsqueeze(-1) + trailing_shape = full.shape[2:] + if trailing_shape: + gather_indices = gather_indices.expand(*positions.shape, *trailing_shape) + padding_mask = padding_mask.expand_as(gather_indices) + out = full.gather(1, gather_indices) + return out.masked_fill(padding_mask, fill) if layout.input_row_shape is not None: if layout.original_seq_len is not None: full = full.narrow(seq_dim, 0, layout.original_seq_len) diff --git a/nemo_automodel/components/distributed/context_parallel/utils.py b/nemo_automodel/components/distributed/context_parallel/utils.py index 5e77187818..beddad4176 100644 --- a/nemo_automodel/components/distributed/context_parallel/utils.py +++ b/nemo_automodel/components/distributed/context_parallel/utils.py @@ -559,12 +559,18 @@ def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_ if is_thd: # The THD partition is data-dependent (cu_seqlens), so shard_batch # installs the index map it just computed on the sharder for the token - # verbs (chunked streams carry none). The BSHD->THD flatten is a pure - # reshape, so the pre-flatten row shape is the caller's coordinate - # system and the stream length is rows x cols. + # verbs. Chunked streams retain one independent index map per pipeline + # microbatch. The BSHD->THD flatten is a pure reshape, so the + # pre-flatten row shape is the caller's coordinate system and the stream + # length is rows x cols. def _shard_batch_te(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id=0): - input_ids = batch.get("input_ids") - row_shape = tuple(input_ids.shape[:2]) if input_ids is not None and input_ids.dim() >= 2 else None + final_thd = "cu_seqlens" in batch and "seq_lens" not in batch and "seq_lens_padded" not in batch + primary = batch.get("inputs_embeds", batch.get("input_ids")) + row_shape = ( + tuple(primary.shape[:2]) + if isinstance(primary, torch.Tensor) and primary.dim() >= 2 and not final_thd + else None + ) prepped, local_indices = make_cp_batch_for_te( cp_mesh, batch, @@ -575,10 +581,24 @@ def _shard_batch_te(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id return_local_indices=True, ) layout = None - if local_indices is not None: + if isinstance(local_indices, tuple): + cp_size = cp_mesh.size() if cp_mesh is not None else 1 + layout = ShardLayout( + chunk_layouts=tuple( + ShardLayout( + local_token_global_indices=indices, + original_seq_len=indices.numel() * cp_size, + padded_seq_len=indices.numel() * cp_size, + ) + for indices in local_indices + ) + ) + elif local_indices is not None: + stream_len = local_indices.numel() * (cp_mesh.size() if cp_mesh is not None else 1) layout = ShardLayout( local_token_global_indices=local_indices, - padded_seq_len=row_shape[0] * row_shape[1] if row_shape is not None else None, + original_seq_len=stream_len if final_thd else None, + padded_seq_len=stream_len if final_thd or row_shape is not None else None, input_row_shape=row_shape, ) return contextlib.nullcontext, prepped, layout @@ -731,13 +751,15 @@ def make_cp_batch_for_te( seq_lens/seq_lens_padded tensors (default: -1000) return_local_indices (bool): Also return this rank's local-token global index map (the ``thd_get_partitioned_indices`` partition; an - identity arange when CP is inactive; None in chunked mode, where - each chunk is its own token space). Used by the THD ContextParallelSharder's - token verbs. + identity arange when CP is inactive). Chunked mode returns one map + per chunk because each chunk has its own token coordinate space. + Used by the THD ContextParallelSharder's token verbs. Returns: - dict: Processed batch in THD format (or ``(dict, LongTensor | None)`` - when ``return_local_indices``) with the following keys: + dict: Processed batch in THD format. With ``return_local_indices``, the + return value is ``(dict, LongTensor)`` for one stream or + ``(dict, tuple[LongTensor, ...])`` for chunked streams. The batch has + the following keys: - input_ids: Sharded input token IDs [total_tokens] or [num_chunks, chunk_tokens] - labels: Sharded labels [total_tokens] or [num_chunks, chunk_tokens] - position_ids: Generated and sharded position IDs [total_tokens] or [num_chunks, chunk_tokens] @@ -794,14 +816,18 @@ def make_cp_batch_for_te( if cp_mesh is None or cp_mesh.size() <= 1: if not return_local_indices: return batch - # Unsharded THD stream: identity index map. Chunked streams are - # per-chunk token spaces with no single step-wide map -> None. + # Unsharded THD streams use identity maps. Chunked streams retain one + # independent identity map per chunk rather than pretending there is a + # single step-wide token space. primary = batch.get("inputs_embeds", batch.get("input_ids")) if not isinstance(primary, torch.Tensor): raise ValueError("THD batch requires tensor input_ids or inputs_embeds") - local_indices = ( - torch.arange(primary.shape[0], device=primary.device, dtype=torch.long) if num_chunks <= 1 else None - ) + if num_chunks <= 1: + local_indices = torch.arange(primary.shape[0], device=primary.device, dtype=torch.long) + else: + local_indices = tuple( + torch.arange(primary.shape[1], device=primary.device, dtype=torch.long) for _ in range(num_chunks) + ) return batch, local_indices if num_chunks <= 1: @@ -812,6 +838,7 @@ def make_cp_batch_for_te( # Extract each chunk from the batched result and shard it chunks = [] + chunk_local_indices = [] for i in range(num_chunks): chunk_batch = { key: value[i] @@ -819,14 +846,16 @@ def make_cp_batch_for_te( else value for key, value in batch.items() } - chunks.append( - _shard_thd_chunk_for_te(chunk_batch, cp_mesh, qkv_format, seq_lens_padding_value, padding_token_id)[0] + chunk, local_indices = _shard_thd_chunk_for_te( + chunk_batch, cp_mesh, qkv_format, seq_lens_padding_value, padding_token_id ) + chunks.append(chunk) + chunk_local_indices.append(local_indices) return_dict = stack_thd_chunks(chunks, seq_lens_padding_value) - # Chunked mode: each chunk is its own token space, so there is no single - # step-wide local-token index map to expose. - return (return_dict, None) if return_local_indices else return_dict + # Chunked mode: each chunk is its own token space, so expose one partition + # map per chunk rather than flattening them into a misleading step-wide map. + return (return_dict, tuple(chunk_local_indices)) if return_local_indices else return_dict def _shard_thd_chunk_for_te( diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index fe3805366f..14fbdb84e9 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -16,9 +16,12 @@ from __future__ import annotations +import hashlib +import pickle from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass +from math import prod from typing import Any, TypeVar import torch @@ -47,6 +50,7 @@ scale_grads_and_clip_grad_norm, ) from nemo_automodel.components.utils.model_utils import filter_forward_kwargs +from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput CollateFn = Callable[ [list[Datum]], @@ -56,14 +60,22 @@ LossInputs = dict[str, LossInputValue] LossFn = Callable[ [Any, LossInputs], - torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], + torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]] | LossFnOutputBatch], ] +ParsedLossOutputs = list[dict[str, Any]] | LossFnOutputBatch | None _LOSS_FIELD_PREFIX = "__engine_loss__" _LOSS_METADATA = ("cu_seqlens", "cu_seqlens_padded", "max_seqlen", "padding_mask") _T = TypeVar("_T") -__all__ = ["Engine", "ForwardResult", "OptimStepResult", "collate_prebatched"] +__all__ = [ + "Engine", + "ForwardResult", + "LossFnOutputBatch", + "OptimStepResult", + "PerTokenOutput", + "collate_prebatched", +] def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextManager[Any]: @@ -83,7 +95,10 @@ def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLos The Datum represents the whole prebatched item. Consequently, one optional ``loss_fn_output`` mapping also describes that whole batch, not - each sample inside it. + each sample inside it. A typed token output likewise remains one record + containing the full inner ``[B, S, ...]`` (or flat THD) tensor; the Engine + cannot split hidden inner samples. Output records are rejected when PP + divides that prebatched Datum into multiple inner microbatches. Args: datums: A one-item list whose model and loss tensor fields already have @@ -119,6 +134,18 @@ class _LossBatchLayout: unresolved_fields: frozenset[str] = frozenset() +@dataclass(frozen=True) +class _OutputRestorePlan: + """How to turn one CP-local callback token stream back into Datums.""" + + sharder: ContextParallelSharder + is_thd: bool + item_to_datum: tuple[int, ...] | None + real_lengths: tuple[int, ...] | None + padded_lengths: tuple[int, ...] | None + token_mask: torch.Tensor | None + + @dataclass(frozen=True) class ForwardResult: """Forward-only loss statistics and per-Datum outputs. @@ -129,8 +156,10 @@ class ForwardResult: at the end of an evaluation epoch. Distributed model wrappers may still communicate during forward and therefore retain their own call-alignment requirements. ``loss_fn_outputs`` remains local to the replica's input - Datums; PP stages receive identical detached copies, while any CP-local - tensor layout inside those mappings remains caller defined. + Datums. PP stages receive identical detached copies. Legacy output mappings + remain opaque and therefore CP-local; fields declared through + :class:`LossFnOutputBatch` are restored to full token order across CP before + being split back into per-Datum records. Attributes: loss_sum: Detached weighted numerator for this Datum window. @@ -287,12 +316,16 @@ def forward( datums: Flat sequence of Datum items for this forward-only window. loss_fn: Computes a per-element loss tensor or scalar local weighted numerator from the raw model output and CP-local loss - inputs. It may additionally return one output mapping per - Datum. + inputs. It may additionally return opaque mappings per Datum, + or a :class:`LossFnOutputBatch` whose explicitly typed token + streams the Engine restores across CP. Returns: Detached model-parallel-complete loss statistics and per-Datum - outputs. The sums are local to one data-parallel replica. + outputs. The sums and outputs are local to one data-parallel + replica. A prebatched Datum produces one outer record, not one + record per hidden inner sample, and cannot produce records when PP + splits it into multiple inner microbatches. """ microbatches = self._group_datums(datums) self._validate_execution_parallelism() @@ -310,11 +343,11 @@ def forward( returns_outputs: bool | None = None for batch_datums in microbatches: - cp_context, model_inputs, loss_inputs, loss_batch_layout = self._prepare_batch( + cp_context, model_inputs, loss_inputs, loss_batch_layout, output_restore_plan = self._prepare_batch( batch_datums, inner_microbatches ) if self.pipeline is not None: - batch_returns_outputs, batch_outputs = self._pipeline_execute( + batch_returns_outputs, batch_outputs, batch_error = self._pipeline_execute( model_inputs, loss_inputs, batch_datums, @@ -322,9 +355,12 @@ def forward( local_loss_sum, cp_context, loss_batch_layout, + output_restore_plan, backward_scale=None, zero_weight_sum=zero_weight_sum, ) + if batch_error is not None: + raise batch_error if batch_returns_outputs is not None: if returns_outputs is None: returns_outputs = batch_returns_outputs @@ -337,14 +373,31 @@ def forward( with self.context_fn(model_inputs), cp_context(): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - numerator, outputs = _parse_loss_result(loss_fn(output, loss_inputs), loss_inputs["weights"]) - returns_outputs = _update_output_mode(returns_outputs, outputs) - if outputs is not None: - if len(outputs) != len(batch_datums): - raise ValueError("loss_fn outputs must contain one mapping per Datum") - loss_fn_outputs.extend(outputs) + numerator, parsed_outputs, output_parse_error = _parse_loss_result( + loss_fn(output, loss_inputs), loss_inputs["weights"] + ) + self._validate_loss_fn_outputs_across_cp( + parsed_outputs, + loss_inputs.get("weights"), + expected_records=len(batch_datums), + local_error=output_parse_error, + restore_plan=output_restore_plan, + datum_indices=tuple(range(len(batch_datums))), + ) + returns_outputs = _update_output_mode(returns_outputs, parsed_outputs) if zero_weight_sum: numerator = numerator * 0 + if parsed_outputs is not None: + outputs = self._restore_loss_fn_outputs( + parsed_outputs, + loss_inputs, + output_restore_plan, + datum_indices=tuple(range(len(batch_datums))), + chunk_index=None, + ) + if len(outputs) != len(batch_datums): + raise ValueError("loss_fn outputs must contain one mapping per Datum") + loss_fn_outputs.extend(outputs) local_loss_sum.add_(numerator.detach().to(torch.float64)) if cp_size > 1: @@ -379,11 +432,14 @@ def forward_backward( ``loss_fn_inputs["weights"]``, or a scalar local weighted-sum numerator. For a scalar, the callback must apply weights and masks; the Engine will only apply global normalization. The callback may - also return one output mapping per Datum. - Those mappings are detached and preserved in input order; the Engine - deliberately does not interpret or reduce them. Under pipeline - parallelism the last stage computes the mappings and broadcasts them - to every stage in the pipeline group. + also return one opaque output mapping per Datum, or a + :class:`LossFnOutputBatch`. Explicit token streams in that envelope are + detached, restored from CP-local to full token order, and then split + back into per-Datum records; ordinary mappings remain opaque. Under + pipeline parallelism the last stage performs CP restoration before it + broadcasts the records to every stage in the pipeline group. + Values in pipeline output records must therefore be pickle-compatible; + large records also incur CPU serialization and PP broadcast cost. A prebatched Datum intentionally hides its inner sample boundaries. It can therefore return a single output mapping only when a pipeline @@ -403,10 +459,11 @@ def forward_backward( ``(loss, loss_fn_outputs)``. ``loss`` is a detached scalar reduced over the DP-CP gradient group and, for pipeline execution, synchronized across PP stages. ``loss_fn_outputs`` contains - per-Datum mappings in window order; pipeline execution returns the - same mappings on every physical stage rank. Model parameters are - unchanged, but their gradients contain the complete window's - globally normalized backward result. + mappings for this DP replica's outer Datums in window order; + pipeline execution returns the same mappings on every physical + stage rank in that replica. Model parameters are unchanged, but + their gradients contain the complete window's globally normalized + backward result. """ microbatches = self._group_datums(datums) self._validate_parallelism() @@ -430,13 +487,16 @@ def forward_backward( local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) loss_fn_outputs: list[dict[str, Any]] = [] returns_outputs: bool | None = None + output_error: Exception | None = None for index, datums in enumerate(microbatches): is_last = index == len(microbatches) - 1 if is_last: prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) - cp_context, model_inputs, loss_inputs, loss_batch_layout = self._prepare_batch(datums, inner_microbatches) + cp_context, model_inputs, loss_inputs, loss_batch_layout, output_restore_plan = self._prepare_batch( + datums, inner_microbatches + ) if self.pipeline is not None: backward_scale = ( @@ -444,7 +504,7 @@ def forward_backward( if zero_denominator else safe_denominator.new_tensor(grad_group_size) / safe_denominator ) - batch_returns_outputs, batch_outputs = self._pipeline_execute( + batch_returns_outputs, batch_outputs, batch_error = self._pipeline_execute( model_inputs, loss_inputs, datums, @@ -452,15 +512,25 @@ def forward_backward( local_loss_sum, cp_context, loss_batch_layout, + output_restore_plan, backward_scale=backward_scale, zero_weight_sum=zero_denominator, ) - if batch_returns_outputs is not None: - if returns_outputs is None: - returns_outputs = batch_returns_outputs - elif returns_outputs != batch_returns_outputs: - raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") - loss_fn_outputs.extend(batch_outputs) + if output_error is None: + if batch_error is not None: + output_error = batch_error + else: + try: + if batch_returns_outputs is not None: + if returns_outputs is None: + returns_outputs = batch_returns_outputs + elif returns_outputs != batch_returns_outputs: + raise ValueError( + "loss_fn must return per-Datum outputs for every microbatch or none of them" + ) + loss_fn_outputs.extend(batch_outputs) + except Exception as error: + output_error = error else: loss_inputs = _with_loss_metadata(model_inputs, loss_inputs) with ( @@ -470,27 +540,67 @@ def forward_backward( ): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - numerator, outputs = _parse_loss_result(loss_fn(output, loss_inputs), loss_inputs["weights"]) - returns_outputs = _update_output_mode(returns_outputs, outputs) - if outputs is not None: - if len(outputs) != len(datums): - raise ValueError("loss_fn outputs must contain one mapping per Datum") - loss_fn_outputs.extend(outputs) + numerator, parsed_outputs, output_parse_error = _parse_loss_result( + loss_fn(output, loss_inputs), loss_inputs["weights"] + ) + if output_error is None and dp_size <= 1: + self._validate_loss_fn_outputs_across_cp( + parsed_outputs, + loss_inputs.get("weights"), + expected_records=len(datums), + local_error=output_parse_error, + restore_plan=output_restore_plan, + datum_indices=tuple(range(len(datums))), + ) + returns_outputs = _update_output_mode(returns_outputs, parsed_outputs) if zero_denominator: numerator = numerator * 0 (numerator * (grad_group_size / safe_denominator)).backward() + if output_error is None: + try: + if dp_size > 1: + self._validate_loss_fn_outputs_across_cp( + parsed_outputs, + loss_inputs.get("weights"), + expected_records=len(datums), + local_error=output_parse_error, + restore_plan=output_restore_plan, + datum_indices=tuple(range(len(datums))), + ) + returns_outputs = _update_output_mode(returns_outputs, parsed_outputs) + if parsed_outputs is not None: + outputs = self._restore_loss_fn_outputs( + parsed_outputs, + loss_inputs, + output_restore_plan, + datum_indices=tuple(range(len(datums))), + chunk_index=None, + ) + if len(outputs) != len(datums): + raise ValueError("loss_fn outputs must contain one mapping per Datum") + loss_fn_outputs.extend(outputs) + except Exception as error: + output_error = error local_loss_sum.add_(numerator.detach().to(torch.float64)) if index == 0: prepare_after_first_microbatch() + # Piggyback the output-error bit on the existing end-of-window loss + # reductions. Every gradient rank therefore finishes backward before + # any replica raises a data-local output-routing error. + step_state = torch.stack((local_loss_sum, local_loss_sum.new_tensor(int(output_error is not None)))) if grad_group_size > 1: - dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=grad_group) + dist.all_reduce(step_state, op=dist.ReduceOp.SUM, group=grad_group) pp_group, pp_size = self._pp_group_and_size() if pp_size > 1: - dist.all_reduce(local_loss_sum, op=dist.ReduceOp.SUM, group=pp_group) + dist.all_reduce(step_state, op=dist.ReduceOp.SUM, group=pp_group) + if bool(step_state[1] > 0): + if output_error is not None: + raise output_error + raise RuntimeError("another data-parallel replica failed while restoring loss_fn outputs") - loss = (local_loss_sum / safe_denominator).detach() + loss = (step_state[0] / safe_denominator).detach() return loss, loss_fn_outputs @torch.no_grad() @@ -659,7 +769,13 @@ def _prepare_batch( self, datums: list[Datum], num_pipeline_microbatches: int, - ) -> tuple[Callable[[], AbstractContextManager[Any]], dict[str, Any], LossInputs, _LossBatchLayout]: + ) -> tuple[ + Callable[[], AbstractContextManager[Any]], + dict[str, Any], + LossInputs, + _LossBatchLayout, + _OutputRestorePlan, + ]: """Collate, move, and CP-shard one outer batch. Args: @@ -670,7 +786,8 @@ def _prepare_batch( Returns: The CP context factory, CP-local model inputs, CP-local loss - inputs, and collated field-layout metadata. Token-aligned model + inputs, collated field-layout metadata, and an internal plan for + restoring explicitly declared token outputs. Token-aligned model and loss tensors use the same padded, packed THD, Magi, or model-owned local sequence layout. """ @@ -691,6 +808,17 @@ def _prepare_batch( loss_inputs, loss_batch_layout.fields, ) + output_routing = ( + _output_sequence_lengths( + datums, + model_inputs, + loss_inputs, + loss_batch_layout.item_to_datum, + is_thd=model_inputs.get("qkv_format") == "thd", + ) + if loss_seq_dim is not None and weight_layout is LossInputLayout.PER_TOKEN + else (None, None, None) + ) model_inputs = _to_device(model_inputs, self.device) loss_inputs = _to_device(loss_inputs, self.device) token_reference = ( @@ -801,7 +929,21 @@ def _prepare_batch( item_to_datum=loss_batch_layout.item_to_datum, unresolved_fields=loss_batch_layout.unresolved_fields, ) - return cp_context, model_inputs, loss_inputs, loss_batch_layout + real_lengths, padded_lengths, token_mask = output_routing + return ( + cp_context, + model_inputs, + loss_inputs, + loss_batch_layout, + _OutputRestorePlan( + sharder=sharder, + is_thd=is_thd, + item_to_datum=loss_batch_layout.item_to_datum, + real_lengths=real_lengths, + padded_lengths=padded_lengths, + token_mask=token_mask, + ), + ) def _prepare_mtp_cp_inputs(self, batch: dict[str, Any]) -> MTPContextParallelInputs | None: """Prepare global MTP future-token tensors before CP shards the batch. @@ -893,10 +1035,11 @@ def _pipeline_execute( local_loss_sum: torch.Tensor, cp_context: Callable[[], AbstractContextManager[Any]], loss_batch_layout: _LossBatchLayout, + output_restore_plan: _OutputRestorePlan, *, backward_scale: torch.Tensor | None, zero_weight_sum: bool, - ) -> tuple[bool | None, list[dict[str, Any]]]: + ) -> tuple[bool | None, list[dict[str, Any]], Exception | None]: """Run prepared pipeline microbatches in training or forward-only mode. Args: @@ -908,17 +1051,23 @@ def _pipeline_execute( cp_context: Context covering the complete pipeline schedule. loss_batch_layout: Semantic layout of every loss field plus the collater's logical item-to-Datum routing, when available. + output_restore_plan: Captured full-token routing and CP sharder for + explicit per-token callback outputs. backward_scale: Multiplier returned to the training schedule for backward, or ``None`` to run the forward-only schedule. zero_weight_sum: Whether reporting numerators must be forced to graph-connected zero. Returns: - Whether the callback returned outputs, and its detached outputs in - logical Datum order. + Whether the callback returned outputs, its detached outputs in + logical Datum order, and any post-schedule output-restoration + error synchronized across PP stages. """ outputs_by_microbatch: list[list[dict[str, Any]] | None] = [None] * self.pipeline.num_microbatches outputs_by_datum: list[dict[str, Any] | None] = [None] * len(datums) + parsed_outputs_by_microbatch: list[ParsedLossOutputs] = [None] * self.pipeline.num_microbatches + output_parse_errors_by_microbatch: list[Exception | None] = [None] * self.pipeline.num_microbatches + loss_called_by_microbatch = [False] * self.pipeline.num_microbatches returns_outputs: bool | None = None with cp_context(): @@ -952,37 +1101,15 @@ def _pipeline_execute( ) def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: - nonlocal returns_outputs loss_inputs_mb = loss_microbatches[microbatch_index] - numerator, batch_outputs = _parse_loss_result( + numerator, batch_outputs, output_parse_error = _parse_loss_result( loss_fn(output, loss_inputs_mb), loss_inputs_mb["weights"] ) - returns_outputs = _update_output_mode(returns_outputs, batch_outputs) - if batch_outputs is not None: - datum_indices = ( - None - if datum_indices_by_microbatch is None - else datum_indices_by_microbatch[microbatch_index] - ) - if datum_indices is None and len(datums) == 1 and self.pipeline.num_microbatches > 1: - raise ValueError( - "a prebatched Datum may return outputs only when num_microbatches=1 because " - "its inner sample boundaries are not part of the Datum contract" - ) - if datum_indices is not None: - if len(batch_outputs) != len(datum_indices): - raise ValueError( - f"pipeline loss_fn returned {len(batch_outputs)} outputs for microbatch " - f"{microbatch_index}, expected {len(datum_indices)} from its Datum mapping" - ) - for datum_index, item in zip(datum_indices, batch_outputs): - if outputs_by_datum[datum_index] is not None: - raise RuntimeError( - f"pipeline returned more than one output for Datum {datum_index}" - ) - outputs_by_datum[datum_index] = item - else: - outputs_by_microbatch[microbatch_index] = batch_outputs + if loss_called_by_microbatch[microbatch_index]: + raise RuntimeError(f"pipeline evaluated loss_fn twice for microbatch {microbatch_index}") + loss_called_by_microbatch[microbatch_index] = True + parsed_outputs_by_microbatch[microbatch_index] = batch_outputs + output_parse_errors_by_microbatch[microbatch_index] = output_parse_error if zero_weight_sum: numerator = numerator * 0 local_loss_sum.add_(numerator.detach().to(torch.float64)) @@ -1000,24 +1127,261 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: ) outputs: list[dict[str, Any]] = [] - if self.pipeline.info.has_last_stage and returns_outputs: - if datum_indices_by_microbatch is not None: - if any(item is None for item in outputs_by_datum): - raise RuntimeError("pipeline schedule did not return exactly one output for every Datum") - outputs = [item for item in outputs_by_datum if item is not None] - else: - if any(items is None for items in outputs_by_microbatch): - raise RuntimeError("pipeline schedule did not evaluate loss_fn for every logical microbatch") - outputs = [item for items in outputs_by_microbatch if items is not None for item in items] - if len(outputs) != len(datums): - raise ValueError( - f"pipeline loss_fn returned {len(outputs)} outputs across the outer batch, " - f"expected one for each of its {len(datums)} Datums" + serialized_outputs: bytes | None = None + output_error: Exception | None = None + if self.pipeline.info.has_last_stage: + try: + for microbatch_index, batch_outputs in enumerate(parsed_outputs_by_microbatch): + local_error = output_parse_errors_by_microbatch[microbatch_index] + if not loss_called_by_microbatch[microbatch_index]: + local_error = RuntimeError( + f"pipeline did not evaluate loss_fn for microbatch {microbatch_index}" + ) + weights = loss_microbatches[microbatch_index].get("weights") + datum_indices = ( + None if datum_indices_by_microbatch is None else datum_indices_by_microbatch[microbatch_index] + ) + expected_records = ( + len(datum_indices) + if datum_indices is not None + else (1 if len(datums) == 1 and self.pipeline.num_microbatches == 1 else None) + ) + self._validate_loss_fn_outputs_across_cp( + batch_outputs, + weights, + expected_records=expected_records, + microbatch_index=microbatch_index, + local_error=local_error, + restore_plan=output_restore_plan, + datum_indices=datum_indices, + chunk_index=(microbatch_index if is_thd and self.pipeline.num_microbatches > 1 else None), ) - outputs = self._broadcast_pipeline_outputs(outputs) - return bool(outputs), outputs + returns_outputs = _update_output_mode(returns_outputs, batch_outputs) + if batch_outputs is None: + continue + restored = ( + batch_outputs + if isinstance(batch_outputs, list) + else self._restore_loss_fn_outputs( + batch_outputs, + loss_microbatches[microbatch_index], + output_restore_plan, + datum_indices=datum_indices, + chunk_index=(microbatch_index if is_thd and self.pipeline.num_microbatches > 1 else None), + ) + ) + if datum_indices is None: + outputs_by_microbatch[microbatch_index] = restored + continue + if len(restored) != len(datum_indices): + raise RuntimeError("restored token outputs do not match their pipeline Datum routing") + for datum_index, item in zip(datum_indices, restored): + if outputs_by_datum[datum_index] is not None: + raise RuntimeError(f"pipeline returned more than one output for Datum {datum_index}") + outputs_by_datum[datum_index] = item + if returns_outputs: + if datum_indices_by_microbatch is not None: + if any(item is None for item in outputs_by_datum): + raise RuntimeError("pipeline schedule did not return exactly one output for every Datum") + outputs = [item for item in outputs_by_datum if item is not None] + else: + if any(items is None for items in outputs_by_microbatch): + raise RuntimeError( + "pipeline schedule did not evaluate loss_fn for every logical microbatch" + ) + outputs = [item for items in outputs_by_microbatch if items is not None for item in items] + if len(outputs) != len(datums): + raise ValueError( + f"pipeline loss_fn returned {len(outputs)} outputs across the outer batch, " + f"expected one for each of its {len(datums)} Datums" + ) + if outputs: + # broadcast_object_list serializes only on the source + # stage. Preflight while errors can still be propagated to + # every PP stage instead of leaving peers in the broadcast. + # Broadcast these already validated bytes so arbitrary + # record objects are never pickled again inside a PP + # collective. + serialized_outputs = pickle.dumps(_to_device(outputs, torch.device("cpu"))) + except Exception as error: + output_error = error + output_error = self._synchronize_pipeline_output_error(output_error) + if output_error is None: + outputs = self._broadcast_pipeline_outputs(outputs, serialized_outputs=serialized_outputs) + return bool(outputs), outputs, output_error + + def _restore_loss_fn_outputs( + self, + outputs: list[dict[str, Any]] | LossFnOutputBatch, + loss_inputs: LossInputs, + plan: _OutputRestorePlan, + *, + datum_indices: tuple[int, ...] | None, + chunk_index: int | None, + ) -> list[dict[str, Any]]: + """Restore explicitly typed token streams and merge per-Datum records.""" + weights = loss_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("loss_fn token outputs require Tensor loss weights") + if isinstance(outputs, list): + return outputs + + if datum_indices is None: + if plan.item_to_datum is not None: + raise RuntimeError("loss_fn output routing is missing the collater's Datum mapping") + if chunk_index is not None: + raise NotImplementedError( + "a prebatched Datum cannot restore token outputs across multiple pipeline microbatches" + ) + datum_indices = (0,) + + source_records = ({},) * len(datum_indices) if outputs.per_datum is None else outputs.per_datum + records = [dict(record) for record in source_records] + if len(records) != len(datum_indices): + raise ValueError( + f"LossFnOutputBatch.per_datum contains {len(records)} records, expected {len(datum_indices)}" + ) + record_keys = {key for record in records for key in record} + collisions = record_keys & set(outputs.per_token) + if collisions: + raise ValueError(f"per-token output keys collide with per-Datum records: {sorted(collisions)}") + + restored_fields: list[tuple[str, torch.Tensor, int]] = [] + for name in sorted(outputs.per_token): + spec = outputs.per_token[name] + tensor = spec.tensor.detach() + if tensor.shape[: weights.ndim] != weights.shape: + raise ValueError( + f"per-token output {name!r} must start with the loss weight shape {tuple(weights.shape)}, " + f"got {tuple(tensor.shape)}" + ) + if tensor.device != weights.device: + raise ValueError( + f"per-token output {name!r} must be on the loss weight device {weights.device}, got {tensor.device}" + ) + seq_dim = _loss_sequence_dim_from_weights(weights) + shard_layout = plan.sharder.shard_layout + selected_layout = shard_layout + if chunk_index is not None: + if shard_layout is None or shard_layout.chunk_layouts is None: + if self._cp_size() > 1: + raise NotImplementedError( + "the active context-parallel backend does not report reversible per-pipeline-chunk " + "token layouts; typed per-token outputs are unavailable for this packed PP+CP batch" + ) + selected_layout = None + else: + if chunk_index < 0 or chunk_index >= len(shard_layout.chunk_layouts): + raise IndexError( + f"pipeline chunk {chunk_index} is out of range for " + f"{len(shard_layout.chunk_layouts)} reported CP layouts" + ) + selected_layout = shard_layout.chunk_layouts[chunk_index] + if self._cp_size() > 1 and selected_layout is None: + raise NotImplementedError( + "the active context-parallel backend does not report a reversible token layout; " + "typed per-token outputs are unavailable for this batch" + ) + if selected_layout is not None: + tensor = plan.sharder.gather_token_tensor( + tensor, + seq_dim=seq_dim, + trim=True, + fill=spec.fill_value, + chunk_index=chunk_index, + ) + if selected_layout.input_row_shape is not None: + seq_dim = len(selected_layout.input_row_shape) - 1 + restored_fields.append((name, tensor, seq_dim)) + + # Complete every field's CP collective before Datum-local splitting. + # A routing error can then no longer leave a peer entering the next + # field gather while this rank exits early. + for name, tensor, seq_dim in restored_fields: + pieces = _split_restored_token_output( + tensor, + plan, + datum_indices=datum_indices, + seq_dim=seq_dim, + ) + if len(pieces) != len(records): + raise RuntimeError(f"restored per-token output {name!r} does not match its Datum records") + for record, piece in zip(records, pieces): + record[name] = piece.detach() + return records + + def _validate_loss_fn_outputs_across_cp( + self, + outputs: ParsedLossOutputs, + weights: Any, + *, + expected_records: int | None, + microbatch_index: int | None = None, + local_error: Exception | None = None, + restore_plan: _OutputRestorePlan | None = None, + datum_indices: tuple[int, ...] | None = None, + chunk_index: int | None = None, + ) -> None: + """Validate output routing and reach CP consensus before token gathers.""" + if local_error is None: + error, schema = _loss_fn_output_contract( + outputs, + weights, + expected_records=expected_records, + microbatch_index=microbatch_index, + ) + else: + error, schema = str(local_error), ("invalid-output", type(local_error).__name__) + if error is None and isinstance(outputs, LossFnOutputBatch) and restore_plan is not None: + restore_error, restore_schema = _loss_fn_output_restore_contract( + outputs, + weights, + restore_plan, + datum_indices=datum_indices, + chunk_index=chunk_index, + cp_size=self._cp_size(), + ) + error = restore_error + schema = (*schema, restore_schema) + cp_group, cp_size = self._cp_group_and_size() + if cp_size <= 1 or not (dist.is_available() and dist.is_initialized()): + if error is not None: + raise ValueError(error) + return + digest = int.from_bytes(hashlib.sha256(repr(schema).encode()).digest()[:8], "little") & ((1 << 63) - 1) + local = torch.tensor([int(error is not None), digest], dtype=torch.int64, device=self.device) + gathered = torch.empty(cp_size * 2, dtype=torch.int64, device=self.device) + dist.all_gather_into_tensor(gathered, local, group=cp_group) + gathered = gathered.view(cp_size, 2) + if bool((gathered[:, 0] != 0).any()): + detail = f": {error}" if error is not None else "" + raise ValueError(f"invalid loss_fn outputs on one or more context-parallel ranks{detail}") + if not bool((gathered[:, 1] == gathered[0, 1]).all()): + raise ValueError("context-parallel ranks returned different loss output schemas") + + def _synchronize_pipeline_output_error(self, error: Exception | None) -> Exception | None: + """Propagate output failures across CP lanes and then PP stages.""" + failed = torch.tensor(int(error is not None), dtype=torch.int64, device=self.device) + cp_group, cp_size = self._cp_group_and_size() + if cp_size > 1: + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=cp_group) + if bool(failed.item()) and error is None: + error = RuntimeError("another context-parallel rank failed while restoring loss_fn outputs") + + pp_group, pp_size = self._pp_group_and_size() + if pp_size <= 1: + return error + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=pp_group) + if not bool(failed.item()): + return None + return error or RuntimeError("pipeline last stage failed while restoring loss_fn outputs") - def _broadcast_pipeline_outputs(self, outputs: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _broadcast_pipeline_outputs( + self, + outputs: list[dict[str, Any]], + *, + serialized_outputs: bytes | None, + ) -> list[dict[str, Any]]: """Broadcast last-stage per-Datum outputs to every pipeline stage.""" pp_group, pp_size = self._pp_group_and_size() if pp_size <= 1: @@ -1038,11 +1402,14 @@ def _broadcast_pipeline_outputs(self, outputs: list[dict[str, Any]]) -> list[dic if not bool(stage_states[source_group_rank, 1]): return [] source_global_rank = dist.get_global_rank(pp_group, source_group_rank) - object_list: list[Any] = [ - _to_device(outputs, torch.device("cpu")) if dist.get_rank(group=pp_group) == source_group_rank else None - ] + object_list: list[Any] = [serialized_outputs if dist.get_rank(group=pp_group) == source_group_rank else None] dist.broadcast_object_list(object_list, src=source_global_rank, group=pp_group, device=self.device) - received = object_list[0] + payload = object_list[0] + if not isinstance(payload, bytes): + raise RuntimeError("pipeline output synchronization received an invalid serialized payload") + # The payload was serialized by this Engine's trusted PP source rank + # immediately above; this is not an external deserialization boundary. + received = pickle.loads(payload) # noqa: S301 if not isinstance(received, list) or not all(isinstance(item, dict) for item in received): raise RuntimeError("pipeline output synchronization received invalid per-Datum outputs") return _to_device(received, self.device) @@ -1718,24 +2085,440 @@ def _weighted_numerator(losses: Any, weights: torch.Tensor) -> torch.Tensor: return numerator +def _output_sequence_lengths( + datums: Sequence[Datum], + model_inputs: Mapping[str, Any], + loss_inputs: Mapping[str, LossInputValue], + item_to_datum: tuple[int, ...] | None, + *, + is_thd: bool, +) -> tuple[ + tuple[int, ...] | None, + tuple[int, ...] | None, + torch.Tensor | None, +]: + """Capture real and padded sequence lengths before CP mutates the batch.""" + if item_to_datum is None: + return None, None, None + if item_to_datum != tuple(range(len(datums))): + raise ValueError("token output restoration requires collater items to preserve Datum order") + + weights = loss_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("token output restoration requires Tensor loss weights") + if is_thd: + if isinstance(model_inputs.get("seq_lens"), torch.Tensor): + real_lengths = _valid_row_values(model_inputs["seq_lens"]) + padded_source = model_inputs.get("seq_lens_padded", model_inputs["seq_lens"]) + if not isinstance(padded_source, torch.Tensor): + raise ValueError("packed THD seq_lens_padded must be a Tensor") + padded_lengths = _valid_row_values(padded_source) + else: + cu_seqlens = model_inputs.get("cu_seqlens") + if not isinstance(cu_seqlens, torch.Tensor): + raise ValueError("packed token outputs require seq_lens or cu_seqlens metadata") + real_lengths = _lengths_from_cu_seqlens(cu_seqlens) + padded_cu = model_inputs.get("cu_seqlens_padded", cu_seqlens) + if not isinstance(padded_cu, torch.Tensor): + raise ValueError("packed THD cu_seqlens_padded must be a Tensor") + padded_lengths = _lengths_from_cu_seqlens(padded_cu) + if len(real_lengths) != len(item_to_datum) or len(padded_lengths) != len(item_to_datum): + raise ValueError( + "collater item_to_datum does not match packed token metadata; expected one real and padded " + "length per Datum, " + f"got real={real_lengths}, padded={padded_lengths}, Datums={len(item_to_datum)}" + ) + if any(real < 0 or padded < 0 or real > padded for real, padded in zip(real_lengths, padded_lengths)): + raise ValueError( + "packed token metadata requires 0 <= real_length <= padded_length for every Datum; " + f"got real={real_lengths}, padded={padded_lengths}" + ) + if sum(padded_lengths) != weights.numel(): + raise ValueError( + f"packed padded lengths sum to {sum(padded_lengths)}, but loss weights have token width " + f"{weights.numel()}" + ) + return real_lengths, padded_lengths, None + + if weights.ndim < 2 or weights.shape[0] != len(item_to_datum): + raise ValueError("padded token outputs require weights with one row per Datum") + width = int(weights.shape[1]) + attention_mask = model_inputs.get("attention_mask") + if ( + isinstance(attention_mask, torch.Tensor) + and attention_mask.ndim == 2 + and tuple(attention_mask.shape) == tuple(weights.shape[:2]) + ): + token_mask = attention_mask.to(torch.bool) + # Keep token routing compact. Turning every active position into a + # Python int is prohibitively expensive for long-context batches and + # would penalize ordinary training that never returns token outputs. + real_lengths = tuple(int(length) for length in token_mask.sum(dim=1).tolist()) + else: + inferred: list[int] = [] + for datum in datums: + try: + inferred.append(datum.seq_len) + except ValueError: + inferred.append(width) + real_lengths = tuple(inferred) + token_mask = None + if any(length < 0 or length > width for length in real_lengths): + raise ValueError(f"Datum token lengths must be within padded width {width}; got {real_lengths}") + return real_lengths, (width,) * len(item_to_datum), token_mask + + +def _valid_row_values(values: torch.Tensor, sentinel: int = -1000) -> tuple[int, ...]: + rows = values.reshape(1, -1) if values.ndim == 1 else values.reshape(values.shape[0], -1) + result: list[int] = [] + for row in rows: + result.extend(int(value) for value in row.tolist() if int(value) != sentinel) + return tuple(result) + + +def _lengths_from_cu_seqlens(cu_seqlens: torch.Tensor, sentinel: int = -1000) -> tuple[int, ...]: + rows = cu_seqlens.reshape(1, -1) if cu_seqlens.ndim == 1 else cu_seqlens.reshape(cu_seqlens.shape[0], -1) + result: list[int] = [] + for row in rows: + valid = row[row != sentinel] + if valid.numel() < 2: + continue + lengths = valid[1:] - valid[:-1] + if bool((lengths < 0).any()): + raise ValueError("cu_seqlens must be monotonically non-decreasing within each row") + result.extend(int(value) for value in lengths.tolist()) + return tuple(result) + + +def _loss_sequence_dim_from_weights(weights: torch.Tensor) -> int: + if weights.ndim == 1: + return 0 + if weights.ndim == 2: + return 1 + raise ValueError(f"per-token output weights must be one- or two-dimensional, got {tuple(weights.shape)}") + + +def _split_restored_token_output( + tensor: torch.Tensor, + plan: _OutputRestorePlan, + *, + datum_indices: tuple[int, ...], + seq_dim: int, +) -> list[torch.Tensor]: + """Split one restored collated token tensor into input-Datum coordinates.""" + if plan.item_to_datum is None: + if datum_indices != (0,): + raise NotImplementedError("prebatched token outputs can only produce their single outer Datum record") + return [tensor] + if plan.real_lengths is None or plan.padded_lengths is None: + raise RuntimeError("token output routing metadata is incomplete") + + if plan.is_thd: + expected_width = sum(plan.padded_lengths[index] for index in datum_indices) + if tensor.shape[seq_dim] != expected_width and seq_dim == 1 and tensor.ndim >= 2: + # Some THD backends restore the caller's pre-flatten [B, S] + # coordinates. Sequence metadata is a row-major flat stream, so + # collapse those two token axes before routing individual Datums. + if tensor.shape[0] * tensor.shape[1] == expected_width: + tensor = tensor.flatten(0, 1) + seq_dim = 0 + if tensor.shape[seq_dim] != expected_width: + raise ValueError( + f"restored packed token output has width {tensor.shape[seq_dim]}, expected {expected_width}" + ) + pieces: list[torch.Tensor] = [] + start = 0 + for datum_index in datum_indices: + real_length = plan.real_lengths[datum_index] + padded_length = plan.padded_lengths[datum_index] + piece = tensor.narrow(seq_dim, start, real_length) + if seq_dim == 1 and piece.shape[0] == 1: + piece = piece.squeeze(0) + pieces.append(piece.contiguous()) + start += padded_length + return pieces + + if tensor.ndim < 2 or tensor.shape[0] != len(datum_indices): + raise ValueError(f"restored padded token output must have {len(datum_indices)} rows, got {tuple(tensor.shape)}") + pieces = [] + for row_index, datum_index in enumerate(datum_indices): + row = tensor.select(0, row_index) + if plan.token_mask is not None: + mask = plan.token_mask[datum_index].to(device=row.device) + pieces.append(row[mask].contiguous()) + else: + pieces.append(row.narrow(0, 0, plan.real_lengths[datum_index]).contiguous()) + return pieces + + +def _loss_fn_output_contract( + outputs: ParsedLossOutputs, + weights: Any, + *, + expected_records: int | None, + microbatch_index: int | None, +) -> tuple[str | None, tuple[Any, ...]]: + """Return a local validation error and a rank-comparable output schema.""" + if outputs is None: + return None, ("none",) + if expected_records is None: + return ( + "a prebatched Datum may return outputs only when num_microbatches=1 because its inner sample " + "boundaries are not part of the Datum contract", + ("unsupported", type(outputs).__name__), + ) + if not isinstance(weights, torch.Tensor): + return "loss_fn outputs require Tensor loss weights", ("invalid-weights", type(weights).__name__) + + if isinstance(outputs, list): + error = None + if len(outputs) != expected_records: + error = ( + f"pipeline loss_fn returned {len(outputs)} outputs for microbatch {microbatch_index}, " + f"expected {expected_records} from its Datum mapping" + if microbatch_index is not None + else f"loss_fn outputs must contain one mapping per Datum; got {len(outputs)} for {expected_records}" + ) + # Legacy records are deliberately opaque and may contain rank-local + # fields. Only their presence and routing count participate in CP + # consensus; typed outputs opt into stronger schema agreement. + return error, ("legacy", len(outputs)) + + error = None + if weights.ndim not in (1, 2): + error = f"per-token output weights must be one- or two-dimensional, got {tuple(weights.shape)}" + records = outputs.per_datum + if records is not None and len(records) != expected_records: + error = ( + f"LossFnOutputBatch.per_datum contains {len(records)} records, expected {expected_records}" + if error is None + else error + ) + record_keys = () if records is None else tuple(tuple(sorted(record)) for record in records) + field_schema: list[tuple[Any, ...]] = [] + for name, spec in sorted(outputs.per_token.items()): + tensor = spec.tensor + if tensor.shape[: weights.ndim] != weights.shape and error is None: + error = ( + f"per-token output {name!r} must start with the loss weight shape {tuple(weights.shape)}, " + f"got {tuple(tensor.shape)}" + ) + if tensor.device != weights.device and error is None: + error = f"per-token output {name!r} must be on the loss weight device {weights.device}, got {tensor.device}" + field_schema.append( + ( + name, + str(tensor.dtype), + tensor.device.type, + tuple(tensor.shape), + type(spec.fill_value).__name__, + repr(spec.fill_value), + ) + ) + return error, ("typed", len(records) if records is not None else None, record_keys, tuple(field_schema)) + + +def _token_mask_schema(token_mask: torch.Tensor | None) -> tuple[tuple[int, ...], str] | None: + """Return compact, rank-comparable routing metadata for a padded token mask.""" + if token_mask is None: + return None + compact = token_mask.detach().to(device="cpu", dtype=torch.uint8).contiguous() + digest = hashlib.sha256(compact.numpy().tobytes()).hexdigest() + return tuple(compact.shape), digest + + +def _loss_fn_output_restore_contract( + outputs: LossFnOutputBatch, + weights: Any, + plan: _OutputRestorePlan, + *, + datum_indices: tuple[int, ...] | None, + chunk_index: int | None, + cp_size: int, +) -> tuple[str | None, tuple[Any, ...]]: + """Preflight one typed restore before any field enters a CP collective.""" + try: + if not isinstance(weights, torch.Tensor): + raise ValueError("loss_fn token outputs require Tensor loss weights") + seq_dim = _loss_sequence_dim_from_weights(weights) + if datum_indices is None: + if plan.item_to_datum is not None: + raise RuntimeError("loss_fn output routing is missing the collater's Datum mapping") + if chunk_index is not None: + raise NotImplementedError( + "a prebatched Datum cannot restore token outputs across multiple pipeline microbatches" + ) + datum_indices = (0,) + + if any(index < 0 for index in datum_indices): + raise ValueError(f"loss_fn output Datum indices must be non-negative, got {datum_indices}") + if plan.item_to_datum is not None: + if plan.real_lengths is None or plan.padded_lengths is None: + raise RuntimeError("token output routing metadata is incomplete") + if any(index >= len(plan.real_lengths) or index >= len(plan.padded_lengths) for index in datum_indices): + raise ValueError(f"loss_fn output Datum indices are out of range: {datum_indices}") + + layout = plan.sharder.shard_layout + selected_layout = layout + if chunk_index is not None: + if layout is None or layout.chunk_layouts is None: + if cp_size > 1: + raise NotImplementedError( + "the active context-parallel backend does not report reversible per-pipeline-chunk " + "token layouts; typed per-token outputs are unavailable for this packed PP+CP batch" + ) + selected_layout = None + else: + if chunk_index < 0 or chunk_index >= len(layout.chunk_layouts): + raise IndexError( + f"pipeline chunk {chunk_index} is out of range for {len(layout.chunk_layouts)} reported CP layouts" + ) + selected_layout = layout.chunk_layouts[chunk_index] + if cp_size > 1 and selected_layout is None: + raise NotImplementedError( + "the active context-parallel backend does not report a reversible token layout; " + "typed per-token outputs are unavailable for this batch" + ) + + local_width = int(weights.shape[seq_dim]) + global_width = local_width * cp_size + if selected_layout is not None: + captured = selected_layout.local_token_global_indices + if captured is not None and captured.numel() != local_width: + raise ValueError( + f"the reported CP layout has {captured.numel()} local token indices, " + f"but loss outputs have local width {local_width}" + ) + if selected_layout.padded_seq_len is not None and selected_layout.padded_seq_len != global_width: + raise ValueError( + f"the reported CP layout has padded width {selected_layout.padded_seq_len}, " + f"but loss outputs imply global width {global_width}" + ) + + restored_width = global_width + if selected_layout is not None: + if selected_layout.input_token_stream_positions is not None: + positions = selected_layout.input_token_stream_positions + if bool((positions < -1).any()) or bool((positions >= global_width).any()): + raise ValueError(f"the reported CP position map must contain -1 or indices below {global_width}") + restored_width = positions.numel() if plan.is_thd else int(positions.shape[1]) + elif selected_layout.input_row_shape is not None: + restored_width = prod(selected_layout.input_row_shape) + elif selected_layout.original_seq_len is not None: + restored_width = selected_layout.original_seq_len + if ( + plan.is_thd + and weights.ndim == 2 + and ( + selected_layout is None + or (selected_layout.input_token_stream_positions is None and selected_layout.input_row_shape is None) + ) + ): + # Some model-owned THD paths retain caller [B, S] rows instead of + # reporting an explicit input_row_shape. Routing flattens those + # token axes row-major after restoration. + restored_width *= int(weights.shape[0]) + + if plan.item_to_datum is not None: + assert plan.real_lengths is not None and plan.padded_lengths is not None + if plan.is_thd: + expected_width = sum(plan.padded_lengths[index] for index in datum_indices) + if restored_width != expected_width: + raise ValueError( + f"restored packed token output has width {restored_width}, expected {expected_width}" + ) + else: + if weights.ndim != 2 or weights.shape[0] != len(datum_indices): + raise ValueError( + f"restored padded token output must have {len(datum_indices)} rows, " + f"got local loss weights {tuple(weights.shape)}" + ) + expected_widths = {plan.padded_lengths[index] for index in datum_indices} + if expected_widths != {restored_width}: + raise ValueError( + f"restored padded token output has width {restored_width}, expected {sorted(expected_widths)}" + ) + positions = selected_layout.input_token_stream_positions if selected_layout is not None else None + if positions is not None and positions.shape[0] != len(datum_indices): + raise ValueError( + f"the reported CP position map has {positions.shape[0]} rows, expected {len(datum_indices)}" + ) + + layout_schema = None + if selected_layout is not None: + layout_schema = ( + selected_layout.original_seq_len, + selected_layout.padded_seq_len, + selected_layout.input_row_shape, + ( + None + if selected_layout.input_token_stream_positions is None + else tuple(selected_layout.input_token_stream_positions.shape) + ), + ( + None + if selected_layout.local_token_global_indices is None + else selected_layout.local_token_global_indices.numel() + ), + ) + schema = ( + "restore", + plan.is_thd, + datum_indices, + plan.real_lengths, + plan.padded_lengths, + _token_mask_schema(plan.token_mask), + seq_dim, + tuple(weights.shape), + layout_schema, + tuple(sorted(outputs.per_token)), + ) + return None, schema + except Exception as error: + return str(error), ("invalid-restore", type(error).__name__) + + def _parse_loss_result( - result: torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]]], + result: torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]] | LossFnOutputBatch], weights: torch.Tensor, -) -> tuple[torch.Tensor, list[dict[str, Any]] | None]: - """Normalize one loss callback result without applying Datum routing.""" +) -> tuple[torch.Tensor, ParsedLossOutputs, Exception | None]: + """Normalize one loss callback result without applying Datum routing. + + Output-only contract errors are returned separately so distributed peers + can finish the same backward schedule and reach a common error decision. + Loss-tensor errors still raise immediately because no valid backward value + exists in that case. + """ if not isinstance(result, tuple): - return _weighted_numerator(result, weights), None - losses, outputs = result - if ( - not isinstance(outputs, Sequence) - or isinstance(outputs, (str, bytes)) - or not all(isinstance(item, Mapping) for item in outputs) - ): - raise ValueError("loss_fn outputs must be a sequence of mappings") - return _weighted_numerator(losses, weights), [_detach(dict(item)) for item in outputs] + return _weighted_numerator(result, weights), None, None + if not result: + raise ValueError("loss_fn returned an empty tuple instead of a loss Tensor") + numerator = _weighted_numerator(result[0], weights) + if len(result) != 2: + return numerator, None, ValueError("loss_fn must return a loss Tensor or a two-item (loss, outputs) tuple") + outputs = result[1] + try: + if isinstance(outputs, LossFnOutputBatch): + detached = LossFnOutputBatch( + per_token={ + name: PerTokenOutput(spec.tensor.detach(), fill_value=spec.fill_value) + for name, spec in outputs.per_token.items() + }, + per_datum=(None if outputs.per_datum is None else [_detach(record) for record in outputs.per_datum]), + ) + return numerator, detached, None + if ( + not isinstance(outputs, Sequence) + or isinstance(outputs, (str, bytes)) + or not all(isinstance(item, Mapping) for item in outputs) + ): + raise ValueError("loss_fn outputs must be a sequence of mappings") + return numerator, [_detach(dict(item)) for item in outputs], None + except Exception as error: + return numerator, None, error -def _update_output_mode(previous: bool | None, outputs: list[dict[str, Any]] | None) -> bool: +def _update_output_mode(previous: bool | None, outputs: ParsedLossOutputs) -> bool: current = outputs is not None if previous is not None and previous != current: raise ValueError("loss_fn must return per-Datum outputs for every microbatch or none of them") diff --git a/nemo_automodel/engine/outputs.py b/nemo_automodel/engine/outputs.py new file mode 100644 index 0000000000..58a4c0b6c0 --- /dev/null +++ b/nemo_automodel/engine/outputs.py @@ -0,0 +1,149 @@ +# 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. + +"""Structured output records returned by Engine loss callbacks. + +The types in this module describe output layout only. The Engine owns +context-parallel restoration and per-Datum routing; constructing these records +does not detach, clone, pad, gather, or otherwise transform tensors. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +import torch + +__all__ = ["LossFnOutputBatch", "PerTokenOutput"] + + +def _validate_field_name(name: Any, *, container: str) -> None: + if not isinstance(name, str) or not name: + raise TypeError(f"{container} field names must be non-empty strings, got {name!r}") + + +@dataclass(frozen=True, eq=False) +class PerTokenOutput: + """One callback tensor aligned with the current local token stream. + + The tensor's leading dimensions must exactly match the current CP-local + loss ``weights`` shape (``[B, S]`` or ``[T]``); only trailing feature + dimensions may be added. The tensor must also be on the same device as the + weights. The Engine therefore knows the token axis without guessing from + arbitrary shapes. + ``fill_value`` is used only when restoring a caller position for which a + backend did not retain a token (for example, a dropped slot in a repacked + sequence). + + The tensor is deliberately retained by reference rather than cloned. This + preserves its device, autograd relationship, and avoids copying potentially + large token outputs; the Engine detaches returned records at its API + boundary. + + Args: + tensor: Tensor covering the complete CP-local token stream for the + current loss-callback invocation. + fill_value: Scalar value used for positions introduced by layout + restoration. It must be representable by ``tensor.dtype`` when the + Engine performs that restoration. + """ + + tensor: torch.Tensor + fill_value: float | int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.tensor, torch.Tensor): + raise TypeError(f"PerTokenOutput.tensor must be a torch.Tensor, got {type(self.tensor).__name__}") + if self.tensor.ndim == 0: + raise ValueError("PerTokenOutput.tensor must have at least one dimension") + if not isinstance(self.fill_value, (int, float)): + raise TypeError(f"PerTokenOutput.fill_value must be an int or float, got {type(self.fill_value).__name__}") + + +@dataclass(frozen=True, eq=False) +class LossFnOutputBatch: + """Layout-aware outputs from one loss-callback invocation. + + ``per_token`` holds batch-level local token streams. Callbacks should not + split those streams into per-Datum tensors themselves: under packed context + parallelism, different ranks can own different numbers of tokens from one + Datum. The Engine can restore each complete stream first and then route it + to the final per-Datum records. + + Restoration requires the active context-parallel backend to report a + reversible token layout. Backends that do not report one fail explicitly + when this typed output is requested; legacy output records remain usable. + + ``per_datum`` optionally supplies the ordinary output records that the + restored token fields will augment. It follows the existing callback + convention of one mapping per logical Datum in the current microbatch. + These records are opaque and must already be identical on CP ranks (for + example, sample IDs copied from a PER_DATUM loss input). A token-derived + CP-local value belongs in ``per_token`` so the Engine can restore it before + producing records. + + Input mappings, the record sequence, and each record mapping are copied and + exposed as read-only views. Tensor and other nested values are intentionally + retained by reference so the envelope is cheap and preserves autograd until + the Engine consumes it. + + Args: + per_token: Non-empty mapping from output field name to its local token + tensor and restoration fill value. Tensor leading dimensions and + device must match the callback's loss weights. + per_datum: Optional sequence of existing per-Datum output mappings. + These mappings may not already contain a field named by + ``per_token`` because the Engine will add those fields after token + restoration. + """ + + per_token: Mapping[str, PerTokenOutput] + per_datum: Sequence[Mapping[str, Any]] | None = None + + def __post_init__(self) -> None: + if not isinstance(self.per_token, Mapping): + raise TypeError(f"LossFnOutputBatch.per_token must be a mapping, got {type(self.per_token).__name__}") + copied_per_token = dict(self.per_token) + if not copied_per_token: + raise ValueError("LossFnOutputBatch.per_token cannot be empty") + for name, output in copied_per_token.items(): + _validate_field_name(name, container="per_token") + if not isinstance(output, PerTokenOutput): + raise TypeError(f"per_token field {name!r} must be a PerTokenOutput, got {type(output).__name__}") + + object.__setattr__(self, "per_token", MappingProxyType(copied_per_token)) + if self.per_datum is None: + return + if not isinstance(self.per_datum, Sequence) or isinstance(self.per_datum, (str, bytes)): + raise TypeError( + "LossFnOutputBatch.per_datum must be a sequence of mappings or None, " + f"got {type(self.per_datum).__name__}" + ) + + copied_records: list[Mapping[str, Any]] = [] + token_names = set(copied_per_token) + for index, record in enumerate(self.per_datum): + if not isinstance(record, Mapping): + raise TypeError(f"per_datum record {index} must be a mapping, got {type(record).__name__}") + copied_record = dict(record) + for name in copied_record: + _validate_field_name(name, container=f"per_datum record {index}") + conflicts = token_names.intersection(copied_record) + if conflicts: + raise ValueError(f"per_datum record {index} conflicts with per_token fields {sorted(conflicts)}") + copied_records.append(MappingProxyType(copied_record)) + object.__setattr__(self, "per_datum", tuple(copied_records)) diff --git a/tests/functional_tests/context_parallel/L2_PP_CP_Dense_Packed_Test.sh b/tests/functional_tests/context_parallel/L2_PP_CP_Dense_Packed_Test.sh new file mode 100755 index 0000000000..80c60b7db6 --- /dev/null +++ b/tests/functional_tests/context_parallel/L2_PP_CP_Dense_Packed_Test.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# 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. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1,2,3" +export CP_SIZE=2 + +python -m torch.distributed.run --nproc_per_node=4 --nnodes=1 -m coverage run \ + tests/functional_tests/context_parallel/run_packed_pp.py diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index 1af694da16..41eb425151 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -20,17 +20,18 @@ same pipeline. Evaluation must match a native eager Llama in summed loss and weight statistics without creating gradients; both surrounding training calls must match eager loss plus every local-stage gradient. Additional flat-Datum -forwards cover explicit per-token, per-Datum, and replicated loss layouts for -raw 2+2 and final-THD 3+1 document splits. A final padded two-Datum update -verifies that Engine broadcasts callback mappings to both pipeline ranks in -logical input order and that the same pipeline can return to training. +runs cover explicit loss layouts and typed per-token callback outputs for raw +2+2 and final-THD 3+1 document splits. Under CP2 they run both forward and +forward/backward, use ragged internally padded sequences, and require every PP +and CP rank to receive the same restored Datum records. A final padded +two-Datum update verifies legacy callback mappings in logical input order. Run with:: torchrun --standalone --nproc-per-node=2 run_packed_pp.py -Set ``CP_SIZE=2`` and use four ranks to run the forward-only PP2 x CP2 -explicit-layout checks:: +Set ``CP_SIZE=2`` and use four ranks to run the PP2 x CP2 output-restoration +matrix:: CP_SIZE=2 torchrun --standalone --nproc-per-node=4 run_packed_pp.py """ @@ -39,6 +40,7 @@ import os import warnings +from functools import partial import torch import torch.distributed as dist @@ -61,6 +63,7 @@ from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.llama.model import LlamaForCausalLM from nemo_automodel.engine import Engine, collate_prebatched +from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput VOCAB_SIZE = 64 SEQ_LEN = 8 @@ -248,9 +251,78 @@ def _final_explicit_layout_collate(datums: list[Datum]): ) +def _cp_padded_explicit_layout_collate(datums: list[Datum], *, final_thd: bool): + """Pack Datums with per-sequence CP padding, retaining their real boundaries.""" + real_lengths = torch.tensor( + [datum.seq_len for datum in datums], dtype=torch.int32, device=datums[0].input_ids.device + ) + padded_lengths = ((real_lengths + 3) // 4) * 4 # TE CP2 requires every sequence slot to divide by 2*CP. + total_tokens = int(padded_lengths.sum().item()) + device = datums[0].input_ids.device + + input_ids = torch.zeros((1, total_tokens), dtype=torch.long, device=device) + position_ids = torch.zeros_like(input_ids) + padding_mask = torch.ones_like(input_ids, dtype=torch.bool) + token_fields = { + "labels": torch.full((1, total_tokens), -100, dtype=torch.long, device=device), + "weights": torch.zeros((1, total_tokens), dtype=torch.float32, device=device), + "advantages": torch.zeros((1, total_tokens), dtype=torch.float32, device=device), + "old_logprobs": torch.zeros((1, total_tokens), dtype=torch.float32, device=device), + } + offset = 0 + for datum, real_length, padded_length in zip(datums, real_lengths.tolist(), padded_lengths.tolist()): + token_slice = slice(offset, offset + real_length) + input_ids[0, token_slice] = datum.input_ids + position_ids[0, token_slice] = torch.arange(real_length, device=device) + padding_mask[0, token_slice] = False + for name in token_fields: + token_fields[name][0, token_slice] = datum.loss_fn_inputs[name] + offset += padded_length + + model_inputs: dict[str, object] + if final_thd: + model_inputs = { + "input_ids": input_ids.reshape(-1), + "position_ids": position_ids.reshape(-1), + "padding_mask": padding_mask.reshape(-1), + "cu_seqlens": F.pad(real_lengths.cumsum(0), (1, 0)).to(torch.int32), + "cu_seqlens_padded": F.pad(padded_lengths.cumsum(0), (1, 0)).to(torch.int32), + "max_seqlen": real_lengths.max(), + "qkv_format": "thd", + } + token_fields = {name: value.reshape(-1) for name, value in token_fields.items()} + else: + model_inputs = { + "input_ids": input_ids, + "position_ids": position_ids, + "padding_mask": padding_mask, + "seq_lens": real_lengths.unsqueeze(0), + "seq_lens_padded": padded_lengths.unsqueeze(0), + "qkv_format": "thd", + } + + loss_inputs = CollatedLossInputs( + { + **token_fields, + "sample_id": torch.stack([datum.loss_fn_inputs["sample_id"] for datum in datums]), + "global_coefficients": datums[0].loss_fn_inputs["global_coefficients"], + }, + layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + "advantages": LossInputLayout.PER_TOKEN, + "old_logprobs": LossInputLayout.PER_TOKEN, + "sample_id": LossInputLayout.PER_DATUM, + "global_coefficients": LossInputLayout.REPLICATED, + }, + item_to_datum=tuple(range(len(datums))), + ) + return model_inputs, loss_inputs + + def _explicit_layout_losses(output, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: """Use both RL-style token fields so their routing affects the numerator.""" - return _token_losses(output, loss_inputs) + 0.01 * loss_inputs["advantages"] + 0.02 * loss_inputs["old_logprobs"] + return _token_losses(output, loss_inputs) + 0.1 * loss_inputs["advantages"] + 0.05 * loss_inputs["old_logprobs"] def _explicit_layout_eager_reference( @@ -438,92 +510,110 @@ def _run_thd_layout( return pipeline -def _run_explicit_loss_layout_forward( +def _run_explicit_loss_layout( pipeline: AutoPipeline, layout: str, + execution: str, device: torch.device, mesh_context: MeshContext, ) -> None: - """Validate semantic loss routing for one real two-stage packed pipeline.""" + """Validate loss routing plus typed token-output restoration for one layout.""" if layout == "raw": - lengths = [4, 4, 4, 4] if mesh_context.cp_size > 1 else [2, 2, 2, 2] - collate_fn = _raw_explicit_layout_collate + lengths = [3, 5, 7, 3] if mesh_context.cp_size > 1 else [2, 2, 2, 2] + collate_fn = ( + partial(_cp_padded_explicit_layout_collate, final_thd=False) + if mesh_context.cp_size > 1 + else _raw_explicit_layout_collate + ) expected_microbatch_ids = {(11, 22), (33, 44)} elif layout == "final": - # Keep the uneven 3+1 Datum split while making every sequence length - # divisible by TE's 2*CP head/tail partition count under CP2. - lengths = [4, 4, 4, 12] if mesh_context.cp_size > 1 else [1, 1, 2, 4] - collate_fn = _final_explicit_layout_collate + lengths = [3, 3, 3, 9] if mesh_context.cp_size > 1 else [1, 1, 2, 4] + collate_fn = ( + partial(_cp_padded_explicit_layout_collate, final_thd=True) + if mesh_context.cp_size > 1 + else _final_explicit_layout_collate + ) expected_microbatch_ids = {(11, 22, 33), (44,)} else: raise ValueError(f"unknown explicit loss layout: {layout}") + if execution not in {"forward", "forward_backward"}: + raise ValueError(f"unknown Engine execution: {execution}") for part in pipeline.parts: part.zero_grad(set_to_none=True) datums = _explicit_layout_datums(device, lengths) reference_loss_sum, reference_weight_sum = _explicit_layout_eager_reference(device, datums, collate_fn) - expected_tokens_by_ids = {} - for ids in expected_microbatch_ids: - indices = [sample_id // 11 - 1 for sample_id in ids] - expected_tokens_by_ids[ids] = { - "advantages": torch.cat([datums[index].loss_fn_inputs["advantages"] for index in indices]), - "old_logprobs": torch.cat([datums[index].loss_fn_inputs["old_logprobs"] for index in indices]), - } def loss_with_outputs(output, loss_inputs): sample_ids = tuple(int(value) for value in loss_inputs["sample_id"].tolist()) if sample_ids not in expected_microbatch_ids: raise AssertionError(f"PP2 {layout} routed unexpected sample IDs {sample_ids}") - expected_tokens = expected_tokens_by_ids[sample_ids] valid_cu_seqlens = loss_inputs["cu_seqlens"].reshape(-1) valid_cu_seqlens = valid_cu_seqlens[valid_cu_seqlens >= 0] assert valid_cu_seqlens.numel() - 1 == len(sample_ids) if mesh_context.cp_size > 1: - import transformer_engine_torch as tex - - cp_mesh = mesh_context.device_mesh["cp"] - local_indices = tex.thd_get_partitioned_indices( - valid_cu_seqlens.to(torch.int32), - int(valid_cu_seqlens[-1].item()), - cp_mesh.size(), - cp_mesh.get_local_rank(), - ).to(device=device, dtype=torch.long) - expected_tokens = {name: value.index_select(0, local_indices) for name, value in expected_tokens.items()} - torch.testing.assert_close(loss_inputs["advantages"].reshape(-1), expected_tokens["advantages"]) - torch.testing.assert_close(loss_inputs["old_logprobs"].reshape(-1), expected_tokens["old_logprobs"]) + # Each full pipeline chunk reserves 12 padded slots and CP2 owns six. + assert loss_inputs["weights"].numel() == 6 torch.testing.assert_close( loss_inputs["global_coefficients"], torch.tensor([0.25, 0.75], device=device), ) assert loss_inputs["global_coefficients"].shape == (2,) losses = _explicit_layout_losses(output, loss_inputs) - return losses, [{"sample_id": value} for value in loss_inputs["sample_id"]] + rl_probe = torch.stack((loss_inputs["advantages"], loss_inputs["old_logprobs"]), dim=-1) + return losses, LossFnOutputBatch( + per_token={"rl_probe": PerTokenOutput(rl_probe)}, + per_datum=[{"sample_id": value} for value in loss_inputs["sample_id"]], + ) - result = Engine( + engine = Engine( pipeline, device=device, mesh_context=mesh_context, microbatch_size=4, collate_fn=collate_fn, - ).forward(datums, loss_with_outputs) + ) + result = getattr(engine, execution)(datums, loss_with_outputs) + + if execution == "forward": + torch.testing.assert_close(result.loss_sum.float(), reference_loss_sum.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(result.weight_sum.float(), reference_weight_sum.float(), atol=0, rtol=0) + outputs = result.loss_fn_outputs + if any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): + raise AssertionError(f"PP2 x CP{mesh_context.cp_size} {layout} forward unexpectedly created gradients") + else: + loss, outputs = result + reference_loss = reference_loss_sum / reference_weight_sum + torch.testing.assert_close(loss.float(), reference_loss.float(), atol=4e-2, rtol=2e-3) + if not any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): + raise AssertionError(f"PP2 x CP{mesh_context.cp_size} {layout} backward created no gradients") - torch.testing.assert_close(result.loss_sum.float(), reference_loss_sum.float(), atol=4e-2, rtol=2e-3) - torch.testing.assert_close(result.weight_sum.float(), reference_weight_sum.float(), atol=0, rtol=0) - output_ids = torch.stack([item["sample_id"] for item in result.loss_fn_outputs]).to(torch.long) + output_ids = torch.stack([item["sample_id"] for item in outputs]).to(torch.long) expected_ids = torch.tensor([11, 22, 33, 44], device=device) torch.testing.assert_close(output_ids, expected_ids) + restored_probe = [] + for datum, item in zip(datums, outputs): + expected_probe = torch.stack( + (datum.loss_fn_inputs["advantages"], datum.loss_fn_inputs["old_logprobs"]), + dim=-1, + ) + torch.testing.assert_close(item["rl_probe"], expected_probe) + assert not item["rl_probe"].requires_grad + restored_probe.append(item["rl_probe"]) + flat_probe = torch.cat(restored_probe) gathered = [torch.empty_like(output_ids) for _ in range(dist.get_world_size())] dist.all_gather(gathered, output_ids) assert all(torch.equal(ids, expected_ids) for ids in gathered) - if any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): - raise AssertionError(f"PP2 {layout} explicit-layout forward unexpectedly created gradients") + gathered_probe = [torch.empty_like(flat_probe) for _ in range(dist.get_world_size())] + dist.all_gather(gathered_probe, flat_probe) + assert all(torch.equal(probe, flat_probe) for probe in gathered_probe) if dist.get_rank() == 0: datum_counts = "2+2" if layout == "raw" else "3+1" + reported_loss = result.loss_sum.item() if execution == "forward" else result[0].item() print( - f"PP2 x CP{mesh_context.cp_size} {layout} explicit " - f"PER_TOKEN/PER_DATUM/REPLICATED routing passed ({datum_counts} Datums; " - f"loss_sum={result.loss_sum.item():.6f}, weight_sum={result.weight_sum.item():.1f})" + f"PP2 x CP{mesh_context.cp_size} {layout} {execution} explicit loss/output routing passed " + f"({datum_counts} Datums; loss={reported_loss:.6f})" ) @@ -594,10 +684,10 @@ def main() -> None: try: if cp_size > 1: pipeline = _build_pipeline(device, mesh_context) - _run_explicit_loss_layout_forward(pipeline, "raw", device, mesh_context) - dist.barrier() - _run_explicit_loss_layout_forward(pipeline, "final", device, mesh_context) - dist.barrier() + for layout in ("raw", "final"): + for execution in ("forward", "forward_backward"): + _run_explicit_loss_layout(pipeline, layout, execution, device, mesh_context) + dist.barrier() return ( @@ -635,9 +725,9 @@ def main() -> None: weights, ) dist.barrier() - _run_explicit_loss_layout_forward(final_pipeline, "raw", device, mesh_context) + _run_explicit_loss_layout(final_pipeline, "raw", "forward", device, mesh_context) dist.barrier() - _run_explicit_loss_layout_forward(final_pipeline, "final", device, mesh_context) + _run_explicit_loss_layout(final_pipeline, "final", "forward", device, mesh_context) dist.barrier() _run_padded_output_broadcast(final_pipeline, device, mesh_context) dist.barrier() diff --git a/tests/functional_tests/context_parallel/test_context_parallel.py b/tests/functional_tests/context_parallel/test_context_parallel.py index 99d4e286af..f85d2bc081 100644 --- a/tests/functional_tests/context_parallel/test_context_parallel.py +++ b/tests/functional_tests/context_parallel/test_context_parallel.py @@ -33,7 +33,9 @@ CP_DENSE_PACKED_TEST_FILENAME = "L2_CP_Dense_Packed_Test.sh" TP_DENSE_PACKED_TEST_FILENAME = "L2_TP_Dense_Packed_Test.sh" PP_DENSE_PACKED_TEST_FILENAME = "L2_PP_Dense_Packed_Test.sh" +PP_CP_DENSE_PACKED_TEST_FILENAME = "L2_PP_CP_Dense_Packed_Test.sh" TP_CP_DENSE_PACKED_TEST_FILENAME = "L2_TP_CP_Dense_Packed_Test.sh" +PP_CP_DENSE_PACKED_REQUIRED_GPUS = 4 TP_CP_DENSE_PACKED_REQUIRED_GPUS = 4 @@ -76,6 +78,14 @@ def test_pp_dense_packed(self): """Test Engine PP=2 raw/final THD loss and gradient parity for dense Llama.""" run_test_script(TEST_FOLDER, PP_DENSE_PACKED_TEST_FILENAME) + @pytest.mark.skipif( + torch.cuda.device_count() < PP_CP_DENSE_PACKED_REQUIRED_GPUS, + reason="requires 4 GPUs for PP=2 x CP=2", + ) + def test_pp_cp_dense_packed(self): + """Test Engine PP=2 x CP=2 raw/final THD token-output restoration.""" + run_test_script(TEST_FOLDER, PP_CP_DENSE_PACKED_TEST_FILENAME) + @pytest.mark.skipif( torch.cuda.device_count() < TP_CP_DENSE_PACKED_REQUIRED_GPUS, reason="requires 4 GPUs for TP=2 x CP=2; remove once context_parallel CI runs on a 4-GPU runner", diff --git a/tests/unit_tests/distributed/test_cp_sharder.py b/tests/unit_tests/distributed/test_cp_sharder.py index a6c15d6761..5b1e7bf525 100644 --- a/tests/unit_tests/distributed/test_cp_sharder.py +++ b/tests/unit_tests/distributed/test_cp_sharder.py @@ -208,6 +208,13 @@ def test_sharder_repositioned_layout_round_trips_input_coordinates(): ) out = gather_sharder.gather_token_tensor(full_rows, trim=True, fill=-5.0) assert torch.equal(out, torch.tensor([[10.0, 20.0, -5.0]])) + + full_features = torch.tensor([[[10.0, 110.0], [20.0, 120.0], [7.0, 8.0], [9.0, 10.0]]]) + feature_out = gather_sharder.gather_token_tensor(full_features, trim=True, fill=-5.0) + assert torch.equal( + feature_out, + torch.tensor([[[10.0, 110.0], [20.0, 120.0], [-5.0, -5.0]]]), + ) with pytest.raises(ValueError, match="fill"): gather_sharder.gather_token_tensor(full_rows, trim=True) @@ -222,6 +229,65 @@ def test_gather_trim_raises_without_captured_facts(): sharder.gather_token_tensor(torch.zeros(1, 4), trim=True) +def test_gather_token_tensor_selects_one_chunk_layout(monkeypatch): + chunk_layouts = ( + cs.ShardLayout( + local_token_global_indices=torch.tensor([0, 3]), + original_seq_len=4, + padded_seq_len=4, + ), + cs.ShardLayout( + local_token_global_indices=torch.tensor([1, 2]), + original_seq_len=4, + padded_seq_len=4, + ), + ) + sharder = cs.ContextParallelSharder( + device_mesh=_FakeDeviceMesh(_FakeMesh(2)), + shard_batch=cs.shard_batch_identity, + shard_layout=cs.ShardLayout(chunk_layouts=chunk_layouts), + ) + seen_indices = [] + + def fake_gather(_mesh, _tensor, local_indices, seq_dim=1): + assert seq_dim == 0 + seen_indices.append(local_indices.clone()) + return torch.arange(4.0) + + monkeypatch.setattr(cs, "gather_token_tensor_by_indices", fake_gather) + + gathered = sharder.gather_token_tensor( + torch.tensor([10.0, 40.0]), + seq_dim=0, + trim=True, + chunk_index=0, + ) + + assert torch.equal(gathered, torch.arange(4.0)) + assert len(seen_indices) == 1 + assert torch.equal(seen_indices[0], torch.tensor([0, 3])) + + +def test_gather_token_tensor_requires_a_valid_chunk_index(): + sharder = cs.ContextParallelSharder( + device_mesh=_FakeDeviceMesh(_FakeMesh(1)), + shard_batch=cs.shard_batch_identity, + local_token_global_indices=cs.identity_local_indices, + shard_layout=cs.ShardLayout( + chunk_layouts=(cs.ShardLayout(original_seq_len=2, padded_seq_len=2),), + ), + ) + + with pytest.raises(ValueError, match="chunk_index is required"): + sharder.gather_token_tensor(torch.zeros(2), seq_dim=0) + with pytest.raises(IndexError, match="out of range"): + sharder.gather_token_tensor(torch.zeros(2), seq_dim=0, chunk_index=1) + + sharder.shard_layout = cs.ShardLayout(original_seq_len=2, padded_seq_len=2) + with pytest.raises(ValueError, match="only valid for a chunked"): + sharder.gather_token_tensor(torch.zeros(2), seq_dim=0, chunk_index=0) + + def test_reported_indices_validate_stream_length(): # Reported index maps flatten + cast to long, and reject a padded_seq_len # that does not match the partition the shard reported. diff --git a/tests/unit_tests/distributed/test_cp_utils.py b/tests/unit_tests/distributed/test_cp_utils.py index be2a62001a..a0df3d3749 100644 --- a/tests/unit_tests/distributed/test_cp_utils.py +++ b/tests/unit_tests/distributed/test_cp_utils.py @@ -1065,6 +1065,28 @@ def test_make_cp_batch_for_te_identity_indices_without_cp(): assert torch.equal(local_indices, torch.arange(out["input_ids"].shape[-1])) +def test_make_cp_batch_for_te_retains_identity_indices_for_each_final_thd_chunk(): + batch = { + "input_ids": torch.arange(8), + "labels": torch.arange(8), + "position_ids": torch.arange(8), + "cu_seqlens": torch.tensor([0, 4, 8], dtype=torch.int32), + "max_seqlen": torch.tensor(4, dtype=torch.int32), + "qkv_format": "thd", + } + + out, local_indices = _cu.make_cp_batch_for_te( + None, + batch, + num_chunks=2, + return_local_indices=True, + ) + + assert out["input_ids"].shape == (2, 4) + assert isinstance(local_indices, tuple) and len(local_indices) == 2 + assert all(torch.equal(indices, torch.arange(4)) for indices in local_indices) + + def test_round_robin_sharder_captures_lengths_and_pads_token_tensors(monkeypatch): """The generic sharder captures original/padded lengths at shard time so the token verbs accept caller-coordinate tensors: unpadded down (with explicit @@ -1127,6 +1149,72 @@ def fake_make_cp_batch_for_te(cp_mesh, batch, *, return_local_indices=False, **k assert torch.equal(sharder.gather_token_tensor(torch.tensor([1.0, 2.0, 3.0, 4.0]), seq_dim=0, trim=True), rows) +def test_te_sharder_captures_final_thd_stream_length_for_trim(monkeypatch): + local_indices = torch.arange(4) + + def fake_make_cp_batch_for_te(cp_mesh, batch, *, return_local_indices=False, **kwargs): + return (dict(batch), local_indices) if return_local_indices else dict(batch) + + monkeypatch.setattr(_cu, "make_cp_batch_for_te", fake_make_cp_batch_for_te) + + strategy = _cu._resolve_cp_sharder( + _DummySubMesh(1), + None, + magi=None, + is_thd=True, + num_chunks=1, + seq_lens_padding_value=-1000, + model=None, + ) + sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=1, tp_size=1)) + batch = { + "input_ids": torch.arange(4), + "cu_seqlens": torch.tensor([0, 2, 4], dtype=torch.int32), + "qkv_format": "thd", + } + sharder.shard(batch) + + layout = sharder.shard_layout + assert layout.input_row_shape is None + assert (layout.original_seq_len, layout.padded_seq_len) == (4, 4) + values = torch.arange(4.0) + assert torch.equal(sharder.gather_token_tensor(values, seq_dim=0, trim=True), values) + + +def test_te_sharder_captures_one_layout_per_pipeline_chunk(monkeypatch): + chunk_indices = (torch.tensor([0, 3]), torch.tensor([1, 2])) + + def fake_make_cp_batch_for_te(cp_mesh, batch, *, return_local_indices=False, **kwargs): + prepped = {"input_ids": torch.tensor([[0, 3], [1, 2]]), "qkv_format": "thd"} + return (prepped, chunk_indices) if return_local_indices else prepped + + monkeypatch.setattr(_cu, "make_cp_batch_for_te", fake_make_cp_batch_for_te) + + strategy = _cu._resolve_cp_sharder( + _DummySubMesh(2), + None, + magi=None, + is_thd=True, + num_chunks=2, + seq_lens_padding_value=-1000, + model=None, + ) + sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=2, tp_size=1)) + sharder.shard( + { + "input_ids": torch.arange(8).reshape(2, 4), + "seq_lens": torch.tensor([[4], [4]]), + "seq_lens_padded": torch.tensor([[4], [4]]), + } + ) + + layout = sharder.shard_layout + assert layout.chunk_layouts is not None and len(layout.chunk_layouts) == 2 + assert torch.equal(layout.chunk_layouts[0].local_token_global_indices, chunk_indices[0]) + assert torch.equal(layout.chunk_layouts[1].local_token_global_indices, chunk_indices[1]) + assert all((child.original_seq_len, child.padded_seq_len) == (4, 4) for child in layout.chunk_layouts) + + def test_resolve_cp_sharder_layers(): """Resolution order: model-owned > magi > TE > generic round-robin > none.""" from nemo_automodel.components.distributed.context_parallel.sharder import round_robin_local_indices diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 4333b0fb9f..7feaf964c8 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -14,8 +14,10 @@ from __future__ import annotations +import pickle import sys from contextlib import contextmanager +from datetime import timedelta from functools import partial from types import SimpleNamespace @@ -50,6 +52,7 @@ from nemo_automodel.components.models.common.mtp import prepare_mtp_context_parallel_inputs from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler from nemo_automodel.engine import Engine, ForwardResult, OptimStepResult, collate_prebatched +from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput class ScaleModel(nn.Module): @@ -855,6 +858,248 @@ def loss_with_outputs(output, _loss_inputs): assert all(not item["model_value"].requires_grad for item in outputs) +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +def test_typed_batch_outputs_split_token_fields_into_datum_records(execution): + """The callback describes one collated token tensor; Engine owns Datum splitting.""" + datums = [_datum([1, 2]), _datum([3, 4, 5])] + + def loss_with_outputs(output, _loss_inputs): + token_probe = torch.stack((output, output + 100), dim=-1) + return output, LossFnOutputBatch( + per_token={"token_probe": PerTokenOutput(token_probe, fill_value=-1.0)}, + per_datum=[{"sample_id": torch.tensor(11)}, {"sample_id": torch.tensor(22)}], + ) + + result = getattr(Engine(ScaleModel(), device="cpu", microbatch_size=2), execution)(datums, loss_with_outputs) + outputs = result.loss_fn_outputs if execution == "forward" else result[1] + + assert [item["sample_id"].item() for item in outputs] == [11, 22] + torch.testing.assert_close(outputs[0]["token_probe"], torch.tensor([[1.0, 101.0], [2.0, 102.0]])) + torch.testing.assert_close( + outputs[1]["token_probe"], + torch.tensor([[3.0, 103.0], [4.0, 104.0], [5.0, 105.0]]), + ) + assert all(not item["token_probe"].requires_grad for item in outputs) + + +def test_raw_thd_multirow_outputs_restore_trailing_features_in_datum_order(): + """Raw [B, S] THD coordinates flatten back to one routed token stream.""" + datums = [_datum([1, 2, 3]), _datum([4, 5])] + + def raw_thd_collate(_datums): + token_rows = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]) + weights = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) + return ( + { + "input_ids": token_rows, + "position_ids": torch.arange(4).expand(2, -1), + "seq_lens": torch.tensor([[3], [2]], dtype=torch.int32), + "seq_lens_padded": torch.tensor([[4], [4]], dtype=torch.int32), + "qkv_format": "thd", + }, + CollatedLossInputs( + {"weights": weights}, + layouts={"weights": LossInputLayout.PER_TOKEN}, + item_to_datum=(0, 1), + ), + ) + + def loss_with_outputs(output, loss_inputs): + assert output.shape == loss_inputs["weights"].shape == (8,) + probe = torch.stack((output, output + 100), dim=-1) + return output, LossFnOutputBatch(per_token={"probe": PerTokenOutput(probe)}) + + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + result = Engine( + model, + device="cpu", + microbatch_size=2, + collate_fn=raw_thd_collate, + ).forward(datums, loss_with_outputs) + + torch.testing.assert_close( + result.loss_fn_outputs[0]["probe"], torch.tensor([[1.0, 101.0], [2.0, 102.0], [3.0, 103.0]]) + ) + torch.testing.assert_close(result.loss_fn_outputs[1]["probe"], torch.tensor([[4.0, 104.0], [5.0, 105.0]])) + + +def test_typed_batch_outputs_follow_left_and_noncontiguous_attention_masks(): + datums = [_datum([10, 11]), _datum([20, 21, 22])] + + def masked_collate(_datums): + input_ids = torch.tensor( + [ + [99, 99, 10, 11, 99], + [20, 99, 21, 99, 22], + ] + ) + attention_mask = torch.tensor( + [ + [False, False, True, True, False], + [True, False, True, False, True], + ] + ) + return ( + {"input_ids": input_ids, "attention_mask": attention_mask}, + CollatedLossInputs( + {"weights": attention_mask.to(torch.float32)}, + layouts={"weights": LossInputLayout.PER_TOKEN}, + item_to_datum=(0, 1), + ), + ) + + def loss_with_outputs(output, _loss_inputs): + probe = torch.stack((output, output + 100), dim=-1) + return output, LossFnOutputBatch(per_token={"probe": PerTokenOutput(probe)}) + + result = Engine( + ScaleModel(), + device="cpu", + microbatch_size=2, + collate_fn=masked_collate, + ).forward(datums, loss_with_outputs) + + torch.testing.assert_close(result.loss_fn_outputs[0]["probe"], torch.tensor([[10.0, 110.0], [11.0, 111.0]])) + torch.testing.assert_close( + result.loss_fn_outputs[1]["probe"], + torch.tensor([[20.0, 120.0], [21.0, 121.0], [22.0, 122.0]]), + ) + + +def test_padded_output_routing_keeps_attention_mask_compact(): + mask = torch.tensor([[False, True, True], [True, False, True]]) + real_lengths, padded_lengths, token_mask = engine_module._output_sequence_lengths( + [_datum([1, 2]), _datum([3, 4])], + {"attention_mask": mask}, + {"weights": mask.to(torch.float32)}, + (0, 1), + is_thd=False, + ) + + assert real_lengths == (2, 2) + assert padded_lengths == (3, 3) + assert isinstance(token_mask, torch.Tensor) + assert token_mask.dtype is torch.bool + torch.testing.assert_close(token_mask, mask) + + +def test_legacy_per_datum_output_tensors_remain_opaque(): + """A vector whose length resembles a token axis must not opt into restoration by shape.""" + opaque = torch.tensor([91.0, 92.0, 93.0]) + + result = Engine(ScaleModel(), device="cpu", microbatch_size=2).forward( + [_datum([1, 2]), _datum([3, 4, 5])], + lambda output, _inputs: (output, [{"opaque": opaque.clone()} for _ in output]), + ) + + assert len(result.loss_fn_outputs[0]["opaque"]) == 3 + assert len(result.loss_fn_outputs[1]["opaque"]) == 3 + assert all(torch.equal(item["opaque"], opaque) for item in result.loss_fn_outputs) + + +def test_typed_batch_output_rejects_a_mismatched_token_shape(): + def loss_with_bad_shape(output, _loss_inputs): + return output, LossFnOutputBatch( + per_token={"bad_probe": PerTokenOutput(output[:, :-1])}, + per_datum=[{}, {}], + ) + + with pytest.raises(ValueError, match=r"bad_probe.*shape|shape.*bad_probe"): + Engine(ScaleModel(), device="cpu", microbatch_size=2).forward( + [_datum([1, 2]), _datum([3, 4])], + loss_with_bad_shape, + ) + + +def test_output_only_parse_error_is_reported_before_single_rank_backward(): + model = ScaleModel() + + with pytest.raises(ValueError, match="outputs must be a sequence of mappings"): + Engine(model, device="cpu").forward_backward( + [_datum([1, 2])], + lambda output, _inputs: (output, "not-records"), + ) + + assert model.weight.grad is None + + +def test_typed_batch_output_rejects_explicit_empty_per_datum_records(): + def loss_with_empty_records(output, _loss_inputs): + return output, LossFnOutputBatch( + per_token={"probe": PerTokenOutput(output)}, + per_datum=[], + ) + + with pytest.raises(ValueError, match=r"per_datum contains 0 records, expected 2"): + Engine(ScaleModel(), device="cpu", microbatch_size=2).forward( + [_datum([1]), _datum([2])], + loss_with_empty_records, + ) + + +@pytest.mark.parametrize( + ("cu_seqlens", "cu_seqlens_padded", "match"), + [ + ( + torch.tensor([0, 4], dtype=torch.int32), + torch.tensor([0, 3], dtype=torch.int32), + "0 <= real_length <= padded_length", + ), + ( + torch.tensor([0, 4], dtype=torch.int32), + torch.tensor([0, 5], dtype=torch.int32), + r"padded lengths sum to 5.*loss weights.*4", + ), + ], + ids=("real-exceeds-padded", "padded-sum-mismatch"), +) +def test_typed_batch_output_rejects_invalid_final_thd_lengths_before_forward( + cu_seqlens, + cu_seqlens_padded, + match, +): + model = ScaleModel() + + def bad_thd_collate(_datums): + return ( + { + "input_ids": torch.tensor([1, 2, 3, 4]), + "position_ids": torch.arange(4), + "cu_seqlens": cu_seqlens, + "cu_seqlens_padded": cu_seqlens_padded, + "max_seqlen": torch.tensor(4, dtype=torch.int32), + "qkv_format": "thd", + }, + CollatedLossInputs( + {"weights": torch.ones(4)}, + layouts={"weights": LossInputLayout.PER_TOKEN}, + item_to_datum=(0,), + ), + ) + + with pytest.raises(ValueError, match=match): + Engine(model, device="cpu", collate_fn=bad_thd_collate).forward( + [_datum([1, 2, 3, 4])], + _identity_loss, + ) + assert model.forward_calls == 0 + + +def test_typed_batch_output_rejects_per_token_and_per_datum_key_collision(): + def loss_with_duplicate_key(output, _loss_inputs): + return output, LossFnOutputBatch( + per_token={"score": PerTokenOutput(output)}, + per_datum=[{"score": torch.tensor(1.0)}, {"score": torch.tensor(2.0)}], + ) + + with pytest.raises(ValueError, match=r"score.*(?:both|conflict|duplicate)|(?:both|conflict|duplicate).*score"): + Engine(ScaleModel(), device="cpu", microbatch_size=2).forward( + [_datum([1, 2]), _datum([3, 4])], + loss_with_duplicate_key, + ) + + def test_loss_fn_outputs_must_align_with_datums(): model = ScaleModel() with pytest.raises(ValueError, match="one mapping per Datum"): @@ -1258,6 +1503,50 @@ def _final_thd_layout_collate(datums: list[Datum]): ) +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +@pytest.mark.parametrize( + ("lengths", "collate_fn"), + [ + pytest.param([2, 2, 2, 2], partial(collate_datums, packed=True), id="raw-thd-2-plus-2"), + pytest.param([1, 1, 2, 4], _final_thd_layout_collate, id="final-thd-3-plus-1"), + ], +) +def test_pipeline_typed_batch_outputs_split_packed_stream_in_datum_order(execution, lengths, collate_fn): + """Packed callbacks return one stream; Engine restores records after logical PP ordering.""" + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = _packed_layout_datums(lengths) + + def loss_with_outputs(output, loss_inputs): + probe = torch.stack((loss_inputs["advantages"], loss_inputs["old_logprobs"]), dim=-1) + return output, LossFnOutputBatch( + per_token={"rl_probe": PerTokenOutput(probe, fill_value=-999.0)}, + per_datum=[{"sample_id": sample_id} for sample_id in loss_inputs["sample_id"]], + ) + + result = getattr( + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + microbatch_size=4, + collate_fn=collate_fn, + ), + execution, + )(datums, loss_with_outputs) + outputs = result.loss_fn_outputs if execution == "forward" else result[1] + + assert [item["sample_id"].item() for item in outputs] == [11, 22, 33, 44] + for datum, item in zip(datums, outputs): + expected = torch.stack( + (datum.loss_fn_inputs["advantages"], datum.loss_fn_inputs["old_logprobs"]), + dim=-1, + ) + torch.testing.assert_close(item["rl_probe"], expected) + assert not item["rl_probe"].requires_grad + + @pytest.mark.parametrize("execution", ["forward", "forward_backward"]) def test_pipeline_raw_thd_routes_explicit_loss_layouts_and_outputs_in_datum_order(execution): model = ScaleModel() @@ -1361,7 +1650,7 @@ def test_pipeline_final_thd_rejects_wrong_outputs_even_when_window_total_matches pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) datums = _packed_layout_datums([1, 1, 2, 4]) - with pytest.raises(ValueError, match="returned 2 outputs for microbatch 1, expected 1"): + with pytest.raises(ValueError, match="returned 2 outputs for microbatch 0, expected 3"): Engine( pipeline, device="cpu", @@ -1453,14 +1742,14 @@ def fake_broadcast_object_list(objects, *, src, group, device): assert src == 11 assert group is pp_group assert device == torch.device("cpu") - objects[0] = expected + objects[0] = pickle.dumps(expected) monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) monkeypatch.setattr(dist, "get_rank", lambda *, group: 1) monkeypatch.setattr(dist, "get_global_rank", lambda group, group_rank: 11 if group_rank == 0 else 12) monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast_object_list) - result = engine._broadcast_pipeline_outputs([]) + result = engine._broadcast_pipeline_outputs([], serialized_outputs=None) assert result == expected @@ -1483,7 +1772,7 @@ def fake_all_gather_into_tensor(gathered, local, *, group): lambda *_args, **_kwargs: pytest.fail("empty pipeline outputs must not use an object collective"), ) - assert engine._broadcast_pipeline_outputs([]) == [] + assert engine._broadcast_pipeline_outputs([], serialized_outputs=None) == [] def test_pipeline_prebatched_outputs_require_one_inner_microbatch(): @@ -1974,6 +2263,100 @@ def _mismatched_context_parallel_weights_worker(rank: int, world_size: int, init dist.destroy_process_group() +def _context_parallel_output_consensus_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + engine = Engine( + ScaleModel(), + device="cpu", + mesh_context=SimpleNamespace( + pp_size=1, + cp_size=world_size, + device_mesh=_CPMesh(size=world_size, rank=rank), + process_group=None, + ), + ) + weights = torch.ones(2) + valid = LossFnOutputBatch( + per_token={"probe": PerTokenOutput(torch.tensor([1.0, 2.0]))}, + per_datum=[{}], + ) + cases = [ + (None if rank == 0 else valid, "different loss output schemas"), + ( + valid + if rank == 0 + else LossFnOutputBatch( + per_token={"probe": PerTokenOutput(torch.tensor([1.0]))}, + per_datum=[{}], + ), + "invalid loss_fn outputs", + ), + ( + valid + if rank == 0 + else LossFnOutputBatch( + per_token={"probe": PerTokenOutput(torch.tensor([1.0, 2.0]))}, + per_datum=[{}, {}], + ), + "invalid loss_fn outputs", + ), + ] + for outputs, match in cases: + with pytest.raises(ValueError, match=match): + engine._validate_loss_fn_outputs_across_cp( + outputs, + weights, + expected_records=1, + ) + dist.barrier() + + with pytest.raises(ValueError, match="invalid loss_fn outputs"): + engine._validate_loss_fn_outputs_across_cp( + valid, + weights, + expected_records=1, + local_error=ValueError("rank-local output parse failure") if rank == 0 else None, + ) + finally: + dist.destroy_process_group() + + +def _data_parallel_output_error_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + model = nn.parallel.DistributedDataParallel(ScaleModel()) + + def loss_with_rank_local_output_error(output, _loss_inputs): + records = [{}, {}] if rank == 0 else [{}] + return output, LossFnOutputBatch( + per_token={"probe": PerTokenOutput(output)}, + per_datum=records, + ) + + expected = "per_datum contains 2 records" if rank == 0 else "another data-parallel replica" + with pytest.raises((ValueError, RuntimeError), match=expected): + Engine(model, device="cpu").forward_backward( + [_datum([1, 2])], + loss_with_rank_local_output_error, + ) + assert model.module.weight.grad is not None + finally: + dist.destroy_process_group() + + def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): mp.spawn( _distributed_worker, @@ -2008,3 +2391,21 @@ def test_context_parallel_replicas_require_the_same_full_sequence_weights(tmp_pa nprocs=2, join=True, ) + + +def test_context_parallel_output_consensus_rejects_rank_local_contracts_without_hanging(tmp_path): + mp.spawn( + _context_parallel_output_consensus_worker, + args=(2, str(tmp_path / "engine_cp_output_consensus_init")), + nprocs=2, + join=True, + ) + + +def test_data_parallel_output_errors_propagate_after_backward_without_hanging(tmp_path): + mp.spawn( + _data_parallel_output_error_worker, + args=(2, str(tmp_path / "engine_dp_output_error_init")), + nprocs=2, + join=True, + ) diff --git a/tests/unit_tests/test_engine_outputs.py b/tests/unit_tests/test_engine_outputs.py new file mode 100644 index 0000000000..a15f6775ab --- /dev/null +++ b/tests/unit_tests/test_engine_outputs.py @@ -0,0 +1,109 @@ +# 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. + +from __future__ import annotations + +import pytest +import torch + +from nemo_automodel import LossFnOutputBatch as PublicLossFnOutputBatch +from nemo_automodel import PerTokenOutput as PublicPerTokenOutput +from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput + + +def test_output_types_are_public_lazy_exports() -> None: + assert PublicPerTokenOutput is PerTokenOutput + assert PublicLossFnOutputBatch is LossFnOutputBatch + + +def test_loss_fn_output_batch_copies_containers_but_retains_values() -> None: + logprobs = torch.randn(5, requires_grad=True) + sample_id = torch.tensor(7) + token_output = PerTokenOutput(logprobs, fill_value=-1.0) + token_fields = {"logprobs": token_output} + record = {"sample_id": sample_id} + records = [record] + + output = LossFnOutputBatch(per_token=token_fields, per_datum=records) + + token_fields["entropy"] = PerTokenOutput(torch.ones(5)) + record["new_field"] = 3 + records.append({"sample_id": torch.tensor(8)}) + + assert tuple(output.per_token) == ("logprobs",) + assert output.per_token["logprobs"] is token_output + assert output.per_token["logprobs"].tensor is logprobs + assert output.per_token["logprobs"].fill_value == -1.0 + assert output.per_datum is not None + assert len(output.per_datum) == 1 + assert dict(output.per_datum[0]) == {"sample_id": sample_id} + assert output.per_datum[0]["sample_id"] is sample_id + + with pytest.raises(TypeError): + output.per_token["another"] = PerTokenOutput(torch.ones(5)) # type: ignore[index] + with pytest.raises(TypeError): + output.per_datum[0]["another"] = 1 # type: ignore[index] + + +def test_loss_fn_output_batch_preserves_none_per_datum() -> None: + output = LossFnOutputBatch(per_token={"logprobs": PerTokenOutput(torch.ones(3))}) + + assert output.per_datum is None + + +@pytest.mark.parametrize( + ("kwargs", "error", "match"), + [ + ({"tensor": [1.0]}, TypeError, "tensor must be a torch.Tensor"), + ({"tensor": torch.tensor(1.0)}, ValueError, "at least one dimension"), + ({"tensor": torch.ones(1), "fill_value": "zero"}, TypeError, "fill_value must be an int or float"), + ], +) +def test_per_token_output_rejects_invalid_values(kwargs, error, match) -> None: + with pytest.raises(error, match=match): + PerTokenOutput(**kwargs) + + +@pytest.mark.parametrize( + ("per_token", "error", "match"), + [ + ([], TypeError, "per_token must be a mapping"), + ({}, ValueError, "per_token cannot be empty"), + ({"": PerTokenOutput(torch.ones(1))}, TypeError, "field names must be non-empty strings"), + ({1: PerTokenOutput(torch.ones(1))}, TypeError, "field names must be non-empty strings"), + ({"logprobs": torch.ones(1)}, TypeError, "must be a PerTokenOutput"), + ], +) +def test_loss_fn_output_batch_rejects_invalid_token_fields(per_token, error, match) -> None: + with pytest.raises(error, match=match): + LossFnOutputBatch(per_token=per_token) + + +@pytest.mark.parametrize( + ("per_datum", "error", "match"), + [ + ({"sample_id": 1}, TypeError, "per_datum must be a sequence of mappings"), + ("record", TypeError, "per_datum must be a sequence of mappings"), + ([1], TypeError, "record 0 must be a mapping"), + ([{"": 1}], TypeError, "field names must be non-empty strings"), + ([{1: "sample"}], TypeError, "field names must be non-empty strings"), + ([{"logprobs": torch.ones(1)}], ValueError, "conflicts with per_token fields"), + ], +) +def test_loss_fn_output_batch_rejects_invalid_per_datum_records(per_datum, error, match) -> None: + with pytest.raises(error, match=match): + LossFnOutputBatch( + per_token={"logprobs": PerTokenOutput(torch.ones(1))}, + per_datum=per_datum, + ) From c53349461001b0fa67fe547a9276d08ae9c9cceb Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 01:39:48 -0700 Subject: [PATCH 14/34] feat(engine): return structured training results Signed-off-by: HuiyingLi --- nemo_automodel/engine/__init__.py | 58 +++++++-- nemo_automodel/recipes/llm/train_ft.py | 9 +- nemo_automodel/recipes/vlm/finetune.py | 24 ++-- .../context_parallel/run_cp_pp_image_sink.py | 4 +- .../context_parallel/run_cp_pp_layer2_sink.py | 4 +- .../context_parallel/run_dense_packed_cp.py | 10 +- .../context_parallel/run_packed_pp.py | 35 +++--- .../moe/test_experts_ep_tp_grad_parity.py | 4 +- .../recipes/test_finetune_vlm_helpers.py | 42 ++++++- tests/unit_tests/recipes/test_train_ft.py | 23 +++- tests/unit_tests/test_engine.py | 115 ++++++++++-------- .../test_engine_recipe_integration.py | 2 +- 12 files changed, 222 insertions(+), 108 deletions(-) diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 14fbdb84e9..4319f178e9 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -70,6 +70,7 @@ __all__ = [ "Engine", + "ForwardBackwardResult", "ForwardResult", "LossFnOutputBatch", "OptimStepResult", @@ -172,6 +173,33 @@ class ForwardResult: loss_fn_outputs: list[dict[str, Any]] +@dataclass(frozen=True) +class ForwardBackwardResult: + """Training-window loss statistics and per-Datum callback outputs. + + The numerator is summed across the DP-CP gradient group. The full-sequence + denominator is summed across DP only because CP ranks begin with replicated + weights; both are synchronized across PP stages. ``loss`` is + ``loss_sum / weight_sum`` when the denominator is nonzero and zero + otherwise. ``loss_fn_outputs`` remains local to one data-parallel replica, + is restored to full token order across CP when explicitly typed, and is + identical on every PP stage in that replica. + + Attributes: + loss: Detached weighted mean for the complete optimizer window. + loss_sum: Detached numerator summed across DP and CP, then synchronized + across PP stages. + weight_sum: Detached full-window denominator summed across DP, but not + CP, then synchronized across PP stages. + loss_fn_outputs: Detached per-Datum mappings in input order. + """ + + loss: torch.Tensor + loss_sum: torch.Tensor + weight_sum: torch.Tensor + loss_fn_outputs: list[dict[str, Any]] + + @dataclass(frozen=True) class OptimStepResult: """Statistics from one completed optimizer step. @@ -416,7 +444,7 @@ def forward_backward( self, datums: Sequence[Datum], loss_fn: LossFn, - ) -> tuple[torch.Tensor, list[dict[str, Any]]]: + ) -> ForwardBackwardResult: """Accumulate gradients for a complete optimizer window. ``datums`` is a flat optimizer accumulation window. The Engine groups @@ -456,14 +484,16 @@ def forward_backward( collated loss inputs. Returns: - ``(loss, loss_fn_outputs)``. ``loss`` is a detached scalar reduced - over the DP-CP gradient group and, for pipeline execution, - synchronized across PP stages. ``loss_fn_outputs`` contains - mappings for this DP replica's outer Datums in window order; - pipeline execution returns the same mappings on every physical - stage rank in that replica. Model parameters are unchanged, but - their gradients contain the complete window's globally normalized - backward result. + Structured loss statistics and callback outputs. ``loss_sum`` is + reduced across the DP-CP gradient group, while ``weight_sum`` is + reduced across DP only so replicated CP weights are not counted + twice. Both are synchronized across PP stages, and ``loss`` is + their safe quotient. + ``loss_fn_outputs`` contains mappings for this DP replica's outer + Datums in window order; pipeline execution returns the same + mappings on every physical stage rank in that replica. Model + parameters are unchanged, but their gradients contain the complete + window's globally normalized backward result. """ microbatches = self._group_datums(datums) self._validate_parallelism() @@ -600,8 +630,14 @@ def forward_backward( raise output_error raise RuntimeError("another data-parallel replica failed while restoring loss_fn outputs") - loss = (step_state[0] / safe_denominator).detach() - return loss, loss_fn_outputs + loss_sum = step_state[0].detach() + loss = (loss_sum / safe_denominator).detach() + return ForwardBackwardResult( + loss=loss, + loss_sum=loss_sum, + weight_sum=denominator.detach(), + loss_fn_outputs=loss_fn_outputs, + ) @torch.no_grad() def optim_step( diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 6aa46d0b99..afc577d8d6 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -1216,11 +1216,6 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: Metrics for the completed optimizer step. """ - num_label_tokens = torch.tensor( - sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long - ) - num_label_tokens = self._dp_allreduce(num_label_tokens).item() - # number of tokens in the batch, excluding any tail padding. num_tokens_in_batch = torch.tensor( sum(batch["labels"].numel() - count_tail_padding(batch["labels"]) for batch in batches), @@ -1228,10 +1223,12 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - reporting_loss, _ = self.engine.forward_backward( + forward_backward_result = self.engine.forward_backward( [self._make_engine_datum(batch) for batch in batches], self._engine_loss_fn, ) + reporting_loss = forward_backward_result.loss + num_label_tokens = int(forward_backward_result.weight_sum.item()) step_result = self.engine.optim_step(before_optimizer_step=self.checkpointer.maybe_wait_for_staging) # Precompute FP8 scales diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index eacd5582fb..a1cefc8580 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -552,6 +552,7 @@ def setup(self): cfg_quantization=self.cfg.get("quantization", None), ) capability_model = model.parts[0] if isinstance(model, AutoPipeline) else model + self._has_joint_drafter = getattr(capability_model, "drafter", None) is not None _validate_cp_vision_frame_sharding_support(capability_model, self.cp_vision_frame_sharding) apply_te_patches() optimizer = self.cfg.optimizer.build(model, device_mesh=self.device_mesh, is_peft=self.peft_config is not None) @@ -1062,11 +1063,6 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: Returns: Metrics for the completed optimizer step. """ - num_label_tokens = torch.tensor( - sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long - ) - num_label_tokens = self._dp_allreduce(num_label_tokens).item() - # number of tokens in the batch, excluding any tail padding. num_tokens_in_batch = torch.tensor( sum(batch["labels"].numel() - count_tail_padding(batch["labels"]) for batch in batches), @@ -1074,7 +1070,17 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: ) num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item() - log_drafter = self.step_scheduler.is_remote_logging_step + log_drafter = self.step_scheduler.is_remote_logging_step and self._has_joint_drafter + log_denominator = None + if log_drafter: + # The optional drafter breakdown is emitted from inside the first + # loss callback, before forward_backward can return its weight_sum. + # Preserve its exact global mean only on logging steps; ordinary + # steps avoid this otherwise-duplicate collective. + local_label_tokens = torch.tensor( + sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long + ) + log_denominator = max(self._dp_allreduce(local_label_tokens).item(), 1) def engine_loss_fn( out: Any, @@ -1087,13 +1093,15 @@ def engine_loss_fn( out, loss_inputs, log_drafter=should_log, - log_denominator=max(num_label_tokens, 1), + log_denominator=log_denominator, ) - reporting_loss, _ = self.engine.forward_backward( + forward_backward_result = self.engine.forward_backward( [self._make_engine_datum(batch) for batch in batches], engine_loss_fn, ) + reporting_loss = forward_backward_result.loss + num_label_tokens = int(forward_backward_result.weight_sum.item()) step_result = self.engine.optim_step( before_optimizer_step=self.checkpointer.maybe_wait_for_staging, diff --git a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py index 72e90ab4e4..9c9da7e743 100644 --- a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py +++ b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py @@ -241,8 +241,8 @@ def engine_context(model_inputs): model_inputs=batch, loss_fn_inputs={"labels": labels, "weights": torch.ones_like(labels, dtype=torch.float32)}, ) - local, _ = engine.forward_backward([datum], loss_fn) - losses.append(float(local.detach())) + result = engine.forward_backward([datum], loss_fn) + losses.append(float(result.loss.detach())) embed = model_part0.get_input_embeddings() egrad = ( diff --git a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py index 1f24e015f7..7e2ad78edd 100644 --- a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py +++ b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py @@ -279,8 +279,8 @@ def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name= model_inputs={"input_ids": input_ids.clone(), "position_ids": pos.clone()}, loss_fn_inputs={"labels": labels, "weights": torch.ones_like(labels, dtype=torch.float32)}, ) - local, _ = engine.forward_backward([datum], loss_fn) - losses.append(float(local.detach())) + result = engine.forward_backward([datum], loss_fn) + losses.append(float(result.loss.detach())) # (3) embeddings receive gradients -- only the first PP stage owns embed_tokens. embed = model_part0.get_input_embeddings() diff --git a/tests/functional_tests/context_parallel/run_dense_packed_cp.py b/tests/functional_tests/context_parallel/run_dense_packed_cp.py index d590574b43..e87e119994 100644 --- a/tests/functional_tests/context_parallel/run_dense_packed_cp.py +++ b/tests/functional_tests/context_parallel/run_dense_packed_cp.py @@ -206,7 +206,7 @@ def loss_fn(output, loss_inputs): observed["labels"] = loss_inputs["labels"].detach() return local_logits.float().square().mean(dim=-1) - engine_loss, _ = Engine( + forward_backward_result = Engine( cp_model, device=device, mesh_context=mesh_context, @@ -230,16 +230,16 @@ def loss_fn(output, loss_inputs): # The Engine must normalize the exact distributed logits tightly. The # baseline comparison is looser because independently executed BF16 # attention paths can differ by roughly one ULP, which doubles under x**2. - torch.testing.assert_close(engine_loss.float(), cp_loss, atol=1e-6, rtol=1e-5) - torch.testing.assert_close(engine_loss.float(), baseline_loss.detach(), atol=1e-6, rtol=1e-2) + torch.testing.assert_close(forward_backward_result.loss.float(), cp_loss, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(forward_backward_result.loss.float(), baseline_loss.detach(), atol=1e-6, rtol=1e-2) torch.testing.assert_close(cp_logits, baseline_logits, atol=3e-2, rtol=3e-2) for cp_grad, baseline_grad in zip(cp_grads, baseline_grads): torch.testing.assert_close(cp_grad, baseline_grad, atol=1e-3, rtol=5e-2) assert torch.isfinite(cp_logits).all() assert all(torch.isfinite(cp_grad).all() for cp_grad in cp_grads) if dist.get_rank() == 0: - loss_diff = (engine_loss.float() - baseline_loss.detach()).abs().item() - normalization_diff = (engine_loss.float() - cp_loss).abs().item() + loss_diff = (forward_backward_result.loss.float() - baseline_loss.detach()).abs().item() + normalization_diff = (forward_backward_result.loss.float() - cp_loss).abs().item() output_diff = (cp_logits.float() - baseline_logits.float()).abs().max().item() grad_diff = max( (cp_grad - baseline_grad).abs().max().item() for cp_grad, baseline_grad in zip(cp_grads, baseline_grads) diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index 41eb425151..e0b39513e0 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -464,9 +464,11 @@ def _run_thd_layout( collate_fn=collate_prebatched, ) - pre_eval_loss, pre_eval_outputs = engine.forward_backward([datum], _token_losses) - torch.testing.assert_close(pre_eval_loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) - assert pre_eval_outputs == [] + pre_eval_result = engine.forward_backward([datum], _token_losses) + torch.testing.assert_close(pre_eval_result.loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) + torch.testing.assert_close(pre_eval_result.loss_sum.float(), reference_eval_loss_sum.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(pre_eval_result.weight_sum.float(), reference_eval_weight_sum.float(), atol=0, rtol=0) + assert pre_eval_result.loss_fn_outputs == [] pre_eval_grad_diff = _assert_local_grad_parity(pipeline, reference_grads) for part in pipeline.parts: part.zero_grad(set_to_none=True) @@ -494,17 +496,19 @@ def _run_thd_layout( # Reuse the exact pipeline immediately. Together with the training call # above, this proves both train->eval and eval->train schedule transitions # restore temporary split/loss callbacks and backward state. - loss, outputs = engine.forward_backward([datum], _token_losses) + train_result = engine.forward_backward([datum], _token_losses) - torch.testing.assert_close(loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) - assert outputs == [] + torch.testing.assert_close(train_result.loss.float(), reference_loss.float(), atol=2e-3, rtol=2e-3) + torch.testing.assert_close(train_result.loss_sum.float(), reference_eval_loss_sum.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(train_result.weight_sum.float(), reference_eval_weight_sum.float(), atol=0, rtol=0) + assert train_result.loss_fn_outputs == [] grad_diff = _assert_local_grad_parity(pipeline, reference_grads) if dist.get_rank() == 0: print( f"PP2 {layout} THD forward+backward parity passed " f"(eval loss-sum diff=" f"{(forward_result.loss_sum.float() - reference_eval_loss_sum.float()).abs().item():.6f}, " - f"train loss diff={(loss.float() - reference_loss.float()).abs().item():.6f}, " + f"train loss diff={(train_result.loss.float() - reference_loss.float()).abs().item():.6f}, " f"pre/post-eval grad max={pre_eval_grad_diff:.6f}/{grad_diff:.6f})" ) return pipeline @@ -579,16 +583,17 @@ def loss_with_outputs(output, loss_inputs): if execution == "forward": torch.testing.assert_close(result.loss_sum.float(), reference_loss_sum.float(), atol=4e-2, rtol=2e-3) torch.testing.assert_close(result.weight_sum.float(), reference_weight_sum.float(), atol=0, rtol=0) - outputs = result.loss_fn_outputs if any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): raise AssertionError(f"PP2 x CP{mesh_context.cp_size} {layout} forward unexpectedly created gradients") else: - loss, outputs = result reference_loss = reference_loss_sum / reference_weight_sum - torch.testing.assert_close(loss.float(), reference_loss.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(result.loss.float(), reference_loss.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(result.loss_sum.float(), reference_loss_sum.float(), atol=4e-2, rtol=2e-3) + torch.testing.assert_close(result.weight_sum.float(), reference_weight_sum.float(), atol=0, rtol=0) if not any(parameter.grad is not None for part in pipeline.parts for parameter in part.parameters()): raise AssertionError(f"PP2 x CP{mesh_context.cp_size} {layout} backward created no gradients") + outputs = result.loss_fn_outputs output_ids = torch.stack([item["sample_id"] for item in outputs]).to(torch.long) expected_ids = torch.tensor([11, 22, 33, 44], device=device) torch.testing.assert_close(output_ids, expected_ids) @@ -610,7 +615,7 @@ def loss_with_outputs(output, loss_inputs): assert all(torch.equal(probe, flat_probe) for probe in gathered_probe) if dist.get_rank() == 0: datum_counts = "2+2" if layout == "raw" else "3+1" - reported_loss = result.loss_sum.item() if execution == "forward" else result[0].item() + reported_loss = result.loss_sum.item() if execution == "forward" else result.loss.item() print( f"PP2 x CP{mesh_context.cp_size} {layout} {execution} explicit loss/output routing passed " f"({datum_counts} Datums; loss={reported_loss:.6f})" @@ -642,7 +647,7 @@ def loss_with_output(output, loss_inputs): losses = _token_losses(output, loss_inputs) return losses, [{"sample_id": loss_inputs["sample_id"][0], "score": logits.float().mean()}] - loss, outputs = Engine( + result = Engine( pipeline, device=device, mesh_context=mesh_context, @@ -650,13 +655,13 @@ def loss_with_output(output, loss_inputs): ).forward_backward(datums, loss_with_output) expected_ids = torch.tensor([17, 29], device=device) - output_ids = torch.stack([item["sample_id"] for item in outputs]).to(device=device, dtype=torch.long) + output_ids = torch.stack([item["sample_id"] for item in result.loss_fn_outputs]).to(device=device, dtype=torch.long) torch.testing.assert_close(output_ids, expected_ids) gathered = [torch.empty_like(output_ids) for _ in range(dist.get_world_size())] dist.all_gather(gathered, output_ids) assert all(torch.equal(ids, expected_ids) for ids in gathered) - assert torch.isfinite(loss) - assert all(torch.isfinite(item["score"]) for item in outputs) + assert torch.isfinite(result.loss) + assert all(torch.isfinite(item["score"]) for item in result.loss_fn_outputs) if dist.get_rank() == 0: print("PP2 padded per-Datum outputs passed (logical order [17, 29] synchronized on both ranks)") diff --git a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py index 482f87c37f..6bac72c669 100644 --- a/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py +++ b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py @@ -216,9 +216,9 @@ def loss_fn(output: torch.Tensor, loss_inputs: dict[str, torch.Tensor]) -> torch optimizers=optimizer, max_grad_norm=1e6, ) - engine_loss, _ = engine.forward_backward([datum], loss_fn) + forward_backward_result = engine.forward_backward([datum], loss_fn) torch.testing.assert_close(observed["output"], y_ref, rtol=1e-4, atol=1e-5) - torch.testing.assert_close(engine_loss, loss_ref.to(torch.float64), rtol=1e-5, atol=1e-7) + torch.testing.assert_close(forward_backward_result.loss, loss_ref.to(torch.float64), rtol=1e-5, atol=1e-7) n_local_experts = _N_EXPERTS // world_size start = rank * n_local_experts diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 54d42bcf9b..12233fc004 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -34,6 +34,7 @@ from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.optim.optimizer import LRSchedulerConfig, build_optimizer_config from nemo_automodel.components.training.step_scheduler import StepSchedulerConfig +from nemo_automodel.engine import ForwardBackwardResult from nemo_automodel.recipes._typed_config import ( _STEP_SCHEDULER_RUNTIME_KEYS, _as_dict, @@ -392,6 +393,7 @@ def _build_engine_recipe_for_optim_step(*, pp_enabled: bool = False): recipe.moe_mesh = None recipe.loss_fn = object() recipe.model_parts = [_TensorModel()] + recipe._has_joint_drafter = False recipe.pp_enabled = pp_enabled if pp_enabled: recipe.pp = SimpleNamespace(info=SimpleNamespace(has_first_stage=True)) @@ -406,7 +408,12 @@ def _build_engine_recipe_for_optim_step(*, pp_enabled: bool = False): recipe._get_dp_group_size = lambda include_cp=True: 1 recipe._get_cp_group_size = lambda: 1 recipe.engine = MagicMock() - recipe.engine.forward_backward.return_value = (torch.tensor(0.25), []) + recipe.engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.25), + loss_sum=torch.tensor(1.0), + weight_sum=torch.tensor(4.0), + loss_fn_outputs=[], + ) recipe.engine.optim_step.return_value = SimpleNamespace(grad_norm=2.5, learning_rates=(0.01,)) return recipe @@ -447,6 +454,7 @@ def test_run_train_step_passes_flat_prebatched_datums_to_engine(): @pytest.mark.cuda(False) def test_train_step_logs_joint_drafter_only_on_first_engine_loss_call(): recipe = _build_engine_recipe_for_optim_step() + recipe._has_joint_drafter = True recipe.step_scheduler.is_remote_logging_step = True batches = [ {"labels": torch.tensor([[1, -100, 2]]), "input_ids": torch.tensor([[1, 2, 3]])}, @@ -457,7 +465,12 @@ def test_train_step_logs_joint_drafter_only_on_first_engine_loss_call(): def forward_backward(datums, loss_fn): for datum in datums: loss_fn(object(), datum.loss_fn_inputs) - return torch.tensor(0.25), [] + return ForwardBackwardResult( + loss=torch.tensor(0.25), + loss_sum=torch.tensor(1.0), + weight_sum=torch.tensor(4.0), + loss_fn_outputs=[], + ) recipe.engine.forward_backward.side_effect = forward_backward @@ -473,10 +486,33 @@ def forward_backward(datums, loss_fn): assert second_call.kwargs["log_denominator"] == 4 +@pytest.mark.cuda(False) +def test_train_step_does_not_reduce_drafter_denominator_for_regular_model(): + recipe = _build_engine_recipe_for_optim_step() + recipe.step_scheduler.is_remote_logging_step = True + reductions = [] + + def allreduce(tensor, include_cp=False): + reductions.append(tensor.clone()) + return tensor + + recipe._dp_allreduce = allreduce + batch = {"labels": torch.tensor([[1, -100, 2]]), "input_ids": torch.tensor([[1, 2, 3]])} + + recipe._run_train_optim_step([batch]) + + assert len(reductions) == 1 # Throughput tokens only; Engine owns the loss denominator. + + @pytest.mark.cuda(False) def test_run_train_step_uses_engine_for_empty_supervision(): recipe = _build_engine_recipe_for_optim_step() - recipe.engine.forward_backward.return_value = (torch.tensor(0.0), []) + recipe.engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.0), + loss_sum=torch.tensor(0.0), + weight_sum=torch.tensor(0.0), + loss_fn_outputs=[], + ) batch = {"labels": torch.full((1, 4), -100), "input_ids": torch.arange(4).reshape(1, 4)} metrics = recipe._run_train_optim_step([batch]) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 6108d8e11d..61004ae690 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -40,6 +40,7 @@ from nemo_automodel.components.loss.mtp import PipelineCausalLMLoss from nemo_automodel.components.models.deepseek_v4.cp import dsv4_cp_local_seq_multiple from nemo_automodel.components.optim.optimizer import build_optimizer_config +from nemo_automodel.engine import ForwardBackwardResult from nemo_automodel.recipes._typed_config import RecipeConfig, _as_dict, _callable_and_kwargs from nemo_automodel.recipes.llm.train_ft import ( TrainFinetuneRecipeForNextTokenPrediction, @@ -2079,7 +2080,12 @@ def _make_recipe( object.__setattr__(recipe, "checkpointer", SimpleNamespace(maybe_wait_for_staging=lambda: None)) object.__setattr__(recipe, "loss_fn", object()) engine = MagicMock() - engine.forward_backward.return_value = (torch.tensor(0.5), []) + engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.5), + loss_sum=torch.tensor(2.0), + weight_sum=torch.tensor(4.0), + loss_fn_outputs=[], + ) engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) object.__setattr__(recipe, "engine", engine) object.__setattr__(recipe, "timestamp", 0.0) @@ -2096,7 +2102,12 @@ def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypat make_datum = MagicMock(side_effect=datums) monkeypatch.setattr(recipe, "_make_engine_datum", make_datum) engine = MagicMock() - engine.forward_backward.return_value = (torch.tensor(0.25), []) + engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.25), + loss_sum=torch.tensor(1.75), + weight_sum=torch.tensor(7.0), + loss_fn_outputs=[], + ) engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.02,)) object.__setattr__(recipe, "engine", engine) optimizer = SimpleNamespace( @@ -2115,6 +2126,7 @@ def test_pp_engine_owns_forward_backward_and_token_normalization(self, monkeypat assert metrics.metrics["loss"] == pytest.approx(0.25) assert metrics.metrics["grad_norm"] == pytest.approx(1.0) assert metrics.metrics["lr"] == pytest.approx(0.02) + assert metrics.metrics["num_label_tokens"] == 7 def test_pp_thd_batch_uses_engine(self, monkeypatch): recipe = self._make_recipe(monkeypatch, pp_enabled=True) @@ -2124,7 +2136,12 @@ def test_pp_thd_batch_uses_engine(self, monkeypatch): "qkv_format": "thd", } engine = MagicMock() - engine.forward_backward.return_value = (torch.tensor(0.5), []) + engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.5), + loss_sum=torch.tensor(1.0), + weight_sum=torch.tensor(2.0), + loss_fn_outputs=[], + ) engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) object.__setattr__(recipe, "engine", engine) diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 7feaf964c8..74ca76fb00 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -51,7 +51,7 @@ from nemo_automodel.components.distributed.pipelining import AutoPipeline from nemo_automodel.components.models.common.mtp import prepare_mtp_context_parallel_inputs from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler -from nemo_automodel.engine import Engine, ForwardResult, OptimStepResult, collate_prebatched +from nemo_automodel.engine import Engine, ForwardBackwardResult, ForwardResult, OptimStepResult, collate_prebatched from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput @@ -482,13 +482,16 @@ def test_forward_backward_uses_one_denominator_for_the_window(): initial_weight = model.weight.detach().clone() engine = Engine(model, device="cpu") - loss, outputs = engine.forward_backward( + result = engine.forward_backward( [_datum([1, 2]), _datum([3])], _identity_loss, ) - assert loss.item() == pytest.approx(2.0) - assert outputs == [] + assert isinstance(result, ForwardBackwardResult) + assert result.loss.item() == pytest.approx(2.0) + assert result.loss_sum.item() == pytest.approx(6.0) + assert result.weight_sum.item() == pytest.approx(3.0) + assert result.loss_fn_outputs == [] assert model.weight.grad.item() == pytest.approx(2.0) assert torch.equal(model.weight, initial_weight) assert model.forward_calls == 2 @@ -502,7 +505,7 @@ def recording_collate(datums): return collate_datums(datums) model = ScaleModel() - loss, _ = Engine( + result = Engine( model, device="cpu", microbatch_size=2, @@ -511,7 +514,7 @@ def recording_collate(datums): assert group_sizes == [2, 2, 1] assert model.forward_calls == 3 - assert loss.item() == pytest.approx(3.0) + assert result.loss.item() == pytest.approx(3.0) def test_raw_thd_packed_collater_is_prepared_by_context_parallel_sharder(): @@ -524,14 +527,14 @@ def loss_fn(output, inputs): assert inputs["weights"].shape == output.shape == (3,) return output - loss, _ = Engine( + result = Engine( model, device="cpu", microbatch_size=2, collate_fn=partial(collate_datums, packed=True), ).forward_backward([_datum([1, 2]), _datum([3])], loss_fn) - assert loss.item() == pytest.approx(2.0) + assert result.loss.item() == pytest.approx(2.0) assert "seq_lens" not in seen assert "seq_lens_padded" not in seen assert seen["cu_seqlens"].tolist() == [0, 2, 3] @@ -615,9 +618,9 @@ def loss_fn(output, inputs): torch.testing.assert_close(inputs["advantages"], torch.tensor([[0.5, 0.6, 0.0, 0.0]])) return output - loss, _ = engine.forward_backward([datum], loss_fn) + result = engine.forward_backward([datum], loss_fn) - assert loss.item() == pytest.approx(11 / 6) + assert result.loss.item() == pytest.approx(11 / 6) assert model.weight.grad.item() == pytest.approx(11 / 6) assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(2.0) assert not cp_context_active @@ -776,8 +779,9 @@ def loss_fn(output, loss_inputs): assert result.weight_sum.item() == pytest.approx(8) assert model.weight.grad is None else: - loss, _ = result - assert loss.item() == pytest.approx(46 / 8) + assert result.loss.item() == pytest.approx(46 / 8) + assert result.loss_sum.item() == pytest.approx(46) + assert result.weight_sum.item() == pytest.approx(8) assert model.weight.grad.item() == pytest.approx(46 / 8) @@ -823,23 +827,27 @@ def prepare_mtp_inputs_for_cp(self, batch, *, ignore_index): def test_weights_mask_loss_and_denominator(): model = ScaleModel() - loss, _ = Engine(model, device="cpu").forward_backward( + result = Engine(model, device="cpu").forward_backward( [_datum([1, 100], [1.0, 0.0]), _datum([3, 5], [0.5, 1.0])], _identity_loss, ) - assert loss.item() == pytest.approx(3.0) + assert result.loss.item() == pytest.approx(3.0) + assert result.loss_sum.item() == pytest.approx(7.5) + assert result.weight_sum.item() == pytest.approx(2.5) assert model.weight.grad.item() == pytest.approx(3.0) def test_fractional_weight_sum_below_one_is_not_clamped(): model = ScaleModel() - loss, _ = Engine(model, device="cpu").forward_backward( + result = Engine(model, device="cpu").forward_backward( [_datum([2, 4], [0.2, 0.3])], _identity_loss, ) - assert loss.item() == pytest.approx(3.2) + assert result.loss.item() == pytest.approx(3.2) + assert result.loss_sum.item() == pytest.approx(1.6) + assert result.weight_sum.item() == pytest.approx(0.5) assert model.weight.grad.item() == pytest.approx(3.2) @@ -849,13 +857,13 @@ def test_loss_fn_outputs_follow_datum_order_and_are_detached(): def loss_with_outputs(output, _loss_inputs): return output, [{"first_token": row.flatten()[0], "model_value": row.sum()} for row in output] - _, outputs = Engine(model, device="cpu", microbatch_size=2).forward_backward( + result = Engine(model, device="cpu", microbatch_size=2).forward_backward( [_datum([1, 2]), _datum([3]), _datum([4])], loss_with_outputs, ) - assert [item["first_token"].item() for item in outputs] == [1, 3, 4] - assert all(not item["model_value"].requires_grad for item in outputs) + assert [item["first_token"].item() for item in result.loss_fn_outputs] == [1, 3, 4] + assert all(not item["model_value"].requires_grad for item in result.loss_fn_outputs) @pytest.mark.parametrize("execution", ["forward", "forward_backward"]) @@ -871,7 +879,7 @@ def loss_with_outputs(output, _loss_inputs): ) result = getattr(Engine(ScaleModel(), device="cpu", microbatch_size=2), execution)(datums, loss_with_outputs) - outputs = result.loss_fn_outputs if execution == "forward" else result[1] + outputs = result.loss_fn_outputs assert [item["sample_id"].item() for item in outputs] == [11, 22] torch.testing.assert_close(outputs[0]["token_probe"], torch.tensor([[1.0, 101.0], [2.0, 102.0]])) @@ -1159,11 +1167,11 @@ def policy_loss(logits, inputs): losses = -(ratio * inputs["advantages"]) return losses, [{"policy_sum": (losses * inputs["weights"]).sum()}] - loss, outputs = Engine(model, device="cpu").forward_backward([datum], policy_loss) + result = Engine(model, device="cpu").forward_backward([datum], policy_loss) - assert torch.isfinite(loss) - assert torch.isfinite(outputs[0]["policy_sum"]) - assert not outputs[0]["policy_sum"].requires_grad + assert torch.isfinite(result.loss) + assert torch.isfinite(result.loss_fn_outputs[0]["policy_sum"]) + assert not result.loss_fn_outputs[0]["policy_sum"].requires_grad assert model.embedding.weight.grad is not None assert model.output.weight.grad is not None @@ -1230,7 +1238,7 @@ def loss_fn(output, inputs): assert output.shape == inputs["weights"].shape == (1, 2) return output - loss, outputs = Engine( + result = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), @@ -1244,9 +1252,9 @@ def loss_fn(output, inputs): loss_fn, ) - assert loss.item() == pytest.approx(4.5) + assert result.loss.item() == pytest.approx(4.5) assert model.weight.grad.item() == pytest.approx(4.5) - assert outputs == [] + assert result.loss_fn_outputs == [] assert pipeline.step_calls == 2 # The fake schedule performs and counts every backward, then returns None. # A second Engine-owned backward would either fail or change these counts. @@ -1434,7 +1442,7 @@ def test_pipeline_outputs_follow_logical_microbatch_order(): model = ScaleModel() pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) - _, outputs = Engine( + result = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), @@ -1446,7 +1454,7 @@ def test_pipeline_outputs_follow_logical_microbatch_order(): assert pipeline.step_calls == 1 assert pipeline.backward_calls == 2 - assert [item["metric"].item() for item in outputs] == [1.0, 2.0] + assert [item["metric"].item() for item in result.loss_fn_outputs] == [1.0, 2.0] def _packed_layout_datums(lengths: list[int]) -> list[Datum]: @@ -1535,7 +1543,7 @@ def loss_with_outputs(output, loss_inputs): ), execution, )(datums, loss_with_outputs) - outputs = result.loss_fn_outputs if execution == "forward" else result[1] + outputs = result.loss_fn_outputs assert [item["sample_id"].item() for item in outputs] == [11, 22, 33, 44] for datum, item in zip(datums, outputs): @@ -1577,18 +1585,19 @@ def loss_with_outputs(output, loss_inputs): if execution == "forward": assert result.loss_sum.item() == pytest.approx(36.0) assert result.weight_sum.item() == pytest.approx(8.0) - outputs = result.loss_fn_outputs assert pipeline.eval_calls == 1 assert pipeline.step_calls == 0 assert pipeline.backward_calls == 0 assert model.weight.grad is None else: - loss, outputs = result - assert loss.item() == pytest.approx(4.5) + assert result.loss.item() == pytest.approx(4.5) + assert result.loss_sum.item() == pytest.approx(36.0) + assert result.weight_sum.item() == pytest.approx(8.0) assert pipeline.eval_calls == 0 assert pipeline.step_calls == 1 assert pipeline.backward_calls == 2 assert model.weight.grad.item() == pytest.approx(4.5) + outputs = result.loss_fn_outputs assert [item["sample_id"].item() for item in outputs] == [11, 22, 33, 44] @@ -1911,9 +1920,9 @@ def loss_fn(output, inputs): assert output.shape == inputs["weights"].shape == (1, 2) return output - loss, _ = engine.forward_backward([datum], loss_fn) + result = engine.forward_backward([datum], loss_fn) - assert loss.item() == pytest.approx(18 / 8) + assert result.loss.item() == pytest.approx(18 / 8) assert len(pipeline.prepared_inputs) == 1 assert [item["input_ids"].shape for item in pipeline.prepared_inputs[0]] == [(1, 2), (1, 2)] torch.testing.assert_close(seen[0][0], torch.tensor([[1.0, 4.0]])) @@ -1942,14 +1951,14 @@ def test_pipeline_groups_multiple_flat_datums_into_one_outer_batch(): model = ScaleModel() pipeline = _FakeAutoPipeline(model) - loss, _ = Engine( + result = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), microbatch_size=2, ).forward_backward([_datum([1]), _datum([2])], _identity_loss) - assert loss.item() == pytest.approx(1.5) + assert result.loss.item() == pytest.approx(1.5) assert pipeline.step_calls == 1 assert model.forward_calls == 2 @@ -2056,9 +2065,9 @@ def test_prebatched_datum_keeps_existing_recipe_batch_layout(): loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0], [1.0, 0.0]])}, ) - loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) + result = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) - assert loss.item() == pytest.approx(2.0) + assert result.loss.item() == pytest.approx(2.0) assert model.weight.grad.item() == pytest.approx(2.0) @@ -2072,9 +2081,9 @@ def test_prebatched_datum_keeps_vlm_media_layout(): loss_fn_inputs={"weights": torch.ones(2)}, ) - loss, _ = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) + result = Engine(model, device="cpu", collate_fn=collate_prebatched).forward_backward([datum], _identity_loss) - assert torch.isfinite(loss) + assert torch.isfinite(result.loss) assert model.text.weight.grad is not None assert model.vision.weight.grad is not None @@ -2082,22 +2091,24 @@ def test_prebatched_datum_keeps_vlm_media_layout(): def test_scalar_loss_is_a_local_weighted_sum_numerator(): model = ScaleModel() - loss, _ = Engine(model, device="cpu").forward_backward( + result = Engine(model, device="cpu").forward_backward( [_datum([1, 100], [1.0, 0.0]), _datum([3, 5], [0.5, 1.0])], lambda output, inputs: (output * inputs["weights"]).sum(), ) - assert loss.item() == pytest.approx(3.0) + assert result.loss.item() == pytest.approx(3.0) assert model.weight.grad.item() == pytest.approx(3.0) def test_zero_weights_run_graph_connected_zero_backward(): model = ScaleModel() - loss, _ = Engine(model, device="cpu").forward_backward( + result = Engine(model, device="cpu").forward_backward( [_datum([1, 2], [0.0, 0.0])], _identity_loss, ) - assert loss.item() == 0 + assert result.loss.item() == 0 + assert result.loss_sum.item() == 0 + assert result.weight_sum.item() == 0 assert model.forward_calls == 1 assert model.weight.grad.item() == 0 @@ -2170,9 +2181,11 @@ def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: assert model.module.forward_calls == 0 window = [_datum([1, 2]), _datum([3])] if rank == 0 else [_datum([4]), _datum([5, 6])] - loss, outputs = Engine(model, device="cpu").forward_backward(window, _identity_loss) - assert loss.item() == pytest.approx(3.5) - assert outputs == [] + result = Engine(model, device="cpu").forward_backward(window, _identity_loss) + assert result.loss.item() == pytest.approx(3.5) + assert result.loss_sum.item() == pytest.approx(21.0) + assert result.weight_sum.item() == pytest.approx(6.0) + assert result.loss_fn_outputs == [] assert model.module.weight.grad.item() == pytest.approx(3.5) finally: dist.destroy_process_group() @@ -2209,14 +2222,16 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size ) ] - loss, _ = Engine( + result = Engine( model, device="cpu", mesh_context=mesh_context, collate_fn=collate_prebatched, ).forward_backward(window, _identity_loss) - assert loss.item() == pytest.approx(4.5) + assert result.loss.item() == pytest.approx(4.5) + assert result.loss_sum.item() == pytest.approx(18.0 if dp_size == 1 else 36.0) + assert result.weight_sum.item() == pytest.approx(4.0 if dp_size == 1 else 8.0) assert model.module.weight.grad.item() == pytest.approx(4.5) model.module.weight.grad = None diff --git a/tests/unit_tests/test_engine_recipe_integration.py b/tests/unit_tests/test_engine_recipe_integration.py index 29c8f79413..c9f5c9035d 100644 --- a/tests/unit_tests/test_engine_recipe_integration.py +++ b/tests/unit_tests/test_engine_recipe_integration.py @@ -148,6 +148,6 @@ def local_reduce(value, include_cp=False): assert model.forward_calls == 2 assert optimizer.step_calls == 1 assert optimizer.zero_calls == 1 - assert reductions == 2 # token counters only; Engine already reduced the loss + assert reductions == 1 # throughput tokens only; Engine returns the loss denominator for actual, expected in zip(model.parameters(), reference.parameters()): torch.testing.assert_close(actual, expected) From 99349fe2d6170ef891c08f3564c5b9432f3d8e5f Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 09:57:59 -0700 Subject: [PATCH 15/34] feat(engine): support planned multi-call accumulation Signed-off-by: HuiyingLi --- nemo_automodel/engine/__init__.py | 669 +++++++++++++++++++++++++--- tests/unit_tests/test_engine.py | 711 +++++++++++++++++++++++++++++- 2 files changed, 1314 insertions(+), 66 deletions(-) diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 4319f178e9..6f4ba155b0 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -91,6 +91,14 @@ def _as_tuple(value: _T | Sequence[_T] | None) -> tuple[_T, ...]: return (value,) +def _tensor_version(tensor: torch.Tensor) -> int: + """Return the in-place mutation counter when the Tensor exposes one.""" + try: + return int(tensor._version) + except RuntimeError: + return -1 + + def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLossInputs | dict[str, torch.Tensor]]: """Return one already-collated Datum without changing its layout. @@ -175,7 +183,7 @@ class ForwardResult: @dataclass(frozen=True) class ForwardBackwardResult: - """Training-window loss statistics and per-Datum callback outputs. + """One backward call's loss statistics and per-Datum callback outputs. The numerator is summed across the DP-CP gradient group. The full-sequence denominator is summed across DP only because CP ranks begin with replicated @@ -186,11 +194,11 @@ class ForwardBackwardResult: identical on every PP stage in that replica. Attributes: - loss: Detached weighted mean for the complete optimizer window. + loss: Detached weighted mean for this call's Datum window. loss_sum: Detached numerator summed across DP and CP, then synchronized across PP stages. - weight_sum: Detached full-window denominator summed across DP, but not - CP, then synchronized across PP stages. + weight_sum: Detached denominator for this call, summed across DP but + not CP, then synchronized across PP stages. loss_fn_outputs: Detached per-Datum mappings in input order. """ @@ -216,6 +224,28 @@ class OptimStepResult: learning_rates: tuple[float, ...] +@dataclass(frozen=True) +class _PlannedDatum: + datum: Datum + weights: torch.Tensor + weights_version: int + weights_shape: torch.Size + weights_dtype: torch.dtype + weights_device: torch.device + weights_sum: float + + +@dataclass +class _AccumulationState: + windows: tuple[tuple[_PlannedDatum, ...], ...] + weight_sums: tuple[torch.Tensor, ...] + total_weight_sum: torch.Tensor + total_microbatches: int + microbatch_size: int + next_window: int = 0 + status: str = "active" + + class Engine: """Run model forward or forward/backward over Datum windows. @@ -226,10 +256,11 @@ class Engine: gradient-accumulation synchronization, and backward. When optimizers are provided, :meth:`optim_step` owns distributed gradient finalization, clipping, parameter updates, gradient clearing, model post-step hooks, and - LR-scheduler advancement. One :meth:`forward_backward` call represents the - complete optimizer accumulation window whose gradients :meth:`optim_step` - consumes. Dynamic loss scaling and overflow-skipped updates are not part of - this contract. + LR-scheduler advancement. By default, one :meth:`forward_backward` call is + the complete optimizer accumulation window consumed by + :meth:`optim_step`. Call :meth:`begin_accumulation` first when that window + must be split across multiple ``forward_backward`` calls. Dynamic loss + scaling and overflow-skipped updates are not part of this contract. Args: model: An already configured and distributed model, or a built @@ -317,6 +348,183 @@ def __init__( self.optimizers = _as_tuple(optimizers) self.lr_schedulers = _as_tuple(lr_schedulers) self.max_grad_norm = max_grad_norm + self._accumulation_state: _AccumulationState | None = None + self._optim_step_consumed = False + self._grads_finalized = False + self._finalized_grad_norm: torch.Tensor | float | None = None + self._optim_step_in_progress = False + self._implicit_backward_status = "idle" + + def begin_accumulation(self, windows: Sequence[Sequence[Datum]]) -> None: + """Plan one optimizer window split across multiple backward calls. + + The complete plan is required before the first backward because the + supervised loss and MoE auxiliary loss use different global + denominators. Each subsequent :meth:`forward_backward` call must pass + the exact planned Datum objects for the next window, in order. A + :class:`ForwardBackwardResult` continues to describe only that call; + callers combine results with ``sum(loss_sum) / sum(weight_sum)``. + + The first version supports eager/DDP/FSDP execution. Pipeline schedules + currently finalize gradients at the end of every schedule invocation, + so planned multi-call accumulation with PP fails explicitly instead of + silently treating each call as a complete optimizer window. If a + planned backward call fails after execution starts, partial distributed + state cannot be rolled back safely: the plan becomes broken and the + Engine must not be stepped or reused. To make rank-local loss-callback + failures fail together before backward, an explicit plan performs one + small control consensus per outer microbatch; the ordinary one-call + path adds no such collective. + + Args: + windows: Non-empty sequence of non-empty Datum windows in their + future call order. The same Datum objects and weight tensors + must be passed unchanged to :meth:`forward_backward`. Every + non-final window must end on an Engine outer-microbatch + boundary. + + Raises: + RuntimeError: If no optimizer is configured, another plan is + active, or gradients from an earlier optimizer window have not + been cleared. + NotImplementedError: If pipeline parallelism is enabled. + ValueError: If the plan is empty, malformed, or splits an outer + microbatch across calls. + """ + if self.pipeline is not None: + raise NotImplementedError( + "planned multi-call accumulation is not supported with pipeline parallelism; " + "pipeline schedules currently finalize gradients after every schedule call" + ) + if self._accumulation_state is not None: + raise RuntimeError("an Engine accumulation plan is already active") + + self._validate_parallelism() + dp_group, dp_size = self._dp_group_and_size() + control_group, control_group_size = self._accumulation_control_group_and_size() + local_error: Exception | None = None + planned_windows: list[tuple[_PlannedDatum, ...]] = [] + microbatch_counts: list[int] = [] + local_weight_sums: list[float] = [] + try: + if not self.optimizers: + raise RuntimeError("Engine.begin_accumulation requires at least one optimizer") + if not isinstance(windows, Sequence) or isinstance(windows, (str, bytes)) or not windows: + raise ValueError("begin_accumulation requires a non-empty sequence of Datum windows") + if ( + self._implicit_backward_status != "idle" + or self._grads_finalized + or any(parameter.grad is not None for part in self.model_parts for parameter in part.parameters()) + ): + raise RuntimeError("begin_accumulation requires cleared gradients") + + for window_index, window in enumerate(windows): + microbatches = self._group_datums(window) + if window_index < len(windows) - 1 and len(window) % self.microbatch_size != 0: + raise ValueError("every non-final accumulation window must end on an outer-microbatch boundary") + + planned_window: list[_PlannedDatum] = [] + local_weight_sum = 0.0 + for datum in window: + weights = datum.loss_fn_inputs.get("weights") + if not isinstance(weights, torch.Tensor): + raise ValueError("every Datum must contain a Tensor loss_fn_inputs['weights']") + if weights.numel() == 0 or not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()): + raise ValueError("Datum weights must be non-empty, finite, and non-negative") + weight_sum = float(weights.to(torch.float64).sum()) + local_weight_sum += weight_sum + planned_window.append( + _PlannedDatum( + datum=datum, + weights=weights, + weights_version=_tensor_version(weights), + weights_shape=weights.shape, + weights_dtype=weights.dtype, + weights_device=weights.device, + weights_sum=weight_sum, + ) + ) + planned_windows.append(tuple(planned_window)) + microbatch_counts.append(len(microbatches)) + local_weight_sums.append(local_weight_sum) + except Exception as error: + local_error = error + + self._synchronize_accumulation_error( + local_error, + control_group, + control_group_size, + peer_message="another model-parallel rank rejected the accumulation plan", + ) + self._validate_window_size_across_group(len(planned_windows), control_group, control_group_size) + for count in microbatch_counts: + self._validate_window_size_across_group(count, control_group, control_group_size) + + local_denominators = torch.tensor(local_weight_sums, dtype=torch.float64, device=self.device) + cp_group, cp_size = self._cp_group_and_size() + cp_error: Exception | None = None + if cp_size > 1: + gathered_cp_denominators = torch.empty( + cp_size * len(local_weight_sums), + dtype=local_denominators.dtype, + device=local_denominators.device, + ) + dist.all_gather_into_tensor(gathered_cp_denominators, local_denominators, group=cp_group) + gathered_cp_denominators = gathered_cp_denominators.view(cp_size, len(local_weight_sums)) + if not torch.allclose( + gathered_cp_denominators, + gathered_cp_denominators[0].expand_as(gathered_cp_denominators), + rtol=1e-8, + atol=1e-12, + ): + cp_error = ValueError( + "context-parallel ranks must plan identical full-sequence weights; " + f"got {gathered_cp_denominators.tolist()}" + ) + self._synchronize_accumulation_error( + cp_error, + control_group, + control_group_size, + peer_message="another context-parallel group rejected the planned weight sums", + ) + + global_denominators = local_denominators.clone() + if dp_size > 1: + dist.all_reduce(global_denominators, op=dist.ReduceOp.SUM, group=dp_group) + weight_sums = list(global_denominators.detach().unbind()) + + if control_group_size > 1: + local_denominators = torch.stack(weight_sums) + gathered_denominators = torch.empty( + control_group_size * len(weight_sums), + dtype=local_denominators.dtype, + device=local_denominators.device, + ) + dist.all_gather_into_tensor(gathered_denominators, local_denominators, group=control_group) + gathered_denominators = gathered_denominators.view(control_group_size, len(weight_sums)) + if not torch.allclose( + gathered_denominators, + gathered_denominators[0].expand_as(gathered_denominators), + rtol=1e-8, + atol=1e-12, + ): + raise ValueError( + "every model-parallel rank must plan the same DP-global weight sums; " + f"got {gathered_denominators.tolist()}" + ) + + total_weight_sum = torch.stack(weight_sums).sum() + self._accumulation_state = _AccumulationState( + windows=tuple(planned_windows), + weight_sums=tuple(weight_sums), + total_weight_sum=total_weight_sum, + total_microbatches=sum(microbatch_counts), + microbatch_size=self.microbatch_size, + ) + self._optim_step_consumed = False + self._grads_finalized = False + self._finalized_grad_norm = None + self._implicit_backward_status = "idle" @torch.no_grad() def forward( @@ -355,6 +563,8 @@ def forward( record per hidden inner sample, and cannot produce records when PP splits it into multiple inner microbatches. """ + if self._accumulation_state is not None: + raise RuntimeError("Engine.forward cannot run while a backward accumulation plan is active") microbatches = self._group_datums(datums) self._validate_execution_parallelism() cp_group, cp_size = self._cp_group_and_size() @@ -445,7 +655,95 @@ def forward_backward( datums: Sequence[Datum], loss_fn: LossFn, ) -> ForwardBackwardResult: - """Accumulate gradients for a complete optimizer window. + """Run one backward window, optionally inside a predeclared accumulation plan. + + Without :meth:`begin_accumulation`, this call remains a complete + optimizer window. With an active plan, calls must consume its Datum + windows in order. Every returned result is call-local even though all + gradients use the plan's full-step normalization. + """ + state = self._accumulation_state + if state is None: + if self.optimizers: + if self._implicit_backward_status == "ready": + raise RuntimeError( + "forward_backward already produced the current optimizer window; " + "call optim_step, or declare multiple calls up front with begin_accumulation" + ) + if self._implicit_backward_status == "broken": + raise RuntimeError("the previous implicit backward window failed and this Engine cannot be reused") + if self._implicit_backward_status == "running": + raise RuntimeError("an implicit forward_backward call is already running") + if self._grads_finalized: + raise RuntimeError( + "gradients were already finalized; retry optim_step before another forward_backward call" + ) + if self.optimizers: + self._implicit_backward_status = "running" + try: + result = self._forward_backward_window(datums, loss_fn) + except BaseException: + if self.optimizers: + self._implicit_backward_status = "broken" + raise + self._optim_step_consumed = False + self._grads_finalized = False + self._finalized_grad_norm = None + if self.optimizers: + self._implicit_backward_status = "ready" + return result + + if state.status == "broken": + raise RuntimeError("the active accumulation plan is broken and this Engine cannot be reused") + if state.status == "running": + raise RuntimeError("an accumulation forward_backward call is already running") + if state.status == "ready": + raise RuntimeError("the accumulation plan is complete; call optim_step before another backward") + if state.next_window >= len(state.windows): + raise RuntimeError("the accumulation plan has no remaining backward windows") + + try: + self._validate_planned_window( + datums, + state.windows[state.next_window], + microbatch_size=state.microbatch_size, + ) + is_first_window = state.next_window == 0 + is_final_window = state.next_window == len(state.windows) - 1 + state.status = "running" + result = self._forward_backward_window( + datums, + loss_fn, + result_denominator=state.weight_sums[state.next_window], + backward_denominator=state.total_weight_sum, + total_microbatches=state.total_microbatches, + is_first_window=is_first_window, + is_final_window=is_final_window, + ) + except BaseException: + state.status = "broken" + raise + + state.next_window += 1 + state.status = "ready" if state.next_window == len(state.windows) else "active" + self._optim_step_consumed = False + self._grads_finalized = False + self._finalized_grad_norm = None + self._implicit_backward_status = "idle" + return result + + def _forward_backward_window( + self, + datums: Sequence[Datum], + loss_fn: LossFn, + *, + result_denominator: torch.Tensor | None = None, + backward_denominator: torch.Tensor | None = None, + total_microbatches: int | None = None, + is_first_window: bool = True, + is_final_window: bool = True, + ) -> ForwardBackwardResult: + """Accumulate gradients for one implicit or explicitly planned window. ``datums`` is a flat optimizer accumulation window. The Engine groups it into outer batches of ``microbatch_size`` and invokes ``collate_fn`` @@ -475,10 +773,10 @@ def forward_backward( sample use ordinary flat Datums. Args: - datums: Flat sequence of Datum items in the complete optimizer - accumulation window. A Datum's token weights may have shape - [tokens] or the custom collater's batched token layout; the - loss tensor must use the identical shape. + datums: Flat sequence of Datum items in this call's window. A + Datum's token weights may have shape [tokens] or the custom + collater's batched token layout; the loss tensor must use the + identical shape. loss_fn: Computes either that per-token loss tensor or a scalar local weighted-sum numerator from the raw model output and collated loss inputs. @@ -492,35 +790,59 @@ def forward_backward( ``loss_fn_outputs`` contains mappings for this DP replica's outer Datums in window order; pipeline execution returns the same mappings on every physical stage rank in that replica. Model - parameters are unchanged, but their gradients contain the complete - window's globally normalized backward result. + parameters are unchanged. Without an explicit accumulation plan, + gradients contain this call's globally normalized result. With a + plan, they accumulate using the complete plan's denominator even + though the returned statistics remain call-local. """ microbatches = self._group_datums(datums) self._validate_parallelism() dp_group, dp_size = self._dp_group_and_size() grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) - denominator = self._global_weight_sum(microbatches, dp_group, dp_size) + denominator = ( + self._global_weight_sum(microbatches, dp_group, dp_size) + if result_denominator is None + else result_denominator + ) + gradient_denominator = denominator if backward_denominator is None else backward_denominator + planned_accumulation = result_denominator is not None + plan_control_group, plan_control_group_size = ( + self._accumulation_control_group_and_size() if planned_accumulation else (None, 1) + ) zero_denominator = bool(denominator == 0) + zero_gradient_denominator = bool(gradient_denominator == 0) safe_denominator = torch.where(denominator > 0, denominator, torch.ones_like(denominator)) + safe_gradient_denominator = torch.where( + gradient_denominator > 0, + gradient_denominator, + torch.ones_like(gradient_denominator), + ) self._validate_pipeline_window(len(microbatches), denominator) pp_enabled = self.pipeline is not None for part in self.model_parts: part.train() - prepare_for_grad_accumulation(self.model_parts, pp_enabled=pp_enabled) + if is_first_window: + prepare_for_grad_accumulation(self.model_parts, pp_enabled=pp_enabled) inner_microbatches = self.pipeline.num_microbatches if self.pipeline is not None else 1 - MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( - self._cp_size() / (len(microbatches) * inner_microbatches) + effective_total_microbatches = ( + len(microbatches) * inner_microbatches if total_microbatches is None else total_microbatches ) + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(self._cp_size() / effective_total_microbatches) local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) loss_fn_outputs: list[dict[str, Any]] = [] returns_outputs: bool | None = None output_error: Exception | None = None + backward_scale = ( + safe_gradient_denominator.new_zeros(()) + if zero_denominator or zero_gradient_denominator + else safe_gradient_denominator.new_tensor(grad_group_size) / safe_gradient_denominator + ) for index, datums in enumerate(microbatches): - is_last = index == len(microbatches) - 1 + is_last = is_final_window and index == len(microbatches) - 1 if is_last: prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) @@ -529,11 +851,6 @@ def forward_backward( ) if self.pipeline is not None: - backward_scale = ( - safe_denominator.new_zeros(()) - if zero_denominator - else safe_denominator.new_tensor(grad_group_size) / safe_denominator - ) batch_returns_outputs, batch_outputs, batch_error = self._pipeline_execute( model_inputs, loss_inputs, @@ -570,10 +887,27 @@ def forward_backward( ): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - numerator, parsed_outputs, output_parse_error = _parse_loss_result( - loss_fn(output, loss_inputs), loss_inputs["weights"] - ) - if output_error is None and dp_size <= 1: + loss_error: Exception | None = None + numerator: torch.Tensor | None = None + parsed_outputs: ParsedLossOutputs = None + output_parse_error: Exception | None = None + try: + numerator, parsed_outputs, output_parse_error = _parse_loss_result( + loss_fn(output, loss_inputs), loss_inputs["weights"] + ) + except Exception as error: + loss_error = error + if planned_accumulation: + self._synchronize_accumulation_error( + loss_error, + plan_control_group, + plan_control_group_size, + peer_message="another model-parallel rank failed in the planned loss callback", + ) + elif loss_error is not None: + raise loss_error + assert numerator is not None + if output_error is None and dp_size <= 1 and not planned_accumulation: self._validate_loss_fn_outputs_across_cp( parsed_outputs, loss_inputs.get("weights"), @@ -585,11 +919,11 @@ def forward_backward( returns_outputs = _update_output_mode(returns_outputs, parsed_outputs) if zero_denominator: numerator = numerator * 0 - (numerator * (grad_group_size / safe_denominator)).backward() + (numerator * backward_scale).backward() if output_error is None: try: - if dp_size > 1: + if dp_size > 1 or planned_accumulation: self._validate_loss_fn_outputs_across_cp( parsed_outputs, loss_inputs.get("weights"), @@ -613,7 +947,7 @@ def forward_backward( except Exception as error: output_error = error local_loss_sum.add_(numerator.detach().to(torch.float64)) - if index == 0: + if is_first_window and index == 0: prepare_after_first_microbatch() # Piggyback the output-error bit on the existing end-of-window loss @@ -625,10 +959,15 @@ def forward_backward( pp_group, pp_size = self._pp_group_and_size() if pp_size > 1: dist.all_reduce(step_state, op=dist.ReduceOp.SUM, group=pp_group) + if planned_accumulation: + if plan_control_group_size > grad_group_size: + control_error = step_state[1].clamp(max=1) + dist.all_reduce(control_error, op=dist.ReduceOp.MAX, group=plan_control_group) + step_state[1].copy_(control_error) if bool(step_state[1] > 0): if output_error is not None: raise output_error - raise RuntimeError("another data-parallel replica failed while restoring loss_fn outputs") + raise RuntimeError("another model-parallel rank failed while restoring loss_fn outputs") loss_sum = step_state[0].detach() loss = (loss_sum / safe_denominator).detach() @@ -654,12 +993,18 @@ def optim_step( checkpointers use that fence to preserve ``finalize/clip -> wait -> step`` overlap. + Once the first optimizer mutation begins, failures from an optimizer, + model post-step hook, or scheduler cannot be rolled back and are + process-fatal in distributed execution. + Args: before_optimizer_step: Optional callback invoked exactly once after gradient finalization and clipping, but before the first optimizer step. If it raises, parameters, optimizer state, model post-step state, and schedulers remain untouched; the - finalized gradients remain available. + finalized gradients remain available. Outside an explicit + accumulation plan, distributed ranks must invoke this callback + consistently; rank-local execution failures are process-fatal. Returns: Gradient norm and post-scheduler learning rates for the completed @@ -668,47 +1013,145 @@ def optim_step( Raises: RuntimeError: If this Engine was constructed without optimizers. """ - if not self.optimizers: - raise RuntimeError("Engine.optim_step requires at least one optimizer") - if before_optimizer_step is not None and not callable(before_optimizer_step): - raise TypeError("before_optimizer_step must be callable or None") + state = self._accumulation_state + local_preflight_error: Exception | None = None + try: + if not self.optimizers: + raise RuntimeError("Engine.optim_step requires at least one optimizer") + if before_optimizer_step is not None and not callable(before_optimizer_step): + raise TypeError("before_optimizer_step must be callable or None") + if self._optim_step_in_progress: + raise RuntimeError("Engine.optim_step is already running") + if state is not None: + if state.status == "broken": + raise RuntimeError("the active accumulation plan is broken and cannot be optimized") + if state.status != "ready": + raise RuntimeError( + "the active accumulation plan must finish every forward_backward call before optim_step" + ) + else: + if self._implicit_backward_status == "broken": + raise RuntimeError( + "the previous implicit backward window failed and this Engine cannot be optimized" + ) + if self._implicit_backward_status == "running": + raise RuntimeError("an implicit forward_backward call is still running") + if self._optim_step_consumed: + raise RuntimeError("optim_step already consumed the current gradients; run forward_backward first") + except Exception as error: + local_preflight_error = error + + control_group: dist.ProcessGroup | None = None + control_group_size = 1 + # Explicit plans pay the small control collectives needed to fail + # together before mutation. Keep the ordinary one-call path free of + # new per-step collectives; its distributed callback contract remains + # the same as before planned accumulation was introduced. + synchronize_update = state is not None + if synchronize_update: + control_group, control_group_size = self._accumulation_control_group_and_size() + if synchronize_update: + self._synchronize_accumulation_error( + local_preflight_error, + control_group, + control_group_size, + peer_message="another model-parallel rank rejected optim_step", + ) + elif local_preflight_error is not None: + raise local_preflight_error device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None moe_mesh = self.mesh_context.moe_mesh if self.mesh_context is not None else None dp_group, dp_size = self._dp_group_and_size() _, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) pp_enabled = self.pipeline is not None - grad_norm = scale_grads_and_clip_grad_norm( - max_grad_norm=self.max_grad_norm, - model_parts=self.model_parts, - norm_type=2.0, - pp_enabled=pp_enabled, - device_mesh=device_mesh, - moe_mesh=moe_mesh, - ep_axis_name="ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None, - pp_axis_name="pp" if pp_enabled else None, - foreach=True, - num_label_tokens=None, - dp_group_size=grad_group_size, - expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, device_mesh), - ) + self._optim_step_in_progress = True + mutation_started = False + try: + if not self._grads_finalized: + finalization_error: Exception | None = None + try: + self._finalized_grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=self.max_grad_norm, + model_parts=self.model_parts, + norm_type=2.0, + pp_enabled=pp_enabled, + device_mesh=device_mesh, + moe_mesh=moe_mesh, + ep_axis_name=( + "ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None + ), + pp_axis_name="pp" if pp_enabled else None, + foreach=True, + num_label_tokens=None, + dp_group_size=grad_group_size, + expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, device_mesh), + ) + if self._finalized_grad_norm is None: + raise RuntimeError("gradient finalization did not return a gradient norm") + except Exception as error: + finalization_error = error + if synchronize_update: + self._synchronize_accumulation_error( + finalization_error, + control_group, + control_group_size, + peer_message="another model-parallel rank failed while finalizing gradients", + ) + elif finalization_error is not None: + raise finalization_error + self._grads_finalized = True + grad_norm = self._finalized_grad_norm + assert grad_norm is not None + + fence_error: Exception | None = None + if before_optimizer_step is not None: + try: + before_optimizer_step() + except Exception as error: + fence_error = error + if synchronize_update: + self._synchronize_accumulation_error( + fence_error, + control_group, + control_group_size, + peer_message="another model-parallel rank failed before the optimizer mutation fence", + ) + elif fence_error is not None: + raise fence_error - if before_optimizer_step is not None: - before_optimizer_step() + mutation_started = True + for optimizer in self.optimizers: + optimizer.step() + optimizer.zero_grad(set_to_none=True) - for optimizer in self.optimizers: - optimizer.step() - optimizer.zero_grad(set_to_none=True) + for part in self.model_parts: + update_moe_gate_bias = getattr(part, "update_moe_gate_bias", None) + if callable(update_moe_gate_bias): + update_moe_gate_bias() - for part in self.model_parts: - update_moe_gate_bias = getattr(part, "update_moe_gate_bias", None) - if callable(update_moe_gate_bias): - update_moe_gate_bias() + for scheduler in self.lr_schedulers: + scheduler.step(1) - for scheduler in self.lr_schedulers: - scheduler.step(1) - - learning_rates = tuple(float(group["lr"]) for optimizer in self.optimizers for group in optimizer.param_groups) + learning_rates = tuple( + float(group["lr"]) for optimizer in self.optimizers for group in optimizer.param_groups + ) + except Exception: + if mutation_started or not self._grads_finalized: + self._optim_step_consumed = True + if state is not None: + state.status = "broken" + else: + self._implicit_backward_status = "broken" + raise + finally: + self._optim_step_in_progress = False + + self._accumulation_state = None + self._optim_step_consumed = True + self._grads_finalized = False + self._finalized_grad_norm = None + self._implicit_backward_status = "idle" return OptimStepResult(grad_norm=grad_norm, learning_rates=learning_rates) def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: @@ -1644,6 +2087,33 @@ def _gradient_group_and_size( size = int(dp_cp_mesh.size()) return (dp_cp_mesh.get_group() if size > 1 else None), size + def _accumulation_control_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: + """Return the full non-PP model group used to keep plan control flow aligned.""" + if not dist.is_available() or not dist.is_initialized(): + return None, 1 + if self.mesh_context is None: + return None, dist.get_world_size() + if (group := getattr(self.mesh_context, "process_group", None)) is not None: + return group, dist.get_world_size(group=group) + if (device_mesh := getattr(self.mesh_context, "device_mesh", None)) is None: + return None, dist.get_world_size() + + root_mesh = device_mesh._get_root_mesh() if hasattr(device_mesh, "_get_root_mesh") else device_mesh + size = int(root_mesh.size()) + if size <= 1: + return None, 1 + if root_mesh.ndim == 1: + return root_mesh.get_group(), size + if size == dist.get_world_size(): + return None, size + for flat_mesh in getattr(root_mesh, "_flatten_mapping", {}).values(): + if int(flat_mesh.size()) == size: + return flat_mesh.get_group(), size + raise NotImplementedError( + "planned multi-call accumulation on a rank-subset multi-axis DeviceMesh requires " + "MeshContext.process_group for the complete model" + ) + def _pp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: if self.pipeline is None or not dist.is_available() or not dist.is_initialized(): return None, 1 @@ -1689,6 +2159,75 @@ def _local_weight_sum(self, microbatches: list[list[Datum]]) -> torch.Tensor: self._validate_weight_sum_across_cp(denominator) return denominator + def _synchronize_accumulation_error( + self, + local_error: Exception | None, + group: dist.ProcessGroup | None, + group_size: int, + *, + peer_message: str, + ) -> None: + """Make every gradient rank reject a bad accumulation contract together.""" + if group_size <= 1: + if local_error is not None: + raise local_error + return + + failed = torch.tensor(int(local_error is not None), dtype=torch.int64, device=self.device) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group) + if bool(failed): + if local_error is not None: + raise local_error + raise RuntimeError(peer_message) + + def _validate_planned_window( + self, + datums: Sequence[Datum], + planned: tuple[_PlannedDatum, ...], + *, + microbatch_size: int, + ) -> None: + """Validate an accumulation-plan slot before any model collective.""" + local_error: Exception | None = None + try: + if self.microbatch_size != microbatch_size: + raise RuntimeError( + "Engine.microbatch_size changed after begin_accumulation; " + f"planned {microbatch_size}, found {self.microbatch_size}" + ) + if not isinstance(datums, Sequence) or isinstance(datums, (str, bytes)): + raise TypeError("planned forward_backward input must be a sequence of Datum") + if len(datums) != len(planned): + raise ValueError( + f"planned forward_backward window has {len(planned)} Datums, but received {len(datums)}" + ) + for index, (datum, expected) in enumerate(zip(datums, planned)): + if datum is not expected.datum: + raise ValueError( + f"planned forward_backward window Datum {index} is not the object declared to begin_accumulation" + ) + weights = datum.loss_fn_inputs.get("weights") + if weights is not expected.weights: + raise ValueError(f"planned Datum {index} replaced its weights Tensor") + if ( + _tensor_version(weights) != expected.weights_version + or weights.shape != expected.weights_shape + or weights.dtype != expected.weights_dtype + or weights.device != expected.weights_device + or float(weights.to(torch.float64).sum()) != expected.weights_sum + ): + raise ValueError(f"planned Datum {index} weights changed after begin_accumulation") + except Exception as error: + local_error = error + + control_group, control_group_size = self._accumulation_control_group_and_size() + self._synchronize_accumulation_error( + local_error, + control_group, + control_group_size, + peer_message="another model-parallel rank changed its planned Datum window", + ) + def _validate_window_size_across_group( self, size: int, diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 74ca76fb00..6166577684 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -66,6 +66,21 @@ def forward(self, input_ids: torch.Tensor, **_) -> torch.Tensor: return input_ids.to(torch.float32) * self.weight +class _MainAndAuxScaleModel(nn.Module): + """Small model whose main and auto-scaled auxiliary gradients are separable.""" + + def __init__(self) -> None: + super().__init__() + self.main_weight = nn.Parameter(torch.tensor(1.0)) + self.aux_weight = nn.Parameter(torch.tensor(1.0)) + self.forward_calls = 0 + + def forward(self, input_ids: torch.Tensor, **_) -> torch.Tensor: + self.forward_calls += 1 + output = input_ids.to(torch.float32) * self.main_weight + return MoEAuxLossAutoScaler.apply(output, self.aux_weight) + + class _SubMesh: def __init__(self, size, rank=0): self._size = size @@ -389,6 +404,14 @@ def fail_before_step(): torch.testing.assert_close(model.weight, torch.tensor(1.0)) torch.testing.assert_close(model.weight.grad, torch.tensor(2.0)) + # Retrying the mutation fence must consume the already-finalized gradients; + # running finalization/scaling twice would silently change the update. + result = engine.optim_step() + assert events == ["finalize", "before-step", "step", "zero", "gate", "scheduler"] + torch.testing.assert_close(model.weight, torch.tensor(0.8)) + assert model.weight.grad is None + assert result.learning_rates == (0.1,) + def test_optim_step_requires_an_optimizer(monkeypatch): monkeypatch.setattr( @@ -497,6 +520,348 @@ def test_forward_backward_uses_one_denominator_for_the_window(): assert model.forward_calls == 2 +def test_planned_multi_call_matches_one_window_with_unequal_denominators(): + window_a = [_datum([2, 100], [1.0, 0.0])] + window_b = [_datum([4, 8], [0.5, 1.5])] + + reference_model = ScaleModel() + reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) + reference_engine = Engine( + reference_model, + device="cpu", + optimizers=reference_optimizer, + max_grad_norm=None, + ) + reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) + reference_grad = reference_model.weight.grad.detach().clone() + reference_step = reference_engine.optim_step() + + planned_model = ScaleModel() + planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) + planned_engine = Engine( + planned_model, + device="cpu", + optimizers=planned_optimizer, + max_grad_norm=None, + ) + planned_engine.begin_accumulation([window_a, window_b]) + result_a = planned_engine.forward_backward(window_a, _identity_loss) + result_b = planned_engine.forward_backward(window_b, _identity_loss) + planned_grad = planned_model.weight.grad.detach().clone() + planned_step = planned_engine.optim_step() + + assert result_a.loss_sum.item() == pytest.approx(2.0) + assert result_a.weight_sum.item() == pytest.approx(1.0) + assert result_a.loss.item() == pytest.approx(2.0) + assert result_b.loss_sum.item() == pytest.approx(14.0) + assert result_b.weight_sum.item() == pytest.approx(2.0) + assert result_b.loss.item() == pytest.approx(7.0) + assert reference_result.loss_sum.item() == pytest.approx(16.0) + assert reference_result.weight_sum.item() == pytest.approx(3.0) + assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) + torch.testing.assert_close(planned_grad, reference_grad) + torch.testing.assert_close(planned_step.grad_norm, reference_step.grad_norm) + torch.testing.assert_close(planned_model.weight, reference_model.weight) + + +def test_explicit_one_window_accumulation_matches_implicit_call(): + implicit_window = [_datum([1, 9], [0.25, 0.75])] + explicit_window = [_datum([1, 9], [0.25, 0.75])] + + implicit_model = ScaleModel() + implicit_optimizer = torch.optim.SGD(implicit_model.parameters(), lr=0.1) + implicit_engine = Engine( + implicit_model, + device="cpu", + optimizers=implicit_optimizer, + max_grad_norm=None, + ) + implicit_result = implicit_engine.forward_backward(implicit_window, _identity_loss) + + explicit_model = ScaleModel() + explicit_optimizer = torch.optim.SGD(explicit_model.parameters(), lr=0.1) + explicit_engine = Engine( + explicit_model, + device="cpu", + optimizers=explicit_optimizer, + max_grad_norm=None, + ) + explicit_engine.begin_accumulation([explicit_window]) + explicit_result = explicit_engine.forward_backward(explicit_window, _identity_loss) + + torch.testing.assert_close(explicit_result.loss, implicit_result.loss) + torch.testing.assert_close(explicit_result.loss_sum, implicit_result.loss_sum) + torch.testing.assert_close(explicit_result.weight_sum, implicit_result.weight_sum) + torch.testing.assert_close(explicit_model.weight.grad, implicit_model.weight.grad) + explicit_engine.optim_step() + implicit_engine.optim_step() + torch.testing.assert_close(explicit_model.weight, implicit_model.weight) + + +def test_begin_accumulation_requires_an_optimizer_before_forward(): + model = ScaleModel() + engine = Engine(model, device="cpu") + + with pytest.raises(RuntimeError, match="optimizer"): + engine.begin_accumulation([[_datum([1])]]) + + assert model.forward_calls == 0 + assert model.weight.grad is None + + +def test_multiple_backward_calls_with_an_optimizer_require_an_explicit_plan(): + model = ScaleModel() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=None) + + engine.forward_backward([_datum([2])], _identity_loss) + with pytest.raises(RuntimeError, match="begin_accumulation"): + engine.forward_backward([_datum([6])], _identity_loss) + + torch.testing.assert_close(model.weight.grad, torch.tensor(2.0)) + engine.optim_step() + torch.testing.assert_close(model.weight, torch.tensor(0.8)) + + +def test_failed_implicit_backward_poisoned_gradients_cannot_be_reused(): + model = ScaleModel() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=None) + calls = 0 + + def fail_on_second_microbatch(output, _loss_inputs): + nonlocal calls + calls += 1 + if calls == 2: + raise ValueError("second microbatch failed") + return output + + with pytest.raises(ValueError, match="second microbatch failed"): + engine.forward_backward([_datum([2]), _datum([6])], fail_on_second_microbatch) + + # The first microbatch already produced a partial gradient. It must never + # be consumed or silently combined with a later optimizer window. + torch.testing.assert_close(model.weight.grad, torch.tensor(1.0)) + with pytest.raises(RuntimeError, match="failed"): + engine.forward_backward([_datum([4])], _identity_loss) + with pytest.raises(RuntimeError, match="failed"): + engine.optim_step() + with pytest.raises(RuntimeError, match="cleared gradients"): + engine.begin_accumulation([[_datum([4])]]) + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + + +def test_planned_multi_call_uses_one_lifecycle_and_whole_step_moe_scale(monkeypatch): + events = [] + + @contextmanager + def recording_sync_ctx(_model, is_optim_step, _defer_fsdp_grad_sync): + events.append(f"sync-{is_optim_step}") + yield + + monkeypatch.setattr( + engine_module, + "prepare_for_grad_accumulation", + lambda *_args, **_kwargs: events.append("prepare"), + ) + monkeypatch.setattr( + engine_module, + "prepare_for_final_backward", + lambda *_args, **_kwargs: events.append("final"), + ) + monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after-first")) + monkeypatch.setattr(engine_module, "get_sync_ctx", recording_sync_ctx) + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) + + model = _MainAndAuxScaleModel() + engine = Engine(model, device="cpu", optimizers=torch.optim.SGD(model.parameters(), lr=0.1)) + # A zero-weight call still counts toward the whole-step MoE microbatch + # average even though it contributes no main-loss numerator. + window_a = [_datum([100], [0.0])] + window_b = [_datum([3]), _datum([5])] + + engine.begin_accumulation([window_a, window_b]) + result_a = engine.forward_backward(window_a, _identity_loss) + result_b = engine.forward_backward(window_b, _identity_loss) + + assert events == [ + "prepare", + "sync-False", + "after-first", + "sync-False", + "final", + "sync-True", + ] + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(1.0 / 3.0) + assert result_a.loss_sum.item() == pytest.approx(0.0) + assert result_a.weight_sum.item() == pytest.approx(0.0) + assert result_a.loss.item() == pytest.approx(0.0) + assert result_b.loss.item() == pytest.approx(4.0) + assert model.main_weight.grad.item() == pytest.approx(4.0) + assert model.aux_weight.grad.item() == pytest.approx(1.0) + assert model.forward_calls == 3 + + +def test_planned_accumulation_rejects_unfinished_extra_and_double_steps(monkeypatch): + events = [] + + class RecordingSGD(torch.optim.SGD): + def step(self, closure=None): + events.append("step") + return super().step(closure) + + def zero_grad(self, set_to_none=True): + events.append("zero") + return super().zero_grad(set_to_none=set_to_none) + + class RecordingScheduler: + def step(self, increment): + assert increment == 1 + events.append("scheduler") + + model = ScaleModel() + optimizer = RecordingSGD(model.parameters(), lr=0.1) + engine = Engine( + model, + device="cpu", + optimizers=optimizer, + lr_schedulers=RecordingScheduler(), + max_grad_norm=None, + ) + real_finalize = engine_module.scale_grads_and_clip_grad_norm + + def recording_finalize(**kwargs): + events.append("finalize") + return real_finalize(**kwargs) + + monkeypatch.setattr(engine_module, "scale_grads_and_clip_grad_norm", recording_finalize) + window_a = [_datum([1])] + window_b = [_datum([3])] + engine.begin_accumulation([window_a, window_b]) + engine.forward_backward(window_a, _identity_loss) + + with pytest.raises(RuntimeError): + engine.optim_step() + with pytest.raises(RuntimeError): + engine.begin_accumulation([window_a, window_b]) + assert events == [] + + engine.forward_backward(window_b, _identity_loss) + with pytest.raises(RuntimeError): + engine.forward_backward(window_b, _identity_loss) + result = engine.optim_step() + assert events == ["finalize", "step", "zero", "scheduler"] + assert result.learning_rates == (0.1,) + + with pytest.raises(RuntimeError): + engine.optim_step() + assert events == ["finalize", "step", "zero", "scheduler"] + + # A new successful implicit window starts a new optimizer step. + engine.forward_backward(window_a, _identity_loss) + engine.optim_step() + assert events == [ + "finalize", + "step", + "zero", + "scheduler", + "finalize", + "step", + "zero", + "scheduler", + ] + + +def test_planned_accumulation_failure_breaks_engine_before_optimizer_step(): + events = [] + + class RecordingSGD(torch.optim.SGD): + def step(self, closure=None): + events.append("step") + return super().step(closure) + + model = ScaleModel() + optimizer = RecordingSGD(model.parameters(), lr=0.1) + engine = Engine(model, device="cpu", optimizers=optimizer) + window_a = [_datum([1])] + window_b = [_datum([3])] + engine.begin_accumulation([window_a, window_b]) + engine.forward_backward(window_a, _identity_loss) + + def failing_loss(_output, _loss_inputs): + raise ValueError("loss callback failed") + + with pytest.raises(ValueError, match="loss callback failed"): + engine.forward_backward(window_b, failing_loss) + + with pytest.raises(RuntimeError): + engine.optim_step() + with pytest.raises(RuntimeError): + engine.begin_accumulation([window_a]) + with pytest.raises(RuntimeError): + engine.forward_backward(window_b, _identity_loss) + assert events == [] + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + + +@pytest.mark.parametrize("mismatch", ["datum_reference", "weights"]) +def test_planned_accumulation_validates_declared_datums_before_forward(mismatch): + model = ScaleModel() + engine = Engine( + model, + device="cpu", + microbatch_size=2, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([1], [1.0]), _datum([2], [1.0])] + engine.begin_accumulation([window]) + + actual_window = window + if mismatch == "datum_reference": + actual_window = [_datum([1], [1.0]), window[1]] + else: + window[0].loss_fn_inputs["weights"].mul_(2.0) + + with pytest.raises((RuntimeError, ValueError)): + engine.forward_backward(actual_window, _identity_loss) + assert model.forward_calls == 0 + + +def test_planned_accumulation_rejects_microbatch_size_changes_before_forward(): + model = ScaleModel() + engine = Engine( + model, + device="cpu", + microbatch_size=2, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([1]), _datum([2])] + engine.begin_accumulation([window]) + engine.microbatch_size = 1 + + with pytest.raises(RuntimeError, match="microbatch_size changed"): + engine.forward_backward(window, _identity_loss) + + assert model.forward_calls == 0 + + +def test_planned_accumulation_rejects_nonfinal_partial_outer_microbatch(): + model = ScaleModel() + engine = Engine( + model, + device="cpu", + microbatch_size=2, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window_a = [_datum([1])] + window_b = [_datum([2])] + + with pytest.raises(ValueError, match="microbatch|aligned|divisible"): + engine.begin_accumulation([window_a, window_b]) + + assert model.forward_calls == 0 + + def test_forward_backward_groups_flat_datums_by_microbatch_size(): group_sizes = [] @@ -1438,6 +1803,26 @@ def prepare_final(parts, *, pp_enabled): assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) +def test_pipeline_rejects_planned_multi_call_before_forward(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + engine = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + ) + window_a = [_datum([1, 2])] + window_b = [_datum([3, 4])] + + with pytest.raises(NotImplementedError, match="pipeline|PP"): + engine.begin_accumulation([window_a, window_b]) + + assert pipeline.step_calls == 0 + assert pipeline.backward_calls == 0 + assert model.forward_calls == 0 + assert model.weight.grad is None + + def test_pipeline_outputs_follow_logical_microbatch_order(): model = ScaleModel() pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) @@ -2247,6 +2632,82 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size assert forward_result.loss_sum.item() == pytest.approx(expected_sum) assert forward_result.weight_sum.item() == pytest.approx(4.0) assert model.module.weight.grad is None + + # Planned accumulation must use the sum of the two DP-only + # denominators while letting CP ranks contribute disjoint numerators. + # The two calls deliberately have different weight sums, and DP ranks + # deliberately own different amounts of supervision when dp_size=2. + if dp_size == 1: + window_a = [ + Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 0.0, 0.0, 0.0]])}, + ) + ] + window_b = [ + Datum( + model_inputs={"input_ids": torch.tensor([[5, 6, 7, 8]])}, + loss_fn_inputs={"weights": torch.tensor([[0.0, 1.0, 1.0, 1.0]])}, + ) + ] + expected_a_sum, expected_a_weight = 1.0, 1.0 + expected_b_sum, expected_b_weight = 21.0, 3.0 + else: + dp_rank = get_flat_mesh(mesh_context.device_mesh, "dp").get_local_rank() + if dp_rank == 0: + window_a = [ + Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 0.0, 0.0, 0.0]])}, + ) + ] + window_b = [ + Datum( + model_inputs={"input_ids": torch.tensor([[5, 6, 7, 8]])}, + loss_fn_inputs={"weights": torch.tensor([[0.0, 1.0, 1.0, 1.0]])}, + ) + ] + else: + window_a = [ + Datum( + model_inputs={"input_ids": torch.tensor([[9, 10, 11, 12]])}, + loss_fn_inputs={"weights": torch.tensor([[1.0, 1.0, 0.0, 0.0]])}, + ) + ] + window_b = [ + Datum( + model_inputs={"input_ids": torch.tensor([[13, 14, 15, 16]])}, + loss_fn_inputs={"weights": torch.tensor([[0.0, 0.0, 0.0, 1.0]])}, + ) + ] + expected_a_sum, expected_a_weight = 20.0, 3.0 + expected_b_sum, expected_b_weight = 37.0, 4.0 + + planned_model = _DDPWithCP(_DistributedCPModel()) + planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) + planned_engine = Engine( + planned_model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + optimizers=planned_optimizer, + max_grad_norm=None, + ) + planned_engine.begin_accumulation([window_a, window_b]) + result_a = planned_engine.forward_backward(window_a, _identity_loss) + result_b = planned_engine.forward_backward(window_b, _identity_loss) + + assert result_a.loss_sum.item() == pytest.approx(expected_a_sum) + assert result_a.weight_sum.item() == pytest.approx(expected_a_weight) + assert result_a.loss.item() == pytest.approx(expected_a_sum / expected_a_weight) + assert result_b.loss_sum.item() == pytest.approx(expected_b_sum) + assert result_b.weight_sum.item() == pytest.approx(expected_b_weight) + assert result_b.loss.item() == pytest.approx(expected_b_sum / expected_b_weight) + expected_global_mean = (expected_a_sum + expected_b_sum) / (expected_a_weight + expected_b_weight) + assert planned_model.module.weight.grad.item() == pytest.approx(expected_global_mean) + + planned_engine.optim_step() + assert planned_model.module.weight.item() == pytest.approx(1.0 - 0.1 * expected_global_mean) finally: dist.destroy_process_group() @@ -2361,7 +2822,7 @@ def loss_with_rank_local_output_error(output, _loss_inputs): per_datum=records, ) - expected = "per_datum contains 2 records" if rank == 0 else "another data-parallel replica" + expected = "per_datum contains 2 records" if rank == 0 else "another model-parallel rank" with pytest.raises((ValueError, RuntimeError), match=expected): Engine(model, device="cpu").forward_backward( [_datum([1, 2])], @@ -2372,6 +2833,218 @@ def loss_with_rank_local_output_error(output, _loss_inputs): dist.destroy_process_group() +def _planned_accumulation_validation_consensus_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + model = nn.parallel.DistributedDataParallel(ScaleModel()) + engine = Engine( + model, + device="cpu", + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([rank + 1])] + planned_windows = [window] if rank == 0 else [window, [_datum([rank + 2])]] + + with pytest.raises((RuntimeError, ValueError)): + engine.begin_accumulation(planned_windows) + assert model.module.forward_calls == 0 + dist.barrier() + + model = nn.parallel.DistributedDataParallel(ScaleModel()) + engine = Engine( + model, + device="cpu", + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([rank + 1])] + engine.begin_accumulation([window]) + if rank == 0: + window[0].loss_fn_inputs["weights"].mul_(2.0) + + # Rank 1 must observe rank 0's local plan-validation failure instead of + # entering DDP forward and hanging on a different collective. + with pytest.raises((RuntimeError, ValueError)): + engine.forward_backward(window, _identity_loss) + assert model.module.forward_calls == 0 + finally: + dist.destroy_process_group() + + +def _model_parallel_planned_output_error_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=1, tp_size=world_size), + world_size=world_size, + ) + model = ScaleModel() + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([1, 2])] + engine.begin_accumulation([window]) + + def loss_with_rank_local_output_error(output, _loss_inputs): + return output, ([{}, {}] if rank == 0 else [{}]) + + expected = "one mapping per Datum" if rank == 0 else "another model-parallel rank" + with pytest.raises((ValueError, RuntimeError), match=expected): + engine.forward_backward(window, loss_with_rank_local_output_error) + + # Both model-parallel ranks complete backward, then enter the same + # terminal state even though only rank 0 owns the bad output contract. + assert model.weight.grad is not None + assert engine._accumulation_state is not None + assert engine._accumulation_state.status == "broken" + with pytest.raises(RuntimeError, match="broken"): + engine.optim_step() + + model = ScaleModel() + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + window = [_datum([1, 2])] + engine.begin_accumulation([window]) + + def rank_local_loss_failure(output, _loss_inputs): + if rank == 0: + raise ValueError("rank-local loss failure") + return output + + expected = "rank-local loss failure" if rank == 0 else "another model-parallel rank" + with pytest.raises((ValueError, RuntimeError), match=expected): + engine.forward_backward(window, rank_local_loss_failure) + + # Loss/shape errors are agreed before backward, so no peer enters a + # TP/EP/DP backward collective while another exits locally. + assert model.weight.grad is None + assert engine._accumulation_state is not None + assert engine._accumulation_state.status == "broken" + finally: + dist.destroy_process_group() + + +def _model_parallel_optimizer_fence_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + real_finalize = engine_module.scale_grads_and_clip_grad_norm + try: + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=1, tp_size=world_size), + world_size=world_size, + ) + steps = [] + + class RecordingSGD(torch.optim.SGD): + def step(self, closure=None): + steps.append("step") + return super().step(closure) + + model = ScaleModel() + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + optimizers=RecordingSGD(model.parameters(), lr=0.1), + max_grad_norm=None, + ) + window = [_datum([1, 2])] + engine.begin_accumulation([window]) + engine.forward_backward(window, _identity_loss) + + finalize_calls = [] + + def recording_finalize(**_kwargs): + finalize_calls.append("finalize") + return torch.tensor(1.5) + + engine_module.scale_grads_and_clip_grad_norm = recording_finalize + + def rank_local_fence(): + if rank == 0: + raise ValueError("rank-local fence failure") + + expected = "rank-local fence failure" if rank == 0 else "another model-parallel rank" + with pytest.raises((ValueError, RuntimeError), match=expected): + engine.optim_step(before_optimizer_step=rank_local_fence) + + assert finalize_calls == ["finalize"] + assert steps == [] + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + assert model.weight.grad is not None + + result = engine.optim_step() + assert finalize_calls == ["finalize"] + assert steps == ["step"] + torch.testing.assert_close(result.grad_norm, torch.tensor(1.5)) + torch.testing.assert_close(model.weight, torch.tensor(0.85)) + finally: + engine_module.scale_grads_and_clip_grad_norm = real_finalize + dist.destroy_process_group() + + +def _planned_cp_preflight_consensus_worker(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=2, cp_size=2), + world_size=world_size, + ) + model = ScaleModel() + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + ) + # Only one CP subgroup disagrees. The full-control error reduction must + # stop all four ranks before any rank enters the later DP all-reduce. + weight = 2.0 if rank == 0 else 1.0 + window = [_datum([rank + 1], [weight])] + + with pytest.raises((ValueError, RuntimeError), match="context-parallel"): + engine.begin_accumulation([window]) + + assert engine._accumulation_state is None + assert model.forward_calls == 0 + assert model.weight.grad is None + dist.barrier() + finally: + dist.destroy_process_group() + + def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): mp.spawn( _distributed_worker, @@ -2424,3 +3097,39 @@ def test_data_parallel_output_errors_propagate_after_backward_without_hanging(tm nprocs=2, join=True, ) + + +def test_planned_accumulation_validation_errors_reach_every_data_rank(tmp_path): + mp.spawn( + _planned_accumulation_validation_consensus_worker, + args=(2, str(tmp_path / "engine_planned_validation_init")), + nprocs=2, + join=True, + ) + + +def test_planned_output_error_breaks_every_model_parallel_rank_after_backward(tmp_path): + mp.spawn( + _model_parallel_planned_output_error_worker, + args=(2, str(tmp_path / "engine_model_parallel_output_init")), + nprocs=2, + join=True, + ) + + +def test_optimizer_fence_failure_reaches_every_model_parallel_rank_and_can_retry(tmp_path): + mp.spawn( + _model_parallel_optimizer_fence_worker, + args=(2, str(tmp_path / "engine_model_parallel_fence_init")), + nprocs=2, + join=True, + ) + + +def test_planned_cp_preflight_failure_reaches_all_dp_cp_ranks(tmp_path): + mp.spawn( + _planned_cp_preflight_consensus_worker, + args=(4, str(tmp_path / "engine_planned_cp_preflight_init")), + nprocs=4, + join=True, + ) From ea2308a715e76f2686e97ba459da4204e3e7599b Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 11:02:47 -0700 Subject: [PATCH 16/34] fix(engine): centralize FP8 scale precompute Signed-off-by: HuiyingLi --- nemo_automodel/components/quantization/fp8.py | 13 +- nemo_automodel/engine/__init__.py | 46 ++++++- nemo_automodel/recipes/llm/train_ft.py | 13 -- nemo_automodel/recipes/vlm/finetune.py | 12 -- tests/unit_tests/quantization/test_fp8.py | 20 ++- .../recipes/test_finetune_vlm_helpers.py | 16 +-- tests/unit_tests/recipes/test_train_ft.py | 36 ++--- tests/unit_tests/test_engine.py | 127 ++++++++++++++++++ 8 files changed, 214 insertions(+), 69 deletions(-) diff --git a/nemo_automodel/components/quantization/fp8.py b/nemo_automodel/components/quantization/fp8.py index ceac1cf186..8aca5b2e41 100644 --- a/nemo_automodel/components/quantization/fp8.py +++ b/nemo_automodel/components/quantization/fp8.py @@ -190,13 +190,6 @@ def apply_fp8_to_model( if not HAVE_TORCHAO: raise ImportError(MISSING_TORCHAO_MSG) - # Set precompute attribute on model - model.precompute_float8_dynamic_scale_for_fsdp = ( - fp8_config.precompute_float8_dynamic_scale_for_fsdp - and fp8_config.recipe_name == "tensorwise" - and fp8_config.enable_fsdp_float8_all_gather - ) - # Handle config creation or recipe-based configuration if fp8_config.recipe_name is not None and fp8_config.recipe_name != "tensorwise": torchao_config = Float8LinearConfig.from_recipe_name(fp8_config.recipe_name) @@ -215,6 +208,12 @@ def apply_fp8_to_model( ) logger.info("Using FP8 tensorwise scaling") + # Record the resolved torchao capability before distributed partitioning. + # Pipeline parts inherit it when the model is copied and split. + model.precompute_float8_dynamic_scale_for_fsdp = fp8_config.precompute_float8_dynamic_scale_for_fsdp and getattr( + torchao_config, "enable_fsdp_float8_all_gather", False + ) + # Check hardware capability if not using emulation config_emulate = getattr(torchao_config, "emulate", fp8_config.emulate) if not _has_cuda_capability(8, 9) and not config_emulate: diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 6f4ba155b0..ce53311fbe 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -51,6 +51,7 @@ ) from nemo_automodel.components.utils.model_utils import filter_forward_kwargs from nemo_automodel.engine.outputs import LossFnOutputBatch, PerTokenOutput +from nemo_automodel.shared.import_utils import MISSING_TORCHAO_MSG, safe_import_from CollateFn = Callable[ [list[Datum]], @@ -99,6 +100,40 @@ def _tensor_version(tensor: torch.Tensor) -> int: return -1 +def _resolve_fp8_scale_precompute( + model_parts: Sequence[nn.Module], +) -> tuple[tuple[nn.Module, ...], Callable[[nn.Module], None] | None]: + """Resolve the torchao post-step function only for opted-in model parts. + + Args: + model_parts: Local eager model or pipeline parts after FP8 conversion + and distributed model partitioning. + + Returns: + Opted-in local model parts and the torchao precompute function. Both + are empty when no part requests FP8 FSDP scale precomputation. + + Raises: + ImportError: If an opted-in part requires a torchao API that is not + available. Resolution happens during Engine construction, before + any optimizer mutation. + """ + capable_parts = tuple( + part for part in model_parts if getattr(part, "precompute_float8_dynamic_scale_for_fsdp", False) + ) + if not capable_parts: + return (), None + + available, precompute = safe_import_from( + "torchao.float8", + "precompute_float8_dynamic_scale_for_fsdp", + msg=MISSING_TORCHAO_MSG, + ) + if not available: + raise ImportError(MISSING_TORCHAO_MSG) + return capable_parts, precompute + + def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLossInputs | dict[str, torch.Tensor]]: """Return one already-collated Datum without changing its layout. @@ -337,6 +372,9 @@ def __init__( self.pipeline = model if isinstance(model, AutoPipeline) else None self.model_parts = model.parts if self.pipeline is not None else [model] self.model = self.model_parts[0] + self._fp8_scale_precompute_parts, self._fp8_scale_precompute_fn = _resolve_fp8_scale_precompute( + self.model_parts + ) self.device = torch.device(device) self.mesh_context = mesh_context self.microbatch_size = microbatch_size @@ -991,7 +1029,9 @@ def optim_step( expert-gradient correction and global clipping once, then invokes an optional mutation fence before any optimizer changes. Async checkpointers use that fence to preserve ``finalize/clip -> wait -> - step`` overlap. + step`` overlap. After updating weights, it runs model maintenance for + MoE gate bias and opted-in FP8 FSDP scale precomputation before + advancing LR schedulers. Once the first optimizer mutation begins, failures from an optimizer, model post-step hook, or scheduler cannot be rolled back and are @@ -1130,6 +1170,10 @@ def optim_step( if callable(update_moe_gate_bias): update_moe_gate_bias() + if self._fp8_scale_precompute_fn is not None: + for part in self._fp8_scale_precompute_parts: + self._fp8_scale_precompute_fn(part) + for scheduler in self.lr_schedulers: scheduler.step(1) diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index afc577d8d6..5fdd269928 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -39,7 +39,6 @@ import torch.nn as nn import wandb from huggingface_hub import constants as hf_constants -from torchao.float8 import precompute_float8_dynamic_scale_for_fsdp from transformers import AutoConfig from nemo_automodel._transformers import ( @@ -1231,18 +1230,6 @@ def _run_train_optim_step(self, batches: list[dict[str, Any]]) -> MetricsSample: num_label_tokens = int(forward_backward_result.weight_sum.item()) step_result = self.engine.optim_step(before_optimizer_step=self.checkpointer.maybe_wait_for_staging) - # Precompute FP8 scales - fp8_config = self.cfg.get("fp8", None) - if ( - fp8_config is not None - and fp8_config.get("enabled", False) - and fp8_config.get("precompute_float8_dynamic_scale_for_fsdp", False) - and not self.pp_enabled - and self.device_mesh is not None - and self.device_mesh["dp_shard"].size() > 1 - ): - precompute_float8_dynamic_scale_for_fsdp(self.model_parts[0]) - # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api # self.model_parts[0].install_optimized_model_weights() # self.model_parts[0].zero_grad_buffer() diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index a1cefc8580..9a8ab317b5 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -37,7 +37,6 @@ import torch.nn as nn import wandb from torch.utils.data import DataLoader -from torchao.float8 import precompute_float8_dynamic_scale_for_fsdp from transformers.processing_utils import ProcessorMixin from nemo_automodel._transformers import ( @@ -1107,17 +1106,6 @@ def engine_loss_fn( before_optimizer_step=self.checkpointer.maybe_wait_for_staging, ) - # Precompute FP8 scales - fp8_config = self.cfg.get("fp8", None) - if ( - fp8_config is not None - and fp8_config.get("enabled", False) - and fp8_config.get("precompute_float8_dynamic_scale_for_fsdp", False) - and self.device_mesh is not None - and self.device_mesh["dp_shard"].size() > 1 - ): - precompute_float8_dynamic_scale_for_fsdp(self.model_parts[0]) - # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api # self.model.install_optimized_model_weights() # self.model.zero_grad_buffer() diff --git a/tests/unit_tests/quantization/test_fp8.py b/tests/unit_tests/quantization/test_fp8.py index 582037ecdb..91d1b1429c 100644 --- a/tests/unit_tests/quantization/test_fp8.py +++ b/tests/unit_tests/quantization/test_fp8.py @@ -210,6 +210,24 @@ def test_apply_fp8_to_model_disabled(self): # Should return the same model instance when disabled assert result is model + @patch("nemo_automodel.components.quantization.fp8.convert_to_float8_training") + def test_default_recipe_enables_tensorwise_fsdp_scale_precompute_capability(self, mock_convert): + """``recipe_name=None`` selects the tensorwise path and must expose the same capability.""" + model = nn.Linear(32, 64) + config = FP8Config( + enabled=True, + recipe_name=None, + enable_fsdp_float8_all_gather=True, + precompute_float8_dynamic_scale_for_fsdp=True, + emulate=True + ) + + result = apply_fp8_to_model(model, config=config) + + assert result is model + assert result.precompute_float8_dynamic_scale_for_fsdp is True + mock_convert.assert_called_once() + def test_apply_fp8_to_model_with_individual_params(self): """Test apply_fp8_to_model with individual parameters instead of config.""" model = nn.Linear(32, 64) @@ -246,7 +264,7 @@ def test_verify_fp8_conversion_with_mock_fp8(self): """Test verification with mock FP8 modules.""" # This test requires torchao to work properly try: - from torchao.float8.float8_linear import Float8Linear + from torchao.float8.float8_linear import Float8Linear # noqa: F401 except ImportError: pytest.skip("torchao not available") diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 12233fc004..f533b850c6 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -524,7 +524,7 @@ def test_run_train_step_uses_engine_for_empty_supervision(): @pytest.mark.cuda(False) -def test_run_train_step_keeps_fp8_precompute_after_engine_optim_step(monkeypatch): +def test_train_step_leaves_fp8_post_step_work_to_engine(monkeypatch): recipe = _build_engine_recipe_for_optim_step() recipe.cfg = _Cfg( fp8={ @@ -533,22 +533,18 @@ def test_run_train_step_keeps_fp8_precompute_after_engine_optim_step(monkeypatch } ) recipe.device_mesh = {"dp_shard": SimpleNamespace(size=lambda: 2)} - events = [] - - def optim_step(**kwargs): - events.append("optim_step") - return SimpleNamespace(grad_norm=2.5, learning_rates=(0.01,)) - - recipe.engine.optim_step.side_effect = optim_step monkeypatch.setattr( "nemo_automodel.recipes.vlm.finetune.precompute_float8_dynamic_scale_for_fsdp", - lambda model: events.append(("fp8_precompute", model)), + lambda _model: pytest.fail("the VLM recipe must not run FP8 post-step work directly"), + raising=False, ) batch = {"labels": torch.tensor([[1, 2]]), "input_ids": torch.tensor([[3, 4]])} recipe._run_train_optim_step([batch]) - assert events == ["optim_step", ("fp8_precompute", recipe.model_parts[0])] + recipe.engine.optim_step.assert_called_once_with( + before_optimizer_step=recipe.checkpointer.maybe_wait_for_staging, + ) def test_make_engine_datum_filters_raw_media_off_first_pipeline_stage(): diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 61004ae690..216235e5c7 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -2154,7 +2154,7 @@ def test_pp_thd_batch_uses_engine(self, monkeypatch): assert datums[0].model_inputs["qkv_format"] == "thd" assert loss_fn == recipe._engine_loss_fn - def test_fp8_scale_precompute_stays_after_engine_optim_step(self, monkeypatch): + def test_train_step_leaves_fp8_post_step_work_to_engine(self, monkeypatch): recipe = self._make_recipe(monkeypatch, pp_enabled=False) object.__setattr__( recipe, @@ -2168,36 +2168,22 @@ def test_fp8_scale_precompute_stays_after_engine_optim_step(self, monkeypatch): } ), ) - - class _DeviceMesh: - def __getitem__(self, name): - assert name == "dp_shard" - return SimpleNamespace(size=lambda: 2) - - object.__setattr__(recipe, "device_mesh", _DeviceMesh()) - events = [] - recipe.checkpointer.maybe_wait_for_staging = lambda: events.append("checkpoint_wait") - - def optim_step(*, before_optimizer_step): - events.append("optim_step_start") - before_optimizer_step() - events.append("optim_step_done") - return SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) - - recipe.engine.optim_step.side_effect = optim_step + object.__setattr__( + recipe, + "device_mesh", + {"dp_shard": SimpleNamespace(size=lambda: 2)}, + ) monkeypatch.setattr( "nemo_automodel.recipes.llm.train_ft.precompute_float8_dynamic_scale_for_fsdp", - lambda model: events.append(("fp8_precompute", model)), + lambda _model: pytest.fail("the LLM recipe must not run FP8 post-step work directly"), + raising=False, ) recipe._run_train_optim_step([{"input_ids": torch.tensor([[1, 2]]), "labels": torch.tensor([[2, -100]])}]) - assert events == [ - "optim_step_start", - "checkpoint_wait", - "optim_step_done", - ("fp8_precompute", recipe.model_parts[0]), - ] + recipe.engine.optim_step.assert_called_once_with( + before_optimizer_step=recipe.checkpointer.maybe_wait_for_staging, + ) # ----------------------------------------------------------------------------- diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 6166577684..9ae1d1dcfc 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -228,6 +228,64 @@ def test_engine_and_datum_are_lazy_top_level_exports(): assert PublicCollatedLossInputs is CollatedLossInputs +def test_fp8_scale_resolver_skips_import_without_capable_model_parts(monkeypatch): + first = ScaleModel() + second = ScaleModel() + second.precompute_float8_dynamic_scale_for_fsdp = False + monkeypatch.setattr( + engine_module, + "safe_import_from", + lambda *_args, **_kwargs: pytest.fail("non-FP8 Engine construction must not import torchao"), + ) + + parts, precompute = engine_module._resolve_fp8_scale_precompute([first, second]) + + assert parts == () + assert precompute is None + + +def test_fp8_scale_resolver_filters_capable_parts_and_caches_callable(monkeypatch): + first = ScaleModel() + disabled = ScaleModel() + third = ScaleModel() + first.precompute_float8_dynamic_scale_for_fsdp = True + disabled.precompute_float8_dynamic_scale_for_fsdp = False + third.precompute_float8_dynamic_scale_for_fsdp = True + precompute = lambda _part: None + calls = [] + + def resolve(module, symbol, *, msg): + calls.append((module, symbol, msg)) + return True, precompute + + monkeypatch.setattr(engine_module, "safe_import_from", resolve) + + parts, resolved = engine_module._resolve_fp8_scale_precompute([first, disabled, third]) + + assert parts == (first, third) + assert resolved is precompute + assert calls == [ + ( + "torchao.float8", + "precompute_float8_dynamic_scale_for_fsdp", + engine_module.MISSING_TORCHAO_MSG, + ) + ] + + +def test_capable_fp8_model_rejects_missing_torchao_api_during_engine_construction(monkeypatch): + model = ScaleModel() + model.precompute_float8_dynamic_scale_for_fsdp = True + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + monkeypatch.setattr(engine_module, "safe_import_from", lambda *_args, **_kwargs: (False, object())) + + with pytest.raises(ImportError, match="torchao"): + Engine(model, device="cpu", optimizers=optimizer) + + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + assert model.weight.grad is None + + def test_optim_step_clips_updates_and_clears_real_gradients(): model = ScaleModel() optimizer = torch.optim.SGD(model.parameters(), lr=0.1) @@ -293,6 +351,15 @@ def step(self, increment): second_scheduler = RecordingScheduler("second", second_optimizer) device_mesh = _NamedMesh(("cp", "tp"), cp=_SubMesh(2), tp=_SubMesh(4)) moe_mesh = _NamedMesh(("ep_shard", "ep"), ep_shard=_SubMesh(2), ep=_SubMesh(2)) + + def precompute_fp8_scale(part): + events.append(f"fp8-{part.name}") + + def resolve_fp8_scale_precompute(parts): + assert parts == [first, second] + return tuple(parts), precompute_fp8_scale + + monkeypatch.setattr(engine_module, "_resolve_fp8_scale_precompute", resolve_fp8_scale_precompute) engine = Engine( pipeline, device="cpu", @@ -334,6 +401,8 @@ def before_optimizer_step(): "zero-second", "gate-first", "gate-second", + "fp8-first", + "fp8-second", "scheduler-first", "scheduler-second", ] @@ -357,6 +426,64 @@ def before_optimizer_step(): assert second.weight.grad is None +def test_optim_step_fp8_post_step_failure_does_not_advance_scheduler(monkeypatch): + events = [] + + class GateModel(ScaleModel): + def update_moe_gate_bias(self): + events.append("gate") + + class RecordingSGD(torch.optim.SGD): + def step(self, closure=None): + events.append("step") + return super().step(closure) + + def zero_grad(self, set_to_none=True): + events.append("zero") + return super().zero_grad(set_to_none=set_to_none) + + class RecordingScheduler: + def step(self, increment): + events.append("scheduler") + + model = GateModel() + model.weight.grad = torch.tensor(2.0) + optimizer = RecordingSGD(model.parameters(), lr=0.1) + + def fail_fp8_post_step(part): + assert part is model + events.append("fp8") + raise RuntimeError("fp8 scale precompute failed") + + monkeypatch.setattr( + engine_module, + "_resolve_fp8_scale_precompute", + lambda parts: (tuple(parts), fail_fp8_post_step), + ) + engine = Engine( + model, + device="cpu", + optimizers=optimizer, + lr_schedulers=RecordingScheduler(), + ) + monkeypatch.setattr( + engine_module, + "scale_grads_and_clip_grad_norm", + lambda **_kwargs: torch.tensor(2.0), + ) + + with pytest.raises(RuntimeError, match="fp8 scale precompute failed"): + engine.optim_step() + + assert events == ["step", "zero", "gate", "fp8"] + torch.testing.assert_close(model.weight, torch.tensor(0.8)) + assert model.weight.grad is None + + with pytest.raises(RuntimeError, match="already consumed|cannot be optimized"): + engine.optim_step() + assert events == ["step", "zero", "gate", "fp8"] + + def test_optim_step_callback_failure_preserves_optimizer_and_post_step_state(monkeypatch): events = [] From da3713140ba30308757594fdd664c50a3d4a133a Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 12:30:47 -0700 Subject: [PATCH 17/34] fix(benchmark): run training through Engine Signed-off-by: HuiyingLi --- nemo_automodel/recipes/llm/benchmark.py | 82 ++------ .../unit_tests/recipes/llm/test_benchmark.py | 179 +++++++++--------- 2 files changed, 102 insertions(+), 159 deletions(-) diff --git a/nemo_automodel/recipes/llm/benchmark.py b/nemo_automodel/recipes/llm/benchmark.py index ec9633864d..4e354f8f31 100644 --- a/nemo_automodel/recipes/llm/benchmark.py +++ b/nemo_automodel/recipes/llm/benchmark.py @@ -20,11 +20,6 @@ from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.components.training.timers import Timers -from nemo_automodel.components.training.utils import ( - prepare_after_first_microbatch, - prepare_for_final_backward, - prepare_for_grad_accumulation, -) from nemo_automodel.components.utils.flops_utils import calculate_mfu, get_flops_formula_for_hf_config from nemo_automodel.recipes.llm.train_ft import TrainFinetuneRecipeForNextTokenPrediction @@ -104,7 +99,7 @@ class BenchmarkingRecipeForNextTokenPrediction(TrainFinetuneRecipeForNextTokenPr This class extends TrainFinetuneRecipeForNextTokenPrediction to provide a simplified benchmarking-focused training loop with timers and profiling support. - It reuses the setup() and _forward_backward_step() methods from the parent class. + It reuses the parent's setup and Engine execution paths. """ def __init__(self, cfg): @@ -288,7 +283,6 @@ def run_benchmark(self): with timers and profiling support, similar to the original benchmarking script. """ rank = self.dist_env.rank - device = self.dist_env.device # Get benchmarking config steps = self._bench_steps @@ -327,55 +321,27 @@ def run_benchmark(self): if i == nsys_start and rank in nsys_ranks: logger.info(f"Rank {rank} | Starting nsys profiling") torch.cuda.cudart().cudaProfilerStart() - # Per-microbatch NVTX ranges below already delimit the work. + # Per-window NVTX ranges below already delimit the work. # Entering emit_nvtx() without closing it leaks RecordFunction # callbacks into later DTensor/FSDP operations. if rank == 0: logger.info(f"Rank {rank} | Iteration {i}") - # Zero gradients - for opt in self.optimizer: - opt.zero_grad() - # Time the iteration iter_timer = "iteration_warmup" if i < warmup_steps else "iteration" with self.timers(iter_timer, log_level=1): - # Gradient accumulation loop - num_label_tokens = 0 - loss_buffer = [] - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) - - for ga_step_idx in range(ga_steps): - if ga_step_idx == ga_steps - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) - - # Get batch from dataloader - batch = next(dataloader_iter) - torch.cuda.nvtx.range_push(f"iteration_{i}_ga_step_{ga_step_idx}") - - # Accumulate label tokens locally - num_label_tokens += (batch["labels"] != -100).sum().item() - - with self.timers(f"forward_backward_{ga_step_idx}", log_level=2): - self._forward_backward_step( - ga_step_idx, - batch, - loss_buffer=loss_buffer, - num_label_tokens=None, - num_batches=ga_steps, - is_train=True, - ) - + batches = [next(dataloader_iter) for _ in range(ga_steps)] + datums = [self._make_engine_datum(batch) for batch in batches] + torch.cuda.nvtx.range_push(f"iteration_{i}_forward_backward") + try: + with self.timers("forward_backward", log_level=2): + forward_backward_result = self.engine.forward_backward(datums, self._engine_loss_fn) + finally: torch.cuda.nvtx.range_pop() - if ga_step_idx == 0: - prepare_after_first_microbatch() - - # Optimizer step with self.timers("optimizer", log_level=2): - for opt in self.optimizer: - opt.step() + self.engine.optim_step(before_optimizer_step=self.checkpointer.maybe_wait_for_staging) logger.debug("Optimizer step") # Match the training-loop lifecycle: record one complete eager @@ -384,26 +350,8 @@ def run_benchmark(self): self.partial_cuda_graph_manager.capture() self._partial_cuda_graph_capture_pending = False - # Synchronize num_label_tokens across DP ranks - num_label_tokens_tensor = torch.tensor(num_label_tokens, dtype=torch.long, device=device) - num_label_tokens_tensor = self._dp_allreduce(num_label_tokens_tensor) - num_label_tokens = num_label_tokens_tensor.item() - - # Calculate loss - following exact train_ft.py:1059-1071 pattern - reporting_loss = torch.sum(torch.stack(loss_buffer)) - reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True) - reporting_loss = reporting_loss.to(torch.float32) / num_label_tokens - - if self.pp_enabled: - reporting_loss = reporting_loss.to(self.dist_env.device) - # Send loss to first rank if pp group rank is 0 - src_rank = self.device_mesh.mesh.reshape(-1)[-1].item() - if self.dist_env.rank == src_rank: - torch.distributed.send(reporting_loss, dst=0) - elif self.dist_env.is_main: - torch.distributed.recv(reporting_loss, src=src_rank) - - reporting_loss = reporting_loss.cpu().item() + num_label_tokens = int(forward_backward_result.weight_sum.item()) + reporting_loss = forward_backward_result.loss.cpu().item() if rank == 0: print(f"num_label_tokens={num_label_tokens} | loss={reporting_loss:.4f}") @@ -419,7 +367,7 @@ def run_benchmark(self): self._log_moe_metrics(i, self.wandb_run.log) # Calculate and log MFU - self._log_iteration_metrics(iter_timer, ga_steps, peak_tflops, rank, i) + self._log_iteration_metrics(iter_timer, peak_tflops, rank, i) # Stop nsys profiling if configured if i == nsys_end and rank in nsys_ranks: @@ -434,7 +382,7 @@ def run_benchmark(self): # Final summary self._log_benchmark_summary(steps, warmup_steps, peak_tflops, rank) - def _log_iteration_metrics(self, iter_timer, ga_steps, peak_tflops, rank, iteration): + def _log_iteration_metrics(self, iter_timer, peak_tflops, rank, iteration): max_iter_time = self.timers._get_global_min_max_time([iter_timer], reset=False, barrier=False, normalizer=1.0)[ iter_timer ][1] @@ -450,7 +398,7 @@ def _log_iteration_metrics(self, iter_timer, ga_steps, peak_tflops, rank, iterat logger.info(f"MFU: {mfu:.6f}%") # Log detailed timers - timer_names = [iter_timer, "optimizer"] + [f"forward_backward_{ga_step_idx}" for ga_step_idx in range(ga_steps)] + timer_names = [iter_timer, "forward_backward", "optimizer"] # Log timers to wandb if self._wandb_enabled: self.timers.write_to_wandb( diff --git a/tests/unit_tests/recipes/llm/test_benchmark.py b/tests/unit_tests/recipes/llm/test_benchmark.py index 81fae306fb..31ba9fc1ee 100644 --- a/tests/unit_tests/recipes/llm/test_benchmark.py +++ b/tests/unit_tests/recipes/llm/test_benchmark.py @@ -19,6 +19,7 @@ import pytest import torch +from nemo_automodel.engine import ForwardBackwardResult from nemo_automodel.recipes.llm.benchmark import BenchmarkingRecipeForNextTokenPrediction, _infer_vocab_size, main @@ -134,6 +135,7 @@ def mock_recipe(mock_config, monkeypatch): intermediate_size=3072, ) recipe.optimizer = [MagicMock()] + recipe.loss_fn = object() recipe.dataloader = MagicMock() recipe.val_dataloader = None recipe.pp_enabled = False @@ -154,6 +156,15 @@ def mock_recipe(mock_config, monkeypatch): recipe.step_scheduler = SimpleNamespace(step=0, gc_every_steps=None) recipe._dp_allreduce = MagicMock(side_effect=lambda x, include_cp=False: x) recipe.device_mesh = None + recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=MagicMock()) + recipe.engine = MagicMock() + recipe.engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.5), + loss_sum=torch.tensor(12.0), + weight_sum=torch.tensor(24.0), + loss_fn_outputs=[], + ) + recipe.engine.optim_step.return_value = SimpleNamespace(grad_norm=torch.tensor(1.0), learning_rates=(0.01,)) return recipe @@ -339,13 +350,6 @@ class TestBenchmarkingRecipeRunBenchmark: def test_run_benchmark_sets_models_to_train_mode(self, mock_recipe): """Test that run_benchmark sets all models to training mode.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) - - # Mock _forward_backward_step to append loss to loss_buffer - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) # Mock timers to return a dict with the expected structure mock_recipe.timers._get_global_min_max_time = MagicMock( return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} @@ -381,13 +385,6 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): def test_run_benchmark_calculates_gradient_accumulation_steps(self, mock_recipe): """Test that gradient accumulation steps are calculated correctly.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) - - # Mock _forward_backward_step to append loss to loss_buffer - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) # Mock timers to return a dict with the expected structure mock_recipe.timers._get_global_min_max_time = MagicMock( return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} @@ -420,19 +417,22 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): # global_batch_size=256, local_batch_size=4, dp_size=8 # ga_steps = 256 / (4 * 8) = 8 expected_ga_steps = 8 - # Verify forward_backward_step was called expected_ga_steps times per iteration - assert mock_recipe._forward_backward_step.call_count == 30 * expected_ga_steps - - def test_run_benchmark_zero_grads_per_iteration(self, mock_recipe): - """Test that gradients are zeroed at the start of each iteration.""" + assert mock_recipe.engine.forward_backward.call_count == 30 + assert all( + len(call.args[0]) == expected_ga_steps for call in mock_recipe.engine.forward_backward.call_args_list + ) + assert torch.cuda.nvtx.range_push.call_count == 30 + assert torch.cuda.nvtx.range_pop.call_count == 30 + torch.cuda.nvtx.range_push.assert_any_call("iteration_0_forward_backward") + torch.cuda.nvtx.range_push.assert_any_call("iteration_29_forward_backward") + timer_names = [call.args[0] for call in mock_recipe.timers.call_args_list] + assert timer_names.count("iteration_warmup") == 10 + assert timer_names.count("iteration") == 20 + assert timer_names.count("forward_backward") == 30 + + def test_run_benchmark_leaves_gradient_clearing_to_engine(self, mock_recipe): + """Test that the benchmark does not bypass Engine gradient ownership.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) - - # Mock _forward_backward_step to append loss to loss_buffer - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) # Mock timers to return a dict with the expected structure mock_recipe.timers._get_global_min_max_time = MagicMock( return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} @@ -462,22 +462,15 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): with patch("torch.distributed.barrier"): mock_recipe.run_benchmark() - # Should be called 30 times (once per iteration) - assert mock_recipe.optimizer[0].zero_grad.call_count == 30 + mock_recipe.optimizer[0].zero_grad.assert_not_called() + assert mock_recipe.engine.optim_step.call_count == 30 def test_run_benchmark_optimizer_step_per_iteration(self, mock_recipe): - """Test that optimizer step is called once per iteration.""" + """Test that Engine performs one optimizer step per iteration.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) graph_manager = MagicMock() mock_recipe.partial_cuda_graph_manager = graph_manager mock_recipe._partial_cuda_graph_capture_pending = True - - # Mock _forward_backward_step to append loss to loss_buffer - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) # Mock timers to return a dict with the expected structure mock_recipe.timers._get_global_min_max_time = MagicMock( return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} @@ -507,20 +500,16 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): with patch("torch.distributed.barrier"): mock_recipe.run_benchmark() - # Should be called 30 times (once per iteration) - assert mock_recipe.optimizer[0].step.call_count == 30 + assert mock_recipe.engine.optim_step.call_count == 30 + for call in mock_recipe.engine.optim_step.call_args_list: + assert call.kwargs["before_optimizer_step"] is mock_recipe.checkpointer.maybe_wait_for_staging + mock_recipe.optimizer[0].step.assert_not_called() graph_manager.capture.assert_called_once_with() assert mock_recipe._partial_cuda_graph_capture_pending is False def test_run_benchmark_calls_gc_hook_per_iteration(self, mock_recipe): mock_recipe._get_dp_group_size = MagicMock(return_value=8) mock_recipe._maybe_collect_garbage = MagicMock() - - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) mock_recipe.timers._get_global_min_max_time = MagicMock( return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} ) @@ -549,6 +538,33 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): assert mock_recipe._maybe_collect_garbage.call_count == 30 + def test_run_benchmark_balances_nvtx_after_engine_failure(self, mock_recipe): + """Test that the aggregate NVTX range closes when Engine fails.""" + mock_recipe._get_dp_group_size = MagicMock(return_value=8) + mock_recipe._bench_steps = 1 + mock_recipe._bench_warmup_steps = 0 + mock_recipe.dataloader.__iter__ = MagicMock( + return_value=iter( + [ + { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[1, 2, 3]]), + "position_ids": torch.tensor([[0, 1, 2]]), + } + ] + * 8 + ) + ) + + mock_recipe.engine.forward_backward.side_effect = RuntimeError("benchmark backward failed") + + with pytest.raises(RuntimeError, match="benchmark backward failed"): + mock_recipe.run_benchmark() + + torch.cuda.nvtx.range_push.assert_called_once_with("iteration_0_forward_backward") + torch.cuda.nvtx.range_pop.assert_called_once_with() + mock_recipe.engine.optim_step.assert_not_called() + @pytest.mark.usefixtures("patch_torch_distributed_for_benchmark") class TestBenchmarkingRecipeHelpers: @@ -666,17 +682,10 @@ def test_log_iteration_metrics_with_wandb(self, mock_recipe): mock_recipe.timers._get_global_min_max_time = MagicMock(return_value={"iteration": (0.0, 1.5)}) with patch("nemo_automodel.recipes.llm.benchmark.calculate_mfu", return_value=50.0): - mock_recipe._log_iteration_metrics("iteration", ga_steps=4, peak_tflops=989, rank=0, iteration=10) + mock_recipe._log_iteration_metrics("iteration", peak_tflops=989, rank=0, iteration=10) # Verify wandb logging was called with correct parameters - expected_timer_names = [ - "iteration", - "optimizer", - "forward_backward_0", - "forward_backward_1", - "forward_backward_2", - "forward_backward_3", - ] + expected_timer_names = ["iteration", "forward_backward", "optimizer"] mock_recipe.timers.write_to_wandb.assert_called_once_with( names=expected_timer_names, writer=mock_recipe.wandb_run, @@ -870,19 +879,16 @@ def test_wandb_finish_called(self, mock_wandb_finish, mock_recipe): class TestBenchmarkingRecipeLossCalculation: """Test loss calculation and synchronization.""" - def test_loss_buffer_accumulation(self, mock_recipe): - """Test that losses are accumulated in buffer during gradient accumulation.""" + def test_engine_result_drives_reported_loss(self, mock_recipe, capsys): + """Test that Engine's reduced loss and weight sum are reported directly.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) - - # Track loss_buffer contents - captured_loss_buffers = [] - - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5 + ga_step_idx * 0.1)) - captured_loss_buffers.append(list(loss_buffer)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) + mock_recipe.engine.forward_backward.side_effect = None + mock_recipe.engine.forward_backward.return_value = ForwardBackwardResult( + loss=torch.tensor(0.625), + loss_sum=torch.tensor(10.625), + weight_sum=torch.tensor(17.0), + loss_fn_outputs=[], + ) # Mock timers mock_recipe.timers._get_global_min_max_time = MagicMock( @@ -916,27 +922,16 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): with patch("torch.distributed.barrier"): mock_recipe.run_benchmark() - # Verify loss_buffer was populated correctly (8 GA steps) - assert len(captured_loss_buffers[-1]) == 8 + assert "num_label_tokens=17 | loss=0.6250" in capsys.readouterr().out + datums, loss_fn = mock_recipe.engine.forward_backward.call_args.args + assert len(datums) == 8 + assert loss_fn == mock_recipe._engine_loss_fn - def test_dp_allreduce_called_for_loss(self, mock_recipe): - """Test that DP allreduce is called for loss synchronization.""" + def test_pp_engine_result_avoids_duplicate_recipe_collectives(self, mock_recipe): + """Test that the PP recipe trusts Engine-synchronized loss statistics.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) - dp_allreduce_calls = [] - - def track_allreduce(tensor, include_cp=False): - dp_allreduce_calls.append({"tensor_shape": tensor.shape, "include_cp": include_cp}) - return tensor - - # Use __dict__ to bypass state tracking - mock_recipe.__dict__["_dp_allreduce"] = MagicMock(side_effect=track_allreduce) - - # Mock forward_backward_step - def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): - if loss_buffer is not None: - loss_buffer.append(torch.tensor(0.5)) - - mock_recipe._forward_backward_step = MagicMock(side_effect=mock_forward_backward_step) + mock_recipe.pp_enabled = True + mock_recipe.__dict__["_dp_allreduce"] = MagicMock() # Mock timers mock_recipe.timers._get_global_min_max_time = MagicMock( @@ -967,13 +962,13 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): ) ) - with patch("torch.distributed.barrier"): + with ( + patch("torch.distributed.barrier"), + patch("torch.distributed.send") as send, + patch("torch.distributed.recv") as recv, + ): mock_recipe.run_benchmark() - # Verify DP allreduce was called - # Should be called twice per iteration: once for num_label_tokens, once for loss - assert len(dp_allreduce_calls) >= 2 - - # Verify loss reduction includes context parallelism - loss_reduction_calls = [call for call in dp_allreduce_calls if call["include_cp"]] - assert len(loss_reduction_calls) >= 1 + mock_recipe._dp_allreduce.assert_not_called() + send.assert_not_called() + recv.assert_not_called() From 93363eedfdfc3e5543f1e3ca490776492a8ab8c7 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 12:37:54 -0700 Subject: [PATCH 18/34] feat(vlm): enable pipeline validation Signed-off-by: HuiyingLi --- .../gemma4/gemma4_31b_tp4_pp2.yaml | 1 + .../gemma4/gemma4_31b_tp4_pp4.yaml | 1 + .../minimax_m3_vl_lora_pp4ep8_8node.yaml | 1 + .../minimax_m3_vl_sft_cp2_medpix_2k.yaml | 1 + .../minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml | 1 + .../mistral3p5/mistral3p5_128b_medpix.yaml | 1 + .../mistral3p5_128b_medpix_lora.yaml | 1 + .../mistral4/mistral4_medpix.yaml | 1 + .../qwen3_5_moe/qwen3_5_35b_neat_packing.yaml | 1 + .../stepfun/step3p7_medpix_200b_ep32pp4.yaml | 1 + ...step3p7_medpix_200b_lora_pp8ep8_8node.yaml | 1 + nemo_automodel/recipes/vlm/finetune.py | 43 +++++-- .../L2_Parallelism_VLM_Gemma4_PP2_Parity.sh | 16 ++- .../parallelism/compare_parallel_parity.py | 84 +++++++++---- .../recipes/test_finetune_vlm_cp_wiring.py | 118 +++++++++++++++++- .../recipes/test_finetune_vlm_helpers.py | 30 +++++ .../test_compare_parallel_parity.py | 103 +++++++++++++++ 17 files changed, 364 insertions(+), 41 deletions(-) create mode 100644 tests/unit_tests/test_compare_parallel_parity.py diff --git a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml b/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml index ff3bc0d4cf..b1292a433c 100644 --- a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml +++ b/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 diff --git a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml b/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml index 3848cdb375..492f537647 100644 --- a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml index d3b7e5d698..522839d91a 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml index 3d54946f9c..742f31797d 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml index 2aa40543a1..f9103d73d4 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml b/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml index 279c2e8d8e..3cea53752b 100644 --- a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml b/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml index 351a6eac58..952cda63a0 100644 --- a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml b/examples/vlm_finetune/mistral4/mistral4_medpix.yaml index f8033c25e4..ebd2105de3 100644 --- a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml +++ b/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 diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml index 9ace6728c9..6e696eea7c 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml +++ b/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 diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml index f979814fe2..8b0b4206f3 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml index effa21485e..33c1694fb8 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml +++ b/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 diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 9a8ab317b5..6da6e0750b 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -612,8 +612,9 @@ def setup(self): ) from nemo_automodel.components.models.common.packing import configure_packing, get_attn_implementation + model_attn_implementation = get_attn_implementation(self.cfg.model, model=self.model_parts[0]) packing_attn_implementation = dataloader_config.resolve_packing_attn_implementation( - model_attn_implementation=get_attn_implementation(self.cfg.model, model=self.model_parts[0]), + model_attn_implementation=model_attn_implementation, cp_size=self.mesh_context.cp_size, ) if dataloader_config.packing is not None and dataloader_config.packing.packing_format != "thd": @@ -643,6 +644,22 @@ def setup(self): self.val_dataloader = None validation_config = self.cfg.vlm_validation_dataloader if validation_config is not None: + if self.pp_enabled and not validation_config.drop_last: + raise ValueError( + "Pipeline-parallel VLM validation requires validation_dataloader.drop_last=true because " + "AutoPipeline uses a fixed outer batch size. Enable drop_last or remove validation_dataset." + ) + _validate_cp_packing_support( + self.model_parts[0], + packing_enabled=validation_config.packing is not None, + cp_size=self.mesh_context.cp_size, + ) + validation_packing_attn_implementation = validation_config.resolve_packing_attn_implementation( + model_attn_implementation=model_attn_implementation, + cp_size=self.mesh_context.cp_size, + ) + if validation_config.packing is not None and validation_config.packing.packing_format != "thd": + configure_packing(attn_implementation=validation_packing_attn_implementation) validation_build_context = FirstRankPerNode(group=process_group) with ScopedRNG(seed=self.cfg.get("seed", 42), ranked=True): validation_build = validation_config.build( @@ -652,6 +669,8 @@ def setup(self): batch_size=self.cfg.get("step_scheduler.local_batch_size", 1), dataset_build_context=validation_build_context, get_rope_index=get_rope_index, + packing_attn_implementation=validation_packing_attn_implementation, + pp_n_microbatches=pp_n_microbatches, cp_size=self.mesh_context.cp_size, ) self.val_dataloader = validation_build.dataloader @@ -730,12 +749,9 @@ def run_train_validation_loop(self): val_loss = {} if self.step_scheduler.is_val_step and self.val_dataloader is not None: - if self.pp_enabled: - logger.warning("Validation is not supported for pipeline parallelism") - else: - val_log_data = self._run_validation_epoch(self.val_dataloader) - val_loss["val_loss"] = val_log_data.metrics["val_loss"] - self.log_val_metrics(val_log_data) + val_log_data = self._run_validation_epoch(self.val_dataloader) + val_loss["val_loss"] = val_log_data.metrics["val_loss"] + self.log_val_metrics(val_log_data) for mp in self.model_parts: mp.train() @@ -1152,11 +1168,18 @@ def _run_validation_epoch(self, val_dataloader): total_loss += result.loss_sum total_num_label_tokens += result.weight_sum - # Engine.forward has already reconstructed CP shards. Only independent - # DP validation shards remain to combine (VLM PP validation stays disabled). + # Engine.forward has already reconstructed CP shards and synchronized + # PP stages. Only independent DP validation shards remain to combine. total_loss = self._dp_allreduce(total_loss).item() total_num_label_tokens = int(self._dp_allreduce(total_num_label_tokens).item()) - val_loss = total_loss / max(total_num_label_tokens, 1e-8) + if total_num_label_tokens <= 0: + raise ValueError( + "VLM validation produced no supervised label tokens after DP aggregation. " + "With pipeline parallelism, validation_dataloader.drop_last=true may have removed every batch " + "because each DP shard is smaller than the local batch size; otherwise verify that labels are not " + "all masked." + ) + val_loss = total_loss / total_num_label_tokens return MetricsSample( step=self.step_scheduler.step, diff --git a/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh b/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh index 9d3e7b916f..01ad3e0011 100644 --- a/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh +++ b/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh @@ -17,8 +17,9 @@ # # Runs the Gemma4 31B proxy twice with the same seed and data order -- once on a # single rank, once at pp_size=2 -- and asserts both follow the same loss and -# gradient-norm trajectory. `dp_size` is 1 in both runs, so the dataloader yields -# identical batches and any divergence is attributable to the pipeline split. +# gradient-norm trajectory and validation loss. `dp_size` is 1 in both runs, so +# the dataloader yields identical batches and any divergence is attributable to +# the pipeline split. # # Covers the gap from PR #2983 (commit 00f40419). # @@ -55,8 +56,10 @@ COMMON_ARGS=( --validation_dataset.split validation --validation_dataset.limit_dataset_samples 8 --step_scheduler.max_steps 6 + --step_scheduler.val_every_steps 2 --step_scheduler.global_batch_size 2 --step_scheduler.local_batch_size 2 + --validation_dataloader.drop_last true ) # --- Baseline: single rank, no parallelism --- @@ -94,3 +97,12 @@ python tests/functional_tests/parallelism/compare_parallel_parity.py \ --axis pp \ --loss-tol 0.05 \ --grad-norm-rtol 0.20 + +# Both runs must execute recipe-owned validation. The parity helper also rejects +# empty validation logs, so a stale PP validation skip cannot pass this check. +python tests/functional_tests/parallelism/compare_parallel_parity.py \ + "$RUN_DIR/baseline/validation.jsonl" \ + "$RUN_DIR/pp2/validation.jsonl" \ + --axis pp \ + --metric val_loss \ + --loss-tol 0.05 diff --git a/tests/functional_tests/parallelism/compare_parallel_parity.py b/tests/functional_tests/parallelism/compare_parallel_parity.py index 62a680e545..9a7d954111 100644 --- a/tests/functional_tests/parallelism/compare_parallel_parity.py +++ b/tests/functional_tests/parallelism/compare_parallel_parity.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Parallel-vs-single-rank training parity validator. +"""Parallel-vs-single-rank training and validation parity validator. -Compares two ``training.jsonl`` logs produced by the same recipe and seed: a -single-rank baseline and a run with one parallelism axis enabled (TP, PP, CP, or -EP). Both runs must follow the same loss and gradient-norm trajectory. +Compares two metric logs produced by the same recipe and seed: a single-rank +baseline and a run with one parallelism axis enabled (TP, PP, CP, or EP). +Training logs must follow the same loss and gradient-norm trajectory; +validation logs compare their validation loss. This is the generic net for parallelism correctness. A smoke test only fails on a crash or a hang, but wrong stage metadata, a gradient that syncs over the @@ -32,12 +33,14 @@ Usage: python compare_parallel_parity.py baseline.jsonl pp2.jsonl --axis pp + python compare_parallel_parity.py baseline-validation.jsonl pp2-validation.jsonl --axis pp --metric val_loss """ from __future__ import annotations import argparse import json +import math # Both runs share a seed and a data order, so step 1 differs only by floating- # point reduction order. Later steps accumulate that difference through the @@ -50,11 +53,12 @@ DEFAULT_GRAD_NORM_RTOL = 0.05 -def read_metrics(jsonl_path: str) -> dict[int, dict[str, float]]: - """Read per-step training metrics from a ``training.jsonl`` log. +def read_metrics(jsonl_path: str, *, metric: str = "loss") -> dict[int, dict[str, float]]: + """Read per-step loss metrics from a JSONL log. Args: - jsonl_path: Path to a ``training.jsonl`` written by ``MetricLogger``. + jsonl_path: Path to a metric JSONL written by ``MetricLogger``. + metric: Loss field to read: ``loss`` for training or ``val_loss`` for validation. Returns: Mapping of step index to a dict with the ``loss`` key and, when the @@ -67,9 +71,9 @@ def read_metrics(jsonl_path: str) -> dict[int, dict[str, float]]: if not line: continue record = json.loads(line) - if "step" not in record or "loss" not in record: + if "step" not in record or metric not in record: continue - sample: dict[str, float] = {"loss": float(record["loss"])} + sample: dict[str, float] = {"loss": float(record[metric])} grad_norm = record.get("grad_norm") if grad_norm is not None: sample["grad_norm"] = float(grad_norm) @@ -79,10 +83,16 @@ def read_metrics(jsonl_path: str) -> dict[int, dict[str, float]]: def main() -> None: """Compare a single-rank baseline log against a parallel-run log.""" - parser = argparse.ArgumentParser(description="Compare single-rank vs parallel training parity") - parser.add_argument("baseline_jsonl", help="training.jsonl from the single-rank baseline run") - parser.add_argument("parallel_jsonl", help="training.jsonl from the parallel run") + parser = argparse.ArgumentParser(description="Compare single-rank vs parallel training/validation parity") + parser.add_argument("baseline_jsonl", help="Metric JSONL from the single-rank baseline run") + parser.add_argument("parallel_jsonl", help="Metric JSONL from the parallel run") parser.add_argument("--axis", required=True, help="Parallelism axis under test, e.g. pp/tp/cp/ep") + parser.add_argument( + "--metric", + choices=("loss", "val_loss"), + default="loss", + help="Loss field to compare; val_loss performs validation loss-only parity", + ) parser.add_argument("--loss-tol", type=float, default=DEFAULT_LOSS_TOL, help="Absolute per-step loss tolerance") parser.add_argument( "--grad-norm-rtol", @@ -92,11 +102,16 @@ def main() -> None: ) args = parser.parse_args() - baseline = read_metrics(args.baseline_jsonl) - parallel = read_metrics(args.parallel_jsonl) + baseline = read_metrics(args.baseline_jsonl, metric=args.metric) + parallel = read_metrics(args.parallel_jsonl, metric=args.metric) - assert len(baseline) > 0, f"No training records in {args.baseline_jsonl}" - assert len(parallel) > 0, f"No training records in {args.parallel_jsonl}" + assert len(baseline) > 0, f"No {args.metric} records in {args.baseline_jsonl}" + assert len(parallel) > 0, f"No {args.metric} records in {args.parallel_jsonl}" + if args.metric == "val_loss": + assert set(baseline) == set(parallel), ( + f"Validation steps differ between {args.baseline_jsonl} (steps {sorted(baseline)}) " + f"and {args.parallel_jsonl} (steps {sorted(parallel)})" + ) common_steps = sorted(set(baseline) & set(parallel)) assert len(common_steps) > 0, ( @@ -108,21 +123,33 @@ def main() -> None: grad_norm_failures: list[str] = [] compared_grad_norms = 0 - print(f"=== {args.axis} parity: {len(common_steps)} common steps ===") + print(f"=== {args.axis} {args.metric} parity: {len(common_steps)} common steps ===") print(f"{'step':>6} {'baseline':>12} {'parallel':>12} {'delta':>12}") for step in common_steps: base_loss = baseline[step]["loss"] par_loss = parallel[step]["loss"] - delta = abs(base_loss - par_loss) - print(f"{step:>6} {base_loss:>12.6f} {par_loss:>12.6f} {delta:>12.6f}") - if delta > args.loss_tol: - loss_failures.append(f"step {step}: baseline={base_loss:.6f} {args.axis}={par_loss:.6f} delta={delta:.6f}") + if not math.isfinite(base_loss) or not math.isfinite(par_loss): + loss_failures.append( + f"step {step}: non-finite {args.metric}: baseline={base_loss!r} {args.axis}={par_loss!r}" + ) + else: + delta = abs(base_loss - par_loss) + print(f"{step:>6} {base_loss:>12.6f} {par_loss:>12.6f} {delta:>12.6f}") + if delta > args.loss_tol: + loss_failures.append( + f"step {step}: baseline={base_loss:.6f} {args.axis}={par_loss:.6f} delta={delta:.6f}" + ) base_norm = baseline[step].get("grad_norm") par_norm = parallel[step].get("grad_norm") if base_norm is None or par_norm is None: continue compared_grad_norms += 1 + if not math.isfinite(base_norm) or not math.isfinite(par_norm): + grad_norm_failures.append( + f"step {step}: non-finite gradient norm: baseline={base_norm!r} {args.axis}={par_norm!r}" + ) + continue scale = max(abs(base_norm), 1e-8) norm_delta = abs(base_norm - par_norm) / scale if norm_delta > args.grad_norm_rtol: @@ -139,13 +166,16 @@ def main() -> None: f"relative in gradient norm:\n " + "\n ".join(grad_norm_failures) ) - # A log without grad_norm would silently reduce this to a loss-only check. - assert compared_grad_norms > 0, ( - "Neither log reported grad_norm, so the gradient-sync half of this check did not run. " - "Confirm the recipe logs grad_norm to training.jsonl." - ) + if args.metric == "loss": + # A training log without grad_norm would silently reduce this to a loss-only check. + assert compared_grad_norms > 0, ( + "Neither log reported grad_norm, so the gradient-sync half of this check did not run. " + "Confirm the recipe logs grad_norm to training.jsonl." + ) - print(f"{args.axis} parity OK: {len(common_steps)} steps, {compared_grad_norms} gradient norms compared") + print( + f"{args.axis} {args.metric} parity OK: {len(common_steps)} steps, {compared_grad_norms} gradient norms compared" + ) if __name__ == "__main__": diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py index 79541613b0..ff3044e03d 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -24,6 +24,7 @@ from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -151,7 +152,14 @@ class _StageWithoutCPPrepare: pass -def _patch_pp_setup_minimals(monkeypatch, *, cp_size, stage0, dataloader_calls): +def _patch_pp_setup_minimals( + monkeypatch, + *, + cp_size, + stage0, + dataloader_calls, + validation_loader_config=None, +): monkeypatch.setattr(vlm_finetune, "AutoPipeline", _FakePPModel) monkeypatch.setattr("nemo_automodel.engine.AutoPipeline", _FakePPModel) monkeypatch.setattr( @@ -224,7 +232,7 @@ def _build_dataloader(**kwargs): ) monkeypatch.setattr( "nemo_automodel.recipes._typed_config.RecipeConfig.vlm_validation_dataloader", - property(lambda self: None), + property(lambda self: validation_loader_config), ) monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda **kwargs: nullcontext()) monkeypatch.setattr( @@ -303,3 +311,109 @@ def test_setup_always_stages_pp_media_under_pp( assert dataloader_calls[0]["cp_size"] == cp_size assert trainer.engine.pipeline is trainer.pp assert trainer.engine.microbatch_size == 1 + + +def test_setup_stages_pp_validation_media_and_preserves_packing_wiring(monkeypatch): + dataloader_calls = [] + packing_resolutions = [] + configure_packing_calls = [] + + def _resolve_validation_packing(**kwargs): + packing_resolutions.append(kwargs) + return "sdpa" + + validation_loader_config = SimpleNamespace( + drop_last=True, + packing=SimpleNamespace(packing_format="neat"), + resolve_packing_attn_implementation=_resolve_validation_packing, + build=lambda **kwargs: ( + dataloader_calls.append(kwargs) or SimpleNamespace(dataloader="val_dl", processor="processor") + ), + ) + _patch_pp_setup_minimals( + monkeypatch, + cp_size=1, + stage0=_StageWithoutCPPrepare(), + dataloader_calls=dataloader_calls, + validation_loader_config=validation_loader_config, + ) + monkeypatch.setattr( + "nemo_automodel.components.models.common.packing.configure_packing", + lambda **kwargs: configure_packing_calls.append(kwargs), + ) + trainer = FinetuneRecipeForVLM(_minimal_pp_setup_cfg()) + + trainer.setup() + + assert len(dataloader_calls) == 2 + validation_call = dataloader_calls[1] + assert validation_call["pp_n_microbatches"] == 2 + assert validation_call["packing_attn_implementation"] == "sdpa" + assert validation_call["cp_size"] == 1 + assert packing_resolutions == [{"model_attn_implementation": "sdpa", "cp_size": 1}] + assert configure_packing_calls == [{"attn_implementation": "sdpa"}] + assert trainer.val_dataloader == "val_dl" + + +def test_setup_rejects_incomplete_pp_validation_batches(monkeypatch): + dataloader_calls = [] + validation_loader_config = SimpleNamespace( + drop_last=False, + packing=None, + resolve_packing_attn_implementation=lambda **kwargs: None, + build=lambda **kwargs: pytest.fail("validation loader must not build before drop_last validation"), + ) + _patch_pp_setup_minimals( + monkeypatch, + cp_size=1, + stage0=_StageWithoutCPPrepare(), + dataloader_calls=dataloader_calls, + validation_loader_config=validation_loader_config, + ) + trainer = FinetuneRecipeForVLM(_minimal_pp_setup_cfg()) + + with pytest.raises(ValueError, match=r"validation_dataloader\.drop_last=true"): + trainer.setup() + + assert len(dataloader_calls) == 1 + + +def test_train_loop_runs_validation_when_pipeline_is_enabled(): + class _SingleStepScheduler: + epochs = (0,) + step = 1 + epoch = 0 + is_val_step = True + is_ckpt_step = False + sigterm_flag = False + + def set_epoch(self, epoch): + self.epoch = epoch + + def __iter__(self): + yield [object()] + + recipe = object.__new__(FinetuneRecipeForVLM) + model_part = SimpleNamespace(train=MagicMock()) + recipe.model_parts = [model_part] + recipe.step_scheduler = _SingleStepScheduler() + recipe.val_dataloader = object() + recipe.pp_enabled = True + recipe._make_progress_bar = MagicMock(return_value=None) + recipe._run_train_optim_step = MagicMock(return_value=SimpleNamespace(metrics={"loss": 1.0})) + recipe.log_train_metrics = MagicMock() + recipe._update_progress_bar = MagicMock() + validation_metrics = SimpleNamespace(metrics={"val_loss": 0.25}) + recipe._run_validation_epoch = MagicMock(return_value=validation_metrics) + recipe.log_val_metrics = MagicMock() + recipe.save_checkpoint = MagicMock() + recipe._maybe_collect_garbage = MagicMock() + recipe.metric_logger_train = SimpleNamespace(close=MagicMock()) + recipe.metric_logger_valid = SimpleNamespace(close=MagicMock()) + recipe._finalize_and_close_checkpointer = MagicMock() + + recipe.run_train_validation_loop() + + recipe._run_validation_epoch.assert_called_once_with(recipe.val_dataloader) + recipe.log_val_metrics.assert_called_once_with(validation_metrics) + assert model_part.train.call_count == 2 diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index f533b850c6..3fb00bcc07 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -2205,6 +2205,36 @@ def test_vlm_validation_uses_engine_forward_and_aggregates_uneven_batches(monkey assert metrics.metrics["num_label_tokens"] == pytest.approx(5.0) +@pytest.mark.parametrize( + "batches", + [ + [], + [{"input_ids": torch.tensor([[1, 2]]), "labels": torch.tensor([[-100, -100]])}], + ], +) +def test_vlm_validation_rejects_zero_global_denominator(monkeypatch, batches): + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) + recipe.pp_enabled = False + recipe.loss_fn = object() + recipe.engine = MagicMock() + recipe.engine.forward.return_value = SimpleNamespace( + loss_sum=torch.tensor(0.0, dtype=torch.float64), + weight_sum=torch.tensor(0.0, dtype=torch.float64), + loss_fn_outputs=[], + ) + recipe._dp_allreduce = MagicMock(side_effect=lambda tensor, **kwargs: tensor) + monkeypatch.setattr( + "nemo_automodel.recipes.vlm.finetune.ScopedRNG", + lambda **kwargs: nullcontext(), + ) + + with pytest.raises(ValueError, match="no supervised label tokens.*drop_last=true"): + recipe._run_validation_epoch(batches) + + assert recipe._dp_allreduce.call_count == 2 + + def test_vlm_rope_fusion_unchanged_when_cp_eq_1(monkeypatch): """rope_fusion should remain True in VLM setup when cp_size == 1.""" cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=True) diff --git a/tests/unit_tests/test_compare_parallel_parity.py b/tests/unit_tests/test_compare_parallel_parity.py new file mode 100644 index 0000000000..d41d854f6e --- /dev/null +++ b/tests/unit_tests/test_compare_parallel_parity.py @@ -0,0 +1,103 @@ +# 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. + +import json +import sys + +import pytest + +from tests.functional_tests.parallelism import compare_parallel_parity + + +def _write_metrics(path, records): + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def _run_validation_comparison(monkeypatch, baseline_path, parallel_path): + monkeypatch.setattr( + sys, + "argv", + [ + "compare_parallel_parity.py", + str(baseline_path), + str(parallel_path), + "--axis", + "pp", + "--metric", + "val_loss", + ], + ) + compare_parallel_parity.main() + + +def _run_training_comparison(monkeypatch, baseline_path, parallel_path): + monkeypatch.setattr( + sys, + "argv", + [ + "compare_parallel_parity.py", + str(baseline_path), + str(parallel_path), + "--axis", + "pp", + ], + ) + compare_parallel_parity.main() + + +def test_validation_parity_accepts_finite_matching_steps(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics(baseline_path, [{"step": 2, "val_loss": 1.0}, {"step": 4, "val_loss": 0.9}]) + _write_metrics(parallel_path, [{"step": 2, "val_loss": 1.01}, {"step": 4, "val_loss": 0.91}]) + + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf"), float("-inf")]) +def test_validation_parity_rejects_nonfinite_loss(tmp_path, monkeypatch, nonfinite): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics(baseline_path, [{"step": 2, "val_loss": nonfinite}]) + _write_metrics(parallel_path, [{"step": 2, "val_loss": nonfinite}]) + + with pytest.raises(AssertionError, match="non-finite val_loss"): + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_validation_parity_requires_identical_step_sets(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics( + baseline_path, + [ + {"step": 2, "val_loss": 1.0}, + {"step": 4, "val_loss": 0.9}, + ], + ) + _write_metrics(parallel_path, [{"step": 2, "val_loss": 1.0}]) + + with pytest.raises(AssertionError, match="Validation steps differ"): + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_training_parity_rejects_nonfinite_gradient_norm(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + record = {"step": 1, "loss": 1.0, "grad_norm": float("nan")} + _write_metrics(baseline_path, [record]) + _write_metrics(parallel_path, [record]) + + with pytest.raises(AssertionError, match="non-finite gradient norm"): + _run_training_comparison(monkeypatch, baseline_path, parallel_path) From 7fd09131bce29e444006cc99b655ef2b54972c42 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 16:27:49 -0700 Subject: [PATCH 19/34] feat(engine): support planned pipeline accumulation Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 39 +++ .../distributed/pipelining/functional.py | 117 ++++++++- nemo_automodel/components/moe/fsdp_mixin.py | 6 +- nemo_automodel/engine/__init__.py | 63 +++-- .../context_parallel/run_packed_pp.py | 123 +++++++++ .../parallelism/run_pp_grad_accum_parity.py | 236 +++++++++++++++++- .../pipelining/test_autopipeline.py | 189 +++++++++++++- tests/unit_tests/moe/test_fsdp_mixin.py | 44 +++- tests/unit_tests/test_engine.py | 139 ++++++++++- 9 files changed, 902 insertions(+), 54 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index 00bbddbc6b..a85512eba3 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -36,6 +36,7 @@ ) logger = logging.getLogger(__name__) +_MISSING_STAGE_STATE = object() @dataclass @@ -357,6 +358,7 @@ def step_microbatches( loss_fn: Callable[[Any, int], Any], losses: list[torch.Tensor] | None = None, return_outputs: bool = False, + finalize_backward: bool = True, ) -> Any: """Run a schedule step over already prepared model microbatches. @@ -374,16 +376,41 @@ def step_microbatches( losses: Mutable list populated by the schedule on the last stage. return_outputs: Whether the last stage returns merged model outputs when supported by the installed PyTorch version. + finalize_backward: Whether this schedule invocation ends the + optimizer's backward window. ``False`` keeps DDP in no-sync + mode and completes local FSDP post-backward state while + deferring its gradient collective; per-microbatch gradient + reduction remains authoritative when configured. Returns: The value returned by the underlying PyTorch pipeline schedule. + + Raises: + ValueError: If a non-final call would use schedule-local gradient + scaling and rescale gradients accumulated by earlier calls. + RuntimeError: If a non-final call is attempted before + accumulation-aware pipeline stages have been built. """ + if not finalize_backward and self.scale_grads_in_schedule: + raise ValueError( + "planned pipeline accumulation requires scale_grads_in_schedule=False; " + "schedule-local scaling would rescale gradients from earlier calls" + ) + if not finalize_backward and ( + not self._info.stages + or any(not vars(stage).get("_nemo_accumulation_aware", False) for stage in self._info.stages) + ): + raise RuntimeError( + "finalize_backward=False requires accumulation-aware pipeline stages; " + "build the AutoPipeline before running planned accumulation" + ) return self._run_prepared_microbatches( model_inputs, loss_fn=loss_fn, losses=losses, return_outputs=return_outputs, schedule_method="step", + finalize_backward=finalize_backward, ) def eval_microbatches( @@ -425,6 +452,7 @@ def eval_microbatches( losses=losses, return_outputs=return_outputs, schedule_method="eval", + finalize_backward=True, ) def _run_prepared_microbatches( @@ -435,6 +463,7 @@ def _run_prepared_microbatches( losses: list[torch.Tensor] | None, return_outputs: bool, schedule_method: Literal["step", "eval"], + finalize_backward: bool, ) -> Any: """Run one schedule method with an exact prepared-microbatch split.""" schedule = self._info.schedule @@ -480,6 +509,11 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: ) previous_split_inputs = schedule._split_inputs previous_loss_fn = schedule._loss_fn + previous_stage_state: list[tuple[PipelineStage, object]] = [] + if schedule_method == "step": + for stage in self._info.stages or (): + previous_stage_state.append((stage, vars(stage).get("_nemo_finalize_backward", _MISSING_STAGE_STATE))) + stage._nemo_finalize_backward = finalize_backward schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) schedule._loss_fn = indexed_loss try: @@ -491,6 +525,11 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: finally: schedule._loss_fn = previous_loss_fn schedule._split_inputs = previous_split_inputs + for stage, previous in previous_stage_state: + if previous is _MISSING_STAGE_STATE: + del stage._nemo_finalize_backward + else: + stage._nemo_finalize_backward = previous @property def parts(self) -> list[nn.Module]: diff --git a/nemo_automodel/components/distributed/pipelining/functional.py b/nemo_automodel/components/distributed/pipelining/functional.py index 5e0ac1cbef..95ee0519e1 100644 --- a/nemo_automodel/components/distributed/pipelining/functional.py +++ b/nemo_automodel/components/distributed/pipelining/functional.py @@ -20,10 +20,11 @@ import os import time import types -from typing import Callable, Protocol +from typing import Any, Callable, Protocol import torch import torch.nn as nn +from torch.distributed import fsdp as torch_fsdp from torch.distributed.device_mesh import DeviceMesh from torch.distributed.pipelining import PipelineStage from torch.distributed.pipelining.schedules import ( @@ -42,9 +43,20 @@ model_keeps_self_forward, patch_hf_model_for_pp, ) +from nemo_automodel.shared.import_utils import safe_import_from logger = logging.getLogger(__name__) +_HAS_REPLICATE_MODULE, ReplicateModule = safe_import_from( + "torch.distributed._composable.replicate_with_fsdp", + "ReplicateModule", +) +_HAS_REPLICATE_STATE, replicate = safe_import_from( + "torch.distributed._composable.replicate_with_fsdp", + "replicate", +) +_HAS_REPLICATE_WITH_FSDP = _HAS_REPLICATE_MODULE and _HAS_REPLICATE_STATE + def _get_optional_hook(module: object, name: str) -> Callable | None: try: @@ -491,6 +503,99 @@ def _cleanup_preserving_grads(self, *, _cleanup=cleanup) -> None: stage._post_metadata_inference_cleanup = types.MethodType(_cleanup_preserving_grads, stage) +def _accumulation_aware_backward_maybe_with_nosync( + self: PipelineStage, + backward_type: str, + bwd_kwargs: dict[str, Any], + last_backward: bool = False, +) -> tuple[tuple[torch.Tensor | None, ...], list[dict[str, Any]] | None]: + """Keep a schedule-local last backward open across planned Engine windows. + + A configured per-microbatch reduction remains authoritative. Otherwise a + non-final Engine window must not let PyTorch's schedule-local + ``last_backward`` trigger DDP/FSDP gradient synchronization. + + Args: + backward_type: PyTorch stage operation: ``full``, ``input``, or + ``weight`` backward. + bwd_kwargs: Schedule-owned state for one local PP microbatch. ``full`` + and ``input`` carry ``stage_output``, ``output_grads``, and + ``input_values`` tensor trees; ``weight`` carries ``stage_output`` + and the split-backward ``param_groups``. All tensor shapes, dtypes, + devices, and layouts are model- and stage-defined. This wrapper + forwards them unchanged and performs no redistribution. + last_backward: Whether the PyTorch schedule considers this its final + local backward operation. + + Returns: + The original stage backward result unchanged: stage-input gradient + tensors in their model-defined local layouts, plus optional + split-backward parameter-group records. + """ + finalize_backward = self._nemo_finalize_backward or self._reduce_grad_per_microbatch + return self._nemo_original_backward_maybe_with_nosync( + backward_type, + bwd_kwargs, + last_backward=last_backward and finalize_backward, + ) + + +def _accumulation_aware_perform_reduce_grad(self: PipelineStage, grad_scale_factor: int) -> None: + """Complete a PP schedule while deferring its final gradient collective. + + FSDP post-backward is both a communication boundary and required local + lifecycle cleanup. On a non-final planned window, run that cleanup with + gradient synchronization disabled so accumulated unsharded gradients are + preserved for the final window. Schedule-owned gradient scaling still runs + once per schedule. The original implementation is used unchanged for + ordinary calls, final planned windows, and per-microbatch reduction mode. + """ + finalize_backward = self._nemo_finalize_backward or self._reduce_grad_per_microbatch + if finalize_backward: + self._nemo_original_perform_reduce_grad(grad_scale_factor) + return + + if isinstance(self.submod, torch_fsdp.FSDPModule): + fsdp_module = self.submod + fsdp_module.set_is_last_backward(True) + fsdp_module.set_reshard_after_backward(True) + fsdp_module.set_requires_gradient_sync(False) + fsdp_state = ( + replicate.state(fsdp_module) + if _HAS_REPLICATE_WITH_FSDP and isinstance(fsdp_module, ReplicateModule) + else torch_fsdp.fully_shard.state(fsdp_module) # type: ignore[attr-defined] + ) + for state in fsdp_state._state_ctx.all_states: + if state._fsdp_param_group: + state._fsdp_param_group.post_backward() + fsdp_state._root_post_backward_final_callback() + + if grad_scale_factor != 1: + self.scale_grads(grad_scale_factor) + + +def _make_pipeline_stages_accumulation_aware( + stages: list[PipelineStage], + *, + reduce_grad_per_microbatch: bool, +) -> None: + """Install behavior-neutral stage gates used by planned PP accumulation.""" + for stage in stages: + stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch + stage._nemo_finalize_backward = True + if vars(stage).get("_nemo_accumulation_aware", False): + continue + + stage._nemo_original_backward_maybe_with_nosync = stage.backward_maybe_with_nosync + stage.backward_maybe_with_nosync = types.MethodType(_accumulation_aware_backward_maybe_with_nosync, stage) + + perform_reduce_grad = getattr(stage, "perform_reduce_grad", None) + if callable(perform_reduce_grad): + stage._nemo_original_perform_reduce_grad = perform_reduce_grad + stage.perform_reduce_grad = types.MethodType(_accumulation_aware_perform_reduce_grad, stage) + stage._nemo_accumulation_aware = True + + def reset_pp_stage_shapes( schedule: _PipelineSchedule, stages: list[PipelineStage], @@ -1003,13 +1108,21 @@ def pipeline_model( for stage in stages: stage.backward_maybe_with_nosync = types.MethodType(patched_backward_maybe_with_nosync, stage) - stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch logger.info( "Patched pipeline stages with backward_maybe_with_nosync " f"(reduce_grad_per_microbatch={reduce_grad_per_microbatch})" ) + # PyTorch considers every schedule invocation a complete optimizer window: + # its last backward and REDUCE_GRAD action finalize DP/FSDP gradients. Keep + # those exact defaults, but make the boundary controllable when Engine has + # predeclared one optimizer window spanning multiple schedule invocations. + _make_pipeline_stages_accumulation_aware( + stages, + reduce_grad_per_microbatch=reduce_grad_per_microbatch, + ) + # Determine if this rank has first/last stage has_first_stage = False has_last_stage = False diff --git a/nemo_automodel/components/moe/fsdp_mixin.py b/nemo_automodel/components/moe/fsdp_mixin.py index d27320f4aa..766a1ebc0f 100644 --- a/nemo_automodel/components/moe/fsdp_mixin.py +++ b/nemo_automodel/components/moe/fsdp_mixin.py @@ -309,7 +309,11 @@ def run_post_backward(fsdp_module: FSDPModule) -> None: elif isinstance(self.submod, MoEFSDPSyncMixin): _disable_fsdp_for_moe_module(self.submod) result = perform_backward(backward_type)() - if last_backward and get_is_optim_step(): + if hasattr(self, "_nemo_finalize_backward"): + finalize_backward = self._nemo_finalize_backward or getattr(self, "_reduce_grad_per_microbatch", False) + else: + finalize_backward = get_is_optim_step() + if last_backward and finalize_backward: _run_post_backward_for_moe_module(self.submod) else: # Non-DP submodule, regular backward diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index ce53311fbe..06702ef377 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -403,16 +403,19 @@ def begin_accumulation(self, windows: Sequence[Sequence[Datum]]) -> None: :class:`ForwardBackwardResult` continues to describe only that call; callers combine results with ``sum(loss_sum) / sum(weight_sum)``. - The first version supports eager/DDP/FSDP execution. Pipeline schedules - currently finalize gradients at the end of every schedule invocation, - so planned multi-call accumulation with PP fails explicitly instead of - silently treating each call as a complete optimizer window. If a - planned backward call fails after execution starts, partial distributed - state cannot be rolled back safely: the plan becomes broken and the - Engine must not be stepped or reused. To make rank-local loss-callback - failures fail together before backward, an explicit plan performs one - small control consensus per outer microbatch; the ordinary one-call - path adds no such collective. + Under pipeline parallelism the Engine carries the plan's final-window + boundary through AutoPipeline so a schedule-local last backward does + not prematurely synchronize deferred gradients. If a planned backward + call reports an error after execution starts, partial distributed state + cannot be rolled back safely: the plan becomes broken and the Engine + must not be stepped or reused. For non-pipeline execution, an explicit + plan performs one small control consensus per outer microbatch so a + rank-local loss-callback failure is reported together before backward; + the ordinary one-call path adds no such collective. A pipeline loss + callback runs inside PyTorch's distributed schedule, so an exception + from that callback is process-fatal rather than a recoverable + broken-plan error. Output-only callback errors returned through the + normal schedule path are synchronized across pipeline stages. Args: windows: Non-empty sequence of non-empty Datum windows in their @@ -425,15 +428,9 @@ def begin_accumulation(self, windows: Sequence[Sequence[Datum]]) -> None: RuntimeError: If no optimizer is configured, another plan is active, or gradients from an earlier optimizer window have not been cleared. - NotImplementedError: If pipeline parallelism is enabled. ValueError: If the plan is empty, malformed, or splits an outer microbatch across calls. """ - if self.pipeline is not None: - raise NotImplementedError( - "planned multi-call accumulation is not supported with pipeline parallelism; " - "pipeline schedules currently finalize gradients after every schedule call" - ) if self._accumulation_state is not None: raise RuntimeError("an Engine accumulation plan is already active") @@ -556,7 +553,8 @@ def begin_accumulation(self, windows: Sequence[Sequence[Datum]]) -> None: windows=tuple(planned_windows), weight_sums=tuple(weight_sums), total_weight_sum=total_weight_sum, - total_microbatches=sum(microbatch_counts), + total_microbatches=sum(microbatch_counts) + * (self.pipeline.num_microbatches if self.pipeline is not None else 1), microbatch_size=self.microbatch_size, ) self._optim_step_consumed = False @@ -900,6 +898,7 @@ def _forward_backward_window( output_restore_plan, backward_scale=backward_scale, zero_weight_sum=zero_denominator, + finalize_backward=is_last, ) if output_error is None: if batch_error is not None: @@ -1562,6 +1561,7 @@ def _pipeline_execute( *, backward_scale: torch.Tensor | None, zero_weight_sum: bool, + finalize_backward: bool = True, ) -> tuple[bool | None, list[dict[str, Any]], Exception | None]: """Run prepared pipeline microbatches in training or forward-only mode. @@ -1580,6 +1580,9 @@ def _pipeline_execute( backward, or ``None`` to run the forward-only schedule. zero_weight_sum: Whether reporting numerators must be forced to graph-connected zero. + finalize_backward: Whether this pipeline schedule invocation ends + the complete optimizer backward window. Ignored for + forward-only execution. Returns: Whether the callback returned outputs, its detached outputs in @@ -1639,15 +1642,21 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: return numerator if backward_scale is None else numerator * backward_scale losses = [] if self.pipeline.info.has_last_stage else None - run_microbatches = ( - self.pipeline.eval_microbatches if backward_scale is None else self.pipeline.step_microbatches - ) - run_microbatches( - model_microbatches, - loss_fn=pipeline_loss, - losses=losses, - return_outputs=False, - ) + if backward_scale is None: + self.pipeline.eval_microbatches( + model_microbatches, + loss_fn=pipeline_loss, + losses=losses, + return_outputs=False, + ) + else: + self.pipeline.step_microbatches( + model_microbatches, + loss_fn=pipeline_loss, + losses=losses, + return_outputs=False, + finalize_backward=finalize_backward, + ) outputs: list[dict[str, Any]] = [] serialized_outputs: bytes | None = None @@ -2132,7 +2141,7 @@ def _gradient_group_and_size( return (dp_cp_mesh.get_group() if size > 1 else None), size def _accumulation_control_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: - """Return the full non-PP model group used to keep plan control flow aligned.""" + """Return the full model group used to keep plan control flow aligned.""" if not dist.is_available() or not dist.is_initialized(): return None, 1 if self.mesh_context is None: diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index e0b39513e0..2b3fe5a8b5 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -514,6 +514,125 @@ def _run_thd_layout( return pipeline +def _run_planned_accumulation_parity( + device: torch.device, + mesh_context: MeshContext, + raw_inputs: dict[str, object], + labels: torch.Tensor, + weights: torch.Tensor, +) -> None: + """Compare two PP schedule calls in one plan with one complete Engine window. + + Args: + device: CUDA device for this physical pipeline rank. + mesh_context: Runtime PP2 topology; this check runs with CP size one. + raw_inputs: Raw THD mapping whose token and position tensors have shape + ``[batch, sequence]`` before Engine preparation. + labels: Target token IDs shaped ``[batch, sequence]``. + weights: Base token weights shaped ``[batch, sequence]``; the two + planned calls derive unequal denominators from this tensor. + """ + + def make_windows() -> tuple[list[Datum], list[Datum]]: + """Clone two prebatched Datum windows with ``[batch, sequence]`` tensors.""" + weights_a = weights.clone() + weights_a[:, 1::2] = 0 + weights_b = weights.clone().mul(1.5) + weights_b[:, ::2] = 0 + return ( + [ + Datum( + model_inputs=_clone_mapping(raw_inputs), + loss_fn_inputs={"labels": labels.clone(), "weights": weights_a}, + ) + ], + [ + Datum( + model_inputs=_clone_mapping(raw_inputs), + loss_fn_inputs={"labels": labels.clone(), "weights": weights_b}, + ) + ], + ) + + reference_pipeline = _build_pipeline(device, mesh_context) + reference_parameters = [parameter for part in reference_pipeline.parts for parameter in part.parameters()] + reference_optimizer = torch.optim.SGD(reference_parameters, lr=0.05) + reference_engine = Engine( + reference_pipeline, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + optimizers=reference_optimizer, + max_grad_norm=1e6, + ) + reference_window_a, reference_window_b = make_windows() + reference_result = reference_engine.forward_backward( + reference_window_a + reference_window_b, + _token_losses, + ) + reference_grads = [parameter.grad.detach().clone() for parameter in reference_parameters] + reference_step = reference_engine.optim_step() + + planned_pipeline = _build_pipeline(device, mesh_context) + planned_parameters = [parameter for part in planned_pipeline.parts for parameter in part.parameters()] + planned_optimizer = torch.optim.SGD(planned_parameters, lr=0.05) + planned_engine = Engine( + planned_pipeline, + device=device, + mesh_context=mesh_context, + collate_fn=collate_prebatched, + optimizers=planned_optimizer, + max_grad_norm=1e6, + ) + planned_window_a, planned_window_b = make_windows() + planned_engine.begin_accumulation([planned_window_a, planned_window_b]) + planned_result_a = planned_engine.forward_backward(planned_window_a, _token_losses) + planned_result_b = planned_engine.forward_backward(planned_window_b, _token_losses) + + if not reference_parameters or len(reference_parameters) != len(planned_parameters): + raise AssertionError( + f"planned/reference PP parameter mismatch: {len(planned_parameters)} != {len(reference_parameters)}" + ) + if planned_result_a.weight_sum.item() == planned_result_b.weight_sum.item(): + raise AssertionError("planned PP accumulation fixture must use unequal call denominators") + torch.testing.assert_close( + planned_result_a.loss_sum + planned_result_b.loss_sum, + reference_result.loss_sum, + atol=4e-2, + rtol=2e-3, + ) + torch.testing.assert_close( + planned_result_a.weight_sum + planned_result_b.weight_sum, + reference_result.weight_sum, + atol=0, + rtol=0, + ) + combined_loss = (planned_result_a.loss_sum + planned_result_b.loss_sum) / ( + planned_result_a.weight_sum + planned_result_b.weight_sum + ) + torch.testing.assert_close(combined_loss, reference_result.loss, atol=4e-2, rtol=2e-3) + for planned_parameter, reference_grad in zip(planned_parameters, reference_grads): + if planned_parameter.grad is None: + raise AssertionError("planned PP accumulation left a local parameter without a gradient") + torch.testing.assert_close(planned_parameter.grad.float(), reference_grad.float(), atol=5e-3, rtol=5e-2) + + planned_step = planned_engine.optim_step() + torch.testing.assert_close(planned_step.grad_norm.float(), reference_step.grad_norm.float(), atol=5e-3, rtol=5e-2) + for planned_parameter, reference_parameter in zip(planned_parameters, reference_parameters): + torch.testing.assert_close( + planned_parameter.float(), + reference_parameter.float(), + atol=5e-3, + rtol=5e-2, + ) + + if dist.get_rank() == 0: + print( + "PP2 planned two-call accumulation matched one complete window " + f"(weights={planned_result_a.weight_sum.item():.1f}+{planned_result_b.weight_sum.item():.1f})" + ) + + def _run_explicit_loss_layout( pipeline: AutoPipeline, layout: str, @@ -736,6 +855,10 @@ def main() -> None: dist.barrier() _run_padded_output_broadcast(final_pipeline, device, mesh_context) dist.barrier() + del final_pipeline + torch.cuda.empty_cache() + _run_planned_accumulation_parity(device, mesh_context, raw_inputs, labels, weights) + dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py b/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py index c95e0701c1..3b2666a3e7 100644 --- a/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py +++ b/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py @@ -30,8 +30,14 @@ Sequence length changes between windows, which is what triggers the stage reset that caused the loss. +The second phase drives those same varying-length FSDP2 stages through Engine: +one complete ``forward_backward`` call is the reference for an explicit plan +split across two calls. This additionally verifies that a non-final pipeline +schedule completes FSDP post-backward cleanup without synchronizing gradients. + Usage: torchrun --nproc-per-node=2 run_pp_grad_accum_parity.py + torchrun --nproc-per-node=4 run_pp_grad_accum_parity.py # PP2 x DP2 """ import torch @@ -43,6 +49,7 @@ # Two accumulation windows with *different* sequence lengths. The change is the # trigger: equal lengths would skip the stage reset and hide the regression. WINDOW_SEQ_LENS = (32, 48) +LOCAL_WINDOW_WEIGHT_SUM = BATCH * sum(WINDOW_SEQ_LENS) def _build_model(device: torch.device) -> torch.nn.Module: @@ -94,12 +101,14 @@ def _batch(seq_len: int, device: torch.device) -> dict[str, torch.Tensor]: device: Device to place the tensors on. Returns: - Dict with ``input_ids`` and ``labels``, both of shape [BATCH, seq_len]. + Dict with ``input_ids``, ``attention_mask``, and ``labels``, all of + shape [BATCH, seq_len]. The mask is explicitly all ones so direct + schedule and Engine collation exercise identical model inputs. """ torch.manual_seed(seq_len) # same window -> same data on every rank and phase ids = torch.randint(2, VOCAB, (BATCH, seq_len), device=device) dist.broadcast(ids, src=0) - return {"input_ids": ids, "labels": ids.clone()} + return {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids.clone()} def _run_window(pp, seq_len: int, device: torch.device) -> None: @@ -137,21 +146,233 @@ def _zero_grads(model) -> None: param.grad = None +def _scale_aware_grad_close( + actual: torch.Tensor, + expected: torch.Tensor, + *, + relative_max_error: float = 0.03, + absolute_floor: float = 2e-3, + relative_norm_error: float = 0.02, +) -> tuple[bool, float, float, float]: + """Compare equal-shaped local gradient shards with a BF16-aware bound. + + Args: + actual: Float32 local FSDP gradient shard being checked. + expected: Equal-shaped float32 local FSDP reference shard. + relative_max_error: Allowed max-element error relative to the largest + absolute reference element. + absolute_floor: Absolute max-element allowance for values near zero. + relative_norm_error: Allowed relative error in the full shard norm. + + Returns: + Whether both the max-error and norm checks pass, followed by the + observed max error, its allowed bound, and ``||actual||/||expected||``. + """ + if actual.shape != expected.shape: + return False, float("inf"), 0.0, float("inf") + if expected.numel() == 0: + return True, 0.0, absolute_floor, 1.0 + + max_error = float((actual - expected).abs().max()) + error_bound = relative_max_error * float(expected.abs().max()) + absolute_floor + actual_norm = float(actual.norm()) + expected_norm = float(expected.norm()) + if expected_norm == 0.0: + norm_ratio = 1.0 if actual_norm == 0.0 else float("inf") + norm_close = actual_norm <= absolute_floor + else: + norm_ratio = actual_norm / expected_norm + norm_close = abs(norm_ratio - 1.0) <= relative_norm_error + return max_error <= error_bound and norm_close, max_error, error_bound, norm_ratio + + +def _engine_window(seq_len: int, weight_scale: float, device: torch.device): + """Build one flat Datum window with a caller-specific token denominator. + + Args: + seq_len: Token length of every Datum in the window. + weight_scale: Constant value assigned to every token loss weight. + device: Device holding the generated tensors. + + Returns: + ``BATCH`` Datums whose ``input_ids``, ``attention_mask``, ``labels``, + and ``weights`` each have shape ``[sequence]``. The Engine collates + them to ``[BATCH, sequence]`` before PP splits the batch axis. + """ + from nemo_automodel.components.datasets.datum import Datum + + batch = _batch(seq_len, device) + return [ + Datum( + model_inputs={"input_ids": input_ids, "attention_mask": attention_mask}, + loss_fn_inputs={ + "labels": labels, + "weights": torch.full_like(labels, weight_scale, dtype=torch.float32), + }, + ) + for input_ids, attention_mask, labels in zip( + batch["input_ids"], + batch["attention_mask"], + batch["labels"], + ) + ] + + +def _run_engine_planned_parity( + pp, + mesh, + device: torch.device, + rank: int, + direct_normalized_grads: dict[str, torch.Tensor], +) -> None: + """Compare one Engine window with the same FSDP2 PP work split across calls. + + Args: + pp: Built PP2 AutoPipeline whose local stage is FSDP2-wrapped. + mesh: ``[pp, dp]`` device mesh; DP may be one or two. + device: CUDA device for this rank's token tensors. + rank: Global rank used in actionable assertion messages. + direct_normalized_grads: Per-parameter direct-schedule reference. Each + value is the float32 local FSDP shard accumulated over both + varying-length windows, with every microbatch loss divided by the + complete local-window token denominator. + """ + from nemo_automodel.components.distributed.mesh import MeshContext + from nemo_automodel.engine import Engine + + def token_losses(pred, loss_inputs): + """Return unweighted token cross entropy in the Engine loss layout. + + Args: + pred: Pipeline output with logits shaped + ``[pp_microbatch, sequence, vocab]``. + loss_inputs: Mapping containing ``labels`` and ``weights`` shaped + ``[pp_microbatch, sequence]``. + + Returns: + Per-token cross entropy with the same shape as ``weights``. Engine + applies the token weights and complete-window denominator. + """ + logits = pred.logits if hasattr(pred, "logits") else pred + return torch.nn.functional.cross_entropy( + logits.float().flatten(0, 1), + loss_inputs["labels"].flatten(0, 1), + reduction="none", + ).view_as(loss_inputs["weights"]) + + window_a = _engine_window(WINDOW_SEQ_LENS[0], 1.0, device) + window_b = _engine_window(WINDOW_SEQ_LENS[1], 1.0, device) + part = pp.parts[0] + mesh_context = MeshContext.from_meshes(mesh) + + _zero_grads(part) + reference_engine = Engine( + pp, + device=device, + mesh_context=mesh_context, + microbatch_size=BATCH, + optimizers=torch.optim.SGD(part.parameters(), lr=0.01), + max_grad_norm=None, + ) + reference_result = reference_engine.forward_backward(window_a + window_b, token_losses) + reference_grads = _snapshot_grads(part) + + _zero_grads(part) + planned_engine = Engine( + pp, + device=device, + mesh_context=mesh_context, + microbatch_size=BATCH, + optimizers=torch.optim.SGD(part.parameters(), lr=0.01), + max_grad_norm=None, + ) + planned_engine.begin_accumulation([window_a, window_b]) + result_a = planned_engine.forward_backward(window_a, token_losses) + result_b = planned_engine.forward_backward(window_b, token_losses) + planned_grads = _snapshot_grads(part) + + assert result_a.weight_sum.item() != result_b.weight_sum.item(), ( + f"[rank {rank}] Engine fixture must use unequal call denominators" + ) + torch.testing.assert_close(result_a.loss_sum + result_b.loss_sum, reference_result.loss_sum, rtol=2e-2, atol=2e-3) + torch.testing.assert_close( + result_a.weight_sum + result_b.weight_sum, + reference_result.weight_sum, + rtol=0, + atol=0, + ) + combined_loss = (result_a.loss_sum + result_b.loss_sum) / (result_a.weight_sum + result_b.weight_sum) + torch.testing.assert_close(combined_loss, reference_result.loss, rtol=2e-2, atol=2e-3) + + missing = set(reference_grads) ^ set(planned_grads) + assert not missing, f"[rank {rank}] Engine gradient key mismatch: {sorted(missing)[:5]}" + direct_missing = set(direct_normalized_grads) ^ set(planned_grads) + assert not direct_missing, f"[rank {rank}] direct/Engine gradient key mismatch: {sorted(direct_missing)[:5]}" + mismatches = [] + oracle_mismatches = [] + for name in sorted(reference_grads): + got, want = planned_grads[name], reference_grads[name] + if not torch.allclose(got, want, rtol=2e-2, atol=2e-3): + mismatches.append(f"{name}: max|delta|={(got - want).abs().max():.3e}") + direct = direct_normalized_grads[name] + single_close, max_error, error_bound, norm_ratio = _scale_aware_grad_close(want, direct) + if not single_close: + oracle_mismatches.append( + f"{name}/single: max|delta|={max_error:.3e} bound={error_bound:.3e} norm_ratio={norm_ratio:.6f}" + ) + planned_close, max_error, error_bound, norm_ratio = _scale_aware_grad_close(got, direct) + if not planned_close: + oracle_mismatches.append( + f"{name}/planned: max|delta|={max_error:.3e} bound={error_bound:.3e} norm_ratio={norm_ratio:.6f}" + ) + if mismatches: + raise AssertionError( + f"[rank {rank}] Engine planned PP gradients != one complete window " + f"({len(mismatches)}/{len(reference_grads)} parameters differ).\n " + "\n ".join(mismatches[:8]) + ) + if oracle_mismatches: + raise AssertionError( + f"[rank {rank}] Engine normalized gradients disagree with the independent direct normalized oracle " + f"({len(oracle_mismatches)} mismatches; global_weight_sum={reference_result.weight_sum.item():.1f}, " + f"dp_size={mesh['dp'].size()}).\n " + "\n ".join(oracle_mismatches[:8]) + ) + + print(f"[rank {rank}] Engine planned PP/FSDP2 parity OK over {len(reference_grads)} parameters") + + def main() -> None: """Compare accumulated gradients against the sum of per-window gradients.""" dist.init_process_group("nccl") rank, world = dist.get_rank(), dist.get_world_size() + if world not in {2, 4}: + raise ValueError(f"PP grad-accumulation parity requires 2 or 4 ranks, got {world}") torch.cuda.set_device(rank % torch.cuda.device_count()) device = torch.device("cuda", rank % torch.cuda.device_count()) from nemo_automodel.components.distributed.pipelining import AutoPipeline - mesh = init_device_mesh("cuda", (world, 1), mesh_dim_names=("pp", "dp")) + mesh = init_device_mesh("cuda", (2, world // 2), mesh_dim_names=("pp", "dp")) model = _build_model(device) def loss_fn(pred, target): + """Return one PP microbatch's CE normalized by the full local window. + + Args: + pred: Pipeline output with logits shaped + ``[pp_microbatch, sequence, vocab]``. + target: Token labels shaped ``[pp_microbatch, sequence]``. + + Returns: + Scalar summed cross entropy divided by the complete two-window + local token denominator. This matches Engine's backward scale. + """ logits = pred.logits if hasattr(pred, "logits") else pred - return torch.nn.functional.cross_entropy(logits.float().flatten(0, 1), target.flatten(0, 1), reduction="sum") + loss_sum = torch.nn.functional.cross_entropy( + logits.float().flatten(0, 1), + target.flatten(0, 1), + reduction="sum", + ) + return loss_sum / LOCAL_WINDOW_WEIGHT_SUM pp = AutoPipeline( world_mesh=mesh, @@ -191,8 +412,10 @@ def loss_fn(pred, target): _zero_grads(part) _run_window(pp, WINDOW_SEQ_LENS[-1], device) last_only = _snapshot_grads(part) - differs = any(not torch.allclose(last_only[n], reference[n], rtol=1e-3, atol=1e-4) for n in reference) - assert differs, f"[rank {rank}] windows produce identical gradients; test cannot detect a wipe" + catches_dropped_window = any(not _scale_aware_grad_close(last_only[name], reference[name])[0] for name in reference) + assert catches_dropped_window, ( + f"[rank {rank}] last-window-only gradients pass the Engine oracle tolerance; test cannot detect a wipe" + ) mismatches = [] for name in sorted(reference): @@ -211,6 +434,7 @@ def loss_fn(pred, target): ) print(f"[rank {rank}] PP grad-accumulation parity OK over {len(reference)} parameters") + _run_engine_planned_parity(pp, mesh, device, rank, reference) dist.barrier() dist.destroy_process_group() diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 9b93d1e47c..b7518853d6 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -20,6 +20,7 @@ import torch.nn as nn from torch.distributed.pipelining.microbatch import TensorChunkSpec, split_args_kwargs_into_chunks +import nemo_automodel.components.distributed.pipelining.functional as pipeline_functional from nemo_automodel.components.distributed.pipelining.autopipeline import AutoPipeline from nemo_automodel.components.distributed.pipelining.functional import ( generate_hf_model_fqn_per_model_part, @@ -118,6 +119,12 @@ def scale_grads(self, divisor: int): # record the last divisor for verification if needed self._scaled = divisor + def backward_maybe_with_nosync(self, _backward_type, _bwd_kwargs, last_backward=False): + return (), None + + def perform_reduce_grad(self, divisor: int): + self.scale_grads(divisor) + class FakeSchedule: def __init__(self, stages: list[DummyPipelineStage], n_microbatches: int = 1): @@ -359,8 +366,58 @@ class _NoEvalSchedule(_LegacyStepSchedule): eval = None +class _FinalizationStage: + """Record the two schedule-local gradient-finalization signals.""" + + def __init__(self, submod=None): + self.events = [] + self.submod = nn.Module() if submod is None else submod + + def backward_maybe_with_nosync(self, _backward_type, _bwd_kwargs, *, last_backward=False): + self.events.append(("backward", last_backward)) + return (), None + + def perform_reduce_grad(self, divisor): + self.events.append(("reduce", divisor)) + set_requires_gradient_sync = getattr(self.submod, "set_requires_gradient_sync", None) + if callable(set_requires_gradient_sync): + set_requires_gradient_sync(True) + self.scale_grads(divisor) + + def scale_grads(self, divisor): + self.events.append(("scale", divisor)) + + +class _FinalizationSchedule(_KwargsChunkSchedule): + """Model the final backward and reduce calls made by a PyTorch schedule.""" + + def __init__(self, stage): + super().__init__() + self._stage = stage + self._stages = [stage] + + def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): + result = super().step( + *args, + target=target, + losses=losses, + return_outputs=return_outputs, + **kwargs, + ) + self._stage.backward_maybe_with_nosync("full", {}, last_backward=True) + self._stage.perform_reduce_grad(2) + return result + + class TestAutoPipelineKwargsChunkSpec: - def _pipeline_with_parts(self, *parts: nn.Module, schedule=None, has_first_stage: bool = True): + def _pipeline_with_parts( + self, + *parts: nn.Module, + schedule=None, + has_first_stage: bool = True, + defer_fsdp_grad_sync: bool = True, + scale_grads_in_schedule: bool = False, + ): ap = AutoPipeline( world_mesh=FakeDeviceMesh(), pp_axis_name="pp", @@ -368,6 +425,8 @@ def _pipeline_with_parts(self, *parts: nn.Module, schedule=None, has_first_stage pp_microbatch_size=1, pp_batch_size=2, device=torch.device("cpu"), + defer_fsdp_grad_sync=defer_fsdp_grad_sync, + scale_grads_in_schedule=scale_grads_in_schedule, ) ap._info.schedule = schedule or _KwargsChunkSchedule() ap._info.model_parts = list(parts) @@ -463,6 +522,134 @@ def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): assert schedule.args_split == [(), ()] assert all("inputs_embeds" not in kwargs for kwargs in schedule.kwargs_split) + def test_step_microbatches_defers_schedule_finalization_until_the_logical_last_call(self): + stage = _FinalizationStage() + schedule = _FinalizationSchedule(stage) + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + pipeline_functional._make_pipeline_stages_accumulation_aware( + [stage], + reduce_grad_per_microbatch=False, + ) + ap._info.stages = [stage] + model_inputs = [{"input_ids": torch.zeros(1, 8)} for _ in range(2)] + + ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=False) + assert stage.events == [("backward", False), ("scale", 2)] + assert stage._nemo_finalize_backward is True + + ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=True) + assert stage.events == [ + ("backward", False), + ("scale", 2), + ("backward", True), + ("reduce", 2), + ("scale", 2), + ] + + def test_step_microbatches_preserves_requested_per_microbatch_gradient_reduction(self): + stage = _FinalizationStage() + schedule = _FinalizationSchedule(stage) + ap = self._pipeline_with_parts( + nn.Module(), + schedule=schedule, + defer_fsdp_grad_sync=False, + ) + pipeline_functional._make_pipeline_stages_accumulation_aware( + [stage], + reduce_grad_per_microbatch=True, + ) + ap._info.stages = [stage] + + ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + finalize_backward=False, + ) + + assert stage.events == [("backward", True), ("reduce", 2), ("scale", 2)] + + def test_step_microbatches_rejects_cross_call_schedule_gradient_scaling(self): + schedule = _KwargsChunkSchedule() + ap = self._pipeline_with_parts( + nn.Module(), + schedule=schedule, + scale_grads_in_schedule=True, + ) + + with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): + ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + finalize_backward=False, + ) + + assert schedule.step_calls == 0 + + @pytest.mark.parametrize("stage_state", [None, "unwrapped"]) + def test_step_microbatches_requires_accumulation_aware_stages_for_nonfinal_call(self, stage_state): + schedule = _KwargsChunkSchedule() + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + ap._info.stages = None if stage_state is None else [_FinalizationStage()] + + with pytest.raises(RuntimeError, match="accumulation-aware pipeline stages"): + ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + finalize_backward=False, + ) + + assert schedule.step_calls == 0 + + def test_nonfinal_fully_sharded_stage_runs_post_backward_without_sync(self, monkeypatch): + fsdp_events = [] + + class FakeFSDPModule: + def set_is_last_backward(self, value): + fsdp_events.append(("last", value)) + + def set_reshard_after_backward(self, value): + fsdp_events.append(("reshard", value)) + + def set_requires_gradient_sync(self, value): + fsdp_events.append(("sync", value)) + + parameter_group = types.SimpleNamespace(post_backward=lambda: fsdp_events.append(("post", None))) + fsdp_state = types.SimpleNamespace( + _state_ctx=types.SimpleNamespace( + all_states=[types.SimpleNamespace(_fsdp_param_group=parameter_group)], + ), + _root_post_backward_final_callback=lambda: fsdp_events.append(("root", None)), + ) + import torch.distributed.fsdp as torch_fsdp + + monkeypatch.setattr(torch_fsdp, "FSDPModule", FakeFSDPModule) + monkeypatch.setattr(torch_fsdp.fully_shard, "state", lambda _module: fsdp_state) + + stage = _FinalizationStage(FakeFSDPModule()) + pipeline_functional._make_pipeline_stages_accumulation_aware( + [stage], + reduce_grad_per_microbatch=False, + ) + stage._nemo_finalize_backward = False + stage.perform_reduce_grad(2) + + assert fsdp_events == [ + ("last", True), + ("reshard", True), + ("sync", False), + ("post", None), + ("root", None), + ] + assert stage.events == [("scale", 2)] + + fsdp_events.clear() + stage.events.clear() + stage._nemo_finalize_backward = True + stage.perform_reduce_grad(2) + + assert fsdp_events == [("sync", True)] + assert stage.events == [("reduce", 2), ("scale", 2)] + @pytest.mark.parametrize( ("method_name", "schedule_cls"), [ diff --git a/tests/unit_tests/moe/test_fsdp_mixin.py b/tests/unit_tests/moe/test_fsdp_mixin.py index a0da9965ab..0d32e9c039 100644 --- a/tests/unit_tests/moe/test_fsdp_mixin.py +++ b/tests/unit_tests/moe/test_fsdp_mixin.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from unittest.mock import Mock, patch +import pytest + from nemo_automodel.components.models.common.utils import get_is_optim_step, set_is_optim_step from nemo_automodel.components.moe.fsdp_mixin import ( MoEFSDPSyncMixin, @@ -764,10 +767,28 @@ def test_fsdp_module_last_backward(self, mock_fully_shard): assert grads == ((), None) assert param_groups is None + @pytest.mark.parametrize( + ("stage_finalize", "reduce_grad_per_microbatch", "global_optim_step", "expect_post_backward"), + [ + pytest.param(True, False, False, True, id="stage-final-overrides-global-false"), + pytest.param(False, False, True, False, id="stage-nonfinal-overrides-global-true"), + pytest.param(False, True, False, True, id="per-microbatch-reduce-overrides-global-false"), + pytest.param(False, True, True, True, id="per-microbatch-reduce-overrides-global-true"), + pytest.param(None, None, True, True, id="legacy-stage-falls-back-to-global"), + ], + ) @patch("nemo_automodel.components.moe.fsdp_mixin.get_is_optim_step") @patch("nemo_automodel.components.moe.fsdp_mixin.isinstance") - def test_moe_fsdp_mixin_last_backward_with_optim_step(self, mock_isinstance, mock_get_optim): - """Test MoEFSDPSyncMixin path with last_backward=True and IS_OPTIM_STEP=True.""" + def test_moe_fsdp_mixin_pipeline_boundary_is_authoritative( + self, + mock_isinstance, + mock_get_optim, + stage_finalize, + reduce_grad_per_microbatch, + global_optim_step, + expect_post_backward, + ): + """Stage finalization and per-microbatch reduction override the legacy global flag.""" def isinstance_side_effect(obj, cls): if cls == MoEFSDPSyncMixin: @@ -777,12 +798,15 @@ def isinstance_side_effect(obj, cls): return False mock_isinstance.side_effect = isinstance_side_effect - mock_get_optim.return_value = True + mock_get_optim.return_value = global_optim_step - mock_stage = Mock() model = MockFSDPModule() moe_model = MockMoEModel(MockBackend(), model) - mock_stage.submod = moe_model + mock_stage = SimpleNamespace(submod=moe_model) + if stage_finalize is not None: + mock_stage._nemo_finalize_backward = stage_finalize + if reduce_grad_per_microbatch is not None: + mock_stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch bwd_kwargs = { "stage_output": Mock(), @@ -796,8 +820,14 @@ def isinstance_side_effect(obj, cls): result = patched_backward_maybe_with_nosync(mock_stage, "full", bwd_kwargs, last_backward=True) - # Verify post backward was called - mock_run_post.assert_called_once_with(moe_model) + if expect_post_backward: + mock_run_post.assert_called_once_with(moe_model) + else: + mock_run_post.assert_not_called() + if stage_finalize is None: + mock_get_optim.assert_called_once_with() + else: + mock_get_optim.assert_not_called() grads, param_groups = result assert grads == ((), None) assert param_groups is None diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 9ae1d1dcfc..aba755ec9d 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -133,6 +133,7 @@ def __init__( self.step_calls = 0 self.eval_calls = 0 self.backward_calls = 0 + self.finalize_backward_calls = [] self.updated_seq_lens = [] self.updated_microbatch_sizes = [] self.updated_input_shapes = [] @@ -156,11 +157,12 @@ def update_seq_len(self, seq_len, *, microbatch_size=None, input_tensor=None): self.updated_microbatch_sizes.append(microbatch_size) self.updated_input_shapes.append(tuple(input_tensor.shape) if input_tensor is not None else None) - def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): + def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs, finalize_backward=True): assert return_outputs is False assert len(model_inputs) == self.num_microbatches self.prepared_inputs.append(model_inputs) self.step_calls += 1 + self.finalize_backward_calls.append(finalize_backward) if self.events is not None: self.events.append("step") @@ -192,7 +194,7 @@ def eval_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): def _pipeline_mesh_context(): - return SimpleNamespace(pp_size=2, cp_size=1, device_mesh=None, process_group=None) + return SimpleNamespace(pp_size=2, cp_size=1, device_mesh=None, moe_mesh=None, process_group=None) class _DDPWithCP(nn.parallel.DistributedDataParallel): @@ -1748,6 +1750,7 @@ def loss_fn(output, inputs): assert model.weight.grad.item() == pytest.approx(4.5) assert result.loss_fn_outputs == [] assert pipeline.step_calls == 2 + assert pipeline.finalize_backward_calls == [False, True] # The fake schedule performs and counts every backward, then returns None. # A second Engine-owned backward would either fail or change these counts. assert pipeline.backward_calls == backward_calls == 4 @@ -1930,24 +1933,140 @@ def prepare_final(parts, *, pp_enabled): assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) -def test_pipeline_rejects_planned_multi_call_before_forward(): +def test_pipeline_planned_multi_call_matches_one_window_with_unequal_weights(): + window_a = [_datum([[2], [100]], [[1.0], [0.0]])] + window_b = [_datum([[4], [8]], [[0.5], [1.5]])] + + reference_model = ScaleModel() + reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) + reference_pipeline = _FakeAutoPipeline(reference_model, num_microbatches=2) + reference_engine = Engine( + reference_pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + optimizers=reference_optimizer, + max_grad_norm=None, + ) + reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) + reference_grad = reference_model.weight.grad.detach().clone() + reference_step = reference_engine.optim_step() + + planned_model = ScaleModel() + planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) + planned_pipeline = _FakeAutoPipeline(planned_model, num_microbatches=2) + planned_engine = Engine( + planned_pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + optimizers=planned_optimizer, + max_grad_norm=None, + ) + planned_engine.begin_accumulation([window_a, window_b]) + result_a = planned_engine.forward_backward(window_a, _identity_loss) + result_b = planned_engine.forward_backward(window_b, _identity_loss) + planned_grad = planned_model.weight.grad.detach().clone() + planned_step = planned_engine.optim_step() + + assert result_a.loss_sum.item() == pytest.approx(2.0) + assert result_a.weight_sum.item() == pytest.approx(1.0) + assert result_a.loss.item() == pytest.approx(2.0) + assert result_b.loss_sum.item() == pytest.approx(14.0) + assert result_b.weight_sum.item() == pytest.approx(2.0) + assert result_b.loss.item() == pytest.approx(7.0) + assert reference_result.loss_sum.item() == pytest.approx(16.0) + assert reference_result.weight_sum.item() == pytest.approx(3.0) + assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) + assert reference_pipeline.step_calls == planned_pipeline.step_calls == 2 + assert reference_pipeline.finalize_backward_calls == [False, True] + assert planned_pipeline.finalize_backward_calls == [False, True] + assert reference_pipeline.backward_calls == planned_pipeline.backward_calls == 4 + torch.testing.assert_close(planned_grad, reference_grad) + torch.testing.assert_close(planned_step.grad_norm, reference_step.grad_norm) + torch.testing.assert_close(planned_model.weight, reference_model.weight) + + +def test_pipeline_planned_accumulation_uses_one_lifecycle_and_all_inner_microbatches(monkeypatch): + events = [] + + monkeypatch.setattr( + engine_module, + "prepare_for_grad_accumulation", + lambda _parts, *, pp_enabled: events.append(f"prepare:{pp_enabled}"), + ) + monkeypatch.setattr( + engine_module, + "prepare_for_final_backward", + lambda _parts, *, pp_enabled: events.append(f"final:{pp_enabled}"), + ) + monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after_first")) + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) + + model = _MainAndAuxScaleModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2, events=events) + engine = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + max_grad_norm=None, + ) + window_a = [_datum([[100], [200]], [[0.0], [0.0]])] + window_b = [_datum([[3], [5]], [[1.0], [1.0]])] + + engine.begin_accumulation([window_a, window_b]) + result_a = engine.forward_backward(window_a, _identity_loss) + result_b = engine.forward_backward(window_b, _identity_loss) + + assert events == ["prepare:True", "step", "after_first", "final:True", "step"] + assert result_a.loss.item() == pytest.approx(0.0) + assert result_b.loss.item() == pytest.approx(4.0) + assert pipeline.step_calls == 2 + assert pipeline.backward_calls == 4 + assert pipeline.finalize_backward_calls == [False, True] + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) + assert model.main_weight.grad.item() == pytest.approx(4.0) + assert model.aux_weight.grad.item() == pytest.approx(1.0) + + +def test_pipeline_planned_accumulation_enforces_state_and_retries_a_failed_fence(): model = ScaleModel() pipeline = _FakeAutoPipeline(model, num_microbatches=2) engine = Engine( pipeline, device="cpu", mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + optimizers=torch.optim.SGD(model.parameters(), lr=0.1), + max_grad_norm=None, ) - window_a = [_datum([1, 2])] - window_b = [_datum([3, 4])] + window_a = [_datum([[1], [3]])] + window_b = [_datum([[5], [7]])] + engine.begin_accumulation([window_a, window_b]) + engine.forward_backward(window_a, _identity_loss) - with pytest.raises(NotImplementedError, match="pipeline|PP"): - engine.begin_accumulation([window_a, window_b]) + with pytest.raises(RuntimeError, match="finish every forward_backward"): + engine.optim_step() + torch.testing.assert_close(model.weight, torch.tensor(1.0)) - assert pipeline.step_calls == 0 - assert pipeline.backward_calls == 0 - assert model.forward_calls == 0 + engine.forward_backward(window_b, _identity_loss) + expected_grad = model.weight.grad.detach().clone() + + def fail_before_step(): + raise ValueError("checkpoint staging failed") + + with pytest.raises(ValueError, match="checkpoint staging failed"): + engine.optim_step(before_optimizer_step=fail_before_step) + torch.testing.assert_close(model.weight, torch.tensor(1.0)) + torch.testing.assert_close(model.weight.grad, expected_grad) + + engine.optim_step() + torch.testing.assert_close(model.weight, 1.0 - 0.1 * expected_grad) assert model.weight.grad is None + with pytest.raises(RuntimeError, match="already consumed"): + engine.optim_step() def test_pipeline_outputs_follow_logical_microbatch_order(): From 8b4b9320e1de616c846af753544c1cc02ae5db1c Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 16:47:11 -0700 Subject: [PATCH 20/34] feat(engine): support MegatronFSDP summed gradients Signed-off-by: HuiyingLi --- nemo_automodel/engine/__init__.py | 81 ++++++- nemo_automodel/recipes/llm/train_ft.py | 5 - nemo_automodel/recipes/vlm/finetune.py | 5 - ...Parallelism_MegatronFSDP_Per_Token_Loss.sh | 24 ++ .../run_megatron_fsdp_per_token_loss.py | 227 ++++++++++++++++++ .../parallelism/test_parallelism.py | 4 + .../recipes/test_finetune_vlm_helpers.py | 8 +- tests/unit_tests/recipes/test_train_ft.py | 8 +- tests/unit_tests/test_engine.py | 180 +++++++++++++- 9 files changed, 510 insertions(+), 32 deletions(-) create mode 100644 tests/functional_tests/parallelism/L2_Parallelism_MegatronFSDP_Per_Token_Loss.sh create mode 100644 tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 06702ef377..315ee0a7ee 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -92,6 +92,54 @@ def _as_tuple(value: _T | Sequence[_T] | None) -> tuple[_T, ...]: return (value,) +def _resolve_summed_gradient_reduction(model_parts: Sequence[nn.Module]) -> bool: + """Resolve whether every local model part uses summed gradient collectives. + + MegatronFSDP exposes ``calculate_per_token_loss`` on its wrapper. Walking + each part's module tree also finds it below ordinary wrapper layers such as + DDP or compilation wrappers. A part without a declaration uses the normal + averaged-gradient contract. + + Args: + model_parts: Local eager model or pipeline parts after distributed + wrapping. + + Returns: + ``True`` when every part declares summed gradients, otherwise ``False``. + + Raises: + TypeError: If a declaration is not boolean. + ValueError: If modules or model parts disagree on the reduction mode. + """ + part_modes: list[bool] = [] + sentinel = object() + for part_index, part in enumerate(model_parts): + declared_modes: set[bool] = set() + for module in part.modules(): + declared_mode = getattr(module, "calculate_per_token_loss", sentinel) + if declared_mode is sentinel: + continue + if not isinstance(declared_mode, bool): + raise TypeError( + "model calculate_per_token_loss declarations must be boolean; " + f"part {part_index} has {declared_mode!r}" + ) + declared_modes.add(declared_mode) + if len(declared_modes) > 1: + raise ValueError( + "model part mixes calculate_per_token_loss=True and False; " + "Engine requires one gradient-reduction mode per optimizer window" + ) + part_modes.append(next(iter(declared_modes), False)) + + if len(set(part_modes)) > 1: + raise ValueError( + "model parts disagree on calculate_per_token_loss; " + "Engine requires one gradient-reduction mode per optimizer window" + ) + return bool(part_modes and part_modes[0]) + + def _tensor_version(tensor: torch.Tensor) -> int: """Return the in-place mutation counter when the Tensor exposes one.""" try: @@ -372,6 +420,7 @@ def __init__( self.pipeline = model if isinstance(model, AutoPipeline) else None self.model_parts = model.parts if self.pipeline is not None else [model] self.model = self.model_parts[0] + self._summed_gradient_reduction = _resolve_summed_gradient_reduction(self.model_parts) self._fp8_scale_precompute_parts, self._fp8_scale_precompute_fn = _resolve_fp8_scale_precompute( self.model_parts ) @@ -835,6 +884,7 @@ def _forward_backward_window( self._validate_parallelism() dp_group, dp_size = self._dp_group_and_size() grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) + gradient_reduction_multiplier = self._gradient_reduction_multiplier(grad_group_size) self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) denominator = ( self._global_weight_sum(microbatches, dp_group, dp_size) @@ -865,7 +915,9 @@ def _forward_backward_window( effective_total_microbatches = ( len(microbatches) * inner_microbatches if total_microbatches is None else total_microbatches ) - MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(self._cp_size() / effective_total_microbatches) + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( + self._cp_size() * gradient_reduction_multiplier / (grad_group_size * effective_total_microbatches) + ) local_loss_sum = torch.zeros((), dtype=torch.float64, device=self.device) loss_fn_outputs: list[dict[str, Any]] = [] @@ -874,7 +926,7 @@ def _forward_backward_window( backward_scale = ( safe_gradient_denominator.new_zeros(()) if zero_denominator or zero_gradient_denominator - else safe_gradient_denominator.new_tensor(grad_group_size) / safe_gradient_denominator + else safe_gradient_denominator.new_tensor(gradient_reduction_multiplier) / safe_gradient_denominator ) for index, datums in enumerate(microbatches): @@ -2067,18 +2119,25 @@ def _materialize_pipeline_microbatches( def _validate_parallelism(self) -> None: """Validate topology plus backward-specific distributed contracts.""" self._validate_execution_parallelism() - if any( - bool(getattr(module, "calculate_per_token_loss", False)) - for part in self.model_parts - for module in part.modules() - ): - raise NotImplementedError( - "Engine.forward_backward requires averaged distributed gradients; " - "MegatronFSDP calculate_per_token_loss=True uses summed gradients" - ) if self.pipeline is not None and self.pipeline.scale_grads_in_schedule: raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") + def _gradient_reduction_multiplier(self, grad_group_size: int) -> int: + """Return the factor that compensates the backend gradient collective. + + Averaging backends divide by the complete DP-CP gradient group, so the + local loss numerator is multiplied by that group size before backward. + MegatronFSDP ``calculate_per_token_loss=True`` uses SUM collectives and + therefore needs no such compensation. + + Args: + grad_group_size: Size of the complete DP-CP gradient group. + + Returns: + One for summed gradients, otherwise ``grad_group_size``. + """ + return 1 if self._summed_gradient_reduction else grad_group_size + def _validate_execution_parallelism(self) -> None: """Validate model-parallel topology shared by forward and backward.""" if ( diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index 40dbacbe72..db75675233 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -521,11 +521,6 @@ def setup(self): if not self._should_setup_training_components(): return - if getattr(self.distributed_config, "calculate_per_token_loss", False): - raise NotImplementedError( - "Engine-backed finetuning does not support " - "MegatronFSDP calculate_per_token_loss=True; use averaged gradients instead." - ) if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): raise ValueError("Engine-backed finetuning requires distributed.pipeline.scale_grads_in_schedule=False") diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 6da6e0750b..3efbac3b9f 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -451,11 +451,6 @@ def setup(self): if not self._should_setup_training_components(): return - if getattr(self.distributed_config, "calculate_per_token_loss", False): - raise NotImplementedError( - "Engine-backed VLM finetuning does not support " - "MegatronFSDP calculate_per_token_loss=True; use averaged gradients instead." - ) if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): raise ValueError("Engine-backed VLM finetuning requires distributed.pipeline.scale_grads_in_schedule=False") diff --git a/tests/functional_tests/parallelism/L2_Parallelism_MegatronFSDP_Per_Token_Loss.sh b/tests/functional_tests/parallelism/L2_Parallelism_MegatronFSDP_Per_Token_Loss.sh new file mode 100644 index 0000000000..075d5a9be1 --- /dev/null +++ b/tests/functional_tests/parallelism/L2_Parallelism_MegatronFSDP_Per_Token_Loss.sh @@ -0,0 +1,24 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# 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. + +#!/bin/bash +# Tiny real-MegatronFSDP parity for SUM-gradient per-token loss mode. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="0,1" + +torchrun --nproc-per-node=2 --standalone \ + tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py diff --git a/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py new file mode 100644 index 0000000000..1f190d96a0 --- /dev/null +++ b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py @@ -0,0 +1,227 @@ +# 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. + +"""Real MegatronFSDP SUM-vs-average gradient parity through Engine. + +Run with:: + + torchrun --standalone --nproc-per-node=2 run_megatron_fsdp_per_token_loss.py +""" + +from __future__ import annotations + +import math + +import torch +import torch.distributed as dist +from torch import nn + +from nemo_automodel.components.datasets.datum import Datum +from nemo_automodel.components.distributed.config import MegatronFSDPConfig +from nemo_automodel.components.distributed.megatron_fsdp import MegatronFSDPManager +from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes +from nemo_automodel.engine import Engine + + +class TinyBlock(nn.Module): + """One real MegatronFSDP wrapping unit operating on ``[batch, 4]`` tokens.""" + + def __init__(self) -> None: + super().__init__() + self.projection = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + self.projection.weight.copy_(torch.eye(4)) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + """Project float token features from ``[batch, 4]`` to ``[batch, 4]``.""" + return self.projection(tokens) + + +class TinyTokenModel(nn.Module): + """Tiny model whose block class is auto-derived as a MegatronFSDP unit.""" + + _no_split_modules = ["TinyBlock"] + + def __init__(self) -> None: + super().__init__() + self.block = TinyBlock() + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Map numeric token rows ``[batch, 4]`` to outputs of the same shape. + + Args: + input_ids: Numeric token features shaped ``[batch, sequence=4]``. + attention_mask: Engine-generated validity mask with shape + ``[batch, sequence=4]``. This token-independent projection does + not need to consume it. + padding_mask: Engine-generated padding indicator with shape + ``[batch, sequence=4]``; also unused by this projection. + + Returns: + Projected per-token values shaped ``[batch, sequence=4]``. + """ + del attention_mask, padding_mask + return self.block(input_ids.to(torch.float32)) + + +def _windows(rank: int) -> tuple[list[Datum], list[Datum]]: + """Build two rank-local windows with unequal global token denominators. + + Args: + rank: Data-parallel rank, zero or one. + + Returns: + Two one-Datum windows. ``input_ids`` and ``weights`` both have layout + ``[sequence=4]``; Engine collates them to ``[batch=1, sequence=4]``. + Across DP, the first and second windows have weight sums three and + four, respectively. + """ + if rank == 0: + values_a, weights_a = [1, 2, 3, 4], [1.0, 0.0, 0.0, 0.0] + values_b, weights_b = [5, 6, 7, 8], [0.0, 1.0, 1.0, 1.0] + else: + values_a, weights_a = [9, 10, 11, 12], [1.0, 1.0, 0.0, 0.0] + values_b, weights_b = [13, 14, 15, 16], [0.0, 0.0, 0.0, 1.0] + + def datum(values: list[int], weights: list[float]) -> Datum: + """Create one Datum with token values and weights shaped ``[sequence=4]``.""" + return Datum( + model_inputs={"input_ids": torch.tensor(values)}, + loss_fn_inputs={"weights": torch.tensor(weights)}, + ) + + return [datum(values_a, weights_a)], [datum(values_b, weights_b)] + + +def _local_parameters(model: nn.Module) -> dict[str, torch.Tensor]: + """Snapshot float32 local shards of all named model parameters.""" + parameters = {} + for name, parameter in model.named_parameters(): + value = parameter.detach() + value = value.to_local() if hasattr(value, "to_local") else value + parameters[name] = value.float().cpu().clone() + return parameters + + +def _per_token_identity_loss(output: torch.Tensor, loss_inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Return model outputs as Engine-owned weighted per-token losses. + + Args: + output: Model values shaped ``[batch=1, sequence=4]``. + loss_inputs: Collated mapping whose ``weights`` tensor has the same + ``[batch=1, sequence=4]`` layout. + + Returns: + Per-token losses with shape ``[batch=1, sequence=4]``. Engine applies + ``weights`` and the complete planned-window denominator. + """ + assert output.shape == loss_inputs["weights"].shape + return output + + +def _run_mode( + mesh_context: MeshContext, + *, + summed_gradients: bool, +) -> tuple[tuple[float, ...], float, dict[str, torch.Tensor]]: + """Run one planned update through a real MegatronFSDP reduction mode. + + Args: + mesh_context: Two-rank ``[dp=2, cp=1, tp=1]`` CUDA mesh. + summed_gradients: Value passed as ``calculate_per_token_loss``. True + selects SUM gradient collectives; false selects averaged gradients. + + Returns: + Call-local loss sums, weight sums, and normalized losses; the global + gradient norm; and float32 local parameter shards after the update. + """ + torch.manual_seed(1234) + model = TinyTokenModel().cuda() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + config = MegatronFSDPConfig( + zero_dp_strategy=3, + overlap_grad_reduce=False, + overlap_param_gather=False, + check_for_nan_in_grad=False, + disable_bucketing=True, + calculate_per_token_loss=summed_gradients, + ) + model, optimizer = MegatronFSDPManager(config, mesh_context.device_mesh).parallelize(model, optimizer) + engine = Engine( + model, + device=torch.device("cuda", torch.cuda.current_device()), + mesh_context=mesh_context, + optimizers=optimizer, + max_grad_norm=1e9, + ) + expected_multiplier = 1 if summed_gradients else dist.get_world_size() + assert engine._gradient_reduction_multiplier(dist.get_world_size()) == expected_multiplier + + window_a, window_b = _windows(dist.get_rank()) + engine.begin_accumulation([window_a, window_b]) + result_a = engine.forward_backward(window_a, _per_token_identity_loss) + result_b = engine.forward_backward(window_b, _per_token_identity_loss) + step_result = engine.optim_step() + + statistics = ( + result_a.loss_sum.item(), + result_a.weight_sum.item(), + result_a.loss.item(), + result_b.loss_sum.item(), + result_b.weight_sum.item(), + result_b.loss.item(), + ) + return statistics, float(step_result.grad_norm), _local_parameters(model) + + +def main() -> None: + """Assert real MegatronFSDP SUM and average modes produce one update.""" + dist.init_process_group("nccl") + rank = dist.get_rank() + if dist.get_world_size() != 2: + raise ValueError(f"MegatronFSDP per-token parity requires two ranks, got {dist.get_world_size()}") + torch.cuda.set_device(int(torch.distributed.get_rank() % torch.cuda.device_count())) + + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=2, cp_size=1, tp_size=1), + world_size=2, + ) + averaged = _run_mode(mesh_context, summed_gradients=False) + dist.barrier() + summed = _run_mode(mesh_context, summed_gradients=True) + + assert math.isfinite(averaged[1]) and averaged[1] > 0 + assert math.isfinite(summed[1]) and summed[1] > 0 + torch.testing.assert_close(torch.tensor(summed[0]), torch.tensor(averaged[0]), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(torch.tensor(summed[1]), torch.tensor(averaged[1]), rtol=1e-5, atol=1e-5) + assert summed[0][1] == 3.0 + assert summed[0][4] == 4.0 + assert set(summed[2]) == set(averaged[2]) + for name in sorted(averaged[2]): + torch.testing.assert_close(summed[2][name], averaged[2][name], rtol=1e-5, atol=1e-5) + + if rank == 0: + print("MegatronFSDP calculate_per_token_loss SUM gradients match averaged-gradient Engine update") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/functional_tests/parallelism/test_parallelism.py b/tests/functional_tests/parallelism/test_parallelism.py index b7849ff0da..9894eca997 100644 --- a/tests/functional_tests/parallelism/test_parallelism.py +++ b/tests/functional_tests/parallelism/test_parallelism.py @@ -28,6 +28,7 @@ GEMMA4_PP2_PARITY_FILENAME = "L2_Parallelism_VLM_Gemma4_PP2_Parity.sh" GEMMA4_TP2_PARITY_FILENAME = "L2_Parallelism_VLM_Gemma4_TP2_Parity.sh" PP_GRAD_ACCUM_PARITY_FILENAME = "L2_Parallelism_PP_Grad_Accum_Parity.sh" +MEGATRON_FSDP_PER_TOKEN_LOSS_FILENAME = "L2_Parallelism_MegatronFSDP_Per_Token_Loss.sh" class TestParallelismParity: @@ -39,3 +40,6 @@ def test_gemma4_tp2_parity(self): def test_pp_grad_accum_parity(self): run_test_script(TEST_FOLDER, PP_GRAD_ACCUM_PARITY_FILENAME) + + def test_megatron_fsdp_per_token_loss(self): + run_test_script(TEST_FOLDER, MEGATRON_FSDP_PER_TOKEN_LOSS_FILENAME) diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 3fb00bcc07..0fa937ba07 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -1873,14 +1873,16 @@ def _patch_vlm_distributed_setup( ) -def test_vlm_setup_rejects_calculate_per_token_loss(monkeypatch): +def test_vlm_setup_allows_calculate_per_token_loss(monkeypatch): cfg = _minimal_vlm_cfg(cp_size=1, rope_fusion=False) _patch_vlm_setup_minimals(monkeypatch, cp_size=1) _patch_vlm_distributed_setup(monkeypatch, pp_enabled=False, calculate_per_token_loss=True) trainer = FinetuneRecipeForVLM(cfg) - with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): - trainer.setup() + trainer.setup() + + assert trainer.distributed_config.calculate_per_token_loss is True + assert trainer.engine is not None def test_vlm_setup_rejects_pipeline_schedule_gradient_scaling(monkeypatch): diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 47e10a9736..aff67c520c 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -1290,7 +1290,7 @@ class DummyAutoPipeline(SimpleNamespace): assert dataloader_build_kwargs[0]["collate_wrapper"] is (None if dataloader_emits_thd else pp_collate_wrapper) -def test_setup_rejects_per_token_megatron_fsdp(monkeypatch): +def test_setup_allows_per_token_megatron_fsdp(monkeypatch): cfg = _minimal_cfg_with_nvtx(nvtx_value=False) _patch_setup_minimals(monkeypatch, lambda *args, **kwargs: None) monkeypatch.setattr( @@ -1311,8 +1311,10 @@ def test_setup_rejects_per_token_megatron_fsdp(monkeypatch): ) trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) - with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): - trainer.setup() + trainer.setup() + + assert trainer.distributed_config.calculate_per_token_loss is True + assert trainer.engine is not None def test_engine_pipeline_loss_reuses_configured_loss_and_thd_metadata(): diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index aba755ec9d..489d5f2752 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -223,6 +223,13 @@ def _identity_loss(output, _loss_inputs): return output +def _configure_fake_gradient_group(engine: Engine, monkeypatch, *, group_size: int) -> None: + """Expose a logical gradient group without running a real collective.""" + engine._gradient_group_and_size = lambda _group, _size: (None, group_size) + engine._validate_window_size_across_group = lambda *_args, **_kwargs: None + monkeypatch.setattr(engine_module.dist, "all_reduce", lambda *_args, **_kwargs: None) + + def test_engine_and_datum_are_lazy_top_level_exports(): assert PublicEngine is Engine assert PublicDatum is Datum @@ -2757,15 +2764,143 @@ def test_pipeline_parallelism_fails_before_forward(): assert model.forward_calls == 0 -def test_megatron_fsdp_per_token_loss_mode_fails_before_forward(): +@pytest.mark.parametrize( + ("declared_mode", "expected_multiplier", "expected_pre_collective_grad"), + [ + pytest.param(None, 4, 14.0, id="undeclared-averaged"), + pytest.param(False, 4, 14.0, id="explicit-averaged"), + pytest.param(True, 1, 3.5, id="summed"), + ], +) +def test_gradient_reduction_mode_controls_main_loss_scale( + monkeypatch, + declared_mode, + expected_multiplier, + expected_pre_collective_grad, +): + model = ScaleModel() + if declared_mode is not None: + model.calculate_per_token_loss = declared_mode + engine = Engine(model, device="cpu") + _configure_fake_gradient_group(engine, monkeypatch, group_size=4) + + result = engine.forward_backward([_datum([2, 4], [1.0, 3.0])], _identity_loss) + + assert engine._gradient_reduction_multiplier(4) == expected_multiplier + assert result.loss_sum.item() == pytest.approx(14.0) + assert result.weight_sum.item() == pytest.approx(4.0) + assert result.loss.item() == pytest.approx(3.5) + assert model.weight.grad.item() == pytest.approx(expected_pre_collective_grad) + + +def test_gradient_reduction_mode_finds_summed_backend_below_ordinary_wrapper(): + model = ScaleModel() + model.distributed_backend = nn.Identity() + model.distributed_backend.calculate_per_token_loss = True + + engine = Engine(model, device="cpu") + + assert engine._gradient_reduction_multiplier(8) == 1 + + +def test_summed_gradient_mode_planned_accumulation_matches_one_window(monkeypatch): + window_a = [_datum([2, 100], [1.0, 0.0])] + window_b = [_datum([4, 8], [0.5, 1.5])] + + reference_model = ScaleModel() + reference_model.calculate_per_token_loss = True + reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) + reference_engine = Engine( + reference_model, + device="cpu", + optimizers=reference_optimizer, + max_grad_norm=None, + ) + _configure_fake_gradient_group(reference_engine, monkeypatch, group_size=4) + reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) + reference_grad = reference_model.weight.grad.detach().clone() + reference_engine.optim_step() + + planned_model = ScaleModel() + planned_model.calculate_per_token_loss = True + planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) + planned_engine = Engine( + planned_model, + device="cpu", + optimizers=planned_optimizer, + max_grad_norm=None, + ) + _configure_fake_gradient_group(planned_engine, monkeypatch, group_size=4) + planned_engine.begin_accumulation([window_a, window_b]) + result_a = planned_engine.forward_backward(window_a, _identity_loss) + result_b = planned_engine.forward_backward(window_b, _identity_loss) + + assert result_a.weight_sum.item() == pytest.approx(1.0) + assert result_b.weight_sum.item() == pytest.approx(2.0) + assert reference_result.weight_sum.item() == pytest.approx(3.0) + assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) + torch.testing.assert_close(planned_model.weight.grad, reference_grad) + planned_engine.optim_step() + torch.testing.assert_close(planned_model.weight, reference_model.weight) + + +def test_summed_gradient_mode_zero_weight_window_keeps_graph_connected_zero(monkeypatch): model = ScaleModel() model.calculate_per_token_loss = True + engine = Engine(model, device="cpu") + _configure_fake_gradient_group(engine, monkeypatch, group_size=4) - with pytest.raises(NotImplementedError, match="calculate_per_token_loss=True"): - Engine(model, device="cpu").forward_backward([_datum([1])], _identity_loss) + result = engine.forward_backward([_datum([100, 200], [0.0, 0.0])], _identity_loss) - assert model.forward_calls == 0 - assert model.weight.grad is None + assert result.loss.item() == 0 + assert result.loss_sum.item() == 0 + assert result.weight_sum.item() == 0 + assert model.forward_calls == 1 + assert model.weight.grad.item() == 0 + + +@pytest.mark.parametrize( + ("summed_gradients", "expected_scale", "expected_local_aux_grad"), + [ + pytest.param(False, 2.0 / 3.0, 2.0, id="averaged"), + pytest.param(True, 1.0 / 12.0, 0.25, id="summed"), + ], +) +def test_moe_aux_scale_uses_general_gradient_reduction_formula( + monkeypatch, + summed_gradients, + expected_scale, + expected_local_aux_grad, +): + monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) + model = _MainAndAuxScaleModel() + model.calculate_per_token_loss = summed_gradients + engine = Engine(model, device="cpu") + engine._cp_size = lambda: 2 + _configure_fake_gradient_group(engine, monkeypatch, group_size=8) + + engine.forward_backward([_datum([1]), _datum([2]), _datum([3])], _identity_loss) + + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(expected_scale) + assert model.aux_weight.grad.item() == pytest.approx(expected_local_aux_grad) + + +@pytest.mark.parametrize( + ("root_mode", "nested_mode", "error_type", "match"), + [ + (True, False, ValueError, "mixes calculate_per_token_loss"), + ("yes", None, TypeError, "must be boolean"), + ], +) +def test_gradient_reduction_mode_rejects_ambiguous_declarations(root_mode, nested_mode, error_type, match): + model = ScaleModel() + model.calculate_per_token_loss = root_mode + if nested_mode is not None: + model.mode_probe = nn.Identity() + model.mode_probe.calculate_per_token_loss = nested_mode + + with pytest.raises(error_type, match=match): + Engine(model, device="cpu") def test_loss_shape_must_exactly_match_weights(): @@ -2954,6 +3089,41 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size planned_engine.optim_step() assert planned_model.module.weight.item() == pytest.approx(1.0 - 0.1 * expected_global_mean) + + # Mimic MegatronFSDP calculate_per_token_loss=True with a SUM hook over + # the same DP-CP gradient group. Engine must not compensate for an + # average a second time, and the DP-only denominator must still produce + # the identical normalized update under both CP-only and DP+CP meshes. + summed_model = _DistributedCPModel() + summed_model.calculate_per_token_loss = True + gradient_group = get_flat_mesh(mesh_context.device_mesh, "dp_cp").get_group() + + def sum_gradient(gradient): + dist.all_reduce(gradient, op=dist.ReduceOp.SUM, group=gradient_group) + return gradient + + summed_model.weight.register_hook(sum_gradient) + summed_optimizer = torch.optim.SGD(summed_model.parameters(), lr=0.1) + summed_engine = Engine( + summed_model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + optimizers=summed_optimizer, + max_grad_norm=None, + ) + summed_engine.begin_accumulation([window_a, window_b]) + summed_result_a = summed_engine.forward_backward(window_a, _identity_loss) + summed_result_b = summed_engine.forward_backward(window_b, _identity_loss) + + torch.testing.assert_close(summed_result_a.loss_sum, result_a.loss_sum) + torch.testing.assert_close(summed_result_a.weight_sum, result_a.weight_sum) + torch.testing.assert_close(summed_result_b.loss_sum, result_b.loss_sum) + torch.testing.assert_close(summed_result_b.weight_sum, result_b.weight_sum) + assert summed_model.weight.grad.item() == pytest.approx(expected_global_mean) + + summed_engine.optim_step() + torch.testing.assert_close(summed_model.weight, planned_model.module.weight) finally: dist.destroy_process_group() From b02090bc0b48236b2f02bbc2f849cc479d3d6259 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 17:37:58 -0700 Subject: [PATCH 21/34] refactor(engine): simplify update and pipeline finalization Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 4 +- .../distributed/pipelining/functional.py | 6 +- nemo_automodel/components/moe/fsdp_mixin.py | 26 +++---- nemo_automodel/engine/__init__.py | 71 ++++++------------- .../run_megatron_fsdp_per_token_loss.py | 3 - .../pipelining/test_autopipeline.py | 5 ++ tests/unit_tests/moe/test_fsdp_mixin.py | 15 ++-- tests/unit_tests/test_engine.py | 18 ++--- 8 files changed, 56 insertions(+), 92 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index a85512eba3..a12e098354 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -513,7 +513,9 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: if schedule_method == "step": for stage in self._info.stages or (): previous_stage_state.append((stage, vars(stage).get("_nemo_finalize_backward", _MISSING_STAGE_STATE))) - stage._nemo_finalize_backward = finalize_backward + stage._nemo_finalize_backward = finalize_backward or getattr( + stage, "_reduce_grad_per_microbatch", False + ) schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) schedule._loss_fn = indexed_loss try: diff --git a/nemo_automodel/components/distributed/pipelining/functional.py b/nemo_automodel/components/distributed/pipelining/functional.py index 95ee0519e1..cb0557a602 100644 --- a/nemo_automodel/components/distributed/pipelining/functional.py +++ b/nemo_automodel/components/distributed/pipelining/functional.py @@ -532,11 +532,10 @@ def _accumulation_aware_backward_maybe_with_nosync( tensors in their model-defined local layouts, plus optional split-backward parameter-group records. """ - finalize_backward = self._nemo_finalize_backward or self._reduce_grad_per_microbatch return self._nemo_original_backward_maybe_with_nosync( backward_type, bwd_kwargs, - last_backward=last_backward and finalize_backward, + last_backward=last_backward and self._nemo_finalize_backward, ) @@ -550,8 +549,7 @@ def _accumulation_aware_perform_reduce_grad(self: PipelineStage, grad_scale_fact once per schedule. The original implementation is used unchanged for ordinary calls, final planned windows, and per-microbatch reduction mode. """ - finalize_backward = self._nemo_finalize_backward or self._reduce_grad_per_microbatch - if finalize_backward: + if self._nemo_finalize_backward: self._nemo_original_perform_reduce_grad(grad_scale_factor) return diff --git a/nemo_automodel/components/moe/fsdp_mixin.py b/nemo_automodel/components/moe/fsdp_mixin.py index 766a1ebc0f..9725e6bf84 100644 --- a/nemo_automodel/components/moe/fsdp_mixin.py +++ b/nemo_automodel/components/moe/fsdp_mixin.py @@ -290,27 +290,21 @@ def perform_backward( result = perform_backward(backward_type)() if last_backward: # Manually call post backward for FSDP - def run_post_backward(fsdp_module: FSDPModule) -> None: - fsdp_module.set_is_last_backward(True) - fsdp_module.set_reshard_after_backward(True) - fsdp_module.set_requires_gradient_sync(True) - fsdp_state = fully_shard.state(fsdp_module) # type: ignore[attr-defined] - for state in fsdp_state._state_ctx.all_states: - if state._fsdp_param_group: - state._fsdp_param_group.post_backward() - - # it would be much better if pipelining backward invoked .backward so autograd hooks - # worked and modules like DDP/FSDP behaved as expected. Working around this for the time being, - # we need to call this too to ensure FSDP syncs its grad reduction ops back to the default stream. - fsdp_state._root_post_backward_final_callback() - - run_post_backward(self.submod) + _configure_fsdp_module( + self.submod, + is_last_backward=True, + reshard_after_backward=True, + requires_gradient_sync=True, + ) + # Pipeline backward does not invoke ``.backward()``, so manually + # finish FSDP post-backward and synchronize its reduction stream. + _run_post_backward_hooks(self.submod)() # If submod is a MoEFSDPSyncMixin, use the MoE-specific FSDP functions elif isinstance(self.submod, MoEFSDPSyncMixin): _disable_fsdp_for_moe_module(self.submod) result = perform_backward(backward_type)() if hasattr(self, "_nemo_finalize_backward"): - finalize_backward = self._nemo_finalize_backward or getattr(self, "_reduce_grad_per_microbatch", False) + finalize_backward = self._nemo_finalize_backward else: finalize_backward = get_is_optim_step() if last_backward and finalize_backward: diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index 315ee0a7ee..f38d919c1f 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -884,7 +884,7 @@ def _forward_backward_window( self._validate_parallelism() dp_group, dp_size = self._dp_group_and_size() grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) - gradient_reduction_multiplier = self._gradient_reduction_multiplier(grad_group_size) + gradient_reduction_multiplier = 1 if self._summed_gradient_reduction else grad_group_size self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) denominator = ( self._global_weight_sum(microbatches, dp_group, dp_size) @@ -1132,24 +1132,19 @@ def optim_step( except Exception as error: local_preflight_error = error - control_group: dist.ProcessGroup | None = None - control_group_size = 1 # Explicit plans pay the small control collectives needed to fail # together before mutation. Keep the ordinary one-call path free of # new per-step collectives; its distributed callback contract remains # the same as before planned accumulation was introduced. - synchronize_update = state is not None - if synchronize_update: - control_group, control_group_size = self._accumulation_control_group_and_size() - if synchronize_update: - self._synchronize_accumulation_error( - local_preflight_error, - control_group, - control_group_size, - peer_message="another model-parallel rank rejected optim_step", - ) - elif local_preflight_error is not None: - raise local_preflight_error + control_group, control_group_size = ( + self._accumulation_control_group_and_size() if state is not None else (None, 1) + ) + self._synchronize_accumulation_error( + local_preflight_error, + control_group, + control_group_size, + peer_message="another model-parallel rank rejected optim_step", + ) device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None moe_mesh = self.mesh_context.moe_mesh if self.mesh_context is not None else None @@ -1182,15 +1177,12 @@ def optim_step( raise RuntimeError("gradient finalization did not return a gradient norm") except Exception as error: finalization_error = error - if synchronize_update: - self._synchronize_accumulation_error( - finalization_error, - control_group, - control_group_size, - peer_message="another model-parallel rank failed while finalizing gradients", - ) - elif finalization_error is not None: - raise finalization_error + self._synchronize_accumulation_error( + finalization_error, + control_group, + control_group_size, + peer_message="another model-parallel rank failed while finalizing gradients", + ) self._grads_finalized = True grad_norm = self._finalized_grad_norm assert grad_norm is not None @@ -1201,15 +1193,12 @@ def optim_step( before_optimizer_step() except Exception as error: fence_error = error - if synchronize_update: - self._synchronize_accumulation_error( - fence_error, - control_group, - control_group_size, - peer_message="another model-parallel rank failed before the optimizer mutation fence", - ) - elif fence_error is not None: - raise fence_error + self._synchronize_accumulation_error( + fence_error, + control_group, + control_group_size, + peer_message="another model-parallel rank failed before the optimizer mutation fence", + ) mutation_started = True for optimizer in self.optimizers: @@ -2122,22 +2111,6 @@ def _validate_parallelism(self) -> None: if self.pipeline is not None and self.pipeline.scale_grads_in_schedule: raise ValueError("Engine requires AutoPipeline scale_grads_in_schedule=False") - def _gradient_reduction_multiplier(self, grad_group_size: int) -> int: - """Return the factor that compensates the backend gradient collective. - - Averaging backends divide by the complete DP-CP gradient group, so the - local loss numerator is multiplied by that group size before backward. - MegatronFSDP ``calculate_per_token_loss=True`` uses SUM collectives and - therefore needs no such compensation. - - Args: - grad_group_size: Size of the complete DP-CP gradient group. - - Returns: - One for summed gradients, otherwise ``grad_group_size``. - """ - return 1 if self._summed_gradient_reduction else grad_group_size - def _validate_execution_parallelism(self) -> None: """Validate model-parallel topology shared by forward and backward.""" if ( diff --git a/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py index 1f190d96a0..c08da7a55d 100644 --- a/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py +++ b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py @@ -170,9 +170,6 @@ def _run_mode( optimizers=optimizer, max_grad_norm=1e9, ) - expected_multiplier = 1 if summed_gradients else dist.get_world_size() - assert engine._gradient_reduction_multiplier(dist.get_world_size()) == expected_multiplier - window_a, window_b = _windows(dist.get_rank()) engine.begin_accumulation([window_a, window_b]) result_a = engine.forward_backward(window_a, _per_token_identity_loss) diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index b7518853d6..e109422fd6 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -371,9 +371,11 @@ class _FinalizationStage: def __init__(self, submod=None): self.events = [] + self.finalize_backward_states = [] self.submod = nn.Module() if submod is None else submod def backward_maybe_with_nosync(self, _backward_type, _bwd_kwargs, *, last_backward=False): + self.finalize_backward_states.append(getattr(self, "_nemo_finalize_backward", None)) self.events.append(("backward", last_backward)) return (), None @@ -535,6 +537,7 @@ def test_step_microbatches_defers_schedule_finalization_until_the_logical_last_c ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=False) assert stage.events == [("backward", False), ("scale", 2)] + assert stage.finalize_backward_states == [False] assert stage._nemo_finalize_backward is True ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=True) @@ -545,6 +548,7 @@ def test_step_microbatches_defers_schedule_finalization_until_the_logical_last_c ("reduce", 2), ("scale", 2), ] + assert stage.finalize_backward_states == [False, True] def test_step_microbatches_preserves_requested_per_microbatch_gradient_reduction(self): stage = _FinalizationStage() @@ -567,6 +571,7 @@ def test_step_microbatches_preserves_requested_per_microbatch_gradient_reduction ) assert stage.events == [("backward", True), ("reduce", 2), ("scale", 2)] + assert stage.finalize_backward_states == [True] def test_step_microbatches_rejects_cross_call_schedule_gradient_scaling(self): schedule = _KwargsChunkSchedule() diff --git a/tests/unit_tests/moe/test_fsdp_mixin.py b/tests/unit_tests/moe/test_fsdp_mixin.py index 0d32e9c039..d916c407df 100644 --- a/tests/unit_tests/moe/test_fsdp_mixin.py +++ b/tests/unit_tests/moe/test_fsdp_mixin.py @@ -768,13 +768,11 @@ def test_fsdp_module_last_backward(self, mock_fully_shard): assert param_groups is None @pytest.mark.parametrize( - ("stage_finalize", "reduce_grad_per_microbatch", "global_optim_step", "expect_post_backward"), + ("stage_finalize", "global_optim_step", "expect_post_backward"), [ - pytest.param(True, False, False, True, id="stage-final-overrides-global-false"), - pytest.param(False, False, True, False, id="stage-nonfinal-overrides-global-true"), - pytest.param(False, True, False, True, id="per-microbatch-reduce-overrides-global-false"), - pytest.param(False, True, True, True, id="per-microbatch-reduce-overrides-global-true"), - pytest.param(None, None, True, True, id="legacy-stage-falls-back-to-global"), + pytest.param(True, False, True, id="stage-final-overrides-global-false"), + pytest.param(False, True, False, id="stage-nonfinal-overrides-global-true"), + pytest.param(None, True, True, id="legacy-stage-falls-back-to-global"), ], ) @patch("nemo_automodel.components.moe.fsdp_mixin.get_is_optim_step") @@ -784,11 +782,10 @@ def test_moe_fsdp_mixin_pipeline_boundary_is_authoritative( mock_isinstance, mock_get_optim, stage_finalize, - reduce_grad_per_microbatch, global_optim_step, expect_post_backward, ): - """Stage finalization and per-microbatch reduction override the legacy global flag.""" + """Effective stage finalization overrides the legacy global flag.""" def isinstance_side_effect(obj, cls): if cls == MoEFSDPSyncMixin: @@ -805,8 +802,6 @@ def isinstance_side_effect(obj, cls): mock_stage = SimpleNamespace(submod=moe_model) if stage_finalize is not None: mock_stage._nemo_finalize_backward = stage_finalize - if reduce_grad_per_microbatch is not None: - mock_stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch bwd_kwargs = { "stage_output": Mock(), diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 489d5f2752..448a5a884a 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -2765,17 +2765,16 @@ def test_pipeline_parallelism_fails_before_forward(): @pytest.mark.parametrize( - ("declared_mode", "expected_multiplier", "expected_pre_collective_grad"), + ("declared_mode", "expected_pre_collective_grad"), [ - pytest.param(None, 4, 14.0, id="undeclared-averaged"), - pytest.param(False, 4, 14.0, id="explicit-averaged"), - pytest.param(True, 1, 3.5, id="summed"), + pytest.param(None, 14.0, id="undeclared-averaged"), + pytest.param(False, 14.0, id="explicit-averaged"), + pytest.param(True, 3.5, id="summed"), ], ) def test_gradient_reduction_mode_controls_main_loss_scale( monkeypatch, declared_mode, - expected_multiplier, expected_pre_collective_grad, ): model = ScaleModel() @@ -2786,21 +2785,22 @@ def test_gradient_reduction_mode_controls_main_loss_scale( result = engine.forward_backward([_datum([2, 4], [1.0, 3.0])], _identity_loss) - assert engine._gradient_reduction_multiplier(4) == expected_multiplier assert result.loss_sum.item() == pytest.approx(14.0) assert result.weight_sum.item() == pytest.approx(4.0) assert result.loss.item() == pytest.approx(3.5) assert model.weight.grad.item() == pytest.approx(expected_pre_collective_grad) -def test_gradient_reduction_mode_finds_summed_backend_below_ordinary_wrapper(): +def test_gradient_reduction_mode_finds_summed_backend_below_ordinary_wrapper(monkeypatch): model = ScaleModel() model.distributed_backend = nn.Identity() model.distributed_backend.calculate_per_token_loss = True - engine = Engine(model, device="cpu") + _configure_fake_gradient_group(engine, monkeypatch, group_size=8) + + engine.forward_backward([_datum([2, 4], [1.0, 3.0])], _identity_loss) - assert engine._gradient_reduction_multiplier(8) == 1 + assert model.weight.grad.item() == pytest.approx(3.5) def test_summed_gradient_mode_planned_accumulation_matches_one_window(monkeypatch): From f864aadbe57e2c427f4b193264f1988fa24b6e8a Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Fri, 21 Aug 2026 23:18:42 -0700 Subject: [PATCH 22/34] feat(engine): add batch contexts and simplify pipeline execution Expose CP-prepared loss side channels to eager batch contexts for model-scoped router replay. Remove the unused whole-batch AutoPipeline step API and planned multi-call accumulation, making one forward_backward call the complete optimizer window while retaining single-window lifecycle safety. Signed-off-by: HuiyingLi --- nemo_automodel/components/datasets/datum.py | 124 +- .../distributed/context_parallel/sharder.py | 19 +- .../distributed/pipelining/autopipeline.py | 139 -- .../distributed/pipelining/functional.py | 115 +- nemo_automodel/components/moe/fsdp_mixin.py | 6 +- .../components/moe/router_replay.py | 344 +++- nemo_automodel/engine/__init__.py | 726 ++------ nemo_automodel/shared/model_utils.py | 20 +- .../context_parallel/run_packed_pp.py | 123 -- .../run_megatron_fsdp_per_token_loss.py | 27 +- .../parallelism/run_pp_grad_accum_parity.py | 236 +-- tests/unit_tests/datasets/test_datum.py | 149 ++ .../pipelining/test_autopipeline.py | 312 +--- .../unit_tests/distributed/test_cp_sharder.py | 48 + tests/unit_tests/moe/test_fsdp_mixin.py | 39 +- tests/unit_tests/moe/test_router_replay.py | 368 ++++ tests/unit_tests/test_engine.py | 1580 +++++++++-------- 17 files changed, 2058 insertions(+), 2317 deletions(-) diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index d701253a6f..7fb8bff6ca 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -39,8 +39,9 @@ class LossInputLayout(str, Enum): """How one loss input relates to the Datums being collated. - ``PER_TOKEN`` values follow token padding, packing, CP sharding, and PP - microbatching. ``PER_DATUM`` values contain one scalar for each outer + ``PER_TOKEN`` values have a leading token axis and follow token padding, + packing, CP sharding, and PP microbatching; trailing feature axes are + preserved. ``PER_DATUM`` values contain one scalar for each outer :class:`Datum`; CP ranks receive the same scalars, so a loss callback uses them with its CP-local token contribution rather than returning a repeated full-sequence scalar. ``REPLICATED`` values are batch-level metadata copied @@ -61,8 +62,10 @@ class CollatedLossInputs(dict[str, torch.Tensor]): the current Engine accepts the identity mapping (one item per Datum, in input order). It is ``None`` when a collater cannot expose its inner boundaries, as with an already-prebatched batch. ``copy()`` preserves the - side channel; converting this object to a plain ``dict`` intentionally - drops it and opts back into the Engine's conservative legacy inference. + side channel. ``pad_values`` supplies nonzero sentinels for explicitly + padded ``PER_TOKEN`` fields. Converting this object to a plain ``dict`` + intentionally drops it and opts back into the Engine's conservative legacy + inference. """ def __init__( @@ -71,7 +74,19 @@ def __init__( *, layouts: Mapping[str, LossInputLayout], item_to_datum: tuple[int, ...] | None, + pad_values: Mapping[str, float | int | bool] | None = None, ) -> None: + """Attach semantic layout metadata to collated loss tensors. + + Args: + values: Collated tensors. ``PER_TOKEN`` values have shape + ``[batch, sequence, ...]`` or ``[tokens, ...]``; + ``PER_DATUM`` values have leading Datum axis. + layouts: Complete semantic layout mapping for ``values``. + item_to_datum: Optional mapping from collated rows/sequences to + outer Datum indices. + pad_values: Optional scalar fills for padded ``PER_TOKEN`` fields. + """ super().__init__(values) if set(layouts) != set(self): raise ValueError("layouts must contain exactly the CollatedLossInputs keys") @@ -80,10 +95,20 @@ def __init__( resolved_item_to_datum = None if item_to_datum is None else tuple(item_to_datum) if resolved_item_to_datum is not None and not all(isinstance(index, int) for index in resolved_item_to_datum): raise TypeError("item_to_datum must contain integer Datum indices") + resolved_pad_values = dict(pad_values or {}) + unknown_pad_values = set(resolved_pad_values) - set(self) + if unknown_pad_values: + raise ValueError(f"pad_values contains unknown loss inputs: {sorted(unknown_pad_values)}") + if not all(isinstance(value, (bool, int, float)) for value in resolved_pad_values.values()): + raise TypeError("every loss input pad value must be a bool, int, or float") + non_token_pad_values = {name for name in resolved_pad_values if layouts[name] is not LossInputLayout.PER_TOKEN} + if non_token_pad_values: + raise ValueError(f"pad_values are only valid for PER_TOKEN loss inputs: {sorted(non_token_pad_values)}") # Store a normal dict so the public collate result remains pickleable # across DataLoader worker boundaries. Expose only a read-only view. self._layouts = dict(layouts) + self._pad_values = resolved_pad_values self.item_to_datum = resolved_item_to_datum @property @@ -91,9 +116,19 @@ def layouts(self) -> Mapping[str, LossInputLayout]: """Complete, read-only field-layout mapping.""" return MappingProxyType(self._layouts) + @property + def pad_values(self) -> Mapping[str, float | int | bool]: + """Read-only explicit fill values for padded ``PER_TOKEN`` fields.""" + return MappingProxyType(self._pad_values) + def copy(self) -> CollatedLossInputs: """Return a shallow copy that retains the layout side channel.""" - return type(self)(self, layouts=self.layouts, item_to_datum=self.item_to_datum) + return type(self)( + self, + layouts=self.layouts, + item_to_datum=self.item_to_datum, + pad_values=self.pad_values, + ) def __copy__(self) -> CollatedLossInputs: return self.copy() @@ -118,6 +153,9 @@ class Datum: loss_fn_input_layouts: Optional semantic layouts for loss fields. The canonical collater infers omitted fields using its legacy token-aligned-versus-scalar rules. + loss_fn_input_pad_values: Optional nonzero padding sentinels for + ``PER_TOKEN`` loss/side-channel fields. For example, routing replay + uses ``-1`` to mean that a padded token keeps its live route. input_ids: Deprecated convenience spelling for the old text-only API. It cannot be combined with ``model_inputs``. """ @@ -125,6 +163,7 @@ class Datum: model_inputs: dict[str, Any] loss_fn_inputs: dict[str, torch.Tensor] = field(default_factory=dict) loss_fn_input_layouts: dict[str, LossInputLayout] = field(default_factory=dict) + loss_fn_input_pad_values: dict[str, float | int | bool] = field(default_factory=dict) def __init__( self, @@ -133,7 +172,20 @@ def __init__( *, input_ids: torch.Tensor | list[int] | None = None, loss_fn_input_layouts: Mapping[str, LossInputLayout] | None = None, + loss_fn_input_pad_values: Mapping[str, float | int | bool] | None = None, ) -> None: + """Initialize one processor-ready item and its loss side channels. + + Args: + model_inputs: Model keyword mapping, or legacy 1-D ``input_ids``. + Tensor layouts are model-specific. + loss_fn_inputs: Loss/algorithm tensors before collation. A + ``PER_TOKEN`` tensor has shape ``[sequence, ...]``. + input_ids: Legacy 1-D ``[sequence]`` token ids. + loss_fn_input_layouts: Optional semantic layout per loss field. + loss_fn_input_pad_values: Optional scalar padding sentinel per + ``PER_TOKEN`` loss field. + """ # Preserve the old positional ``Datum(input_ids, loss_fn_inputs)`` form # while downstream users move to the model-ready mapping. if model_inputs is not None and not isinstance(model_inputs, dict): @@ -151,6 +203,7 @@ def __init__( self.model_inputs = dict(model_inputs) self.loss_fn_inputs = dict(loss_fn_inputs or {}) self.loss_fn_input_layouts = dict(loss_fn_input_layouts or {}) + self.loss_fn_input_pad_values = dict(loss_fn_input_pad_values or {}) self.__post_init__() def __post_init__(self) -> None: @@ -174,6 +227,11 @@ def __post_init__(self) -> None: raise ValueError(f"loss_fn_input_layouts contains unknown loss inputs: {sorted(unknown_layouts)}") if not all(isinstance(layout, LossInputLayout) for layout in self.loss_fn_input_layouts.values()): raise TypeError("every loss_fn_input_layouts value must be a LossInputLayout") + unknown_pad_values = set(self.loss_fn_input_pad_values) - set(self.loss_fn_inputs) + if unknown_pad_values: + raise ValueError(f"loss_fn_input_pad_values contains unknown loss inputs: {sorted(unknown_pad_values)}") + if not all(isinstance(value, (bool, int, float)) for value in self.loss_fn_input_pad_values.values()): + raise TypeError("every loss_fn_input_pad_values value must be a bool, int, or float") @property def input_ids(self) -> torch.Tensor: @@ -254,9 +312,9 @@ def collate_datums( Returns: ``(model_inputs, loss_fn_inputs)``. The second item remains a ``dict`` and also exposes complete ``layouts`` and ``item_to_datum`` metadata. - Per-token loss inputs have shape ``[B, T]`` in padded mode and - ``[1, total_tokens]`` in packed mode; per-Datum scalar loss inputs have - shape ``[B]``. Replicated inputs retain one copy of their original + Per-token loss inputs have shape ``[B, T, ...]`` in padded mode and + ``[1, total_tokens, ...]`` in packed mode; per-Datum scalar loss inputs + have shape ``[B]``. Replicated inputs retain one copy of their original shape. """ if not datums: @@ -295,6 +353,7 @@ def collate_datums( loss_inputs: dict[str, torch.Tensor] = {} loss_layouts: dict[str, LossInputLayout] = {} + loss_pad_values: dict[str, float | int | bool] = {} for key in sorted(loss_keys): values = [datum.loss_fn_inputs[key] for datum in datums] declared_layouts = [datum.loss_fn_input_layouts[key] for datum in datums if key in datum.loss_fn_input_layouts] @@ -304,12 +363,21 @@ def collate_datums( if len(explicit_layouts) > 1: raise ValueError(f"every Datum must use the same explicit layout for loss input {key!r}") explicit_layout = next(iter(explicit_layouts), None) - - token_aligned = [value.ndim == 1 and value.shape[0] == datum.seq_len for value, datum in zip(values, datums)] + declared_pad_values = [ + datum.loss_fn_input_pad_values[key] for datum in datums if key in datum.loss_fn_input_pad_values + ] + if declared_pad_values and len(declared_pad_values) != len(datums): + raise ValueError(f"every Datum must declare the pad value for loss input {key!r}, or none may declare it") + if len(set(declared_pad_values)) > 1: + raise ValueError(f"every Datum must use the same pad value for loss input {key!r}") + + token_aligned = [value.ndim >= 1 and value.shape[0] == datum.seq_len for value, datum in zip(values, datums)] + legacy_token_aligned = [value.ndim == 1 and aligned for value, aligned in zip(values, token_aligned)] if explicit_layout is LossInputLayout.PER_TOKEN and not all(token_aligned): shapes = [tuple(value.shape) for value in values] raise ValueError( - f"PER_TOKEN loss input {key!r} must be 1-D and match each Datum's token length; got {shapes}" + f"PER_TOKEN loss input {key!r} must have a leading axis matching each Datum's token length; " + f"got {shapes}" ) scalar_per_datum = [value.numel() == 1 for value in values] @@ -319,24 +387,49 @@ def collate_datums( layout = explicit_layout if layout is None: - if all(token_aligned): + if all(legacy_token_aligned): layout = LossInputLayout.PER_TOKEN elif all(scalar_per_datum): layout = LossInputLayout.PER_DATUM else: shapes = [tuple(value.shape) for value in values] raise ValueError( - f"the default collater only supports scalar or 1-D token-aligned loss inputs; {key!r} has {shapes}" + "the default collater infers only scalar or 1-D token-aligned loss inputs; " + f"declare an explicit layout for {key!r} with shapes {shapes}" ) loss_layouts[key] = layout if layout is LossInputLayout.PER_TOKEN: + pad_value = declared_pad_values[0] if declared_pad_values else 0 + if declared_pad_values: + loss_pad_values[key] = pad_value + first = values[0] + if not all( + value.shape[1:] == first.shape[1:] and value.dtype == first.dtype and value.device == first.device + for value in values[1:] + ): + shapes = [tuple(value.shape) for value in values] + raise ValueError( + f"PER_TOKEN loss input {key!r} must use one trailing shape, dtype, and device; got {shapes}" + ) if packed: - loss_inputs[key] = torch.cat(values).unsqueeze(0) + loss_inputs[key] = torch.cat(values, dim=0).unsqueeze(0) else: - loss_inputs[key] = torch.stack([F.pad(value, (0, width - value.shape[0])) for value in values]) + loss_inputs[key] = torch.stack( + [ + F.pad( + value, + (*([0, 0] * (value.ndim - 1)), 0, width - value.shape[0]), + value=pad_value, + ) + for value in values + ] + ) continue + if declared_pad_values: + raise ValueError(f"loss input {key!r} declares a pad value but does not use the PER_TOKEN layout") + if layout is LossInputLayout.PER_DATUM: loss_inputs[key] = torch.stack([value.reshape(()) for value in values]) continue @@ -356,4 +449,5 @@ def collate_datums( loss_inputs, layouts=loss_layouts, item_to_datum=tuple(range(len(datums))), + pad_values=loss_pad_values, ) diff --git a/nemo_automodel/components/distributed/context_parallel/sharder.py b/nemo_automodel/components/distributed/context_parallel/sharder.py index 61929f303d..71d745171c 100644 --- a/nemo_automodel/components/distributed/context_parallel/sharder.py +++ b/nemo_automodel/components/distributed/context_parallel/sharder.py @@ -408,9 +408,10 @@ def shard_token_tensor( caller may pass tensors in its own coordinates and the verb applies the same transform the batch went through: - - ``[B, S_in]`` tensors on a repositioned-row layout (reported position - map, e.g. DSV4 packed repad) are scattered into the padded rows, - ``fill`` filling the pad slots; + - ``[B, S_in, ...]`` tensors on a repositioned-row layout (reported + position map, e.g. DSV4 packed repad) are scattered into the padded + rows, ``fill`` filling the pad slots while trailing feature axes are + preserved; - tensors matching the reported pre-flatten ``input_row_shape`` on a flat-stream (THD) layout are flattened first (the returned shard is in the model's local stream coordinate); @@ -421,15 +422,19 @@ def shard_token_tensor( Any other length raises instead of silently sharding the wrong slice. """ layout = self.shard_layout or _NO_SHARD_LAYOUT - if layout.input_token_stream_positions is not None and tuple(tensor.shape) == tuple( - layout.input_token_stream_positions.shape - ): + position_shape = ( + tuple(layout.input_token_stream_positions.shape) if layout.input_token_stream_positions is not None else () + ) + if position_shape and tuple(tensor.shape[: len(position_shape)]) == position_shape: if fill is None: raise ValueError("sharding an input-coordinate tensor on a repositioned layout requires `fill`") positions = layout.input_token_stream_positions.to(tensor.device) valid = positions >= 0 padded = torch.full( - (tensor.shape[0], layout.padded_seq_len), fill, dtype=tensor.dtype, device=tensor.device + (tensor.shape[0], layout.padded_seq_len, *tensor.shape[len(position_shape) :]), + fill, + dtype=tensor.dtype, + device=tensor.device, ) padded[valid.nonzero(as_tuple=True)[0], positions[valid]] = tensor[valid] tensor, seq_dim = padded, 1 diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index a12e098354..00a3802307 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -20,11 +20,8 @@ import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.pipelining.microbatch import BlockMask, TensorChunkSpec -from torch.distributed.pipelining.microbatch import _Replicate as ReplicateChunkSpec from torch.distributed.pipelining.schedules import _PipelineSchedule from torch.distributed.pipelining.stage import PipelineStage -from torch.utils._pytree import tree_map from nemo_automodel.components.distributed.pipelining.functional import ( ParallelizeFnProtocol, @@ -36,7 +33,6 @@ ) logger = logging.getLogger(__name__) -_MISSING_STAGE_STATE = object() @dataclass @@ -256,101 +252,6 @@ def update_seq_len( effective_microbatch_size, ) - def _get_schedule_kwargs_chunk_spec(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Build pipeline microbatch chunking metadata for keyword inputs. - - PyTorch's default schedule chunking splits every tensor kwarg on dim 0. - Most AutoModel batch tensors are batch-major and should keep that - default, but some model-owned input layouts place batch on another axis. - The canonical local model part can declare those exceptions by implementing - ``get_pipeline_kwargs_chunk_dims(kwargs) -> dict[str, int]``. - - Args: - kwargs: Mapping passed to the pipeline schedule. Tensor values may - have arbitrary model-defined layouts; the model hook identifies - any nonstandard batch axis. - - Returns: - A chunk-spec mapping with the same nested structure as ``kwargs``, - or ``None`` when PyTorch's default chunking applies. - """ - model_parts = self._info.model_parts - if not model_parts: - raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") - - hook = getattr(model_parts[0], "get_pipeline_kwargs_chunk_dims", None) - if hook is None: - return None - - custom_chunk_dims = hook(kwargs) or {} - for key in custom_chunk_dims: - if key not in kwargs: - raise ValueError(f"Model PP chunk hook returned unknown kwarg: {key}") - - if not custom_chunk_dims: - return None - - def default_spec(value): - if isinstance(value, (torch.Tensor, BlockMask)): - return TensorChunkSpec(0) - return ReplicateChunkSpec() - - kwargs_chunk_spec = tree_map(default_spec, kwargs, is_leaf=lambda value: isinstance(value, BlockMask)) - for key, split_dim in custom_chunk_dims.items(): - kwargs_chunk_spec[key] = TensorChunkSpec(split_dim) - return kwargs_chunk_spec - - def step( - self, - model_input: torch.Tensor, - *, - target: torch.Tensor | None = None, - losses: list[torch.Tensor] | None = None, - **kwargs: Any, - ) -> Any: - """Run one pipeline schedule step with model-owned input chunking. - - Args: - model_input: Tensor of shape [batch, ...] containing the first - pipeline stage's input. Ignored on ranks without the first stage. - target: Tensor with a model-defined target layout, or ``None`` on - ranks without the last pipeline stage. - losses: Mutable list populated with scalar loss tensors, or ``None`` - on ranks without the last pipeline stage. - **kwargs: Keyword schedule inputs. Tensor values may have arbitrary - model-defined layouts; model-owned metadata identifies any - nonstandard batch axis. - - Returns: - The value returned by the underlying PyTorch pipeline schedule. - """ - schedule = self._info.schedule - if schedule is None: - raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") - - schedule_args = (model_input,) if self._info.has_first_stage else () - kwargs_chunk_spec = self._get_schedule_kwargs_chunk_spec(kwargs) - - if kwargs_chunk_spec is None: - return schedule.step( - *schedule_args, - target=target, - losses=losses, - **kwargs, - ) - - previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec - schedule._kwargs_chunk_spec = kwargs_chunk_spec - try: - return schedule.step( - *schedule_args, - target=target, - losses=losses, - **kwargs, - ) - finally: - schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec - def step_microbatches( self, model_inputs: list[dict[str, Any]], @@ -358,7 +259,6 @@ def step_microbatches( loss_fn: Callable[[Any, int], Any], losses: list[torch.Tensor] | None = None, return_outputs: bool = False, - finalize_backward: bool = True, ) -> Any: """Run a schedule step over already prepared model microbatches. @@ -376,41 +276,16 @@ def step_microbatches( losses: Mutable list populated by the schedule on the last stage. return_outputs: Whether the last stage returns merged model outputs when supported by the installed PyTorch version. - finalize_backward: Whether this schedule invocation ends the - optimizer's backward window. ``False`` keeps DDP in no-sync - mode and completes local FSDP post-backward state while - deferring its gradient collective; per-microbatch gradient - reduction remains authoritative when configured. Returns: The value returned by the underlying PyTorch pipeline schedule. - - Raises: - ValueError: If a non-final call would use schedule-local gradient - scaling and rescale gradients accumulated by earlier calls. - RuntimeError: If a non-final call is attempted before - accumulation-aware pipeline stages have been built. """ - if not finalize_backward and self.scale_grads_in_schedule: - raise ValueError( - "planned pipeline accumulation requires scale_grads_in_schedule=False; " - "schedule-local scaling would rescale gradients from earlier calls" - ) - if not finalize_backward and ( - not self._info.stages - or any(not vars(stage).get("_nemo_accumulation_aware", False) for stage in self._info.stages) - ): - raise RuntimeError( - "finalize_backward=False requires accumulation-aware pipeline stages; " - "build the AutoPipeline before running planned accumulation" - ) return self._run_prepared_microbatches( model_inputs, loss_fn=loss_fn, losses=losses, return_outputs=return_outputs, schedule_method="step", - finalize_backward=finalize_backward, ) def eval_microbatches( @@ -452,7 +327,6 @@ def eval_microbatches( losses=losses, return_outputs=return_outputs, schedule_method="eval", - finalize_backward=True, ) def _run_prepared_microbatches( @@ -463,7 +337,6 @@ def _run_prepared_microbatches( losses: list[torch.Tensor] | None, return_outputs: bool, schedule_method: Literal["step", "eval"], - finalize_backward: bool, ) -> Any: """Run one schedule method with an exact prepared-microbatch split.""" schedule = self._info.schedule @@ -509,13 +382,6 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: ) previous_split_inputs = schedule._split_inputs previous_loss_fn = schedule._loss_fn - previous_stage_state: list[tuple[PipelineStage, object]] = [] - if schedule_method == "step": - for stage in self._info.stages or (): - previous_stage_state.append((stage, vars(stage).get("_nemo_finalize_backward", _MISSING_STAGE_STATE))) - stage._nemo_finalize_backward = finalize_backward or getattr( - stage, "_reduce_grad_per_microbatch", False - ) schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) schedule._loss_fn = indexed_loss try: @@ -527,11 +393,6 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: finally: schedule._loss_fn = previous_loss_fn schedule._split_inputs = previous_split_inputs - for stage, previous in previous_stage_state: - if previous is _MISSING_STAGE_STATE: - del stage._nemo_finalize_backward - else: - stage._nemo_finalize_backward = previous @property def parts(self) -> list[nn.Module]: diff --git a/nemo_automodel/components/distributed/pipelining/functional.py b/nemo_automodel/components/distributed/pipelining/functional.py index cb0557a602..5e0ac1cbef 100644 --- a/nemo_automodel/components/distributed/pipelining/functional.py +++ b/nemo_automodel/components/distributed/pipelining/functional.py @@ -20,11 +20,10 @@ import os import time import types -from typing import Any, Callable, Protocol +from typing import Callable, Protocol import torch import torch.nn as nn -from torch.distributed import fsdp as torch_fsdp from torch.distributed.device_mesh import DeviceMesh from torch.distributed.pipelining import PipelineStage from torch.distributed.pipelining.schedules import ( @@ -43,20 +42,9 @@ model_keeps_self_forward, patch_hf_model_for_pp, ) -from nemo_automodel.shared.import_utils import safe_import_from logger = logging.getLogger(__name__) -_HAS_REPLICATE_MODULE, ReplicateModule = safe_import_from( - "torch.distributed._composable.replicate_with_fsdp", - "ReplicateModule", -) -_HAS_REPLICATE_STATE, replicate = safe_import_from( - "torch.distributed._composable.replicate_with_fsdp", - "replicate", -) -_HAS_REPLICATE_WITH_FSDP = _HAS_REPLICATE_MODULE and _HAS_REPLICATE_STATE - def _get_optional_hook(module: object, name: str) -> Callable | None: try: @@ -503,97 +491,6 @@ def _cleanup_preserving_grads(self, *, _cleanup=cleanup) -> None: stage._post_metadata_inference_cleanup = types.MethodType(_cleanup_preserving_grads, stage) -def _accumulation_aware_backward_maybe_with_nosync( - self: PipelineStage, - backward_type: str, - bwd_kwargs: dict[str, Any], - last_backward: bool = False, -) -> tuple[tuple[torch.Tensor | None, ...], list[dict[str, Any]] | None]: - """Keep a schedule-local last backward open across planned Engine windows. - - A configured per-microbatch reduction remains authoritative. Otherwise a - non-final Engine window must not let PyTorch's schedule-local - ``last_backward`` trigger DDP/FSDP gradient synchronization. - - Args: - backward_type: PyTorch stage operation: ``full``, ``input``, or - ``weight`` backward. - bwd_kwargs: Schedule-owned state for one local PP microbatch. ``full`` - and ``input`` carry ``stage_output``, ``output_grads``, and - ``input_values`` tensor trees; ``weight`` carries ``stage_output`` - and the split-backward ``param_groups``. All tensor shapes, dtypes, - devices, and layouts are model- and stage-defined. This wrapper - forwards them unchanged and performs no redistribution. - last_backward: Whether the PyTorch schedule considers this its final - local backward operation. - - Returns: - The original stage backward result unchanged: stage-input gradient - tensors in their model-defined local layouts, plus optional - split-backward parameter-group records. - """ - return self._nemo_original_backward_maybe_with_nosync( - backward_type, - bwd_kwargs, - last_backward=last_backward and self._nemo_finalize_backward, - ) - - -def _accumulation_aware_perform_reduce_grad(self: PipelineStage, grad_scale_factor: int) -> None: - """Complete a PP schedule while deferring its final gradient collective. - - FSDP post-backward is both a communication boundary and required local - lifecycle cleanup. On a non-final planned window, run that cleanup with - gradient synchronization disabled so accumulated unsharded gradients are - preserved for the final window. Schedule-owned gradient scaling still runs - once per schedule. The original implementation is used unchanged for - ordinary calls, final planned windows, and per-microbatch reduction mode. - """ - if self._nemo_finalize_backward: - self._nemo_original_perform_reduce_grad(grad_scale_factor) - return - - if isinstance(self.submod, torch_fsdp.FSDPModule): - fsdp_module = self.submod - fsdp_module.set_is_last_backward(True) - fsdp_module.set_reshard_after_backward(True) - fsdp_module.set_requires_gradient_sync(False) - fsdp_state = ( - replicate.state(fsdp_module) - if _HAS_REPLICATE_WITH_FSDP and isinstance(fsdp_module, ReplicateModule) - else torch_fsdp.fully_shard.state(fsdp_module) # type: ignore[attr-defined] - ) - for state in fsdp_state._state_ctx.all_states: - if state._fsdp_param_group: - state._fsdp_param_group.post_backward() - fsdp_state._root_post_backward_final_callback() - - if grad_scale_factor != 1: - self.scale_grads(grad_scale_factor) - - -def _make_pipeline_stages_accumulation_aware( - stages: list[PipelineStage], - *, - reduce_grad_per_microbatch: bool, -) -> None: - """Install behavior-neutral stage gates used by planned PP accumulation.""" - for stage in stages: - stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch - stage._nemo_finalize_backward = True - if vars(stage).get("_nemo_accumulation_aware", False): - continue - - stage._nemo_original_backward_maybe_with_nosync = stage.backward_maybe_with_nosync - stage.backward_maybe_with_nosync = types.MethodType(_accumulation_aware_backward_maybe_with_nosync, stage) - - perform_reduce_grad = getattr(stage, "perform_reduce_grad", None) - if callable(perform_reduce_grad): - stage._nemo_original_perform_reduce_grad = perform_reduce_grad - stage.perform_reduce_grad = types.MethodType(_accumulation_aware_perform_reduce_grad, stage) - stage._nemo_accumulation_aware = True - - def reset_pp_stage_shapes( schedule: _PipelineSchedule, stages: list[PipelineStage], @@ -1106,21 +1003,13 @@ def pipeline_model( for stage in stages: stage.backward_maybe_with_nosync = types.MethodType(patched_backward_maybe_with_nosync, stage) + stage._reduce_grad_per_microbatch = reduce_grad_per_microbatch logger.info( "Patched pipeline stages with backward_maybe_with_nosync " f"(reduce_grad_per_microbatch={reduce_grad_per_microbatch})" ) - # PyTorch considers every schedule invocation a complete optimizer window: - # its last backward and REDUCE_GRAD action finalize DP/FSDP gradients. Keep - # those exact defaults, but make the boundary controllable when Engine has - # predeclared one optimizer window spanning multiple schedule invocations. - _make_pipeline_stages_accumulation_aware( - stages, - reduce_grad_per_microbatch=reduce_grad_per_microbatch, - ) - # Determine if this rank has first/last stage has_first_stage = False has_last_stage = False diff --git a/nemo_automodel/components/moe/fsdp_mixin.py b/nemo_automodel/components/moe/fsdp_mixin.py index 9725e6bf84..78d507554f 100644 --- a/nemo_automodel/components/moe/fsdp_mixin.py +++ b/nemo_automodel/components/moe/fsdp_mixin.py @@ -303,11 +303,7 @@ def perform_backward( elif isinstance(self.submod, MoEFSDPSyncMixin): _disable_fsdp_for_moe_module(self.submod) result = perform_backward(backward_type)() - if hasattr(self, "_nemo_finalize_backward"): - finalize_backward = self._nemo_finalize_backward - else: - finalize_backward = get_is_optim_step() - if last_backward and finalize_backward: + if last_backward and get_is_optim_step(): _run_post_backward_for_moe_module(self.submod) else: # Non-DP submodule, regular backward diff --git a/nemo_automodel/components/moe/router_replay.py b/nemo_automodel/components/moe/router_replay.py index 1f3492fb33..445290fb14 100644 --- a/nemo_automodel/components/moe/router_replay.py +++ b/nemo_automodel/components/moe/router_replay.py @@ -49,15 +49,27 @@ position. This assumes single-threaded model construction (the norm for recipe training); call :meth:`RouterReplay.clear_registry` before building a second model in the same process. + +For rollout-provided routing, :class:`RouterReplayAdapter` is the preferred +eager Engine interface. It maps global decoder-layer ids without using the +registry and consumes ``routed_experts`` only after the Engine has applied its +packing and context-parallel token transform. This adapter intentionally +supports PP=1 only. """ -from contextlib import contextmanager +from collections.abc import Iterator, Mapping +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass from enum import Enum -from typing import Iterator, List +from math import prod +from typing import Any import torch +from torch import nn + +from nemo_automodel.shared.model_utils import iter_transformer_blocks -__all__ = ["RouterReplayMode", "RouterReplay", "replay_selection"] +__all__ = ["RouterReplayMode", "RouterReplay", "RouterReplayAdapter", "replay_selection"] class RouterReplayMode(Enum): @@ -76,14 +88,16 @@ class RouterReplay: ``replay`` context managers). """ - _registry: List["RouterReplay"] = [] + _registry: list["RouterReplay"] = [] - def __init__(self) -> None: - """Create a handle and register it in construction (i.e. layer) order.""" + def __init__(self, *, register: bool = True) -> None: + """Create a handle, optionally registering it for legacy global control.""" self.mode: RouterReplayMode | None = None self.recorded_indices: torch.Tensor | None = None self.target_indices: torch.Tensor | None = None - RouterReplay._registry.append(self) + self._allow_trailing_live_tokens = False + if register: + RouterReplay._registry.append(self) def apply(self, indices: torch.Tensor) -> torch.Tensor: """Record or replay ``indices`` according to the current mode. @@ -95,6 +109,8 @@ def apply(self, indices: torch.Tensor) -> torch.Tensor: Returns: ``indices`` unchanged when no mode is active or while recording; the stored target indices (moved to ``indices.device``) while replaying. + A target row containing ``-1`` keeps that token's complete live + top-k selection, preserving unique expert ids. """ if self.mode == RouterReplayMode.RECORD: # Indices are integer selection ids carrying no gradient; detach so the @@ -107,13 +123,28 @@ def apply(self, indices: torch.Tensor) -> torch.Tensor: "RouterReplay is in REPLAY mode but no target indices were set for this layer. " "Call RouterReplay.replay(indices) / set_replay_indices(...) with one tensor per MoE layer." ) - target = self.target_indices.to(indices.device) - if target.shape != indices.shape: - raise ValueError( - f"Replay indices shape {tuple(target.shape)} does not match the current " - f"selection shape {tuple(indices.shape)}; replay must run on the same tokens and topk." + if self.target_indices.dtype not in {torch.int8, torch.int16, torch.int32, torch.int64}: + raise TypeError( + f"RouterReplay target indices must use a signed integer dtype, got {self.target_indices.dtype}" ) - return target + target = self.target_indices.to(device=indices.device, dtype=indices.dtype) + if target.shape != indices.shape: + if ( + self._allow_trailing_live_tokens + and target.ndim == 2 + and indices.ndim == 2 + and target.shape[1] == indices.shape[1] + and target.shape[0] < indices.shape[0] + ): + trailing = target.new_full((indices.shape[0] - target.shape[0], target.shape[1]), -1) + target = torch.cat((target, trailing), dim=0) + else: + raise ValueError( + f"Replay indices shape {tuple(target.shape)} does not match the current " + f"selection shape {tuple(indices.shape)}; replay must run on the same tokens and topk." + ) + keep_live = (target == -1).any(dim=-1, keepdim=True) + return torch.where(keep_live, indices, target) return indices # -- per-instance state ------------------------------------------------- @@ -130,7 +161,7 @@ def clear(self) -> None: # -- global control over every registered instance --------------------- @staticmethod - def instances() -> List["RouterReplay"]: + def instances() -> list["RouterReplay"]: """Return the registered instances in construction (layer) order.""" return RouterReplay._registry @@ -141,7 +172,7 @@ def set_mode(mode: RouterReplayMode | None) -> None: inst.mode = mode @staticmethod - def set_replay_indices(all_layers_indices: List[torch.Tensor]) -> None: + def set_replay_indices(all_layers_indices: list[torch.Tensor]) -> None: """Distribute one selection tensor per layer to the registered instances. Args: @@ -162,14 +193,14 @@ def set_replay_indices(all_layers_indices: List[torch.Tensor]) -> None: inst.set_target(indices) @staticmethod - def collect() -> List[torch.Tensor]: + def collect() -> list[torch.Tensor]: """Collect the recorded selection from every registered instance, in layer order. Raises: RuntimeError: If any instance has no recorded selection (i.e. a forward pass was not run under :meth:`record`). """ - collected: List[torch.Tensor] = [] + collected: list[torch.Tensor] = [] for layer_idx, inst in enumerate(RouterReplay._registry): if inst.recorded_indices is None: raise RuntimeError( @@ -204,7 +235,7 @@ def record(cls) -> Iterator[None]: @classmethod @contextmanager - def replay(cls, all_layers_indices: List[torch.Tensor]) -> Iterator[None]: + def replay(cls, all_layers_indices: list[torch.Tensor]) -> Iterator[None]: """Replay ``all_layers_indices`` (one tensor per layer) for the duration of the block. Target selections are cleared on exit so a stale replay never leaks into a @@ -220,6 +251,283 @@ def replay(cls, all_layers_indices: List[torch.Tensor]) -> Iterator[None]: inst.target_indices = None +@dataclass(frozen=True) +class _RouterReplayBinding: + """One decoder layer's model-scoped replay handle.""" + + layer_idx: int + replay: RouterReplay + topk: int + num_experts: int | None + + +class RouterReplayAdapter: + """Bind rollout routes to one model's MoE gates for eager Engine execution. + + The adapter is both the model-aware route formatter and the callable passed + as ``Engine(batch_context_fn=...)``. It deliberately ignores the legacy + process-global registry: decoder-layer ids determine the mapping, so sparse + hybrid MoE stacks and multiple models in one process remain unambiguous. + Do not nest the legacy process-global ``record``/``replay`` contexts around + an active adapter context when the same gate handles are registered. + AutoPipeline is not supported by this adapter; the Engine rejects that + combination before execution. + + Args: + model: Complete PP=1 model. Its primary decoder blocks must expose + numeric child ids or a consistent integer ``layer_idx``. A block + may contain at most one module with a ``router_replay`` slot. + """ + + field_name = "routed_experts" + + def __init__(self, model: nn.Module) -> None: + block_root = model + visited_roots: set[int] = set() + blocks: tuple[tuple[nn.Module, str, nn.Module], ...] = () + while id(block_root) not in visited_roots: + visited_roots.add(id(block_root)) + blocks = tuple(iter_transformer_blocks(block_root)) + if blocks: + break + wrapped = getattr(block_root, "module", None) + if not isinstance(wrapped, nn.Module): + break + block_root = wrapped + + bindings: list[_RouterReplayBinding] = [] + seen_replays: set[int] = set() + for _parent, child_name, block in blocks: + slots = [module for module in block.modules() if hasattr(module, "router_replay")] + if not slots: + continue + if len(slots) != 1: + raise ValueError( + f"decoder block {child_name!r} has {len(slots)} router_replay slots; " + "RouterReplayAdapter requires one gate per routed layer" + ) + + declared_ids = { + layer_idx + for module in block.modules() + if isinstance((layer_idx := getattr(module, "layer_idx", None)), int) + and not isinstance(layer_idx, bool) + } + if len(declared_ids) > 1: + raise ValueError( + f"decoder block {child_name!r} contains conflicting layer_idx values {sorted(declared_ids)}" + ) + child_idx = int(child_name) if child_name.isdecimal() else None + declared_idx = next(iter(declared_ids), None) + if child_idx is not None and declared_idx is not None and child_idx != declared_idx: + raise ValueError(f"decoder block key {child_idx} disagrees with its layer_idx {declared_idx}") + layer_idx = declared_idx if declared_idx is not None else child_idx + if layer_idx is None: + raise ValueError(f"cannot resolve the global layer id for routed decoder block {child_name!r}") + if layer_idx < 0: + raise ValueError(f"routed decoder block {child_name!r} has negative layer_idx {layer_idx}") + + gate = slots[0] + topk = getattr(gate, "topk", None) + if not isinstance(topk, int) or isinstance(topk, bool) or topk <= 0: + raise ValueError(f"decoder block {layer_idx} replay gate must expose a positive integer topk") + num_experts = getattr(gate, "n_experts", getattr(gate, "num_experts", None)) + if num_experts is not None and ( + not isinstance(num_experts, int) or isinstance(num_experts, bool) or num_experts <= 0 + ): + raise ValueError(f"decoder block {layer_idx} replay gate has invalid expert count {num_experts!r}") + if getattr(gate, "use_routing_core", False): + raise RuntimeError( + "RouterReplayAdapter is incompatible with partial MoE router CUDA graphs; " + "disable the 'moe_router' graph module before enabling routing replay" + ) + replay = gate.router_replay + if replay is None: + replay = RouterReplay(register=False) + gate.router_replay = replay + if not isinstance(replay, RouterReplay): + raise TypeError( + f"decoder block {layer_idx} router_replay must be RouterReplay or None, got {type(replay).__name__}" + ) + if id(replay) in seen_replays: + raise ValueError("one RouterReplay handle is attached to more than one decoder block") + seen_replays.add(id(replay)) + bindings.append(_RouterReplayBinding(layer_idx, replay, topk, num_experts)) + + if not bindings: + raise ValueError("RouterReplayAdapter found no MoE gate with a router_replay slot in the primary decoder") + bindings.sort(key=lambda binding: binding.layer_idx) + layer_ids = [binding.layer_idx for binding in bindings] + if len(set(layer_ids)) != len(layer_ids): + raise ValueError(f"multiple replay gates map to the same global decoder layer: {layer_ids}") + topks = {binding.topk for binding in bindings} + if len(topks) != 1: + raise ValueError(f"all replay gates must use one topk, got {sorted(topks)}") + self._bindings = tuple(bindings) + self._layer_ids = tuple(layer_ids) + self._topk = next(iter(topks)) + self._topology_tensors: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} + + @property + def layer_ids(self) -> tuple[int, ...]: + """Global decoder-layer ids, in the model's replay order.""" + return self._layer_ids + + def prepare_routed_experts(self, routed_experts: torch.Tensor) -> torch.Tensor: + """Convert rollout routes from sequence-last to Engine PER_TOKEN layout. + + Args: + routed_experts: Signed integer expert ids with shape + ``[batch, global_layers, topk, sequence]``. ``-1`` means that + the training gate should keep the token's complete live top-k; + a missing route row must therefore contain only ``-1``. A + replay row contains unique expert ids. + + Returns: + A contiguous tensor with shape + ``[batch, sequence, global_layers, topk]``. Declare it as a + ``PER_TOKEN`` Datum field with pad value ``-1``. + """ + self._validate_integral_routes(routed_experts) + if routed_experts.ndim != 4: + raise ValueError("sequence-last routed_experts must have shape [batch, global_layers, topk, sequence]") + return routed_experts.permute(0, 3, 1, 2).contiguous() + + def __call__( + self, + model_inputs: Mapping[str, Any], + loss_fn_inputs: Mapping[str, Any], + ) -> AbstractContextManager[None]: + """Create replay context from one CP/packing-prepared Engine batch. + + Args: + model_inputs: CP-local model inputs. ``input_ids`` has token shape + ``[batch, local_sequence]`` or ``[local_tokens]``; + ``inputs_embeds`` adds one trailing hidden axis. + loss_fn_inputs: CP-local loss/side-channel mapping. Optional + ``routed_experts`` has shape ``[batch, local_sequence, + global_layers, topk]`` or ``[local_tokens, global_layers, + topk]`` and signed integer dtype. The token axes are validated + against prepared ``weights`` when present, because models that + shard after embedding may retain full-length ``input_ids``. + + Returns: + A context that replays this model's layer targets through forward, + loss, backward, and activation-checkpoint recomputation. Missing + ``routed_experts`` selects live routing and returns a no-op context. + """ + routed_experts = loss_fn_inputs.get(self.field_name) + if routed_experts is None: + return nullcontext() + if not isinstance(routed_experts, torch.Tensor): + raise TypeError("routed_experts must be a Tensor") + self._validate_integral_routes(routed_experts) + if routed_experts.ndim < 3: + raise ValueError("prepared routed_experts must have token axes followed by [global_layers, topk]") + + weights = loss_fn_inputs.get("weights") + if isinstance(weights, torch.Tensor) and weights.ndim > 0: + token_shape = tuple(weights.shape) + if routed_experts.ndim < weights.ndim + 2 or tuple(routed_experts.shape[: weights.ndim]) != token_shape: + raise ValueError( + f"routed_experts token axes {tuple(routed_experts.shape[:-2])} do not match " + f"the prepared loss weights {token_shape}" + ) + expected_tokens = weights.numel() + else: + primary = model_inputs.get("input_ids") + if isinstance(primary, torch.Tensor): + expected_tokens = primary.numel() + else: + primary = model_inputs.get("inputs_embeds") + if not isinstance(primary, torch.Tensor) or primary.ndim < 2: + raise ValueError( + "loss inputs must contain token-shaped weights, or model inputs must contain input_ids " + "or inputs_embeds" + ) + expected_tokens = prod(primary.shape[:-1]) + route_tokens = prod(routed_experts.shape[:-2]) + if route_tokens != expected_tokens: + raise ValueError( + f"routed_experts describes {route_tokens} tokens but the prepared model batch has {expected_tokens}" + ) + + num_layers, route_topk = routed_experts.shape[-2:] + max_layer_idx = self._bindings[-1].layer_idx + if num_layers <= max_layer_idx: + raise ValueError( + f"routed_experts has {num_layers} global layers but this model requires layer {max_layer_idx}" + ) + per_token = routed_experts.reshape(route_tokens, num_layers, route_topk) + if route_topk != self._topk: + raise ValueError(f"routed_experts topk {route_topk} does not match model topk {self._topk}") + topology = self._topology_tensors.get(per_token.device) + if topology is None: + layer_indices = torch.tensor(self.layer_ids, device=per_token.device, dtype=torch.long) + expert_limits = torch.tensor( + [ + binding.num_experts if binding.num_experts is not None else torch.iinfo(torch.long).max + for binding in self._bindings + ], + device=per_token.device, + dtype=torch.long, + ).view(1, -1, 1) + topology = (layer_indices, expert_limits) + self._topology_tensors[per_token.device] = topology + layer_indices, expert_limits = topology + selected = per_token.index_select(1, layer_indices) + missing_rows = (selected == -1).all(dim=-1) + valid_rows = ((selected >= 0) & (selected < expert_limits)).all(dim=-1) + unique_rows = torch.ones_like(missing_rows) + for offset in range(1, route_topk): + unique_rows &= (selected[..., :-offset] != selected[..., offset:]).all(dim=-1) + torch._assert_async( + (missing_rows | (valid_rows & unique_rows)).all(), + "each routed_experts row must be all -1 or contain unique valid model expert ids", + ) + targets = list(selected.unbind(dim=1)) + return self._activate(targets) + + @staticmethod + def _validate_integral_routes(routed_experts: torch.Tensor) -> None: + """Validate a route tensor without changing its token or layer layout. + + Args: + routed_experts: Signed integer expert ids with arbitrary token axes + followed by ``[global_layers, topk]``. + """ + if routed_experts.dtype not in {torch.int8, torch.int16, torch.int32, torch.int64}: + raise TypeError(f"routed_experts must use a signed integer dtype, got {routed_experts.dtype}") + + @contextmanager + def _activate(self, targets: list[torch.Tensor]) -> Iterator[None]: + """Temporarily install one ``[tokens, topk]`` target per binding. + + Args: + targets: Model-scoped replay targets in ``self._bindings`` order. + Every tensor has shape ``[tokens, topk]``. + + Yields: + ``None`` while replay is active. Previous handle state is restored + on normal exit or exception. + """ + previous = [ + (binding.replay.mode, binding.replay.target_indices, binding.replay._allow_trailing_live_tokens) + for binding in self._bindings + ] + try: + for binding, target in zip(self._bindings, targets): + binding.replay.target_indices = target + binding.replay._allow_trailing_live_tokens = True + binding.replay.mode = RouterReplayMode.REPLAY + yield + finally: + for binding, (mode, target, allow_trailing) in zip(self._bindings, previous): + binding.replay.mode = mode + binding.replay.target_indices = target + binding.replay._allow_trailing_live_tokens = allow_trailing + + def replay_selection(router_replay: RouterReplay | None, indices: torch.Tensor) -> torch.Tensor: """Route ``indices`` through ``router_replay`` when routing replay is enabled. diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index f38d919c1f..bb6a979f73 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -59,6 +59,7 @@ ] LossInputValue = torch.Tensor | tuple[torch.Tensor, ...] LossInputs = dict[str, LossInputValue] +BatchContextFn = Callable[[Mapping[str, Any], Mapping[str, LossInputValue]], AbstractContextManager[Any]] LossFn = Callable[ [Any, LossInputs], torch.Tensor | tuple[torch.Tensor, Sequence[Mapping[str, Any]] | LossFnOutputBatch], @@ -84,6 +85,24 @@ def _nullcontext_for_batch(_model_inputs: dict[str, Any]) -> AbstractContextMana return nullcontext() +def _nullcontext_for_prepared_batch( + _model_inputs: Mapping[str, Any], + _loss_fn_inputs: Mapping[str, LossInputValue], +) -> AbstractContextManager[Any]: + """Return a no-op context for one prepared eager batch. + + Args: + _model_inputs: CP-local model mapping with primary token shape + ``[batch, sequence]`` or ``[tokens]``. + _loss_fn_inputs: CP-local loss tensors with the same leading token + axes for ``PER_TOKEN`` fields. + + Returns: + A no-op context manager. + """ + return nullcontext() + + def _as_tuple(value: _T | Sequence[_T] | None) -> tuple[_T, ...]: if value is None: return () @@ -140,14 +159,6 @@ def _resolve_summed_gradient_reduction(model_parts: Sequence[nn.Module]) -> bool return bool(part_modes and part_modes[0]) -def _tensor_version(tensor: torch.Tensor) -> int: - """Return the in-place mutation counter when the Tensor exposes one.""" - try: - return int(tensor._version) - except RuntimeError: - return -1 - - def _resolve_fp8_scale_precompute( model_parts: Sequence[nn.Module], ) -> tuple[tuple[nn.Module, ...], Callable[[nn.Module], None] | None]: @@ -208,8 +219,11 @@ def collate_prebatched(datums: list[Datum]) -> tuple[dict[str, Any], CollatedLos datum.loss_fn_inputs, layouts=datum.loss_fn_input_layouts, item_to_datum=None, + pad_values=datum.loss_fn_input_pad_values, ) else: + if datum.loss_fn_input_pad_values: + raise ValueError("prebatched loss input pad values require an explicit layout for every loss field") # Source-compatible prebatched callers without complete metadata keep # the legacy inference path. It remains fail-closed for ambiguous # packed per-Datum fields. @@ -223,6 +237,7 @@ class _LossBatchLayout: fields: Mapping[str, LossInputLayout] item_to_datum: tuple[int, ...] | None + pad_values: Mapping[str, float | int | bool] unresolved_fields: frozenset[str] = frozenset() @@ -266,7 +281,7 @@ class ForwardResult: @dataclass(frozen=True) class ForwardBackwardResult: - """One backward call's loss statistics and per-Datum callback outputs. + """One complete optimizer window's loss statistics and callback outputs. The numerator is summed across the DP-CP gradient group. The full-sequence denominator is summed across DP only because CP ranks begin with replicated @@ -277,10 +292,10 @@ class ForwardBackwardResult: identical on every PP stage in that replica. Attributes: - loss: Detached weighted mean for this call's Datum window. + loss: Detached weighted mean for the complete Datum window. loss_sum: Detached numerator summed across DP and CP, then synchronized across PP stages. - weight_sum: Detached denominator for this call, summed across DP but + weight_sum: Detached window denominator, summed across DP but not CP, then synchronized across PP stages. loss_fn_outputs: Detached per-Datum mappings in input order. """ @@ -307,28 +322,6 @@ class OptimStepResult: learning_rates: tuple[float, ...] -@dataclass(frozen=True) -class _PlannedDatum: - datum: Datum - weights: torch.Tensor - weights_version: int - weights_shape: torch.Size - weights_dtype: torch.dtype - weights_device: torch.device - weights_sum: float - - -@dataclass -class _AccumulationState: - windows: tuple[tuple[_PlannedDatum, ...], ...] - weight_sums: tuple[torch.Tensor, ...] - total_weight_sum: torch.Tensor - total_microbatches: int - microbatch_size: int - next_window: int = 0 - status: str = "active" - - class Engine: """Run model forward or forward/backward over Datum windows. @@ -339,10 +332,8 @@ class Engine: gradient-accumulation synchronization, and backward. When optimizers are provided, :meth:`optim_step` owns distributed gradient finalization, clipping, parameter updates, gradient clearing, model post-step hooks, and - LR-scheduler advancement. By default, one :meth:`forward_backward` call is - the complete optimizer accumulation window consumed by - :meth:`optim_step`. Call :meth:`begin_accumulation` first when that window - must be split across multiple ``forward_backward`` calls. Dynamic loss + LR-scheduler advancement. One :meth:`forward_backward` call is the complete + optimizer accumulation window consumed by :meth:`optim_step`. Dynamic loss scaling and overflow-skipped updates are not part of this contract. Args: @@ -367,6 +358,13 @@ class Engine: context_fn: Creates an optional context from the CP-prepared model-input mapping. It covers model forward, loss, and backward, or the full pipeline schedule. Recipes use it for FP8 and model input staging. + batch_context_fn: Creates an optional context from the CP/packing- + prepared model-input and loss/side-channel mappings. It covers + eager model forward, loss, backward, and activation-checkpoint + recomputation. AutoPipeline is intentionally unsupported until it + can select one context payload per inner pipeline microbatch. + Rank-local callback failures are process-fatal in distributed + execution, like failures from ``context_fn`` or model forward. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the final microbatch. optimizers: Already-built optimizer or optimizers for these model parts. @@ -408,6 +406,7 @@ def __init__( padding_token_id: int = 0, mtp_ignore_index: int = -100, context_fn: Callable[[dict[str, Any]], AbstractContextManager[Any]] = _nullcontext_for_batch, + batch_context_fn: BatchContextFn | None = None, defer_fsdp_grad_sync: bool = True, optimizers: torch.optim.Optimizer | Sequence[torch.optim.Optimizer] | None = None, lr_schedulers: OptimizerParamScheduler | Sequence[OptimizerParamScheduler] | None = None, @@ -418,6 +417,11 @@ def __init__( if isinstance(mtp_ignore_index, bool) or not isinstance(mtp_ignore_index, int): raise ValueError(f"mtp_ignore_index must be an integer, got {mtp_ignore_index!r}") self.pipeline = model if isinstance(model, AutoPipeline) else None + if self.pipeline is not None and batch_context_fn is not None: + raise NotImplementedError( + "batch_context_fn currently supports eager PP=1 execution only; " + "AutoPipeline needs per-inner-microbatch context routing" + ) self.model_parts = model.parts if self.pipeline is not None else [model] self.model = self.model_parts[0] self._summed_gradient_reduction = _resolve_summed_gradient_reduction(self.model_parts) @@ -431,185 +435,16 @@ def __init__( self.padding_token_id = padding_token_id self.mtp_ignore_index = mtp_ignore_index self.context_fn = context_fn + self.batch_context_fn = _nullcontext_for_prepared_batch if batch_context_fn is None else batch_context_fn self.defer_fsdp_grad_sync = defer_fsdp_grad_sync self.optimizers = _as_tuple(optimizers) self.lr_schedulers = _as_tuple(lr_schedulers) self.max_grad_norm = max_grad_norm - self._accumulation_state: _AccumulationState | None = None self._optim_step_consumed = False self._grads_finalized = False self._finalized_grad_norm: torch.Tensor | float | None = None self._optim_step_in_progress = False - self._implicit_backward_status = "idle" - - def begin_accumulation(self, windows: Sequence[Sequence[Datum]]) -> None: - """Plan one optimizer window split across multiple backward calls. - - The complete plan is required before the first backward because the - supervised loss and MoE auxiliary loss use different global - denominators. Each subsequent :meth:`forward_backward` call must pass - the exact planned Datum objects for the next window, in order. A - :class:`ForwardBackwardResult` continues to describe only that call; - callers combine results with ``sum(loss_sum) / sum(weight_sum)``. - - Under pipeline parallelism the Engine carries the plan's final-window - boundary through AutoPipeline so a schedule-local last backward does - not prematurely synchronize deferred gradients. If a planned backward - call reports an error after execution starts, partial distributed state - cannot be rolled back safely: the plan becomes broken and the Engine - must not be stepped or reused. For non-pipeline execution, an explicit - plan performs one small control consensus per outer microbatch so a - rank-local loss-callback failure is reported together before backward; - the ordinary one-call path adds no such collective. A pipeline loss - callback runs inside PyTorch's distributed schedule, so an exception - from that callback is process-fatal rather than a recoverable - broken-plan error. Output-only callback errors returned through the - normal schedule path are synchronized across pipeline stages. - - Args: - windows: Non-empty sequence of non-empty Datum windows in their - future call order. The same Datum objects and weight tensors - must be passed unchanged to :meth:`forward_backward`. Every - non-final window must end on an Engine outer-microbatch - boundary. - - Raises: - RuntimeError: If no optimizer is configured, another plan is - active, or gradients from an earlier optimizer window have not - been cleared. - ValueError: If the plan is empty, malformed, or splits an outer - microbatch across calls. - """ - if self._accumulation_state is not None: - raise RuntimeError("an Engine accumulation plan is already active") - - self._validate_parallelism() - dp_group, dp_size = self._dp_group_and_size() - control_group, control_group_size = self._accumulation_control_group_and_size() - local_error: Exception | None = None - planned_windows: list[tuple[_PlannedDatum, ...]] = [] - microbatch_counts: list[int] = [] - local_weight_sums: list[float] = [] - try: - if not self.optimizers: - raise RuntimeError("Engine.begin_accumulation requires at least one optimizer") - if not isinstance(windows, Sequence) or isinstance(windows, (str, bytes)) or not windows: - raise ValueError("begin_accumulation requires a non-empty sequence of Datum windows") - if ( - self._implicit_backward_status != "idle" - or self._grads_finalized - or any(parameter.grad is not None for part in self.model_parts for parameter in part.parameters()) - ): - raise RuntimeError("begin_accumulation requires cleared gradients") - - for window_index, window in enumerate(windows): - microbatches = self._group_datums(window) - if window_index < len(windows) - 1 and len(window) % self.microbatch_size != 0: - raise ValueError("every non-final accumulation window must end on an outer-microbatch boundary") - - planned_window: list[_PlannedDatum] = [] - local_weight_sum = 0.0 - for datum in window: - weights = datum.loss_fn_inputs.get("weights") - if not isinstance(weights, torch.Tensor): - raise ValueError("every Datum must contain a Tensor loss_fn_inputs['weights']") - if weights.numel() == 0 or not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()): - raise ValueError("Datum weights must be non-empty, finite, and non-negative") - weight_sum = float(weights.to(torch.float64).sum()) - local_weight_sum += weight_sum - planned_window.append( - _PlannedDatum( - datum=datum, - weights=weights, - weights_version=_tensor_version(weights), - weights_shape=weights.shape, - weights_dtype=weights.dtype, - weights_device=weights.device, - weights_sum=weight_sum, - ) - ) - planned_windows.append(tuple(planned_window)) - microbatch_counts.append(len(microbatches)) - local_weight_sums.append(local_weight_sum) - except Exception as error: - local_error = error - - self._synchronize_accumulation_error( - local_error, - control_group, - control_group_size, - peer_message="another model-parallel rank rejected the accumulation plan", - ) - self._validate_window_size_across_group(len(planned_windows), control_group, control_group_size) - for count in microbatch_counts: - self._validate_window_size_across_group(count, control_group, control_group_size) - - local_denominators = torch.tensor(local_weight_sums, dtype=torch.float64, device=self.device) - cp_group, cp_size = self._cp_group_and_size() - cp_error: Exception | None = None - if cp_size > 1: - gathered_cp_denominators = torch.empty( - cp_size * len(local_weight_sums), - dtype=local_denominators.dtype, - device=local_denominators.device, - ) - dist.all_gather_into_tensor(gathered_cp_denominators, local_denominators, group=cp_group) - gathered_cp_denominators = gathered_cp_denominators.view(cp_size, len(local_weight_sums)) - if not torch.allclose( - gathered_cp_denominators, - gathered_cp_denominators[0].expand_as(gathered_cp_denominators), - rtol=1e-8, - atol=1e-12, - ): - cp_error = ValueError( - "context-parallel ranks must plan identical full-sequence weights; " - f"got {gathered_cp_denominators.tolist()}" - ) - self._synchronize_accumulation_error( - cp_error, - control_group, - control_group_size, - peer_message="another context-parallel group rejected the planned weight sums", - ) - - global_denominators = local_denominators.clone() - if dp_size > 1: - dist.all_reduce(global_denominators, op=dist.ReduceOp.SUM, group=dp_group) - weight_sums = list(global_denominators.detach().unbind()) - - if control_group_size > 1: - local_denominators = torch.stack(weight_sums) - gathered_denominators = torch.empty( - control_group_size * len(weight_sums), - dtype=local_denominators.dtype, - device=local_denominators.device, - ) - dist.all_gather_into_tensor(gathered_denominators, local_denominators, group=control_group) - gathered_denominators = gathered_denominators.view(control_group_size, len(weight_sums)) - if not torch.allclose( - gathered_denominators, - gathered_denominators[0].expand_as(gathered_denominators), - rtol=1e-8, - atol=1e-12, - ): - raise ValueError( - "every model-parallel rank must plan the same DP-global weight sums; " - f"got {gathered_denominators.tolist()}" - ) - - total_weight_sum = torch.stack(weight_sums).sum() - self._accumulation_state = _AccumulationState( - windows=tuple(planned_windows), - weight_sums=tuple(weight_sums), - total_weight_sum=total_weight_sum, - total_microbatches=sum(microbatch_counts) - * (self.pipeline.num_microbatches if self.pipeline is not None else 1), - microbatch_size=self.microbatch_size, - ) - self._optim_step_consumed = False - self._grads_finalized = False - self._finalized_grad_norm = None - self._implicit_backward_status = "idle" + self._backward_status = "idle" @torch.no_grad() def forward( @@ -648,8 +483,6 @@ def forward( record per hidden inner sample, and cannot produce records when PP splits it into multiple inner microbatches. """ - if self._accumulation_state is not None: - raise RuntimeError("Engine.forward cannot run while a backward accumulation plan is active") microbatches = self._group_datums(datums) self._validate_execution_parallelism() cp_group, cp_size = self._cp_group_and_size() @@ -693,7 +526,7 @@ def forward( continue loss_inputs = _with_loss_metadata(model_inputs, loss_inputs) - with self.context_fn(model_inputs), cp_context(): + with self.context_fn(model_inputs), cp_context(), self.batch_context_fn(model_inputs, loss_inputs): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) numerator, parsed_outputs, output_parse_error = _parse_loss_result( @@ -740,95 +573,44 @@ def forward_backward( datums: Sequence[Datum], loss_fn: LossFn, ) -> ForwardBackwardResult: - """Run one backward window, optionally inside a predeclared accumulation plan. + """Run one complete optimizer accumulation window. - Without :meth:`begin_accumulation`, this call remains a complete - optimizer window. With an active plan, calls must consume its Datum - windows in order. Every returned result is call-local even though all - gradients use the plan's full-step normalization. + A second call with configured optimizers is rejected until + :meth:`optim_step` consumes the first call's gradients. A failed + partial backward poisons the window so its gradients cannot be reused. """ - state = self._accumulation_state - if state is None: - if self.optimizers: - if self._implicit_backward_status == "ready": - raise RuntimeError( - "forward_backward already produced the current optimizer window; " - "call optim_step, or declare multiple calls up front with begin_accumulation" - ) - if self._implicit_backward_status == "broken": - raise RuntimeError("the previous implicit backward window failed and this Engine cannot be reused") - if self._implicit_backward_status == "running": - raise RuntimeError("an implicit forward_backward call is already running") - if self._grads_finalized: + if self.optimizers: + if self._backward_status == "ready": raise RuntimeError( - "gradients were already finalized; retry optim_step before another forward_backward call" + "forward_backward already produced the current optimizer window; call optim_step first" ) - if self.optimizers: - self._implicit_backward_status = "running" - try: - result = self._forward_backward_window(datums, loss_fn) - except BaseException: - if self.optimizers: - self._implicit_backward_status = "broken" - raise - self._optim_step_consumed = False - self._grads_finalized = False - self._finalized_grad_norm = None - if self.optimizers: - self._implicit_backward_status = "ready" - return result - - if state.status == "broken": - raise RuntimeError("the active accumulation plan is broken and this Engine cannot be reused") - if state.status == "running": - raise RuntimeError("an accumulation forward_backward call is already running") - if state.status == "ready": - raise RuntimeError("the accumulation plan is complete; call optim_step before another backward") - if state.next_window >= len(state.windows): - raise RuntimeError("the accumulation plan has no remaining backward windows") - + if self._backward_status == "broken": + raise RuntimeError("the previous backward window failed and this Engine cannot be reused") + if self._backward_status == "running": + raise RuntimeError("a forward_backward call is already running") + if self._grads_finalized: + raise RuntimeError("gradients were already finalized; retry optim_step before forward_backward") + if self.optimizers: + self._backward_status = "running" try: - self._validate_planned_window( - datums, - state.windows[state.next_window], - microbatch_size=state.microbatch_size, - ) - is_first_window = state.next_window == 0 - is_final_window = state.next_window == len(state.windows) - 1 - state.status = "running" - result = self._forward_backward_window( - datums, - loss_fn, - result_denominator=state.weight_sums[state.next_window], - backward_denominator=state.total_weight_sum, - total_microbatches=state.total_microbatches, - is_first_window=is_first_window, - is_final_window=is_final_window, - ) + result = self._forward_backward_window(datums, loss_fn) except BaseException: - state.status = "broken" + if self.optimizers: + self._backward_status = "broken" raise - - state.next_window += 1 - state.status = "ready" if state.next_window == len(state.windows) else "active" self._optim_step_consumed = False self._grads_finalized = False self._finalized_grad_norm = None - self._implicit_backward_status = "idle" + if self.optimizers: + self._backward_status = "ready" return result def _forward_backward_window( self, datums: Sequence[Datum], loss_fn: LossFn, - *, - result_denominator: torch.Tensor | None = None, - backward_denominator: torch.Tensor | None = None, - total_microbatches: int | None = None, - is_first_window: bool = True, - is_final_window: bool = True, ) -> ForwardBackwardResult: - """Accumulate gradients for one implicit or explicitly planned window. + """Accumulate gradients for a complete optimizer window. ``datums`` is a flat optimizer accumulation window. The Engine groups it into outer batches of ``microbatch_size`` and invokes ``collate_fn`` @@ -858,10 +640,10 @@ def _forward_backward_window( sample use ordinary flat Datums. Args: - datums: Flat sequence of Datum items in this call's window. A - Datum's token weights may have shape [tokens] or the custom - collater's batched token layout; the loss tensor must use the - identical shape. + datums: Flat sequence of Datum items in the complete optimizer + accumulation window. A Datum's token weights may have shape + [tokens] or the custom collater's batched token layout; the + loss tensor must use the identical shape. loss_fn: Computes either that per-token loss tensor or a scalar local weighted-sum numerator from the raw model output and collated loss inputs. @@ -875,10 +657,8 @@ def _forward_backward_window( ``loss_fn_outputs`` contains mappings for this DP replica's outer Datums in window order; pipeline execution returns the same mappings on every physical stage rank in that replica. Model - parameters are unchanged. Without an explicit accumulation plan, - gradients contain this call's globally normalized result. With a - plan, they accumulate using the complete plan's denominator even - though the returned statistics remain call-local. + parameters are unchanged, but their gradients contain the complete + window's globally normalized backward result. """ microbatches = self._group_datums(datums) self._validate_parallelism() @@ -886,35 +666,17 @@ def _forward_backward_window( grad_group, grad_group_size = self._gradient_group_and_size(dp_group, dp_size) gradient_reduction_multiplier = 1 if self._summed_gradient_reduction else grad_group_size self._validate_window_size_across_group(len(microbatches), grad_group, grad_group_size) - denominator = ( - self._global_weight_sum(microbatches, dp_group, dp_size) - if result_denominator is None - else result_denominator - ) - gradient_denominator = denominator if backward_denominator is None else backward_denominator - planned_accumulation = result_denominator is not None - plan_control_group, plan_control_group_size = ( - self._accumulation_control_group_and_size() if planned_accumulation else (None, 1) - ) + denominator = self._global_weight_sum(microbatches, dp_group, dp_size) zero_denominator = bool(denominator == 0) - zero_gradient_denominator = bool(gradient_denominator == 0) safe_denominator = torch.where(denominator > 0, denominator, torch.ones_like(denominator)) - safe_gradient_denominator = torch.where( - gradient_denominator > 0, - gradient_denominator, - torch.ones_like(gradient_denominator), - ) self._validate_pipeline_window(len(microbatches), denominator) pp_enabled = self.pipeline is not None for part in self.model_parts: part.train() - if is_first_window: - prepare_for_grad_accumulation(self.model_parts, pp_enabled=pp_enabled) + prepare_for_grad_accumulation(self.model_parts, pp_enabled=pp_enabled) inner_microbatches = self.pipeline.num_microbatches if self.pipeline is not None else 1 - effective_total_microbatches = ( - len(microbatches) * inner_microbatches if total_microbatches is None else total_microbatches - ) + effective_total_microbatches = len(microbatches) * inner_microbatches MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( self._cp_size() * gradient_reduction_multiplier / (grad_group_size * effective_total_microbatches) ) @@ -924,13 +686,13 @@ def _forward_backward_window( returns_outputs: bool | None = None output_error: Exception | None = None backward_scale = ( - safe_gradient_denominator.new_zeros(()) - if zero_denominator or zero_gradient_denominator - else safe_gradient_denominator.new_tensor(gradient_reduction_multiplier) / safe_gradient_denominator + safe_denominator.new_zeros(()) + if zero_denominator + else safe_denominator.new_tensor(gradient_reduction_multiplier) / safe_denominator ) for index, datums in enumerate(microbatches): - is_last = is_final_window and index == len(microbatches) - 1 + is_last = index == len(microbatches) - 1 if is_last: prepare_for_final_backward(self.model_parts, pp_enabled=pp_enabled) @@ -950,7 +712,6 @@ def _forward_backward_window( output_restore_plan, backward_scale=backward_scale, zero_weight_sum=zero_denominator, - finalize_backward=is_last, ) if output_error is None: if batch_error is not None: @@ -973,30 +734,14 @@ def _forward_backward_window( get_sync_ctx(self.model, is_last, self.defer_fsdp_grad_sync), self.context_fn(model_inputs), cp_context(), + self.batch_context_fn(model_inputs, loss_inputs), ): forward_inputs = filter_forward_kwargs(self.model, model_inputs) output = self.model(**forward_inputs) - loss_error: Exception | None = None - numerator: torch.Tensor | None = None - parsed_outputs: ParsedLossOutputs = None - output_parse_error: Exception | None = None - try: - numerator, parsed_outputs, output_parse_error = _parse_loss_result( - loss_fn(output, loss_inputs), loss_inputs["weights"] - ) - except Exception as error: - loss_error = error - if planned_accumulation: - self._synchronize_accumulation_error( - loss_error, - plan_control_group, - plan_control_group_size, - peer_message="another model-parallel rank failed in the planned loss callback", - ) - elif loss_error is not None: - raise loss_error - assert numerator is not None - if output_error is None and dp_size <= 1 and not planned_accumulation: + numerator, parsed_outputs, output_parse_error = _parse_loss_result( + loss_fn(output, loss_inputs), loss_inputs["weights"] + ) + if output_error is None and dp_size <= 1: self._validate_loss_fn_outputs_across_cp( parsed_outputs, loss_inputs.get("weights"), @@ -1012,7 +757,7 @@ def _forward_backward_window( if output_error is None: try: - if dp_size > 1 or planned_accumulation: + if dp_size > 1: self._validate_loss_fn_outputs_across_cp( parsed_outputs, loss_inputs.get("weights"), @@ -1036,7 +781,7 @@ def _forward_backward_window( except Exception as error: output_error = error local_loss_sum.add_(numerator.detach().to(torch.float64)) - if is_first_window and index == 0: + if index == 0: prepare_after_first_microbatch() # Piggyback the output-error bit on the existing end-of-window loss @@ -1048,11 +793,6 @@ def _forward_backward_window( pp_group, pp_size = self._pp_group_and_size() if pp_size > 1: dist.all_reduce(step_state, op=dist.ReduceOp.SUM, group=pp_group) - if planned_accumulation: - if plan_control_group_size > grad_group_size: - control_error = step_state[1].clamp(max=1) - dist.all_reduce(control_error, op=dist.ReduceOp.MAX, group=plan_control_group) - step_state[1].copy_(control_error) if bool(step_state[1] > 0): if output_error is not None: raise output_error @@ -1092,10 +832,7 @@ def optim_step( before_optimizer_step: Optional callback invoked exactly once after gradient finalization and clipping, but before the first optimizer step. If it raises, parameters, optimizer state, - model post-step state, and schedulers remain untouched; the - finalized gradients remain available. Outside an explicit - accumulation plan, distributed ranks must invoke this callback - consistently; rank-local execution failures are process-fatal. + model post-step state, and schedulers remain untouched. Returns: Gradient norm and post-scheduler learning rates for the completed @@ -1104,47 +841,18 @@ def optim_step( Raises: RuntimeError: If this Engine was constructed without optimizers. """ - state = self._accumulation_state - local_preflight_error: Exception | None = None - try: - if not self.optimizers: - raise RuntimeError("Engine.optim_step requires at least one optimizer") - if before_optimizer_step is not None and not callable(before_optimizer_step): - raise TypeError("before_optimizer_step must be callable or None") - if self._optim_step_in_progress: - raise RuntimeError("Engine.optim_step is already running") - if state is not None: - if state.status == "broken": - raise RuntimeError("the active accumulation plan is broken and cannot be optimized") - if state.status != "ready": - raise RuntimeError( - "the active accumulation plan must finish every forward_backward call before optim_step" - ) - else: - if self._implicit_backward_status == "broken": - raise RuntimeError( - "the previous implicit backward window failed and this Engine cannot be optimized" - ) - if self._implicit_backward_status == "running": - raise RuntimeError("an implicit forward_backward call is still running") - if self._optim_step_consumed: - raise RuntimeError("optim_step already consumed the current gradients; run forward_backward first") - except Exception as error: - local_preflight_error = error - - # Explicit plans pay the small control collectives needed to fail - # together before mutation. Keep the ordinary one-call path free of - # new per-step collectives; its distributed callback contract remains - # the same as before planned accumulation was introduced. - control_group, control_group_size = ( - self._accumulation_control_group_and_size() if state is not None else (None, 1) - ) - self._synchronize_accumulation_error( - local_preflight_error, - control_group, - control_group_size, - peer_message="another model-parallel rank rejected optim_step", - ) + if not self.optimizers: + raise RuntimeError("Engine.optim_step requires at least one optimizer") + if before_optimizer_step is not None and not callable(before_optimizer_step): + raise TypeError("before_optimizer_step must be callable or None") + if self._optim_step_in_progress: + raise RuntimeError("Engine.optim_step is already running") + if self._backward_status == "broken": + raise RuntimeError("the previous backward window failed and this Engine cannot be optimized") + if self._backward_status == "running": + raise RuntimeError("a forward_backward call is still running") + if self._optim_step_consumed: + raise RuntimeError("optim_step already consumed the current gradients; run forward_backward first") device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None moe_mesh = self.mesh_context.moe_mesh if self.mesh_context is not None else None @@ -1155,50 +863,28 @@ def optim_step( mutation_started = False try: if not self._grads_finalized: - finalization_error: Exception | None = None - try: - self._finalized_grad_norm = scale_grads_and_clip_grad_norm( - max_grad_norm=self.max_grad_norm, - model_parts=self.model_parts, - norm_type=2.0, - pp_enabled=pp_enabled, - device_mesh=device_mesh, - moe_mesh=moe_mesh, - ep_axis_name=( - "ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None - ), - pp_axis_name="pp" if pp_enabled else None, - foreach=True, - num_label_tokens=None, - dp_group_size=grad_group_size, - expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, device_mesh), - ) - if self._finalized_grad_norm is None: - raise RuntimeError("gradient finalization did not return a gradient norm") - except Exception as error: - finalization_error = error - self._synchronize_accumulation_error( - finalization_error, - control_group, - control_group_size, - peer_message="another model-parallel rank failed while finalizing gradients", + self._finalized_grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=self.max_grad_norm, + model_parts=self.model_parts, + norm_type=2.0, + pp_enabled=pp_enabled, + device_mesh=device_mesh, + moe_mesh=moe_mesh, + ep_axis_name="ep" if moe_mesh is not None and "ep" in (moe_mesh.mesh_dim_names or ()) else None, + pp_axis_name="pp" if pp_enabled else None, + foreach=True, + num_label_tokens=None, + dp_group_size=grad_group_size, + expert_tp_replication_factor=get_expert_tp_replication_factor(self.model_parts, device_mesh), ) + if self._finalized_grad_norm is None: + raise RuntimeError("gradient finalization did not return a gradient norm") self._grads_finalized = True grad_norm = self._finalized_grad_norm assert grad_norm is not None - fence_error: Exception | None = None if before_optimizer_step is not None: - try: - before_optimizer_step() - except Exception as error: - fence_error = error - self._synchronize_accumulation_error( - fence_error, - control_group, - control_group_size, - peer_message="another model-parallel rank failed before the optimizer mutation fence", - ) + before_optimizer_step() mutation_started = True for optimizer in self.optimizers: @@ -1223,19 +909,15 @@ def optim_step( except Exception: if mutation_started or not self._grads_finalized: self._optim_step_consumed = True - if state is not None: - state.status = "broken" - else: - self._implicit_backward_status = "broken" + self._backward_status = "broken" raise finally: self._optim_step_in_progress = False - self._accumulation_state = None self._optim_step_consumed = True self._grads_finalized = False self._finalized_grad_norm = None - self._implicit_backward_status = "idle" + self._backward_status = "idle" return OptimStepResult(grad_norm=grad_norm, learning_rates=learning_rates) def _group_datums(self, datums: Sequence[Datum]) -> list[list[Datum]]: @@ -1254,9 +936,43 @@ def _resolve_loss_batch_layout( loss_inputs: Mapping[str, LossInputValue], ) -> _LossBatchLayout: """Resolve collater metadata without exposing it to the loss callback.""" + datum_layouts: dict[str, LossInputLayout] = {} + for name in sorted({name for datum in datums for name in datum.loss_fn_input_layouts}): + declared = [datum.loss_fn_input_layouts[name] for datum in datums if name in datum.loss_fn_input_layouts] + if len(declared) != len(datums): + raise ValueError(f"every Datum must declare the loss layout for field {name!r}") + if any(layout != declared[0] for layout in declared[1:]): + raise ValueError(f"Datum items disagree on the loss layout for field {name!r}") + datum_layouts[name] = declared[0] + + datum_pad_values: dict[str, float | int | bool] = {} + for name in sorted({name for datum in datums for name in datum.loss_fn_input_pad_values}): + declared = [ + datum.loss_fn_input_pad_values[name] for datum in datums if name in datum.loss_fn_input_pad_values + ] + if len(declared) != len(datums): + raise ValueError(f"every Datum must declare the loss pad value for field {name!r}") + if any(value != declared[0] for value in declared[1:]): + raise ValueError(f"Datum items disagree on the loss pad value for field {name!r}") + datum_pad_values[name] = declared[0] + + missing_fields = (set(datum_layouts) | set(datum_pad_values)) - set(loss_inputs) + if missing_fields: + raise ValueError(f"collate_fn dropped explicitly declared loss fields: {sorted(missing_fields)}") + if isinstance(loss_inputs, CollatedLossInputs): if set(loss_inputs.layouts) != set(loss_inputs): raise ValueError("CollatedLossInputs.layouts must describe every loss field exactly once") + for name, layout in datum_layouts.items(): + if loss_inputs.layouts[name] is not layout: + raise ValueError( + f"CollatedLossInputs layout for field {name!r} disagrees with its Datum declaration" + ) + for name, pad_value in datum_pad_values.items(): + if name not in loss_inputs.pad_values or loss_inputs.pad_values[name] != pad_value: + raise ValueError( + f"CollatedLossInputs pad value for field {name!r} disagrees with its Datum declaration" + ) item_to_datum = loss_inputs.item_to_datum if item_to_datum is not None and item_to_datum != tuple(range(len(datums))): raise ValueError( @@ -1266,6 +982,13 @@ def _resolve_loss_batch_layout( return _LossBatchLayout( fields=dict(loss_inputs.layouts), item_to_datum=item_to_datum, + pad_values=dict(loss_inputs.pad_values), + ) + + if datum_pad_values: + raise ValueError( + "custom collaters must return CollatedLossInputs to preserve " + f"loss_fn_input_pad_values for fields {sorted(datum_pad_values)}" ) weights = loss_inputs.get("weights") @@ -1275,13 +998,8 @@ def _resolve_loss_batch_layout( fields: dict[str, LossInputLayout] = {} unresolved: set[str] = set() for name, value in loss_inputs.items(): - declared = {datum.loss_fn_input_layouts[name] for datum in datums if name in datum.loss_fn_input_layouts} - if len(declared) > 1: - raise ValueError(f"Datum items disagree on the loss layout for field {name!r}") - if declared: - if not all(name in datum.loss_fn_input_layouts for datum in datums): - raise ValueError(f"every Datum must declare the loss layout for field {name!r}") - fields[name] = next(iter(declared)) + if name in datum_layouts: + fields[name] = datum_layouts[name] continue if isinstance(value, torch.Tensor) and _loss_sequence_dim(dict(model_inputs), value) is not None: @@ -1325,6 +1043,7 @@ def _resolve_loss_batch_layout( return _LossBatchLayout( fields=fields, item_to_datum=item_to_datum, + pad_values={}, unresolved_fields=frozenset(unresolved), ) @@ -1409,6 +1128,16 @@ def _prepare_batch( thd_loss_fields: list[str] = [] if is_thd: + nonzero_pad_fields = [ + name + for name, pad_value in loss_batch_layout.pad_values.items() + if pad_value != 0 and loss_batch_layout.fields[name] is LossInputLayout.PER_TOKEN + ] + if num_pipeline_microbatches > 1 and nonzero_pad_fields: + raise NotImplementedError( + "packed pipeline microbatching does not yet preserve nonzero PER_TOKEN pad sentinels for " + f"{nonzero_pad_fields}" + ) if "weights" in loss_batch_layout.unresolved_fields: raise ValueError("packed THD execution requires token-aligned loss weights") unresolved_non_token_fields = [ @@ -1437,7 +1166,8 @@ def _prepare_batch( key = f"{_LOSS_FIELD_PREFIX}{name}" if key in cp_batch: raise ValueError(f"model inputs contain reserved Engine key {key!r}") - cp_batch[key] = value + if loss_batch_layout.pad_values.get(name, 0) == 0: + cp_batch[key] = value thd_loss_fields.append(name) device_mesh = self.mesh_context.device_mesh if self.mesh_context is not None else None @@ -1468,7 +1198,9 @@ def _prepare_batch( local_loss_inputs[name] = candidate else: local_loss_inputs[name] = sharder.shard_token_tensor( - loss_inputs[name], seq_dim=loss_seq_dim or 0, fill=0 + loss_inputs[name], + seq_dim=loss_seq_dim or 0, + fill=loss_batch_layout.pad_values.get(name, 0), ) loss_inputs = local_loss_inputs else: @@ -1477,6 +1209,7 @@ def _prepare_batch( loss_inputs, loss_seq_dim, loss_batch_layout.fields, + loss_batch_layout.pad_values, token_reference, loss_batch_layout.unresolved_fields, ) @@ -1490,6 +1223,7 @@ def _prepare_batch( "mtp_per_depth_targets": LossInputLayout.PER_TOKEN, }, item_to_datum=loss_batch_layout.item_to_datum, + pad_values=loss_batch_layout.pad_values, unresolved_fields=loss_batch_layout.unresolved_fields, ) real_lengths, padded_lengths, token_mask = output_routing @@ -1602,7 +1336,6 @@ def _pipeline_execute( *, backward_scale: torch.Tensor | None, zero_weight_sum: bool, - finalize_backward: bool = True, ) -> tuple[bool | None, list[dict[str, Any]], Exception | None]: """Run prepared pipeline microbatches in training or forward-only mode. @@ -1621,9 +1354,6 @@ def _pipeline_execute( backward, or ``None`` to run the forward-only schedule. zero_weight_sum: Whether reporting numerators must be forced to graph-connected zero. - finalize_backward: Whether this pipeline schedule invocation ends - the complete optimizer backward window. Ignored for - forward-only execution. Returns: Whether the callback returned outputs, its detached outputs in @@ -1696,7 +1426,6 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: loss_fn=pipeline_loss, losses=losses, return_outputs=False, - finalize_backward=finalize_backward, ) outputs: list[dict[str, Any]] = [] @@ -2172,33 +1901,6 @@ def _gradient_group_and_size( size = int(dp_cp_mesh.size()) return (dp_cp_mesh.get_group() if size > 1 else None), size - def _accumulation_control_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: - """Return the full model group used to keep plan control flow aligned.""" - if not dist.is_available() or not dist.is_initialized(): - return None, 1 - if self.mesh_context is None: - return None, dist.get_world_size() - if (group := getattr(self.mesh_context, "process_group", None)) is not None: - return group, dist.get_world_size(group=group) - if (device_mesh := getattr(self.mesh_context, "device_mesh", None)) is None: - return None, dist.get_world_size() - - root_mesh = device_mesh._get_root_mesh() if hasattr(device_mesh, "_get_root_mesh") else device_mesh - size = int(root_mesh.size()) - if size <= 1: - return None, 1 - if root_mesh.ndim == 1: - return root_mesh.get_group(), size - if size == dist.get_world_size(): - return None, size - for flat_mesh in getattr(root_mesh, "_flatten_mapping", {}).values(): - if int(flat_mesh.size()) == size: - return flat_mesh.get_group(), size - raise NotImplementedError( - "planned multi-call accumulation on a rank-subset multi-axis DeviceMesh requires " - "MeshContext.process_group for the complete model" - ) - def _pp_group_and_size(self) -> tuple[dist.ProcessGroup | None, int]: if self.pipeline is None or not dist.is_available() or not dist.is_initialized(): return None, 1 @@ -2244,75 +1946,6 @@ def _local_weight_sum(self, microbatches: list[list[Datum]]) -> torch.Tensor: self._validate_weight_sum_across_cp(denominator) return denominator - def _synchronize_accumulation_error( - self, - local_error: Exception | None, - group: dist.ProcessGroup | None, - group_size: int, - *, - peer_message: str, - ) -> None: - """Make every gradient rank reject a bad accumulation contract together.""" - if group_size <= 1: - if local_error is not None: - raise local_error - return - - failed = torch.tensor(int(local_error is not None), dtype=torch.int64, device=self.device) - dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group) - if bool(failed): - if local_error is not None: - raise local_error - raise RuntimeError(peer_message) - - def _validate_planned_window( - self, - datums: Sequence[Datum], - planned: tuple[_PlannedDatum, ...], - *, - microbatch_size: int, - ) -> None: - """Validate an accumulation-plan slot before any model collective.""" - local_error: Exception | None = None - try: - if self.microbatch_size != microbatch_size: - raise RuntimeError( - "Engine.microbatch_size changed after begin_accumulation; " - f"planned {microbatch_size}, found {self.microbatch_size}" - ) - if not isinstance(datums, Sequence) or isinstance(datums, (str, bytes)): - raise TypeError("planned forward_backward input must be a sequence of Datum") - if len(datums) != len(planned): - raise ValueError( - f"planned forward_backward window has {len(planned)} Datums, but received {len(datums)}" - ) - for index, (datum, expected) in enumerate(zip(datums, planned)): - if datum is not expected.datum: - raise ValueError( - f"planned forward_backward window Datum {index} is not the object declared to begin_accumulation" - ) - weights = datum.loss_fn_inputs.get("weights") - if weights is not expected.weights: - raise ValueError(f"planned Datum {index} replaced its weights Tensor") - if ( - _tensor_version(weights) != expected.weights_version - or weights.shape != expected.weights_shape - or weights.dtype != expected.weights_dtype - or weights.device != expected.weights_device - or float(weights.to(torch.float64).sum()) != expected.weights_sum - ): - raise ValueError(f"planned Datum {index} weights changed after begin_accumulation") - except Exception as error: - local_error = error - - control_group, control_group_size = self._accumulation_control_group_and_size() - self._synchronize_accumulation_error( - local_error, - control_group, - control_group_size, - peer_message="another model-parallel rank changed its planned Datum window", - ) - def _validate_window_size_across_group( self, size: int, @@ -2343,6 +1976,7 @@ def _shard_loss_inputs( loss_inputs: LossInputs, seq_dim: int | None, layouts: Mapping[str, LossInputLayout], + pad_values: Mapping[str, float | int | bool], token_reference: torch.Tensor, unresolved_fields: frozenset[str], ) -> LossInputs: @@ -2357,6 +1991,8 @@ def _shard_loss_inputs( the weights do not follow the model's token axes. layouts: Explicit semantic layout for every loss field. Only ``PER_TOKEN`` fields follow the context-parallel token shard. + pad_values: Explicit fills for ``PER_TOKEN`` fields whose padding + sentinel is not zero. token_reference: Tensor carrying the full collated token axes. unresolved_fields: Legacy fields whose semantics were not declared by their collater. Non-token legacy weights cannot cross a CP @@ -2396,7 +2032,7 @@ def _shard_loss_inputs( token_aligned = _is_token_aligned(value, token_reference) if not token_aligned: raise ValueError(f"per-token loss field {name!r} does not match the collated token layout") - local[name] = sharder.shard_token_tensor(value, seq_dim=seq_dim, fill=0) + local[name] = sharder.shard_token_tensor(value, seq_dim=seq_dim, fill=pad_values.get(name, 0)) return local @staticmethod diff --git a/nemo_automodel/shared/model_utils.py b/nemo_automodel/shared/model_utils.py index 8b5249b63f..cf17fd6cfa 100644 --- a/nemo_automodel/shared/model_utils.py +++ b/nemo_automodel/shared/model_utils.py @@ -21,12 +21,11 @@ _TEXT_MODULE_ATTRS = ("language_model", "text_model", "text_decoder") -def iter_transformer_and_mtp_blocks(model: nn.Module) -> Iterator[tuple[nn.Module, str, nn.Module]]: - """Yield transformer and MTP blocks without depending on a recipe or component. +def iter_transformer_blocks(model: nn.Module) -> Iterator[tuple[nn.Module, str, nn.Module]]: + """Yield primary decoder blocks without depending on a model family. Args: - model: Model root containing a transformer layer collection and optional - multi-token-prediction layers. + model: Model root containing a transformer layer collection. Yields: Tuples containing the parent layer collection, child name, and block. @@ -45,6 +44,19 @@ def iter_transformer_and_mtp_blocks(model: nn.Module) -> Iterator[tuple[nn.Modul for layer_id, block in layers.named_children(): yield layers, layer_id, block + +def iter_transformer_and_mtp_blocks(model: nn.Module) -> Iterator[tuple[nn.Module, str, nn.Module]]: + """Yield primary decoder and MTP blocks without model-family branches. + + Args: + model: Model root containing a transformer layer collection and optional + multi-token-prediction layers. + + Yields: + Tuples containing the parent layer collection, child name, and block. + """ + yield from iter_transformer_blocks(model) + mtp_layers = getattr(getattr(model, "mtp", None), "layers", None) if mtp_layers is not None: for layer_id, block in mtp_layers.named_children(): diff --git a/tests/functional_tests/context_parallel/run_packed_pp.py b/tests/functional_tests/context_parallel/run_packed_pp.py index 2b3fe5a8b5..e0b39513e0 100644 --- a/tests/functional_tests/context_parallel/run_packed_pp.py +++ b/tests/functional_tests/context_parallel/run_packed_pp.py @@ -514,125 +514,6 @@ def _run_thd_layout( return pipeline -def _run_planned_accumulation_parity( - device: torch.device, - mesh_context: MeshContext, - raw_inputs: dict[str, object], - labels: torch.Tensor, - weights: torch.Tensor, -) -> None: - """Compare two PP schedule calls in one plan with one complete Engine window. - - Args: - device: CUDA device for this physical pipeline rank. - mesh_context: Runtime PP2 topology; this check runs with CP size one. - raw_inputs: Raw THD mapping whose token and position tensors have shape - ``[batch, sequence]`` before Engine preparation. - labels: Target token IDs shaped ``[batch, sequence]``. - weights: Base token weights shaped ``[batch, sequence]``; the two - planned calls derive unequal denominators from this tensor. - """ - - def make_windows() -> tuple[list[Datum], list[Datum]]: - """Clone two prebatched Datum windows with ``[batch, sequence]`` tensors.""" - weights_a = weights.clone() - weights_a[:, 1::2] = 0 - weights_b = weights.clone().mul(1.5) - weights_b[:, ::2] = 0 - return ( - [ - Datum( - model_inputs=_clone_mapping(raw_inputs), - loss_fn_inputs={"labels": labels.clone(), "weights": weights_a}, - ) - ], - [ - Datum( - model_inputs=_clone_mapping(raw_inputs), - loss_fn_inputs={"labels": labels.clone(), "weights": weights_b}, - ) - ], - ) - - reference_pipeline = _build_pipeline(device, mesh_context) - reference_parameters = [parameter for part in reference_pipeline.parts for parameter in part.parameters()] - reference_optimizer = torch.optim.SGD(reference_parameters, lr=0.05) - reference_engine = Engine( - reference_pipeline, - device=device, - mesh_context=mesh_context, - collate_fn=collate_prebatched, - optimizers=reference_optimizer, - max_grad_norm=1e6, - ) - reference_window_a, reference_window_b = make_windows() - reference_result = reference_engine.forward_backward( - reference_window_a + reference_window_b, - _token_losses, - ) - reference_grads = [parameter.grad.detach().clone() for parameter in reference_parameters] - reference_step = reference_engine.optim_step() - - planned_pipeline = _build_pipeline(device, mesh_context) - planned_parameters = [parameter for part in planned_pipeline.parts for parameter in part.parameters()] - planned_optimizer = torch.optim.SGD(planned_parameters, lr=0.05) - planned_engine = Engine( - planned_pipeline, - device=device, - mesh_context=mesh_context, - collate_fn=collate_prebatched, - optimizers=planned_optimizer, - max_grad_norm=1e6, - ) - planned_window_a, planned_window_b = make_windows() - planned_engine.begin_accumulation([planned_window_a, planned_window_b]) - planned_result_a = planned_engine.forward_backward(planned_window_a, _token_losses) - planned_result_b = planned_engine.forward_backward(planned_window_b, _token_losses) - - if not reference_parameters or len(reference_parameters) != len(planned_parameters): - raise AssertionError( - f"planned/reference PP parameter mismatch: {len(planned_parameters)} != {len(reference_parameters)}" - ) - if planned_result_a.weight_sum.item() == planned_result_b.weight_sum.item(): - raise AssertionError("planned PP accumulation fixture must use unequal call denominators") - torch.testing.assert_close( - planned_result_a.loss_sum + planned_result_b.loss_sum, - reference_result.loss_sum, - atol=4e-2, - rtol=2e-3, - ) - torch.testing.assert_close( - planned_result_a.weight_sum + planned_result_b.weight_sum, - reference_result.weight_sum, - atol=0, - rtol=0, - ) - combined_loss = (planned_result_a.loss_sum + planned_result_b.loss_sum) / ( - planned_result_a.weight_sum + planned_result_b.weight_sum - ) - torch.testing.assert_close(combined_loss, reference_result.loss, atol=4e-2, rtol=2e-3) - for planned_parameter, reference_grad in zip(planned_parameters, reference_grads): - if planned_parameter.grad is None: - raise AssertionError("planned PP accumulation left a local parameter without a gradient") - torch.testing.assert_close(planned_parameter.grad.float(), reference_grad.float(), atol=5e-3, rtol=5e-2) - - planned_step = planned_engine.optim_step() - torch.testing.assert_close(planned_step.grad_norm.float(), reference_step.grad_norm.float(), atol=5e-3, rtol=5e-2) - for planned_parameter, reference_parameter in zip(planned_parameters, reference_parameters): - torch.testing.assert_close( - planned_parameter.float(), - reference_parameter.float(), - atol=5e-3, - rtol=5e-2, - ) - - if dist.get_rank() == 0: - print( - "PP2 planned two-call accumulation matched one complete window " - f"(weights={planned_result_a.weight_sum.item():.1f}+{planned_result_b.weight_sum.item():.1f})" - ) - - def _run_explicit_loss_layout( pipeline: AutoPipeline, layout: str, @@ -855,10 +736,6 @@ def main() -> None: dist.barrier() _run_padded_output_broadcast(final_pipeline, device, mesh_context) dist.barrier() - del final_pipeline - torch.cuda.empty_cache() - _run_planned_accumulation_parity(device, mesh_context, raw_inputs, labels, weights) - dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py index c08da7a55d..fa42c05039 100644 --- a/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py +++ b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py @@ -129,7 +129,7 @@ def _per_token_identity_loss(output: torch.Tensor, loss_inputs: dict[str, torch. Returns: Per-token losses with shape ``[batch=1, sequence=4]``. Engine applies - ``weights`` and the complete planned-window denominator. + ``weights`` and the complete-window denominator. """ assert output.shape == loss_inputs["weights"].shape return output @@ -139,8 +139,8 @@ def _run_mode( mesh_context: MeshContext, *, summed_gradients: bool, -) -> tuple[tuple[float, ...], float, dict[str, torch.Tensor]]: - """Run one planned update through a real MegatronFSDP reduction mode. +) -> tuple[tuple[float, float, float], float, dict[str, torch.Tensor]]: + """Run one complete update through a real MegatronFSDP reduction mode. Args: mesh_context: Two-rank ``[dp=2, cp=1, tp=1]`` CUDA mesh. @@ -148,8 +148,9 @@ def _run_mode( selects SUM gradient collectives; false selects averaged gradients. Returns: - Call-local loss sums, weight sums, and normalized losses; the global - gradient norm; and float32 local parameter shards after the update. + The complete-window loss sum, weight sum, and normalized loss; the + global gradient norm; and float32 local parameter shards after the + update. """ torch.manual_seed(1234) model = TinyTokenModel().cuda() @@ -171,18 +172,13 @@ def _run_mode( max_grad_norm=1e9, ) window_a, window_b = _windows(dist.get_rank()) - engine.begin_accumulation([window_a, window_b]) - result_a = engine.forward_backward(window_a, _per_token_identity_loss) - result_b = engine.forward_backward(window_b, _per_token_identity_loss) + result = engine.forward_backward(window_a + window_b, _per_token_identity_loss) step_result = engine.optim_step() statistics = ( - result_a.loss_sum.item(), - result_a.weight_sum.item(), - result_a.loss.item(), - result_b.loss_sum.item(), - result_b.weight_sum.item(), - result_b.loss.item(), + result.loss_sum.item(), + result.weight_sum.item(), + result.loss.item(), ) return statistics, float(step_result.grad_norm), _local_parameters(model) @@ -208,8 +204,7 @@ def main() -> None: assert math.isfinite(summed[1]) and summed[1] > 0 torch.testing.assert_close(torch.tensor(summed[0]), torch.tensor(averaged[0]), rtol=1e-5, atol=1e-5) torch.testing.assert_close(torch.tensor(summed[1]), torch.tensor(averaged[1]), rtol=1e-5, atol=1e-5) - assert summed[0][1] == 3.0 - assert summed[0][4] == 4.0 + assert summed[0][1] == 7.0 assert set(summed[2]) == set(averaged[2]) for name in sorted(averaged[2]): torch.testing.assert_close(summed[2][name], averaged[2][name], rtol=1e-5, atol=1e-5) diff --git a/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py b/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py index 3b2666a3e7..c95e0701c1 100644 --- a/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py +++ b/tests/functional_tests/parallelism/run_pp_grad_accum_parity.py @@ -30,14 +30,8 @@ Sequence length changes between windows, which is what triggers the stage reset that caused the loss. -The second phase drives those same varying-length FSDP2 stages through Engine: -one complete ``forward_backward`` call is the reference for an explicit plan -split across two calls. This additionally verifies that a non-final pipeline -schedule completes FSDP post-backward cleanup without synchronizing gradients. - Usage: torchrun --nproc-per-node=2 run_pp_grad_accum_parity.py - torchrun --nproc-per-node=4 run_pp_grad_accum_parity.py # PP2 x DP2 """ import torch @@ -49,7 +43,6 @@ # Two accumulation windows with *different* sequence lengths. The change is the # trigger: equal lengths would skip the stage reset and hide the regression. WINDOW_SEQ_LENS = (32, 48) -LOCAL_WINDOW_WEIGHT_SUM = BATCH * sum(WINDOW_SEQ_LENS) def _build_model(device: torch.device) -> torch.nn.Module: @@ -101,14 +94,12 @@ def _batch(seq_len: int, device: torch.device) -> dict[str, torch.Tensor]: device: Device to place the tensors on. Returns: - Dict with ``input_ids``, ``attention_mask``, and ``labels``, all of - shape [BATCH, seq_len]. The mask is explicitly all ones so direct - schedule and Engine collation exercise identical model inputs. + Dict with ``input_ids`` and ``labels``, both of shape [BATCH, seq_len]. """ torch.manual_seed(seq_len) # same window -> same data on every rank and phase ids = torch.randint(2, VOCAB, (BATCH, seq_len), device=device) dist.broadcast(ids, src=0) - return {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids.clone()} + return {"input_ids": ids, "labels": ids.clone()} def _run_window(pp, seq_len: int, device: torch.device) -> None: @@ -146,233 +137,21 @@ def _zero_grads(model) -> None: param.grad = None -def _scale_aware_grad_close( - actual: torch.Tensor, - expected: torch.Tensor, - *, - relative_max_error: float = 0.03, - absolute_floor: float = 2e-3, - relative_norm_error: float = 0.02, -) -> tuple[bool, float, float, float]: - """Compare equal-shaped local gradient shards with a BF16-aware bound. - - Args: - actual: Float32 local FSDP gradient shard being checked. - expected: Equal-shaped float32 local FSDP reference shard. - relative_max_error: Allowed max-element error relative to the largest - absolute reference element. - absolute_floor: Absolute max-element allowance for values near zero. - relative_norm_error: Allowed relative error in the full shard norm. - - Returns: - Whether both the max-error and norm checks pass, followed by the - observed max error, its allowed bound, and ``||actual||/||expected||``. - """ - if actual.shape != expected.shape: - return False, float("inf"), 0.0, float("inf") - if expected.numel() == 0: - return True, 0.0, absolute_floor, 1.0 - - max_error = float((actual - expected).abs().max()) - error_bound = relative_max_error * float(expected.abs().max()) + absolute_floor - actual_norm = float(actual.norm()) - expected_norm = float(expected.norm()) - if expected_norm == 0.0: - norm_ratio = 1.0 if actual_norm == 0.0 else float("inf") - norm_close = actual_norm <= absolute_floor - else: - norm_ratio = actual_norm / expected_norm - norm_close = abs(norm_ratio - 1.0) <= relative_norm_error - return max_error <= error_bound and norm_close, max_error, error_bound, norm_ratio - - -def _engine_window(seq_len: int, weight_scale: float, device: torch.device): - """Build one flat Datum window with a caller-specific token denominator. - - Args: - seq_len: Token length of every Datum in the window. - weight_scale: Constant value assigned to every token loss weight. - device: Device holding the generated tensors. - - Returns: - ``BATCH`` Datums whose ``input_ids``, ``attention_mask``, ``labels``, - and ``weights`` each have shape ``[sequence]``. The Engine collates - them to ``[BATCH, sequence]`` before PP splits the batch axis. - """ - from nemo_automodel.components.datasets.datum import Datum - - batch = _batch(seq_len, device) - return [ - Datum( - model_inputs={"input_ids": input_ids, "attention_mask": attention_mask}, - loss_fn_inputs={ - "labels": labels, - "weights": torch.full_like(labels, weight_scale, dtype=torch.float32), - }, - ) - for input_ids, attention_mask, labels in zip( - batch["input_ids"], - batch["attention_mask"], - batch["labels"], - ) - ] - - -def _run_engine_planned_parity( - pp, - mesh, - device: torch.device, - rank: int, - direct_normalized_grads: dict[str, torch.Tensor], -) -> None: - """Compare one Engine window with the same FSDP2 PP work split across calls. - - Args: - pp: Built PP2 AutoPipeline whose local stage is FSDP2-wrapped. - mesh: ``[pp, dp]`` device mesh; DP may be one or two. - device: CUDA device for this rank's token tensors. - rank: Global rank used in actionable assertion messages. - direct_normalized_grads: Per-parameter direct-schedule reference. Each - value is the float32 local FSDP shard accumulated over both - varying-length windows, with every microbatch loss divided by the - complete local-window token denominator. - """ - from nemo_automodel.components.distributed.mesh import MeshContext - from nemo_automodel.engine import Engine - - def token_losses(pred, loss_inputs): - """Return unweighted token cross entropy in the Engine loss layout. - - Args: - pred: Pipeline output with logits shaped - ``[pp_microbatch, sequence, vocab]``. - loss_inputs: Mapping containing ``labels`` and ``weights`` shaped - ``[pp_microbatch, sequence]``. - - Returns: - Per-token cross entropy with the same shape as ``weights``. Engine - applies the token weights and complete-window denominator. - """ - logits = pred.logits if hasattr(pred, "logits") else pred - return torch.nn.functional.cross_entropy( - logits.float().flatten(0, 1), - loss_inputs["labels"].flatten(0, 1), - reduction="none", - ).view_as(loss_inputs["weights"]) - - window_a = _engine_window(WINDOW_SEQ_LENS[0], 1.0, device) - window_b = _engine_window(WINDOW_SEQ_LENS[1], 1.0, device) - part = pp.parts[0] - mesh_context = MeshContext.from_meshes(mesh) - - _zero_grads(part) - reference_engine = Engine( - pp, - device=device, - mesh_context=mesh_context, - microbatch_size=BATCH, - optimizers=torch.optim.SGD(part.parameters(), lr=0.01), - max_grad_norm=None, - ) - reference_result = reference_engine.forward_backward(window_a + window_b, token_losses) - reference_grads = _snapshot_grads(part) - - _zero_grads(part) - planned_engine = Engine( - pp, - device=device, - mesh_context=mesh_context, - microbatch_size=BATCH, - optimizers=torch.optim.SGD(part.parameters(), lr=0.01), - max_grad_norm=None, - ) - planned_engine.begin_accumulation([window_a, window_b]) - result_a = planned_engine.forward_backward(window_a, token_losses) - result_b = planned_engine.forward_backward(window_b, token_losses) - planned_grads = _snapshot_grads(part) - - assert result_a.weight_sum.item() != result_b.weight_sum.item(), ( - f"[rank {rank}] Engine fixture must use unequal call denominators" - ) - torch.testing.assert_close(result_a.loss_sum + result_b.loss_sum, reference_result.loss_sum, rtol=2e-2, atol=2e-3) - torch.testing.assert_close( - result_a.weight_sum + result_b.weight_sum, - reference_result.weight_sum, - rtol=0, - atol=0, - ) - combined_loss = (result_a.loss_sum + result_b.loss_sum) / (result_a.weight_sum + result_b.weight_sum) - torch.testing.assert_close(combined_loss, reference_result.loss, rtol=2e-2, atol=2e-3) - - missing = set(reference_grads) ^ set(planned_grads) - assert not missing, f"[rank {rank}] Engine gradient key mismatch: {sorted(missing)[:5]}" - direct_missing = set(direct_normalized_grads) ^ set(planned_grads) - assert not direct_missing, f"[rank {rank}] direct/Engine gradient key mismatch: {sorted(direct_missing)[:5]}" - mismatches = [] - oracle_mismatches = [] - for name in sorted(reference_grads): - got, want = planned_grads[name], reference_grads[name] - if not torch.allclose(got, want, rtol=2e-2, atol=2e-3): - mismatches.append(f"{name}: max|delta|={(got - want).abs().max():.3e}") - direct = direct_normalized_grads[name] - single_close, max_error, error_bound, norm_ratio = _scale_aware_grad_close(want, direct) - if not single_close: - oracle_mismatches.append( - f"{name}/single: max|delta|={max_error:.3e} bound={error_bound:.3e} norm_ratio={norm_ratio:.6f}" - ) - planned_close, max_error, error_bound, norm_ratio = _scale_aware_grad_close(got, direct) - if not planned_close: - oracle_mismatches.append( - f"{name}/planned: max|delta|={max_error:.3e} bound={error_bound:.3e} norm_ratio={norm_ratio:.6f}" - ) - if mismatches: - raise AssertionError( - f"[rank {rank}] Engine planned PP gradients != one complete window " - f"({len(mismatches)}/{len(reference_grads)} parameters differ).\n " + "\n ".join(mismatches[:8]) - ) - if oracle_mismatches: - raise AssertionError( - f"[rank {rank}] Engine normalized gradients disagree with the independent direct normalized oracle " - f"({len(oracle_mismatches)} mismatches; global_weight_sum={reference_result.weight_sum.item():.1f}, " - f"dp_size={mesh['dp'].size()}).\n " + "\n ".join(oracle_mismatches[:8]) - ) - - print(f"[rank {rank}] Engine planned PP/FSDP2 parity OK over {len(reference_grads)} parameters") - - def main() -> None: """Compare accumulated gradients against the sum of per-window gradients.""" dist.init_process_group("nccl") rank, world = dist.get_rank(), dist.get_world_size() - if world not in {2, 4}: - raise ValueError(f"PP grad-accumulation parity requires 2 or 4 ranks, got {world}") torch.cuda.set_device(rank % torch.cuda.device_count()) device = torch.device("cuda", rank % torch.cuda.device_count()) from nemo_automodel.components.distributed.pipelining import AutoPipeline - mesh = init_device_mesh("cuda", (2, world // 2), mesh_dim_names=("pp", "dp")) + mesh = init_device_mesh("cuda", (world, 1), mesh_dim_names=("pp", "dp")) model = _build_model(device) def loss_fn(pred, target): - """Return one PP microbatch's CE normalized by the full local window. - - Args: - pred: Pipeline output with logits shaped - ``[pp_microbatch, sequence, vocab]``. - target: Token labels shaped ``[pp_microbatch, sequence]``. - - Returns: - Scalar summed cross entropy divided by the complete two-window - local token denominator. This matches Engine's backward scale. - """ logits = pred.logits if hasattr(pred, "logits") else pred - loss_sum = torch.nn.functional.cross_entropy( - logits.float().flatten(0, 1), - target.flatten(0, 1), - reduction="sum", - ) - return loss_sum / LOCAL_WINDOW_WEIGHT_SUM + return torch.nn.functional.cross_entropy(logits.float().flatten(0, 1), target.flatten(0, 1), reduction="sum") pp = AutoPipeline( world_mesh=mesh, @@ -412,10 +191,8 @@ def loss_fn(pred, target): _zero_grads(part) _run_window(pp, WINDOW_SEQ_LENS[-1], device) last_only = _snapshot_grads(part) - catches_dropped_window = any(not _scale_aware_grad_close(last_only[name], reference[name])[0] for name in reference) - assert catches_dropped_window, ( - f"[rank {rank}] last-window-only gradients pass the Engine oracle tolerance; test cannot detect a wipe" - ) + differs = any(not torch.allclose(last_only[n], reference[n], rtol=1e-3, atol=1e-4) for n in reference) + assert differs, f"[rank {rank}] windows produce identical gradients; test cannot detect a wipe" mismatches = [] for name in sorted(reference): @@ -434,7 +211,6 @@ def loss_fn(pred, target): ) print(f"[rank {rank}] PP grad-accumulation parity OK over {len(reference)} parameters") - _run_engine_planned_parity(pp, mesh, device, rank, reference) dist.barrier() dist.destroy_process_group() diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index b0a58bc989..8d757e6a2a 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -48,6 +48,31 @@ def _toy_datums(): ] +def _routed_datums(): + """Build ragged Datums with per-token routing metadata. + + Returns: + Two Datums whose routed-expert tensors use the + ``[tokens, layers, topk]`` layout and ``torch.int16`` dtype. + """ + first_routes = torch.arange(3 * 2 * 2, dtype=torch.int16).reshape(3, 2, 2) + second_routes = torch.arange(100, 100 + 2 * 2 * 2, dtype=torch.int16).reshape(2, 2, 2) + return [ + Datum( + input_ids=torch.tensor([10, 11, 12]), + loss_fn_inputs={"weights": torch.ones(3), "routed_experts": first_routes}, + loss_fn_input_layouts={"routed_experts": LossInputLayout.PER_TOKEN}, + loss_fn_input_pad_values={"routed_experts": -1}, + ), + Datum( + input_ids=torch.tensor([20, 21]), + loss_fn_inputs={"weights": torch.ones(2), "routed_experts": second_routes}, + loss_fn_input_layouts={"routed_experts": LossInputLayout.PER_TOKEN}, + loss_fn_input_pad_values={"routed_experts": -1}, + ), + ] + + # ── Datum ───────────────────────────────────────────────────────────────── @@ -106,6 +131,21 @@ def test_datum_rejects_invalid_loss_input_layouts(): ) +def test_datum_validates_loss_input_pad_value_metadata(): + with pytest.raises(ValueError, match="unknown loss inputs"): + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"weights": torch.ones(1)}, + loss_fn_input_pad_values={"missing": -1}, + ) + with pytest.raises(TypeError, match="bool, int, or float"): + Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"weights": torch.ones(1)}, + loss_fn_input_pad_values={"weights": object()}, # type: ignore[dict-item] + ) + + def test_to_features_applies_masking_convention(): feats = _toy_datums()[0].to_features() assert feats["input_ids"] == [10, 11, 12] @@ -173,6 +213,46 @@ def test_collate_packed_side_inputs_ride_the_flat_axis(): assert [t.tolist() for t in split] == [pytest.approx([0.5, 0.5, 0.5]), pytest.approx([0.9, 0.9])] +def test_collate_padded_per_token_trailing_dims_use_the_declared_pad_value(): + datums = _routed_datums() + + _, loss_inputs = collate_datums(datums) + + routes = loss_inputs["routed_experts"] + assert routes.shape == (2, 3, 2, 2) + assert routes.dtype == torch.int16 + torch.testing.assert_close(routes[0], datums[0].loss_fn_inputs["routed_experts"]) + torch.testing.assert_close(routes[1, :2], datums[1].loss_fn_inputs["routed_experts"]) + assert torch.equal(routes[1, 2], torch.full((2, 2), -1, dtype=torch.int16)) + assert loss_inputs.layouts["routed_experts"] is LossInputLayout.PER_TOKEN + assert loss_inputs.pad_values == {"routed_experts": -1} + + +def test_collate_requires_explicit_layout_for_trailing_token_features(): + datums = _routed_datums() + for datum in datums: + datum.loss_fn_input_layouts.clear() + + with pytest.raises(ValueError, match="declare an explicit layout for 'routed_experts'"): + collate_datums(datums) + + +def test_collate_packed_per_token_trailing_dims_preserve_token_order_and_metadata(): + datums = _routed_datums() + + model_inputs, loss_inputs = collate_datums(datums, packed=True) + + routes = loss_inputs["routed_experts"] + assert model_inputs["input_ids"].tolist() == [[10, 11, 12, 20, 21]] + assert routes.shape == (1, 5, 2, 2) + torch.testing.assert_close( + routes[0], + torch.cat([datum.loss_fn_inputs["routed_experts"] for datum in datums]), + ) + assert loss_inputs.layouts["routed_experts"] is LossInputLayout.PER_TOKEN + assert loss_inputs.pad_values == {"routed_experts": -1} + + def test_collate_packed_per_sample_side_input_is_one_per_datum(): datums = [ Datum(model_inputs={"input_ids": torch.tensor([1, 2])}, loss_fn_inputs={"advantages": torch.tensor([0.5])}), @@ -207,6 +287,21 @@ def test_collated_loss_inputs_copy_preserves_side_channel_and_dict_compatibility assert copied.item_to_datum == loss_inputs.item_to_datum +def test_collated_loss_inputs_copy_preserves_read_only_pad_values(): + loss_inputs = collate_datums(_routed_datums())[1] + + for copied in ( + loss_inputs.copy(), + copy(loss_inputs), + deepcopy(loss_inputs), + pickle.loads(pickle.dumps(loss_inputs)), # noqa: S301 - trusted in-process round trip + ): + assert isinstance(copied, CollatedLossInputs) + assert copied.pad_values == {"routed_experts": -1} + with pytest.raises(TypeError): + copied.pad_values["routed_experts"] = 0 # type: ignore[index] + + def test_collated_loss_inputs_requires_complete_read_only_layouts(): with pytest.raises(ValueError, match="exactly"): CollatedLossInputs( @@ -225,6 +320,60 @@ def test_collated_loss_inputs_requires_complete_read_only_layouts(): loss_inputs.layouts["weights"] = LossInputLayout.REPLICATED # type: ignore[index] +@pytest.mark.parametrize( + ("layouts", "pad_values", "error", "message"), + [ + ( + {"weights": LossInputLayout.PER_TOKEN}, + {"missing": -1}, + ValueError, + "unknown loss inputs", + ), + ( + {"weights": LossInputLayout.PER_DATUM}, + {"weights": -1}, + ValueError, + "only valid for PER_TOKEN", + ), + ( + {"weights": LossInputLayout.PER_TOKEN}, + {"weights": object()}, + TypeError, + "bool, int, or float", + ), + ], +) +def test_collated_loss_inputs_validates_pad_value_metadata(layouts, pad_values, error, message): + with pytest.raises(error, match=message): + CollatedLossInputs( + {"weights": torch.ones(1)}, + layouts=layouts, + item_to_datum=(0,), + pad_values=pad_values, + ) + + +def test_collate_requires_consistent_per_token_pad_value_metadata(): + partial = _routed_datums() + partial[1].loss_fn_input_pad_values.clear() + with pytest.raises(ValueError, match="every Datum must declare the pad value"): + collate_datums(partial) + + conflicting = _routed_datums() + conflicting[1].loss_fn_input_pad_values["routed_experts"] = -2 + with pytest.raises(ValueError, match="same pad value"): + collate_datums(conflicting) + + non_token = Datum( + input_ids=torch.tensor([1]), + loss_fn_inputs={"sample_id": torch.tensor(7)}, + loss_fn_input_layouts={"sample_id": LossInputLayout.PER_DATUM}, + loss_fn_input_pad_values={"sample_id": -1}, + ) + with pytest.raises(ValueError, match="does not use the PER_TOKEN layout"): + collate_datums([non_token]) + + def test_collate_explicit_per_datum_overrides_single_token_shape_inference(): datum = Datum( input_ids=torch.tensor([7]), diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index e109422fd6..deae8ac24d 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -18,9 +18,7 @@ import pytest import torch import torch.nn as nn -from torch.distributed.pipelining.microbatch import TensorChunkSpec, split_args_kwargs_into_chunks -import nemo_automodel.components.distributed.pipelining.functional as pipeline_functional from nemo_automodel.components.distributed.pipelining.autopipeline import AutoPipeline from nemo_automodel.components.distributed.pipelining.functional import ( generate_hf_model_fqn_per_model_part, @@ -230,29 +228,13 @@ def test_pp_mesh_extraction(self): assert ap.pp_mesh is not None -class _KwargsChunkHookPart(nn.Module): - def __init__(self, chunk_dims: dict[str, int]): - super().__init__() - self.chunk_dims = chunk_dims - - def get_pipeline_kwargs_chunk_dims(self, kwargs): - return {key: dim for key, dim in self.chunk_dims.items() if key in kwargs} - - -class _UnknownKwargsChunkHookPart(nn.Module): - def get_pipeline_kwargs_chunk_dims(self, kwargs): - return {"unknown": 0} - - -class _KwargsChunkSchedule: +class _PreparedMicrobatchSchedule: def __init__(self, *, fail_on_step: bool = False, invoke_loss: bool = False): - self._kwargs_chunk_spec = None self._loss_fn = Mock(return_value=torch.tensor(0.0)) self.fail_on_step = fail_on_step self.invoke_loss = invoke_loss self.args_during_step = None self.args_split = None - self.kwargs_chunk_spec_during_step = None self.loss_fn_during_step = None self.split_inputs_during_step = None self.kwargs_split = None @@ -265,12 +247,7 @@ def __init__(self, *, fail_on_step: bool = False, invoke_loss: bool = False): self.split_inputs_calls = 0 def _split_inputs(self, args, kwargs=None): - return split_args_kwargs_into_chunks( - args, - kwargs, - 2, - kwargs_chunk_spec=self._kwargs_chunk_spec, - ) + raise AssertionError("prepared microbatch split was not installed") def _run_schedule(self, *args, target=None, losses=None, return_outputs=True, **kwargs): """Split schedule inputs using the chunk spec active during the call. @@ -289,7 +266,6 @@ def _run_schedule(self, *args, target=None, losses=None, return_outputs=True, ** A sentinel string identifying the schedule result. """ self.args_during_step = args - self.kwargs_chunk_spec_during_step = self._kwargs_chunk_spec self.loss_fn_during_step = self._loss_fn self.split_inputs_during_step = self._split_inputs self.target_during_step = target @@ -327,7 +303,7 @@ def eval(self, *args, target=None, losses=None, return_outputs=True, **kwargs): ) -class _LegacyStepSchedule(_KwargsChunkSchedule): +class _LegacyStepSchedule(_PreparedMicrobatchSchedule): """Schedule with the PyTorch 2.6-2.9 step signature.""" def __init__(self): @@ -352,7 +328,7 @@ def eval(self, *args, target=None, losses=None, **kwargs): return self._run_schedule(*args, target=target, losses=losses, **kwargs) -class _ForwardingEvalSchedule(_KwargsChunkSchedule): +class _ForwardingEvalSchedule(_PreparedMicrobatchSchedule): """Current PyTorch shape: eval forwards kwargs to a newer step API.""" def eval(self, *args, target=None, losses=None, **kwargs): @@ -366,59 +342,12 @@ class _NoEvalSchedule(_LegacyStepSchedule): eval = None -class _FinalizationStage: - """Record the two schedule-local gradient-finalization signals.""" - - def __init__(self, submod=None): - self.events = [] - self.finalize_backward_states = [] - self.submod = nn.Module() if submod is None else submod - - def backward_maybe_with_nosync(self, _backward_type, _bwd_kwargs, *, last_backward=False): - self.finalize_backward_states.append(getattr(self, "_nemo_finalize_backward", None)) - self.events.append(("backward", last_backward)) - return (), None - - def perform_reduce_grad(self, divisor): - self.events.append(("reduce", divisor)) - set_requires_gradient_sync = getattr(self.submod, "set_requires_gradient_sync", None) - if callable(set_requires_gradient_sync): - set_requires_gradient_sync(True) - self.scale_grads(divisor) - - def scale_grads(self, divisor): - self.events.append(("scale", divisor)) - - -class _FinalizationSchedule(_KwargsChunkSchedule): - """Model the final backward and reduce calls made by a PyTorch schedule.""" - - def __init__(self, stage): - super().__init__() - self._stage = stage - self._stages = [stage] - - def step(self, *args, target=None, losses=None, return_outputs=True, **kwargs): - result = super().step( - *args, - target=target, - losses=losses, - return_outputs=return_outputs, - **kwargs, - ) - self._stage.backward_maybe_with_nosync("full", {}, last_backward=True) - self._stage.perform_reduce_grad(2) - return result - - -class TestAutoPipelineKwargsChunkSpec: +class TestAutoPipelinePreparedMicrobatches: def _pipeline_with_parts( self, *parts: nn.Module, schedule=None, has_first_stage: bool = True, - defer_fsdp_grad_sync: bool = True, - scale_grads_in_schedule: bool = False, ): ap = AutoPipeline( world_mesh=FakeDeviceMesh(), @@ -427,53 +356,14 @@ def _pipeline_with_parts( pp_microbatch_size=1, pp_batch_size=2, device=torch.device("cpu"), - defer_fsdp_grad_sync=defer_fsdp_grad_sync, - scale_grads_in_schedule=scale_grads_in_schedule, ) - ap._info.schedule = schedule or _KwargsChunkSchedule() + ap._info.schedule = schedule or _PreparedMicrobatchSchedule() ap._info.model_parts = list(parts) ap._info.has_first_stage = has_first_stage return ap - def test_step_splits_mrope_position_ids_on_model_owned_batch_axis(self): - """AutoPipeline.step keeps all mRoPE axes in every microbatch.""" - input_ids = torch.zeros(2, 8, dtype=torch.long) - position_ids = torch.arange(8, dtype=torch.long).view(1, 1, -1).expand(3, 2, -1).clone() - kwargs = { - "position_ids": position_ids, - "attention_mask": torch.ones(2, 8, dtype=torch.bool), - "qkv_format": "thd", - } - - _, default_kwargs_split = split_args_kwargs_into_chunks((input_ids,), kwargs, 2) - assert default_kwargs_split[0]["position_ids"].shape == (2, 2, 8) - - ap = self._pipeline_with_parts(_KwargsChunkHookPart({"position_ids": 1})) - result = ap.step(input_ids, **kwargs) - - fixed_kwargs_split = ap.info.schedule.kwargs_split - assert result == "schedule-result" - assert fixed_kwargs_split[0]["position_ids"].shape == (3, 1, 8) - assert fixed_kwargs_split[1]["position_ids"].shape == (3, 1, 8) - torch.testing.assert_close(fixed_kwargs_split[0]["position_ids"], position_ids[:, :1]) - torch.testing.assert_close(fixed_kwargs_split[1]["position_ids"], position_ids[:, 1:]) - assert fixed_kwargs_split[0]["attention_mask"].shape == (1, 8) - assert fixed_kwargs_split[0]["qkv_format"] == "thd" - assert fixed_kwargs_split[1]["qkv_format"] == "thd" - assert ap.info.schedule.args_during_step == (input_ids,) - assert ap.info.schedule._kwargs_chunk_spec is None - - def test_step_without_model_hook_uses_pytorch_default_chunking(self): - ap = self._pipeline_with_parts(nn.Module()) - - ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) - - assert ap.info.schedule.kwargs_chunk_spec_during_step is None - assert ap.info.schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) - assert ap.info.schedule._kwargs_chunk_spec is None - def test_step_microbatches_passes_prepared_inputs_without_resplitting(self): - schedule = _KwargsChunkSchedule(invoke_loss=True) + schedule = _PreparedMicrobatchSchedule(invoke_loss=True) ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) input_ids = [torch.full((1, 8), index, dtype=torch.long) for index in range(2)] position_ids = [torch.full((3, 1, 8), index, dtype=torch.long) for index in range(2)] @@ -515,7 +405,7 @@ def loss_fn(output, index): assert model_inputs[1]["input_ids"] is input_ids[1] def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): - schedule = _KwargsChunkSchedule() + schedule = _PreparedMicrobatchSchedule() ap = self._pipeline_with_parts(nn.Module(), schedule=schedule, has_first_stage=False) model_inputs = [{"inputs_embeds": torch.zeros(1, 8, 4), "position_ids": torch.zeros(1, 8)} for _ in range(2)] @@ -524,137 +414,6 @@ def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): assert schedule.args_split == [(), ()] assert all("inputs_embeds" not in kwargs for kwargs in schedule.kwargs_split) - def test_step_microbatches_defers_schedule_finalization_until_the_logical_last_call(self): - stage = _FinalizationStage() - schedule = _FinalizationSchedule(stage) - ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - pipeline_functional._make_pipeline_stages_accumulation_aware( - [stage], - reduce_grad_per_microbatch=False, - ) - ap._info.stages = [stage] - model_inputs = [{"input_ids": torch.zeros(1, 8)} for _ in range(2)] - - ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=False) - assert stage.events == [("backward", False), ("scale", 2)] - assert stage.finalize_backward_states == [False] - assert stage._nemo_finalize_backward is True - - ap.step_microbatches(model_inputs, loss_fn=Mock(), finalize_backward=True) - assert stage.events == [ - ("backward", False), - ("scale", 2), - ("backward", True), - ("reduce", 2), - ("scale", 2), - ] - assert stage.finalize_backward_states == [False, True] - - def test_step_microbatches_preserves_requested_per_microbatch_gradient_reduction(self): - stage = _FinalizationStage() - schedule = _FinalizationSchedule(stage) - ap = self._pipeline_with_parts( - nn.Module(), - schedule=schedule, - defer_fsdp_grad_sync=False, - ) - pipeline_functional._make_pipeline_stages_accumulation_aware( - [stage], - reduce_grad_per_microbatch=True, - ) - ap._info.stages = [stage] - - ap.step_microbatches( - [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], - loss_fn=Mock(), - finalize_backward=False, - ) - - assert stage.events == [("backward", True), ("reduce", 2), ("scale", 2)] - assert stage.finalize_backward_states == [True] - - def test_step_microbatches_rejects_cross_call_schedule_gradient_scaling(self): - schedule = _KwargsChunkSchedule() - ap = self._pipeline_with_parts( - nn.Module(), - schedule=schedule, - scale_grads_in_schedule=True, - ) - - with pytest.raises(ValueError, match="scale_grads_in_schedule=False"): - ap.step_microbatches( - [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], - loss_fn=Mock(), - finalize_backward=False, - ) - - assert schedule.step_calls == 0 - - @pytest.mark.parametrize("stage_state", [None, "unwrapped"]) - def test_step_microbatches_requires_accumulation_aware_stages_for_nonfinal_call(self, stage_state): - schedule = _KwargsChunkSchedule() - ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) - ap._info.stages = None if stage_state is None else [_FinalizationStage()] - - with pytest.raises(RuntimeError, match="accumulation-aware pipeline stages"): - ap.step_microbatches( - [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], - loss_fn=Mock(), - finalize_backward=False, - ) - - assert schedule.step_calls == 0 - - def test_nonfinal_fully_sharded_stage_runs_post_backward_without_sync(self, monkeypatch): - fsdp_events = [] - - class FakeFSDPModule: - def set_is_last_backward(self, value): - fsdp_events.append(("last", value)) - - def set_reshard_after_backward(self, value): - fsdp_events.append(("reshard", value)) - - def set_requires_gradient_sync(self, value): - fsdp_events.append(("sync", value)) - - parameter_group = types.SimpleNamespace(post_backward=lambda: fsdp_events.append(("post", None))) - fsdp_state = types.SimpleNamespace( - _state_ctx=types.SimpleNamespace( - all_states=[types.SimpleNamespace(_fsdp_param_group=parameter_group)], - ), - _root_post_backward_final_callback=lambda: fsdp_events.append(("root", None)), - ) - import torch.distributed.fsdp as torch_fsdp - - monkeypatch.setattr(torch_fsdp, "FSDPModule", FakeFSDPModule) - monkeypatch.setattr(torch_fsdp.fully_shard, "state", lambda _module: fsdp_state) - - stage = _FinalizationStage(FakeFSDPModule()) - pipeline_functional._make_pipeline_stages_accumulation_aware( - [stage], - reduce_grad_per_microbatch=False, - ) - stage._nemo_finalize_backward = False - stage.perform_reduce_grad(2) - - assert fsdp_events == [ - ("last", True), - ("reshard", True), - ("sync", False), - ("post", None), - ("root", None), - ] - assert stage.events == [("scale", 2)] - - fsdp_events.clear() - stage.events.clear() - stage._nemo_finalize_backward = True - stage.perform_reduce_grad(2) - - assert fsdp_events == [("sync", True)] - assert stage.events == [("reduce", 2), ("scale", 2)] - @pytest.mark.parametrize( ("method_name", "schedule_cls"), [ @@ -679,7 +438,7 @@ def test_prepared_microbatches_do_not_forward_return_outputs_to_older_pytorch( assert schedule.received_return_outputs is False def test_eval_microbatches_uses_forward_only_schedule_with_exact_prepared_split(self): - schedule = _KwargsChunkSchedule(invoke_loss=True) + schedule = _PreparedMicrobatchSchedule(invoke_loss=True) ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) input_ids = [torch.full((1, 8), index, dtype=torch.long) for index in range(2)] metadata = [object(), object()] @@ -758,7 +517,7 @@ def test_step_microbatches_validates_prepared_inputs(self, model_inputs): @pytest.mark.parametrize("method_name", ["step_microbatches", "eval_microbatches"]) def test_prepared_microbatches_restore_schedule_state_after_failure(self, method_name): - schedule = _KwargsChunkSchedule(fail_on_step=True) + schedule = _PreparedMicrobatchSchedule(fail_on_step=True) original_split_inputs = schedule._split_inputs original_loss_fn = schedule._loss_fn ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) @@ -776,50 +535,14 @@ def test_prepared_microbatches_restore_schedule_state_after_failure(self, method assert schedule._split_inputs == original_split_inputs assert schedule._loss_fn is original_loss_fn - def test_only_canonical_model_part_supplies_chunk_policy(self): - ap = self._pipeline_with_parts( - _KwargsChunkHookPart({"position_ids": 1}), - _KwargsChunkHookPart({"position_ids": 0}), - ) - - ap.step(torch.zeros(2, 8), position_ids=torch.zeros(3, 2, 8)) - - assert ap.info.schedule.kwargs_split[0]["position_ids"].shape == (3, 1, 8) - - def test_nonfirst_stage_ignores_model_input(self): - ap = self._pipeline_with_parts(nn.Module(), has_first_stage=False) - - ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) - - assert ap.info.schedule.args_during_step == () - assert ap.info.schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) - - def test_step_restores_schedule_chunk_spec_after_failure(self): - schedule = _KwargsChunkSchedule(fail_on_step=True) - original_chunk_spec = {"position_ids": TensorChunkSpec(0)} - schedule._kwargs_chunk_spec = original_chunk_spec - ap = self._pipeline_with_parts(_KwargsChunkHookPart({"position_ids": 1}), schedule=schedule) - - with pytest.raises(RuntimeError, match="schedule failed"): - ap.step(torch.zeros(2, 8), position_ids=torch.zeros(3, 2, 8)) - - assert schedule.kwargs_chunk_spec_during_step["position_ids"].split_dim == 1 - assert schedule._kwargs_chunk_spec is original_chunk_spec - - def test_model_hook_cannot_configure_unknown_kwarg(self): - ap = self._pipeline_with_parts(_UnknownKwargsChunkHookPart()) - - with pytest.raises(ValueError, match="unknown kwarg"): - ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) - # ----------------------------- -# Core build/materialize/step tests +# Core build/materialize tests # ----------------------------- -class TestAutoPipelineBuildAndStep: - """Test AutoPipeline build, materialize, and step functionality.""" +class TestAutoPipelineBuild: + """Test AutoPipeline build and materialize functionality.""" def test_autopipeline_basic_creation(self): """Test basic AutoPipeline creation without full build process.""" @@ -839,8 +562,8 @@ def test_autopipeline_basic_creation(self): @pytest.mark.parametrize("pp_size", [2, 4]) @pytest.mark.parametrize("local_rank", [0, 1, 2, 3]) - def test_autopipeline_build_split_materialize_and_step(self, monkeypatch, pp_size, local_rank): - """Test complete AutoPipeline build, materialize, and step workflow.""" + def test_autopipeline_build_split_and_materialize(self, monkeypatch, pp_size, local_rank): + """Test complete AutoPipeline build, split, and materialize workflow.""" _patch_autopipeline_monkey(monkeypatch) if local_rank >= pp_size: pytest.skip("local_rank not part of this pp_size") @@ -955,11 +678,6 @@ def loss_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: assert result is ap assert ap._info.enabled is True - def test_autopipeline_step_workflow(self, monkeypatch): - """Test AutoPipeline step functionality.""" - # Skip this complex test - the step method requires extensive pipeline setup - pytest.skip("Complex step test requires extensive pipeline mocking") - def test_autopipeline_build_assertions(self, monkeypatch): """Test AutoPipeline build method assertion errors.""" _patch_autopipeline_monkey(monkeypatch) diff --git a/tests/unit_tests/distributed/test_cp_sharder.py b/tests/unit_tests/distributed/test_cp_sharder.py index 5b1e7bf525..a1d98dd06b 100644 --- a/tests/unit_tests/distributed/test_cp_sharder.py +++ b/tests/unit_tests/distributed/test_cp_sharder.py @@ -219,6 +219,54 @@ def test_sharder_repositioned_layout_round_trips_input_coordinates(): gather_sharder.gather_token_tensor(full_rows, trim=True) +def test_sharder_repositioned_layout_scatter_gather_preserves_trailing_token_features(): + """Preserve routing features through DSV4-style repositioning. + + The input uses the ``[batch, sequence, layers, topk]`` layout. Only the + first two axes participate in repositioning; both trailing axes are + preserved. Dropped input slots and introduced pad columns retain the fill + sentinel. + """ + positions = torch.tensor( + [ + [2, 0, -1], + [1, 3, -1], + ] + ) + layout = cs.ShardLayout(padded_seq_len=4, input_token_stream_positions=positions) + routes = torch.arange(2 * 3 * 2 * 2, dtype=torch.int16).reshape(2, 3, 2, 2) + expected_padded = torch.full((2, 4, 2, 2), -1, dtype=torch.int16) + expected_padded[0, 2] = routes[0, 0] + expected_padded[0, 0] = routes[0, 1] + expected_padded[1, 1] = routes[1, 0] + expected_padded[1, 3] = routes[1, 1] + + local_parts = [] + for rank in range(2): + sharder = cs.ContextParallelSharder( + device_mesh=_FakeDeviceMesh(_FakeMesh(2, rank)), + shard_batch=cs.shard_batch_identity, + local_token_global_indices=cs.contiguous_local_indices, + shard_layout=layout, + ) + local_parts.append(sharder.shard_token_tensor(routes, fill=-1)) + + assert all(part.shape == (2, 2, 2, 2) for part in local_parts) + torch.testing.assert_close(torch.cat(local_parts, dim=1), expected_padded) + + gather_sharder = cs.ContextParallelSharder( + device_mesh=_FakeDeviceMesh(_FakeMesh(1)), + shard_batch=cs.shard_batch_identity, + local_token_global_indices=cs.contiguous_local_indices, + shard_layout=layout, + ) + restored = gather_sharder.gather_token_tensor(expected_padded, trim=True, fill=-1) + expected_restored = routes.clone() + expected_restored[:, 2] = -1 + assert restored.shape == routes.shape + torch.testing.assert_close(restored, expected_restored) + + def test_gather_trim_raises_without_captured_facts(): sharder = cs.ContextParallelSharder( device_mesh=_FakeDeviceMesh(_FakeMesh(1)), diff --git a/tests/unit_tests/moe/test_fsdp_mixin.py b/tests/unit_tests/moe/test_fsdp_mixin.py index d916c407df..a0da9965ab 100644 --- a/tests/unit_tests/moe/test_fsdp_mixin.py +++ b/tests/unit_tests/moe/test_fsdp_mixin.py @@ -12,11 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from types import SimpleNamespace from unittest.mock import Mock, patch -import pytest - from nemo_automodel.components.models.common.utils import get_is_optim_step, set_is_optim_step from nemo_automodel.components.moe.fsdp_mixin import ( MoEFSDPSyncMixin, @@ -767,25 +764,10 @@ def test_fsdp_module_last_backward(self, mock_fully_shard): assert grads == ((), None) assert param_groups is None - @pytest.mark.parametrize( - ("stage_finalize", "global_optim_step", "expect_post_backward"), - [ - pytest.param(True, False, True, id="stage-final-overrides-global-false"), - pytest.param(False, True, False, id="stage-nonfinal-overrides-global-true"), - pytest.param(None, True, True, id="legacy-stage-falls-back-to-global"), - ], - ) @patch("nemo_automodel.components.moe.fsdp_mixin.get_is_optim_step") @patch("nemo_automodel.components.moe.fsdp_mixin.isinstance") - def test_moe_fsdp_mixin_pipeline_boundary_is_authoritative( - self, - mock_isinstance, - mock_get_optim, - stage_finalize, - global_optim_step, - expect_post_backward, - ): - """Effective stage finalization overrides the legacy global flag.""" + def test_moe_fsdp_mixin_last_backward_with_optim_step(self, mock_isinstance, mock_get_optim): + """Test MoEFSDPSyncMixin path with last_backward=True and IS_OPTIM_STEP=True.""" def isinstance_side_effect(obj, cls): if cls == MoEFSDPSyncMixin: @@ -795,13 +777,12 @@ def isinstance_side_effect(obj, cls): return False mock_isinstance.side_effect = isinstance_side_effect - mock_get_optim.return_value = global_optim_step + mock_get_optim.return_value = True + mock_stage = Mock() model = MockFSDPModule() moe_model = MockMoEModel(MockBackend(), model) - mock_stage = SimpleNamespace(submod=moe_model) - if stage_finalize is not None: - mock_stage._nemo_finalize_backward = stage_finalize + mock_stage.submod = moe_model bwd_kwargs = { "stage_output": Mock(), @@ -815,14 +796,8 @@ def isinstance_side_effect(obj, cls): result = patched_backward_maybe_with_nosync(mock_stage, "full", bwd_kwargs, last_backward=True) - if expect_post_backward: - mock_run_post.assert_called_once_with(moe_model) - else: - mock_run_post.assert_not_called() - if stage_finalize is None: - mock_get_optim.assert_called_once_with() - else: - mock_get_optim.assert_not_called() + # Verify post backward was called + mock_run_post.assert_called_once_with(moe_model) grads, param_groups = result assert grads == ((), None) assert param_groups is None diff --git a/tests/unit_tests/moe/test_router_replay.py b/tests/unit_tests/moe/test_router_replay.py index 27fc1c4f8a..035c763622 100644 --- a/tests/unit_tests/moe/test_router_replay.py +++ b/tests/unit_tests/moe/test_router_replay.py @@ -16,11 +16,13 @@ import pytest import torch +from torch import nn from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.layers import Gate from nemo_automodel.components.moe.router_replay import ( RouterReplay, + RouterReplayAdapter, RouterReplayMode, replay_selection, ) @@ -80,6 +82,59 @@ def run(gate, x): return gate(x, token_mask, None) +class _AdapterGate(nn.Module): + """Small gate exposing the structural contract used by RouterReplayAdapter.""" + + def __init__(self, *, topk=2, num_experts=8): + super().__init__() + self.topk = topk + self.n_experts = num_experts + self.router_replay = None + + def forward(self, live_indices: torch.Tensor) -> torch.Tensor: + """Apply the gate's optional replay selection. + + Args: + live_indices: Naturally selected expert ids with shape + ``[tokens, topk]``. + + Returns: + Expert ids with shape ``[tokens, topk]`` after replay fallback. + """ + return replay_selection(self.router_replay, live_indices) + + +class _AdapterBlock(nn.Module): + def __init__(self, layer_idx, *, routed): + super().__init__() + self.layer_idx = layer_idx + if routed: + self.gate = _AdapterGate() + else: + self.mlp = nn.Identity() + + +class _AdapterDecoder(nn.Module): + def __init__(self, num_layers, routed_layers): + super().__init__() + self.layers = nn.ModuleDict( + { + str(layer_idx): _AdapterBlock(layer_idx, routed=layer_idx in routed_layers) + for layer_idx in range(num_layers) + } + ) + + +class _AdapterModel(nn.Module): + def __init__(self, num_layers=5, routed_layers=(1, 3)): + super().__init__() + self.model = _AdapterDecoder(num_layers, set(routed_layers)) + + +def _adapter_gate(model, layer_idx): + return model.model.layers[str(layer_idx)].gate + + # --------------------------------------------------------------------------- # # Config + helper # --------------------------------------------------------------------------- # @@ -237,6 +292,319 @@ def test_multilayer_distribute_and_collect(): assert torch.equal(rep, rec) +# --------------------------------------------------------------------------- # +# Engine adapter: layout conversion, global mapping, and scoped lifecycle +# --------------------------------------------------------------------------- # + + +def test_adapter_prepares_sequence_last_int16_routes(): + adapter = RouterReplayAdapter(_AdapterModel()) + sequence_last = torch.arange(2 * 5 * 2 * 3, dtype=torch.int16).reshape(2, 5, 2, 3) + + prepared = adapter.prepare_routed_experts(sequence_last) + + assert prepared.shape == (2, 3, 5, 2) + assert prepared.dtype == torch.int16 + assert prepared.is_contiguous() + assert torch.equal(prepared, sequence_last.permute(0, 3, 1, 2)) + + +def test_adapter_maps_sparse_global_layers_and_applies_token_fallback(): + model = _AdapterModel(num_layers=5, routed_layers=(1, 3)) + adapter = RouterReplayAdapter(model) + assert adapter.layer_ids == (1, 3) + + batch, sequence, num_layers, topk = 2, 2, 5, 2 + layer_one = torch.tensor( + [ + [[1, 2], [-1, -1]], + [[-1, -1], [5, 6]], + ], + dtype=torch.int16, + ) + layer_three = torch.tensor( + [ + [[7, 0], [6, 1]], + [[5, 2], [4, 3]], + ], + dtype=torch.int16, + ) + sequence_last = torch.full((batch, num_layers, topk, sequence), -1, dtype=torch.int16) + sequence_last[:, 1] = layer_one.permute(0, 2, 1) + sequence_last[:, 3] = layer_three.permute(0, 2, 1) + prepared = adapter.prepare_routed_experts(sequence_last) + + live_one = torch.tensor([[0, 7], [0, 1], [3, 2], [7, 0]]) + live_three = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 5]]) + gate_one = _adapter_gate(model, 1) + gate_three = _adapter_gate(model, 3) + + with adapter( + {"input_ids": torch.zeros(batch, sequence, dtype=torch.long)}, + {"routed_experts": prepared}, + ): + expected_one_target = layer_one.reshape(-1, topk).long() + expected_three_target = layer_three.reshape(-1, topk).long() + assert torch.equal(gate_one.router_replay.target_indices, expected_one_target) + assert torch.equal(gate_three.router_replay.target_indices, expected_three_target) + keep_live = (expected_one_target == -1).any(dim=-1, keepdim=True) + assert torch.equal(gate_one(live_one), torch.where(keep_live, live_one, expected_one_target)) + assert torch.equal(gate_three(live_three), expected_three_target) + + assert gate_one.router_replay.mode is None + assert gate_one.router_replay.target_indices is None + assert gate_three.router_replay.mode is None + assert gate_three.router_replay.target_indices is None + + +def test_replay_minus_one_row_fallback_casts_int16_target(): + replay = RouterReplay(register=False) + replay.mode = RouterReplayMode.REPLAY + replay.target_indices = torch.tensor([[-1, -1], [3, 4]], dtype=torch.int16) + live = torch.tensor([[1, 2], [5, 6]], dtype=torch.long) + + result = replay.apply(live) + + assert result.dtype == torch.long + assert torch.equal(result, torch.tensor([[1, 2], [3, 4]])) + + +def test_replay_rejects_nonintegral_targets_before_casting(): + replay = RouterReplay(register=False) + replay.mode = RouterReplayMode.REPLAY + replay.target_indices = torch.tensor([[1.9, -1.0]]) + + with pytest.raises(TypeError, match="signed integer dtype"): + replay.apply(torch.tensor([[4, 5]])) + + +def test_adapter_finds_decoder_below_module_wrapper(): + class Wrapper(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + model = _AdapterModel(num_layers=3, routed_layers=(1,)) + adapter = RouterReplayAdapter(Wrapper(model)) + + assert adapter.layer_ids == (1,) + assert isinstance(_adapter_gate(model, 1).router_replay, RouterReplay) + + +def test_adapter_rejects_partial_moe_router_cuda_graph_before_installing_handle(): + model = _AdapterModel(num_layers=3, routed_layers=(1,)) + gate = _adapter_gate(model, 1) + gate.use_routing_core = True + + with pytest.raises(RuntimeError, match="partial MoE router CUDA graphs"): + RouterReplayAdapter(model) + + assert gate.router_replay is None + + +def test_adapter_uses_local_weights_when_model_primary_stays_full_length(): + model = _AdapterModel(num_layers=3, routed_layers=(1,)) + adapter = RouterReplayAdapter(model) + routes = torch.full((2, 3, 2), -1, dtype=torch.int16) + routes[:, 1] = torch.tensor([[1, 2], [3, 4]], dtype=torch.int16) + + with adapter( + {"input_ids": torch.zeros(4, dtype=torch.long)}, + {"weights": torch.ones(2), "routed_experts": routes}, + ): + torch.testing.assert_close( + _adapter_gate(model, 1).router_replay.target_indices, + torch.tensor([[1, 2], [3, 4]], dtype=torch.int16), + ) + + +def test_engine_batch_context_adapter_replays_real_gate_and_preserves_router_gradient(): + from nemo_automodel.components.datasets.datum import Datum, LossInputLayout + from nemo_automodel.engine import Engine, collate_prebatched + + class EngineGateModel(nn.Module): + def __init__(self): + super().__init__() + self.model = _AdapterDecoder(1, {0}) + self.model.layers["0"].gate = make_gate() + self.selected_indices = None + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + """Return selected-router mass for one padded token batch. + + Args: + input_ids: Token ids with shape ``[batch, sequence]``. + + Returns: + Selected probability mass with shape ``[batch, sequence]``. + """ + hidden = torch.nn.functional.one_hot(input_ids.reshape(-1) % 16, num_classes=16).to(torch.float32) + weights, indices, _aux = run(self.model.layers["0"].gate, hidden) + self.selected_indices = indices.detach().clone() + return weights.sum(dim=-1).reshape_as(input_ids) + + model = EngineGateModel() + adapter = RouterReplayAdapter(model) + input_ids = torch.tensor([[1, 2, 3]]) + hidden = torch.nn.functional.one_hot(input_ids.reshape(-1) % 16, num_classes=16).to(torch.float32) + with torch.no_grad(): + _weights, live_indices, _aux = run(model.model.layers["0"].gate, hidden) + target = ((live_indices + 1) % 8).reshape(1, 3, 1, 2).to(torch.int16) + assert not torch.equal(live_indices, target.reshape(-1, 2)) + datum = Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs={"weights": torch.ones_like(input_ids, dtype=torch.float32), "routed_experts": target}, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + def loss_fn(output, _loss_fn_inputs): + """Use selected router mass as one loss per padded token. + + Args: + output: Selected router mass with shape ``[batch, sequence]``. + _loss_fn_inputs: Prepared token side channels with matching leading + ``[batch, sequence]`` axes. + + Returns: + Per-token loss with shape ``[batch, sequence]``. + """ + return output + + Engine( + model, + device="cpu", + collate_fn=collate_prebatched, + batch_context_fn=adapter, + ).forward_backward([datum], loss_fn) + + torch.testing.assert_close(model.selected_indices, target.reshape(-1, 2).long()) + gate_grad = model.model.layers["0"].gate.weight.grad + assert gate_grad is not None + assert torch.count_nonzero(gate_grad) > 0 + + +def test_adapter_auto_installs_model_scoped_handles_and_restores_on_error(): + stale_global = RouterReplay() + stale_target = torch.tensor([[6, 7]]) + stale_global.mode = RouterReplayMode.RECORD + stale_global.target_indices = stale_target + + model_a = _AdapterModel(num_layers=3, routed_layers=(1,)) + model_b = _AdapterModel(num_layers=3, routed_layers=(1,)) + assert _adapter_gate(model_a, 1).router_replay is None + assert _adapter_gate(model_b, 1).router_replay is None + adapter_a = RouterReplayAdapter(model_a) + RouterReplayAdapter(model_b) + replay_a = _adapter_gate(model_a, 1).router_replay + replay_b = _adapter_gate(model_b, 1).router_replay + + assert isinstance(replay_a, RouterReplay) + assert isinstance(replay_b, RouterReplay) + assert RouterReplay.instances() == [stale_global] + + previous_a_target = torch.tensor([[4, 5]]) + previous_b_target = torch.tensor([[2, 3]]) + replay_a.mode = RouterReplayMode.RECORD + replay_a.target_indices = previous_a_target + replay_b.mode = RouterReplayMode.RECORD + replay_b.target_indices = previous_b_target + routes = torch.tensor([[[[-1, -1], [1, 2], [-1, -1]]]], dtype=torch.int16) + + with pytest.raises(RuntimeError, match="forward failed"): + with adapter_a( + {"input_ids": torch.zeros(1, 1, dtype=torch.long)}, + {"routed_experts": routes}, + ): + assert replay_a.mode is RouterReplayMode.REPLAY + assert replay_a.target_indices is not previous_a_target + assert replay_b.mode is RouterReplayMode.RECORD + assert replay_b.target_indices is previous_b_target + assert stale_global.mode is RouterReplayMode.RECORD + assert stale_global.target_indices is stale_target + raise RuntimeError("forward failed") + + assert replay_a.mode is RouterReplayMode.RECORD + assert replay_a.target_indices is previous_a_target + assert replay_b.mode is RouterReplayMode.RECORD + assert replay_b.target_indices is previous_b_target + assert stale_global.mode is RouterReplayMode.RECORD + assert stale_global.target_indices is stale_target + + +def test_adapter_keeps_trailing_unrecorded_tokens_on_live_routing(): + model = _AdapterModel(num_layers=3, routed_layers=(1,)) + adapter = RouterReplayAdapter(model) + gate = _adapter_gate(model, 1) + routes = torch.tensor( + [ + [[-1, -1], [1, 2], [-1, -1]], + [[-1, -1], [3, 4], [-1, -1]], + ], + dtype=torch.int16, + ) + live = torch.tensor([[6, 7], [4, 5], [7, 6], [5, 4]]) + + with adapter( + {"input_ids": torch.zeros(2, dtype=torch.long)}, + {"routed_experts": routes}, + ): + replayed = gate(live) + + assert torch.equal(replayed[:2], torch.tensor([[1, 2], [3, 4]])) + assert torch.equal(replayed[2:], live[2:]) + assert gate.router_replay.mode is None + assert gate.router_replay.target_indices is None + + +@pytest.mark.parametrize( + ("routes", "error", "match"), + [ + (torch.zeros(2, 3, 4, dtype=torch.int16), ValueError, "sequence-last routed_experts"), + (torch.zeros(1, 3, 2, 4), TypeError, "signed integer dtype"), + (torch.zeros(1, 3, 2, 4, dtype=torch.uint8), TypeError, "signed integer dtype"), + ], +) +def test_adapter_prepare_rejects_invalid_rollout_layout(routes, error, match): + adapter = RouterReplayAdapter(_AdapterModel()) + with pytest.raises(error, match=match): + adapter.prepare_routed_experts(routes) + + +@pytest.mark.parametrize( + ("routes", "match"), + [ + (torch.zeros(2, 4, dtype=torch.int16), "token axes"), + (torch.zeros(3, 4, 2, dtype=torch.int16), "describes 3 tokens"), + (torch.zeros(2, 3, 2, dtype=torch.int16), "requires layer 3"), + (torch.zeros(2, 4, 1, dtype=torch.int16), "topk 1"), + ], +) +def test_adapter_rejects_invalid_prepared_layout(routes, match): + adapter = RouterReplayAdapter(_AdapterModel(num_layers=5, routed_layers=(3,))) + with pytest.raises(ValueError, match=match): + adapter( + {"input_ids": torch.zeros(2, dtype=torch.long)}, + {"routed_experts": routes}, + ) + + +@pytest.mark.parametrize("bad_row", [[-1, 2], [-2, -2], [0, 8], [3, 3]]) +def test_adapter_rejects_incomplete_or_invalid_expert_rows(bad_row): + adapter = RouterReplayAdapter(_AdapterModel(num_layers=3, routed_layers=(1,))) + routes = torch.full((2, 3, 2), -1, dtype=torch.int16) + routes[0, 1] = torch.tensor(bad_row, dtype=torch.int16) + + with pytest.raises(RuntimeError, match="all -1 or contain unique valid model expert ids"): + adapter( + {"input_ids": torch.zeros(2, dtype=torch.long)}, + {"weights": torch.ones(2), "routed_experts": routes}, + ) + + # --------------------------------------------------------------------------- # # Error handling + context-manager cleanup # --------------------------------------------------------------------------- # diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 448a5a884a..2db4af4b58 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -27,6 +27,7 @@ import torch.multiprocessing as mp import torch.nn.functional as F from torch import nn +from torch.utils.checkpoint import checkpoint import nemo_automodel.engine as engine_module from nemo_automodel import CollatedLossInputs as PublicCollatedLossInputs @@ -133,7 +134,6 @@ def __init__( self.step_calls = 0 self.eval_calls = 0 self.backward_calls = 0 - self.finalize_backward_calls = [] self.updated_seq_lens = [] self.updated_microbatch_sizes = [] self.updated_input_shapes = [] @@ -157,12 +157,11 @@ def update_seq_len(self, seq_len, *, microbatch_size=None, input_tensor=None): self.updated_microbatch_sizes.append(microbatch_size) self.updated_input_shapes.append(tuple(input_tensor.shape) if input_tensor is not None else None) - def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs, finalize_backward=True): + def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): assert return_outputs is False assert len(model_inputs) == self.num_microbatches self.prepared_inputs.append(model_inputs) self.step_calls += 1 - self.finalize_backward_calls.append(finalize_backward) if self.events is not None: self.events.append("step") @@ -656,102 +655,13 @@ def test_forward_backward_uses_one_denominator_for_the_window(): assert model.forward_calls == 2 -def test_planned_multi_call_matches_one_window_with_unequal_denominators(): - window_a = [_datum([2, 100], [1.0, 0.0])] - window_b = [_datum([4, 8], [0.5, 1.5])] - - reference_model = ScaleModel() - reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) - reference_engine = Engine( - reference_model, - device="cpu", - optimizers=reference_optimizer, - max_grad_norm=None, - ) - reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) - reference_grad = reference_model.weight.grad.detach().clone() - reference_step = reference_engine.optim_step() - - planned_model = ScaleModel() - planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) - planned_engine = Engine( - planned_model, - device="cpu", - optimizers=planned_optimizer, - max_grad_norm=None, - ) - planned_engine.begin_accumulation([window_a, window_b]) - result_a = planned_engine.forward_backward(window_a, _identity_loss) - result_b = planned_engine.forward_backward(window_b, _identity_loss) - planned_grad = planned_model.weight.grad.detach().clone() - planned_step = planned_engine.optim_step() - - assert result_a.loss_sum.item() == pytest.approx(2.0) - assert result_a.weight_sum.item() == pytest.approx(1.0) - assert result_a.loss.item() == pytest.approx(2.0) - assert result_b.loss_sum.item() == pytest.approx(14.0) - assert result_b.weight_sum.item() == pytest.approx(2.0) - assert result_b.loss.item() == pytest.approx(7.0) - assert reference_result.loss_sum.item() == pytest.approx(16.0) - assert reference_result.weight_sum.item() == pytest.approx(3.0) - assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) - torch.testing.assert_close(planned_grad, reference_grad) - torch.testing.assert_close(planned_step.grad_norm, reference_step.grad_norm) - torch.testing.assert_close(planned_model.weight, reference_model.weight) - - -def test_explicit_one_window_accumulation_matches_implicit_call(): - implicit_window = [_datum([1, 9], [0.25, 0.75])] - explicit_window = [_datum([1, 9], [0.25, 0.75])] - - implicit_model = ScaleModel() - implicit_optimizer = torch.optim.SGD(implicit_model.parameters(), lr=0.1) - implicit_engine = Engine( - implicit_model, - device="cpu", - optimizers=implicit_optimizer, - max_grad_norm=None, - ) - implicit_result = implicit_engine.forward_backward(implicit_window, _identity_loss) - - explicit_model = ScaleModel() - explicit_optimizer = torch.optim.SGD(explicit_model.parameters(), lr=0.1) - explicit_engine = Engine( - explicit_model, - device="cpu", - optimizers=explicit_optimizer, - max_grad_norm=None, - ) - explicit_engine.begin_accumulation([explicit_window]) - explicit_result = explicit_engine.forward_backward(explicit_window, _identity_loss) - - torch.testing.assert_close(explicit_result.loss, implicit_result.loss) - torch.testing.assert_close(explicit_result.loss_sum, implicit_result.loss_sum) - torch.testing.assert_close(explicit_result.weight_sum, implicit_result.weight_sum) - torch.testing.assert_close(explicit_model.weight.grad, implicit_model.weight.grad) - explicit_engine.optim_step() - implicit_engine.optim_step() - torch.testing.assert_close(explicit_model.weight, implicit_model.weight) - - -def test_begin_accumulation_requires_an_optimizer_before_forward(): - model = ScaleModel() - engine = Engine(model, device="cpu") - - with pytest.raises(RuntimeError, match="optimizer"): - engine.begin_accumulation([[_datum([1])]]) - - assert model.forward_calls == 0 - assert model.weight.grad is None - - -def test_multiple_backward_calls_with_an_optimizer_require_an_explicit_plan(): +def test_multiple_backward_calls_with_an_optimizer_fail_fast(): model = ScaleModel() optimizer = torch.optim.SGD(model.parameters(), lr=0.1) engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=None) engine.forward_backward([_datum([2])], _identity_loss) - with pytest.raises(RuntimeError, match="begin_accumulation"): + with pytest.raises(RuntimeError, match="optim_step"): engine.forward_backward([_datum([6])], _identity_loss) torch.testing.assert_close(model.weight.grad, torch.tensor(2.0)) @@ -759,7 +669,7 @@ def test_multiple_backward_calls_with_an_optimizer_require_an_explicit_plan(): torch.testing.assert_close(model.weight, torch.tensor(0.8)) -def test_failed_implicit_backward_poisoned_gradients_cannot_be_reused(): +def test_failed_backward_poisoned_gradients_cannot_be_reused(): model = ScaleModel() optimizer = torch.optim.SGD(model.parameters(), lr=0.1) engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=None) @@ -775,227 +685,29 @@ def fail_on_second_microbatch(output, _loss_inputs): with pytest.raises(ValueError, match="second microbatch failed"): engine.forward_backward([_datum([2]), _datum([6])], fail_on_second_microbatch) - # The first microbatch already produced a partial gradient. It must never - # be consumed or silently combined with a later optimizer window. torch.testing.assert_close(model.weight.grad, torch.tensor(1.0)) with pytest.raises(RuntimeError, match="failed"): engine.forward_backward([_datum([4])], _identity_loss) with pytest.raises(RuntimeError, match="failed"): engine.optim_step() - with pytest.raises(RuntimeError, match="cleared gradients"): - engine.begin_accumulation([[_datum([4])]]) torch.testing.assert_close(model.weight, torch.tensor(1.0)) -def test_planned_multi_call_uses_one_lifecycle_and_whole_step_moe_scale(monkeypatch): - events = [] - - @contextmanager - def recording_sync_ctx(_model, is_optim_step, _defer_fsdp_grad_sync): - events.append(f"sync-{is_optim_step}") - yield - - monkeypatch.setattr( - engine_module, - "prepare_for_grad_accumulation", - lambda *_args, **_kwargs: events.append("prepare"), - ) - monkeypatch.setattr( - engine_module, - "prepare_for_final_backward", - lambda *_args, **_kwargs: events.append("final"), - ) - monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after-first")) - monkeypatch.setattr(engine_module, "get_sync_ctx", recording_sync_ctx) - monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) - - model = _MainAndAuxScaleModel() - engine = Engine(model, device="cpu", optimizers=torch.optim.SGD(model.parameters(), lr=0.1)) - # A zero-weight call still counts toward the whole-step MoE microbatch - # average even though it contributes no main-loss numerator. - window_a = [_datum([100], [0.0])] - window_b = [_datum([3]), _datum([5])] - - engine.begin_accumulation([window_a, window_b]) - result_a = engine.forward_backward(window_a, _identity_loss) - result_b = engine.forward_backward(window_b, _identity_loss) - - assert events == [ - "prepare", - "sync-False", - "after-first", - "sync-False", - "final", - "sync-True", - ] - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(1.0 / 3.0) - assert result_a.loss_sum.item() == pytest.approx(0.0) - assert result_a.weight_sum.item() == pytest.approx(0.0) - assert result_a.loss.item() == pytest.approx(0.0) - assert result_b.loss.item() == pytest.approx(4.0) - assert model.main_weight.grad.item() == pytest.approx(4.0) - assert model.aux_weight.grad.item() == pytest.approx(1.0) - assert model.forward_calls == 3 - - -def test_planned_accumulation_rejects_unfinished_extra_and_double_steps(monkeypatch): - events = [] - - class RecordingSGD(torch.optim.SGD): - def step(self, closure=None): - events.append("step") - return super().step(closure) - - def zero_grad(self, set_to_none=True): - events.append("zero") - return super().zero_grad(set_to_none=set_to_none) - - class RecordingScheduler: - def step(self, increment): - assert increment == 1 - events.append("scheduler") - +def test_optim_step_rejects_double_step_and_accepts_a_new_window(): model = ScaleModel() - optimizer = RecordingSGD(model.parameters(), lr=0.1) - engine = Engine( - model, - device="cpu", - optimizers=optimizer, - lr_schedulers=RecordingScheduler(), - max_grad_norm=None, - ) - real_finalize = engine_module.scale_grads_and_clip_grad_norm - - def recording_finalize(**kwargs): - events.append("finalize") - return real_finalize(**kwargs) - - monkeypatch.setattr(engine_module, "scale_grads_and_clip_grad_norm", recording_finalize) - window_a = [_datum([1])] - window_b = [_datum([3])] - engine.begin_accumulation([window_a, window_b]) - engine.forward_backward(window_a, _identity_loss) - - with pytest.raises(RuntimeError): - engine.optim_step() - with pytest.raises(RuntimeError): - engine.begin_accumulation([window_a, window_b]) - assert events == [] - - engine.forward_backward(window_b, _identity_loss) - with pytest.raises(RuntimeError): - engine.forward_backward(window_b, _identity_loss) - result = engine.optim_step() - assert events == ["finalize", "step", "zero", "scheduler"] - assert result.learning_rates == (0.1,) - - with pytest.raises(RuntimeError): - engine.optim_step() - assert events == ["finalize", "step", "zero", "scheduler"] + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + engine = Engine(model, device="cpu", optimizers=optimizer, max_grad_norm=None) - # A new successful implicit window starts a new optimizer step. - engine.forward_backward(window_a, _identity_loss) + engine.forward_backward([_datum([2])], _identity_loss) engine.optim_step() - assert events == [ - "finalize", - "step", - "zero", - "scheduler", - "finalize", - "step", - "zero", - "scheduler", - ] - - -def test_planned_accumulation_failure_breaks_engine_before_optimizer_step(): - events = [] - - class RecordingSGD(torch.optim.SGD): - def step(self, closure=None): - events.append("step") - return super().step(closure) - - model = ScaleModel() - optimizer = RecordingSGD(model.parameters(), lr=0.1) - engine = Engine(model, device="cpu", optimizers=optimizer) - window_a = [_datum([1])] - window_b = [_datum([3])] - engine.begin_accumulation([window_a, window_b]) - engine.forward_backward(window_a, _identity_loss) - - def failing_loss(_output, _loss_inputs): - raise ValueError("loss callback failed") - - with pytest.raises(ValueError, match="loss callback failed"): - engine.forward_backward(window_b, failing_loss) + torch.testing.assert_close(model.weight, torch.tensor(0.8)) - with pytest.raises(RuntimeError): + with pytest.raises(RuntimeError, match="already consumed"): engine.optim_step() - with pytest.raises(RuntimeError): - engine.begin_accumulation([window_a]) - with pytest.raises(RuntimeError): - engine.forward_backward(window_b, _identity_loss) - assert events == [] - torch.testing.assert_close(model.weight, torch.tensor(1.0)) - - -@pytest.mark.parametrize("mismatch", ["datum_reference", "weights"]) -def test_planned_accumulation_validates_declared_datums_before_forward(mismatch): - model = ScaleModel() - engine = Engine( - model, - device="cpu", - microbatch_size=2, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([1], [1.0]), _datum([2], [1.0])] - engine.begin_accumulation([window]) - - actual_window = window - if mismatch == "datum_reference": - actual_window = [_datum([1], [1.0]), window[1]] - else: - window[0].loss_fn_inputs["weights"].mul_(2.0) - - with pytest.raises((RuntimeError, ValueError)): - engine.forward_backward(actual_window, _identity_loss) - assert model.forward_calls == 0 - - -def test_planned_accumulation_rejects_microbatch_size_changes_before_forward(): - model = ScaleModel() - engine = Engine( - model, - device="cpu", - microbatch_size=2, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([1]), _datum([2])] - engine.begin_accumulation([window]) - engine.microbatch_size = 1 - - with pytest.raises(RuntimeError, match="microbatch_size changed"): - engine.forward_backward(window, _identity_loss) - - assert model.forward_calls == 0 - - -def test_planned_accumulation_rejects_nonfinal_partial_outer_microbatch(): - model = ScaleModel() - engine = Engine( - model, - device="cpu", - microbatch_size=2, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window_a = [_datum([1])] - window_b = [_datum([2])] - - with pytest.raises(ValueError, match="microbatch|aligned|divisible"): - engine.begin_accumulation([window_a, window_b]) - assert model.forward_calls == 0 + engine.forward_backward([_datum([5])], _identity_loss) + engine.optim_step() + torch.testing.assert_close(model.weight, torch.tensor(0.3)) def test_forward_backward_groups_flat_datums_by_microbatch_size(): @@ -1022,23 +734,73 @@ def test_raw_thd_packed_collater_is_prepared_by_context_parallel_sharder(): model = ScaleModel() model.backend = SimpleNamespace(attn="te") seen = {} + expected_routes = torch.tensor([[[0], [10]], [[1], [11]], [[2], [12]]], dtype=torch.int16) + + @contextmanager + def batch_context(model_inputs, loss_fn_inputs): + """Check the Engine's final THD ``[tokens, layers, topk]`` route layout. + + Args: + model_inputs: Final THD model mapping with ``input_ids [tokens]``. + loss_fn_inputs: Final THD side channels with + ``routed_experts [tokens, layers, topk]``. + + Yields: + ``None`` while the packed forward and backward execute. + """ + assert model_inputs["input_ids"].shape == (3,) + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_routes) + yield def loss_fn(output, inputs): + """Return one loss per local THD token. + + Args: + output: Model values with shape ``[tokens]``. + inputs: Loss mapping containing ``weights [tokens]`` and + ``routed_experts [tokens, layers, topk]``. + + Returns: + Unreduced token loss with shape ``[tokens]``. + """ seen.update(inputs) assert inputs["weights"].shape == output.shape == (3,) return output + datums = [ + Datum( + input_ids=torch.tensor([1, 2]), + loss_fn_inputs={ + "weights": torch.ones(2), + "routed_experts": expected_routes[:2], + }, + loss_fn_input_layouts={"routed_experts": LossInputLayout.PER_TOKEN}, + loss_fn_input_pad_values={"routed_experts": -1}, + ), + Datum( + input_ids=torch.tensor([3]), + loss_fn_inputs={ + "weights": torch.ones(1), + "routed_experts": expected_routes[2:], + }, + loss_fn_input_layouts={"routed_experts": LossInputLayout.PER_TOKEN}, + loss_fn_input_pad_values={"routed_experts": -1}, + ), + ] + result = Engine( model, device="cpu", microbatch_size=2, collate_fn=partial(collate_datums, packed=True), - ).forward_backward([_datum([1, 2]), _datum([3])], loss_fn) + batch_context_fn=batch_context, + ).forward_backward(datums, loss_fn) assert result.loss.item() == pytest.approx(2.0) assert "seq_lens" not in seen assert "seq_lens_padded" not in seen assert seen["cu_seqlens"].tolist() == [0, 2, 3] + torch.testing.assert_close(seen["routed_experts"], expected_routes) def test_raw_thd_requires_a_thd_capable_context_parallel_sharder(): @@ -1057,6 +819,7 @@ def test_raw_thd_requires_a_thd_capable_context_parallel_sharder(): def test_context_parallel_shards_model_and_rl_loss_inputs_together(): cp_context_active = False + batch_context_active = False @contextmanager def cp_context(): @@ -1082,41 +845,118 @@ def prepare_model_inputs_for_cp(self, batch, *, num_chunks): } def forward(self, input_ids, *args, **kwargs): + """Scale one CP-local token shard while both runtime contexts are active. + + Args: + input_ids: CP-local token ids with shape ``[B, S_local]``. + *args: Additional positional model arguments forwarded to the toy model. + **kwargs: Additional model inputs; loss-only routing data must be absent. + + Returns: + Per-token values with shape ``[B, S_local]``. + """ assert cp_context_active + assert batch_context_active + assert "routed_experts" not in kwargs assert input_ids.tolist() == [[5, 6, 9, 9]] return super().forward(input_ids, *args, **kwargs) model = CPModel() mesh = _CPMesh(size=2, rank=1) mesh_context = SimpleNamespace(pp_size=1, cp_size=2, device_mesh=mesh) + token_ids = torch.arange(6, dtype=torch.int16) + routed_experts = torch.stack((token_ids, token_ids + 10), dim=-1).unsqueeze(0).unsqueeze(-1) + expected_local_routes = torch.tensor( + [[[[4], [14]], [[5], [15]], [[-1], [-1]], [[-1], [-1]]]], + dtype=torch.int16, + ) datum = Datum( model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4, 5, 6]])}, loss_fn_inputs={ "target_tokens": torch.tensor([[11, 12, 13, 14, 15, 16]]), "weights": torch.ones(1, 6), "advantages": torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]]), + "routed_experts": routed_experts, + }, + loss_fn_input_layouts={ + "target_tokens": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + "advantages": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, }, + loss_fn_input_pad_values={"routed_experts": -1}, ) + + @contextmanager + def batch_context(model_inputs, loss_fn_inputs): + """Validate CP-local routing data and mark the batch context active. + + Args: + model_inputs: CP-local model mapping with ``input_ids`` shaped + ``[B, S_local]``. + loss_fn_inputs: CP-local loss mapping with ``routed_experts`` shaped + ``[B, S_local, G, K]``. + + Yields: + None while model forward, loss, and backward execute. + """ + nonlocal batch_context_active + assert cp_context_active + assert not batch_context_active + assert "routed_experts" not in model_inputs + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([[5, 6, 9, 9]])) + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_local_routes) + batch_context_active = True + try: + yield + finally: + batch_context_active = False + engine = Engine( model, device="cpu", mesh_context=mesh_context, collate_fn=collate_prebatched, padding_token_id=9, + batch_context_fn=batch_context, ) # This is a layout-only CPU test with a fake mesh. Distributed CP loss and # gradient scaling are covered separately with a real process group. engine._dp_group_and_size = lambda: (None, 1) engine._gradient_group_and_size = lambda _group, _size: (None, 1) - model.weight.register_hook( - lambda grad: grad if cp_context_active else pytest.fail("CP context ended before backward") - ) + + def assert_backward_context(grad): + """Require both runtime contexts while propagating a scalar parameter gradient. + + Args: + grad: Scalar gradient for the toy model weight. + + Returns: + The unchanged scalar gradient. + """ + if not cp_context_active or not batch_context_active: + pytest.fail("CP or batch context ended before backward") + return grad + + model.weight.register_hook(assert_backward_context) def loss_fn(output, inputs): + """Return CP-local per-token losses after checking aligned side inputs. + + Args: + output: CP-local model values with shape ``[B, S_local]``. + inputs: CP-local loss mapping whose token tensors share + ``[B, S_local]`` and whose routes add trailing ``[G, K]`` axes. + + Returns: + Per-token losses with shape ``[B, S_local]``. + """ assert cp_context_active + assert batch_context_active assert inputs["target_tokens"].tolist() == [[15, 16, 0, 0]] assert inputs["weights"].tolist() == [[1.0, 1.0, 0.0, 0.0]] torch.testing.assert_close(inputs["advantages"], torch.tensor([[0.5, 0.6, 0.0, 0.0]])) + torch.testing.assert_close(inputs["routed_experts"], expected_local_routes) return output result = engine.forward_backward([datum], loss_fn) @@ -1125,6 +965,7 @@ def loss_fn(output, inputs): assert model.weight.grad.item() == pytest.approx(11 / 6) assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(2.0) assert not cp_context_active + assert not batch_context_active def test_context_parallel_rejects_legacy_non_token_weights(): @@ -1757,7 +1598,6 @@ def loss_fn(output, inputs): assert model.weight.grad.item() == pytest.approx(4.5) assert result.loss_fn_outputs == [] assert pipeline.step_calls == 2 - assert pipeline.finalize_backward_calls == [False, True] # The fake schedule performs and counts every backward, then returns None. # A second Engine-owned backward would either fail or change these counts. assert pipeline.backward_calls == backward_calls == 4 @@ -1940,142 +1780,6 @@ def prepare_final(parts, *, pp_enabled): assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) -def test_pipeline_planned_multi_call_matches_one_window_with_unequal_weights(): - window_a = [_datum([[2], [100]], [[1.0], [0.0]])] - window_b = [_datum([[4], [8]], [[0.5], [1.5]])] - - reference_model = ScaleModel() - reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) - reference_pipeline = _FakeAutoPipeline(reference_model, num_microbatches=2) - reference_engine = Engine( - reference_pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - collate_fn=collate_prebatched, - optimizers=reference_optimizer, - max_grad_norm=None, - ) - reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) - reference_grad = reference_model.weight.grad.detach().clone() - reference_step = reference_engine.optim_step() - - planned_model = ScaleModel() - planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) - planned_pipeline = _FakeAutoPipeline(planned_model, num_microbatches=2) - planned_engine = Engine( - planned_pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - collate_fn=collate_prebatched, - optimizers=planned_optimizer, - max_grad_norm=None, - ) - planned_engine.begin_accumulation([window_a, window_b]) - result_a = planned_engine.forward_backward(window_a, _identity_loss) - result_b = planned_engine.forward_backward(window_b, _identity_loss) - planned_grad = planned_model.weight.grad.detach().clone() - planned_step = planned_engine.optim_step() - - assert result_a.loss_sum.item() == pytest.approx(2.0) - assert result_a.weight_sum.item() == pytest.approx(1.0) - assert result_a.loss.item() == pytest.approx(2.0) - assert result_b.loss_sum.item() == pytest.approx(14.0) - assert result_b.weight_sum.item() == pytest.approx(2.0) - assert result_b.loss.item() == pytest.approx(7.0) - assert reference_result.loss_sum.item() == pytest.approx(16.0) - assert reference_result.weight_sum.item() == pytest.approx(3.0) - assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) - assert reference_pipeline.step_calls == planned_pipeline.step_calls == 2 - assert reference_pipeline.finalize_backward_calls == [False, True] - assert planned_pipeline.finalize_backward_calls == [False, True] - assert reference_pipeline.backward_calls == planned_pipeline.backward_calls == 4 - torch.testing.assert_close(planned_grad, reference_grad) - torch.testing.assert_close(planned_step.grad_norm, reference_step.grad_norm) - torch.testing.assert_close(planned_model.weight, reference_model.weight) - - -def test_pipeline_planned_accumulation_uses_one_lifecycle_and_all_inner_microbatches(monkeypatch): - events = [] - - monkeypatch.setattr( - engine_module, - "prepare_for_grad_accumulation", - lambda _parts, *, pp_enabled: events.append(f"prepare:{pp_enabled}"), - ) - monkeypatch.setattr( - engine_module, - "prepare_for_final_backward", - lambda _parts, *, pp_enabled: events.append(f"final:{pp_enabled}"), - ) - monkeypatch.setattr(engine_module, "prepare_after_first_microbatch", lambda: events.append("after_first")) - monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) - - model = _MainAndAuxScaleModel() - pipeline = _FakeAutoPipeline(model, num_microbatches=2, events=events) - engine = Engine( - pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - collate_fn=collate_prebatched, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - max_grad_norm=None, - ) - window_a = [_datum([[100], [200]], [[0.0], [0.0]])] - window_b = [_datum([[3], [5]], [[1.0], [1.0]])] - - engine.begin_accumulation([window_a, window_b]) - result_a = engine.forward_backward(window_a, _identity_loss) - result_b = engine.forward_backward(window_b, _identity_loss) - - assert events == ["prepare:True", "step", "after_first", "final:True", "step"] - assert result_a.loss.item() == pytest.approx(0.0) - assert result_b.loss.item() == pytest.approx(4.0) - assert pipeline.step_calls == 2 - assert pipeline.backward_calls == 4 - assert pipeline.finalize_backward_calls == [False, True] - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.25) - assert model.main_weight.grad.item() == pytest.approx(4.0) - assert model.aux_weight.grad.item() == pytest.approx(1.0) - - -def test_pipeline_planned_accumulation_enforces_state_and_retries_a_failed_fence(): - model = ScaleModel() - pipeline = _FakeAutoPipeline(model, num_microbatches=2) - engine = Engine( - pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - collate_fn=collate_prebatched, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - max_grad_norm=None, - ) - window_a = [_datum([[1], [3]])] - window_b = [_datum([[5], [7]])] - engine.begin_accumulation([window_a, window_b]) - engine.forward_backward(window_a, _identity_loss) - - with pytest.raises(RuntimeError, match="finish every forward_backward"): - engine.optim_step() - torch.testing.assert_close(model.weight, torch.tensor(1.0)) - - engine.forward_backward(window_b, _identity_loss) - expected_grad = model.weight.grad.detach().clone() - - def fail_before_step(): - raise ValueError("checkpoint staging failed") - - with pytest.raises(ValueError, match="checkpoint staging failed"): - engine.optim_step(before_optimizer_step=fail_before_step) - torch.testing.assert_close(model.weight, torch.tensor(1.0)) - torch.testing.assert_close(model.weight.grad, expected_grad) - - engine.optim_step() - torch.testing.assert_close(model.weight, 1.0 - 0.1 * expected_grad) - assert model.weight.grad is None - with pytest.raises(RuntimeError, match="already consumed"): - engine.optim_step() - - def test_pipeline_outputs_follow_logical_microbatch_order(): model = ScaleModel() pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) @@ -2261,18 +1965,102 @@ def strip_layout_metadata(items): assert pipeline.eval_calls == 0 -def test_pipeline_final_thd_routes_three_plus_one_datums_and_restores_output_order(): +def test_custom_collater_cannot_strip_loss_input_pad_values(): model = ScaleModel() - model.backend = SimpleNamespace(attn="te") - pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) - datums = _packed_layout_datums([1, 1, 2, 4]) - seen = [] + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2]])}, + loss_fn_inputs={ + "weights": torch.ones(1, 2), + "routed_experts": torch.tensor([[[[0]], [[-1]]]], dtype=torch.int16), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) - def loss_with_outputs(output, loss_inputs): - torch.testing.assert_close(loss_inputs["global_coefficients"], torch.tensor([701.0, 709.0])) - assert loss_inputs["global_coefficients"].shape == (2,) - torch.testing.assert_close(loss_inputs["advantages"], output * 10) - torch.testing.assert_close(loss_inputs["old_logprobs"], -output) + def strip_pad_metadata(items): + """Return prebatched tensors while intentionally dropping padding metadata. + + Args: + items: One prebatched Datum containing ``input_ids [B, S]`` and + ``routed_experts [B, S, G, K]``. + + Returns: + Model inputs and a plain loss-input mapping without pad metadata. + """ + model_inputs, loss_inputs = collate_prebatched(items) + return model_inputs, dict(loss_inputs) + + with pytest.raises(ValueError, match="CollatedLossInputs.*loss_fn_input_pad_values"): + Engine(model, device="cpu", collate_fn=strip_pad_metadata).forward([datum], _identity_loss) + + assert model.forward_calls == 0 + + +@pytest.mark.parametrize("metadata_error", ["missing_pad_value", "changed_pad_value", "changed_layout"]) +def test_typed_custom_collater_cannot_change_datum_metadata(metadata_error): + model = ScaleModel() + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2]])}, + loss_fn_inputs={ + "weights": torch.ones(1, 2), + "routed_experts": torch.tensor([[[[0]], [[-1]]]], dtype=torch.int16), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + def change_metadata(items): + """Return prebatched tensors with deliberately inconsistent metadata. + + Args: + items: One prebatched Datum containing ``input_ids [B, S]`` and + ``routed_experts [B, S, G, K]``. + + Returns: + Model inputs and typed loss inputs whose layout or padding + metadata disagrees with the Datum declaration. + """ + model_inputs, loss_inputs = collate_prebatched(items) + layouts = dict(loss_inputs.layouts) + pad_values = dict(loss_inputs.pad_values) + if metadata_error == "missing_pad_value": + pad_values.clear() + elif metadata_error == "changed_pad_value": + pad_values["routed_experts"] = 0 + else: + layouts["routed_experts"] = LossInputLayout.REPLICATED + pad_values.clear() + return model_inputs, CollatedLossInputs( + loss_inputs, + layouts=layouts, + item_to_datum=loss_inputs.item_to_datum, + pad_values=pad_values, + ) + + with pytest.raises(ValueError, match="CollatedLossInputs (pad value|layout).+disagrees"): + Engine(model, device="cpu", collate_fn=change_metadata).forward([datum], _identity_loss) + + assert model.forward_calls == 0 + + +def test_pipeline_final_thd_routes_three_plus_one_datums_and_restores_output_order(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) + datums = _packed_layout_datums([1, 1, 2, 4]) + seen = [] + + def loss_with_outputs(output, loss_inputs): + torch.testing.assert_close(loss_inputs["global_coefficients"], torch.tensor([701.0, 709.0])) + assert loss_inputs["global_coefficients"].shape == (2,) + torch.testing.assert_close(loss_inputs["advantages"], output * 10) + torch.testing.assert_close(loss_inputs["old_logprobs"], -output) sample_ids = loss_inputs["sample_id"].clone() seen.append(sample_ids.tolist()) return output, [{"sample_id": sample_id} for sample_id in sample_ids] @@ -2571,6 +2359,124 @@ def loss_fn(output, inputs): torch.testing.assert_close(seen[1][2], torch.tensor([[-5.0, -8.0]])) +def test_eager_te_thd_batch_context_preserves_trailing_routes_through_cp(monkeypatch): + class MockTex: + @staticmethod + def thd_get_partitioned_indices(_cu_seqlens, total_tokens, _cp_size, _cp_rank): + assert total_tokens == 8 + return torch.tensor([0, 3, 4, 7]) + + monkeypatch.setitem(sys.modules, "transformer_engine_torch", MockTex) + monkeypatch.setattr(dist, "get_rank", lambda group=None: 0) + + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + mesh = _CPMesh(size=2, rank=0) + mesh_context = SimpleNamespace(pp_size=1, cp_size=2, device_mesh=mesh, process_group=None) + routes = torch.tensor( + [ + [[[0, 1]], [[-1, -1]], [[4, 5]], [[2, 3]]], + [[[4, 5]], [[0, 1]], [[2, 3]], [[-1, -1]]], + ], + dtype=torch.int16, + ) + expected_routes = routes.reshape(8, 1, 2).index_select(0, torch.tensor([0, 3, 4, 7])) + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]), + "position_ids": torch.arange(4).expand(2, -1), + "seq_lens": torch.tensor([[4], [4]]), + "seq_lens_padded": torch.tensor([[4], [4]]), + "qkv_format": "thd", + }, + loss_fn_inputs={"weights": torch.ones(2, 4), "routed_experts": routes}, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + @contextmanager + def batch_context(model_inputs, loss_fn_inputs): + """Check final THD token order after a fake two-way CP partition. + + Args: + model_inputs: CP-local final THD mapping with ``input_ids [tokens]``. + loss_fn_inputs: CP-local side channels with + ``routed_experts [tokens, layers, topk]``. + + Yields: + ``None`` while the eager forward and backward execute. + """ + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([1, 4, 5, 8])) + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_routes) + yield + + def loss_fn(output, loss_fn_inputs): + """Return one loss per CP-local THD token. + + Args: + output: CP-local model values with shape ``[tokens]``. + loss_fn_inputs: CP-local side channels with leading ``[tokens]``. + + Returns: + Per-token losses with shape ``[tokens]``. + """ + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_routes) + return output + + engine = Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + batch_context_fn=batch_context, + ) + engine._dp_group_and_size = lambda: (None, 1) + engine._gradient_group_and_size = lambda _group, _size: (None, 1) + + result = engine.forward_backward([datum], loss_fn) + + assert result.loss_sum.item() == pytest.approx(18.0) + torch.testing.assert_close(expected_routes[-1], torch.tensor([[-1, -1]], dtype=torch.int16)) + + +def test_pipeline_thd_rejects_nonzero_per_token_pad_sentinel_before_schedule(): + model = ScaleModel() + model.backend = SimpleNamespace(attn="te") + pipeline = _FakeAutoPipeline(model, num_microbatches=2) + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2], [3, 4]]), + "position_ids": torch.arange(2).expand(2, -1), + "seq_lens": torch.tensor([[2], [2]]), + "seq_lens_padded": torch.tensor([[2], [2]]), + "qkv_format": "thd", + }, + loss_fn_inputs={ + "weights": torch.ones(2, 2), + "routed_experts": torch.zeros(2, 2, 1, 1, dtype=torch.int16), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + with pytest.raises(NotImplementedError, match="packed pipeline.*nonzero PER_TOKEN pad sentinels"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + ).forward_backward([datum], _identity_loss) + + assert pipeline.step_calls == 0 + assert model.forward_calls == 0 + + def test_pipeline_rejects_schedule_gradient_scaling_before_forward(): model = ScaleModel() pipeline = _FakeAutoPipeline(model, scale_grads=True) @@ -2641,6 +2547,285 @@ def loss_fn(output, _inputs): assert not active +@pytest.mark.parametrize("execution", ["forward", "forward_backward"]) +def test_batch_context_receives_prepared_side_inputs_without_forwarding_them_to_model(execution): + active = False + legacy_active = False + events = [] + routed_experts = torch.tensor( + [[[[0], [2]], [[1], [-1]], [[3], [0]]]], + dtype=torch.int16, + ) + + @contextmanager + def legacy_context(model_inputs): + """Model-only context remains outermost for FP8-style callers. + + Args: + model_inputs: Prepared model mapping with ``input_ids [B, S]``. + + Yields: + ``None`` while both the batch context and model execution run. + """ + nonlocal legacy_active + assert not legacy_active + assert not active + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([[1, 2, 3]])) + legacy_active = True + events.append("legacy_enter") + try: + yield + finally: + assert not active + events.append("legacy_exit") + legacy_active = False + + @contextmanager + def batch_context(model_inputs, loss_fn_inputs): + """Expose prepared routing tensors only through the batch context. + + Args: + model_inputs: Prepared model mapping with ``input_ids`` shaped ``[B, S]``. + loss_fn_inputs: Prepared loss mapping with ``routed_experts`` shaped + ``[B, S, G, K]``. + + Yields: + None while the eager model, loss, and optional backward execute. + """ + nonlocal active + assert legacy_active + assert not active + assert "routed_experts" not in model_inputs + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([[1, 2, 3]])) + torch.testing.assert_close(loss_fn_inputs["routed_experts"], routed_experts) + active = True + events.append("context_enter") + try: + yield + finally: + events.append("context_exit") + active = False + + class ContextModel(ScaleModel): + def forward(self, input_ids, **kwargs): + """Scale prepared token ids without receiving loss-only routes. + + Args: + input_ids: Prepared token ids with shape ``[B, S]``. + **kwargs: Additional model inputs; ``routed_experts`` must be absent. + + Returns: + Per-token values with shape ``[B, S]``. + """ + assert active + assert "routed_experts" not in kwargs + events.append("model") + return super().forward(input_ids, **kwargs) + + model = ContextModel() + + def record_backward(grad): + """Record backward while preserving the scalar parameter gradient. + + Args: + grad: Scalar gradient for the toy model weight. + + Returns: + The unchanged scalar gradient. + """ + if not active or not legacy_active: + pytest.fail("legacy or batch context ended before backward") + events.append("backward") + return grad + + model.weight.register_hook(record_backward) + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3]])}, + loss_fn_inputs={ + "weights": torch.ones(1, 3), + "routed_experts": routed_experts, + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + def loss_fn(output, loss_fn_inputs): + """Return prepared per-token values as losses after checking routes. + + Args: + output: Model values with shape ``[B, S]``. + loss_fn_inputs: Loss mapping with routes shaped ``[B, S, G, K]``. + + Returns: + Per-token losses with shape ``[B, S]``. + """ + assert active + torch.testing.assert_close(loss_fn_inputs["routed_experts"], routed_experts) + events.append("loss") + return output + + result = getattr( + Engine( + model, + device="cpu", + collate_fn=collate_prebatched, + context_fn=legacy_context, + batch_context_fn=batch_context, + ), + execution, + )([datum], loss_fn) + + assert result.loss_sum.item() == pytest.approx(6.0) + expected_events = ["legacy_enter", "context_enter", "model", "loss"] + if execution == "forward_backward": + expected_events.append("backward") + expected_events.append("context_exit") + expected_events.append("legacy_exit") + assert events == expected_events + assert not active + assert not legacy_active + + +def test_batch_context_covers_activation_checkpoint_recompute(): + active = False + routed_experts = torch.tensor([[[[1]], [[-1]], [[2]]]], dtype=torch.int16) + + @contextmanager + def batch_context(_model_inputs, loss_fn_inputs): + """Keep replay side inputs installed through checkpoint recomputation. + + Args: + _model_inputs: Prepared model mapping with token axes ``[B, S]``. + loss_fn_inputs: Prepared loss mapping with routes shaped ``[B, S, G, K]``. + + Yields: + None while forward, loss, backward, and recomputation execute. + """ + nonlocal active + torch.testing.assert_close(loss_fn_inputs["routed_experts"], routed_experts) + active = True + try: + yield + finally: + active = False + + class CheckpointModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.tensor(1.0)) + self.block_calls = 0 + + def _block(self, value): + """Apply a recomputed nonlinear transform to one token batch. + + Args: + value: Floating-point token values with shape ``[B, S]``. + + Returns: + Nonlinear activations with shape ``[B, S]``. + """ + assert active + self.block_calls += 1 + return torch.sin(value * self.weight) + + def forward(self, input_ids, **kwargs): + """Checkpoint a token transform without forwarding replay side inputs. + + Args: + input_ids: Prepared token ids with shape ``[B, S]``. + **kwargs: Additional model inputs; ``routed_experts`` must be absent. + + Returns: + Checkpointed activations with shape ``[B, S]``. + """ + assert active + assert "routed_experts" not in kwargs + return checkpoint(self._block, input_ids.to(torch.float32), use_reentrant=False) + + model = CheckpointModel() + + def assert_checkpoint_backward_context(grad): + """Require the batch context while propagating a scalar checkpoint gradient. + + Args: + grad: Scalar gradient for the checkpointed model weight. + + Returns: + The unchanged scalar gradient. + """ + if not active: + pytest.fail("context ended before backward") + return grad + + model.weight.register_hook(assert_checkpoint_backward_context) + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3]])}, + loss_fn_inputs={ + "weights": torch.ones(1, 3), + "routed_experts": routed_experts, + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + def loss_fn(output, _loss_fn_inputs): + """Build per-token nonlinear losses that force checkpoint recomputation. + + Args: + output: Checkpointed activations with shape ``[B, S]``. + _loss_fn_inputs: Prepared loss mapping with token axes ``[B, S]``. + + Returns: + Squared per-token losses with shape ``[B, S]``. + """ + return output.square() + + Engine( + model, + device="cpu", + collate_fn=collate_prebatched, + batch_context_fn=batch_context, + ).forward_backward([datum], loss_fn) + + assert model.block_calls == 2 + assert not active + + +def test_batch_context_rejects_auto_pipeline_at_construction(): + model = ScaleModel() + pipeline = _FakeAutoPipeline(model) + + def unused_batch_context(_model_inputs, _loss_fn_inputs): + """Reject any attempted construction of a pipeline batch context. + + Args: + _model_inputs: Pipeline model mapping; this callback must not receive it. + _loss_fn_inputs: Pipeline loss mapping; this callback must not receive it. + + Returns: + No context because AutoPipeline must fail during Engine construction. + """ + pytest.fail("pipeline batch context must not be constructed") + + with pytest.raises(NotImplementedError, match="batch_context_fn.*eager PP=1"): + Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + batch_context_fn=unused_batch_context, + ) + + assert pipeline.step_calls == 0 + assert pipeline.eval_calls == 0 + assert model.forward_calls == 0 + + def test_window_sets_the_same_moe_aux_scale_as_the_recipes(monkeypatch): monkeypatch.setattr(MoEAuxLossAutoScaler, "main_loss_backward_scale", None) @@ -2803,47 +2988,6 @@ def test_gradient_reduction_mode_finds_summed_backend_below_ordinary_wrapper(mon assert model.weight.grad.item() == pytest.approx(3.5) -def test_summed_gradient_mode_planned_accumulation_matches_one_window(monkeypatch): - window_a = [_datum([2, 100], [1.0, 0.0])] - window_b = [_datum([4, 8], [0.5, 1.5])] - - reference_model = ScaleModel() - reference_model.calculate_per_token_loss = True - reference_optimizer = torch.optim.SGD(reference_model.parameters(), lr=0.1) - reference_engine = Engine( - reference_model, - device="cpu", - optimizers=reference_optimizer, - max_grad_norm=None, - ) - _configure_fake_gradient_group(reference_engine, monkeypatch, group_size=4) - reference_result = reference_engine.forward_backward(window_a + window_b, _identity_loss) - reference_grad = reference_model.weight.grad.detach().clone() - reference_engine.optim_step() - - planned_model = ScaleModel() - planned_model.calculate_per_token_loss = True - planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) - planned_engine = Engine( - planned_model, - device="cpu", - optimizers=planned_optimizer, - max_grad_norm=None, - ) - _configure_fake_gradient_group(planned_engine, monkeypatch, group_size=4) - planned_engine.begin_accumulation([window_a, window_b]) - result_a = planned_engine.forward_backward(window_a, _identity_loss) - result_b = planned_engine.forward_backward(window_b, _identity_loss) - - assert result_a.weight_sum.item() == pytest.approx(1.0) - assert result_b.weight_sum.item() == pytest.approx(2.0) - assert reference_result.weight_sum.item() == pytest.approx(3.0) - assert reference_result.loss.item() == pytest.approx(16.0 / 3.0) - torch.testing.assert_close(planned_model.weight.grad, reference_grad) - planned_engine.optim_step() - torch.testing.assert_close(planned_model.weight, reference_model.weight) - - def test_summed_gradient_mode_zero_weight_window_keeps_graph_connected_zero(monkeypatch): model = ScaleModel() model.calculate_per_token_loss = True @@ -2957,6 +3101,145 @@ def _distributed_worker(rank: int, world_size: int, init_file: str) -> None: dist.destroy_process_group() +def _batch_context_cp_worker(rank: int, world_size: int, init_file: str) -> None: + """Validate prepared replay side inputs on a real two-rank Gloo CP mesh. + + Args: + rank: Current Gloo process rank. + world_size: Number of context-parallel ranks. + init_file: Shared file-store path used to initialize the process group. + """ + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=20), + ) + try: + mesh_context = MeshContext.build( + MegatronFSDPConfig(), + ParallelismSizes(dp_size=1, cp_size=world_size), + world_size=world_size, + ) + batch_context_active = False + + class ContextCPModel(_DistributedCPModel): + def forward(self, input_ids, **kwargs): + """Scale one real CP token shard while the batch context is active. + + Args: + input_ids: CP-local token ids with shape ``[B, S_local]``. + **kwargs: Additional model inputs; replay routes must be absent. + + Returns: + Per-token values with shape ``[B, S_local]``. + """ + assert batch_context_active + assert "routed_experts" not in kwargs + return super().forward(input_ids, **kwargs) + + model = _DDPWithCP(ContextCPModel()) + token_ids = torch.arange(6, dtype=torch.int16) + routed_experts = torch.stack((token_ids, token_ids + 10), dim=-1).unsqueeze(0).unsqueeze(-1) + if rank == 0: + expected_input_ids = torch.tensor([[1, 2, 3, 4]]) + expected_routes = routed_experts[:, :4] + else: + expected_input_ids = torch.tensor([[5, 6, 0, 0]]) + expected_routes = torch.cat( + ( + routed_experts[:, 4:], + torch.full((1, 2, 2, 1), -1, dtype=routed_experts.dtype), + ), + dim=1, + ) + + @contextmanager + def batch_context(model_inputs, loss_fn_inputs): + """Check rank-local tokens and routes around forward and backward. + + Args: + model_inputs: CP-local model mapping with ``input_ids`` shaped + ``[B, S_local]``. + loss_fn_inputs: CP-local loss mapping with routes shaped + ``[B, S_local, G, K]``. + + Yields: + None while the DDP-wrapped CP model and loss execute. + """ + nonlocal batch_context_active + assert not batch_context_active + assert "routed_experts" not in model_inputs + torch.testing.assert_close(model_inputs["input_ids"], expected_input_ids) + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_routes) + batch_context_active = True + try: + yield + finally: + batch_context_active = False + + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2, 3, 4, 5, 6]])}, + loss_fn_inputs={ + "weights": torch.ones(1, 6), + "routed_experts": routed_experts, + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "routed_experts": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"routed_experts": -1}, + ) + + def assert_distributed_backward_context(grad): + """Require the batch context while propagating a scalar DDP gradient. + + Args: + grad: Scalar gradient for the DDP-wrapped toy model weight. + + Returns: + The unchanged scalar gradient. + """ + if not batch_context_active: + pytest.fail("batch context ended before backward") + return grad + + model.module.weight.register_hook(assert_distributed_backward_context) + + def loss_fn(output, loss_fn_inputs): + """Return CP-local token losses after checking the route shard. + + Args: + output: CP-local model values with shape ``[B, S_local]``. + loss_fn_inputs: CP-local loss mapping with routes shaped + ``[B, S_local, G, K]``. + + Returns: + Per-token losses with shape ``[B, S_local]``. + """ + assert batch_context_active + torch.testing.assert_close(loss_fn_inputs["routed_experts"], expected_routes) + return output + + result = Engine( + model, + device="cpu", + mesh_context=mesh_context, + collate_fn=collate_prebatched, + batch_context_fn=batch_context, + ).forward_backward([datum], loss_fn) + + assert result.loss.item() == pytest.approx(3.5) + assert result.loss_sum.item() == pytest.approx(21.0) + assert result.weight_sum.item() == pytest.approx(6.0) + assert model.module.weight.grad.item() == pytest.approx(3.5) + assert not batch_context_active + dist.barrier() + finally: + dist.destroy_process_group() + + def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size: int) -> None: dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size) try: @@ -3014,10 +3297,9 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size assert forward_result.weight_sum.item() == pytest.approx(4.0) assert model.module.weight.grad is None - # Planned accumulation must use the sum of the two DP-only - # denominators while letting CP ranks contribute disjoint numerators. - # The two calls deliberately have different weight sums, and DP ranks - # deliberately own different amounts of supervision when dp_size=2. + # One complete window must use its DP-only denominator while letting CP + # ranks contribute disjoint numerators. DP ranks deliberately own + # different amounts of supervision when dp_size=2. if dp_size == 1: window_a = [ Datum( @@ -3064,31 +3346,26 @@ def _context_parallel_worker(rank: int, world_size: int, init_file: str, dp_size expected_a_sum, expected_a_weight = 20.0, 3.0 expected_b_sum, expected_b_weight = 37.0, 4.0 - planned_model = _DDPWithCP(_DistributedCPModel()) - planned_optimizer = torch.optim.SGD(planned_model.parameters(), lr=0.1) - planned_engine = Engine( - planned_model, + averaged_model = _DDPWithCP(_DistributedCPModel()) + averaged_optimizer = torch.optim.SGD(averaged_model.parameters(), lr=0.1) + averaged_engine = Engine( + averaged_model, device="cpu", mesh_context=mesh_context, collate_fn=collate_prebatched, - optimizers=planned_optimizer, + optimizers=averaged_optimizer, max_grad_norm=None, ) - planned_engine.begin_accumulation([window_a, window_b]) - result_a = planned_engine.forward_backward(window_a, _identity_loss) - result_b = planned_engine.forward_backward(window_b, _identity_loss) - - assert result_a.loss_sum.item() == pytest.approx(expected_a_sum) - assert result_a.weight_sum.item() == pytest.approx(expected_a_weight) - assert result_a.loss.item() == pytest.approx(expected_a_sum / expected_a_weight) - assert result_b.loss_sum.item() == pytest.approx(expected_b_sum) - assert result_b.weight_sum.item() == pytest.approx(expected_b_weight) - assert result_b.loss.item() == pytest.approx(expected_b_sum / expected_b_weight) + averaged_result = averaged_engine.forward_backward(window_a + window_b, _identity_loss) + + assert averaged_result.loss_sum.item() == pytest.approx(expected_a_sum + expected_b_sum) + assert averaged_result.weight_sum.item() == pytest.approx(expected_a_weight + expected_b_weight) expected_global_mean = (expected_a_sum + expected_b_sum) / (expected_a_weight + expected_b_weight) - assert planned_model.module.weight.grad.item() == pytest.approx(expected_global_mean) + assert averaged_result.loss.item() == pytest.approx(expected_global_mean) + assert averaged_model.module.weight.grad.item() == pytest.approx(expected_global_mean) - planned_engine.optim_step() - assert planned_model.module.weight.item() == pytest.approx(1.0 - 0.1 * expected_global_mean) + averaged_engine.optim_step() + assert averaged_model.module.weight.item() == pytest.approx(1.0 - 0.1 * expected_global_mean) # Mimic MegatronFSDP calculate_per_token_loss=True with a SUM hook over # the same DP-CP gradient group. Engine must not compensate for an @@ -3112,18 +3389,14 @@ def sum_gradient(gradient): optimizers=summed_optimizer, max_grad_norm=None, ) - summed_engine.begin_accumulation([window_a, window_b]) - summed_result_a = summed_engine.forward_backward(window_a, _identity_loss) - summed_result_b = summed_engine.forward_backward(window_b, _identity_loss) - - torch.testing.assert_close(summed_result_a.loss_sum, result_a.loss_sum) - torch.testing.assert_close(summed_result_a.weight_sum, result_a.weight_sum) - torch.testing.assert_close(summed_result_b.loss_sum, result_b.loss_sum) - torch.testing.assert_close(summed_result_b.weight_sum, result_b.weight_sum) + summed_result = summed_engine.forward_backward(window_a + window_b, _identity_loss) + + torch.testing.assert_close(summed_result.loss_sum, averaged_result.loss_sum) + torch.testing.assert_close(summed_result.weight_sum, averaged_result.weight_sum) assert summed_model.weight.grad.item() == pytest.approx(expected_global_mean) summed_engine.optim_step() - torch.testing.assert_close(summed_model.weight, planned_model.module.weight) + torch.testing.assert_close(summed_model.weight, averaged_model.module.weight) finally: dist.destroy_process_group() @@ -3249,218 +3522,6 @@ def loss_with_rank_local_output_error(output, _loss_inputs): dist.destroy_process_group() -def _planned_accumulation_validation_consensus_worker(rank: int, world_size: int, init_file: str) -> None: - dist.init_process_group( - "gloo", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - timeout=timedelta(seconds=20), - ) - try: - model = nn.parallel.DistributedDataParallel(ScaleModel()) - engine = Engine( - model, - device="cpu", - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([rank + 1])] - planned_windows = [window] if rank == 0 else [window, [_datum([rank + 2])]] - - with pytest.raises((RuntimeError, ValueError)): - engine.begin_accumulation(planned_windows) - assert model.module.forward_calls == 0 - dist.barrier() - - model = nn.parallel.DistributedDataParallel(ScaleModel()) - engine = Engine( - model, - device="cpu", - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([rank + 1])] - engine.begin_accumulation([window]) - if rank == 0: - window[0].loss_fn_inputs["weights"].mul_(2.0) - - # Rank 1 must observe rank 0's local plan-validation failure instead of - # entering DDP forward and hanging on a different collective. - with pytest.raises((RuntimeError, ValueError)): - engine.forward_backward(window, _identity_loss) - assert model.module.forward_calls == 0 - finally: - dist.destroy_process_group() - - -def _model_parallel_planned_output_error_worker(rank: int, world_size: int, init_file: str) -> None: - dist.init_process_group( - "gloo", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - timeout=timedelta(seconds=20), - ) - try: - mesh_context = MeshContext.build( - MegatronFSDPConfig(), - ParallelismSizes(dp_size=1, tp_size=world_size), - world_size=world_size, - ) - model = ScaleModel() - engine = Engine( - model, - device="cpu", - mesh_context=mesh_context, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([1, 2])] - engine.begin_accumulation([window]) - - def loss_with_rank_local_output_error(output, _loss_inputs): - return output, ([{}, {}] if rank == 0 else [{}]) - - expected = "one mapping per Datum" if rank == 0 else "another model-parallel rank" - with pytest.raises((ValueError, RuntimeError), match=expected): - engine.forward_backward(window, loss_with_rank_local_output_error) - - # Both model-parallel ranks complete backward, then enter the same - # terminal state even though only rank 0 owns the bad output contract. - assert model.weight.grad is not None - assert engine._accumulation_state is not None - assert engine._accumulation_state.status == "broken" - with pytest.raises(RuntimeError, match="broken"): - engine.optim_step() - - model = ScaleModel() - engine = Engine( - model, - device="cpu", - mesh_context=mesh_context, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - window = [_datum([1, 2])] - engine.begin_accumulation([window]) - - def rank_local_loss_failure(output, _loss_inputs): - if rank == 0: - raise ValueError("rank-local loss failure") - return output - - expected = "rank-local loss failure" if rank == 0 else "another model-parallel rank" - with pytest.raises((ValueError, RuntimeError), match=expected): - engine.forward_backward(window, rank_local_loss_failure) - - # Loss/shape errors are agreed before backward, so no peer enters a - # TP/EP/DP backward collective while another exits locally. - assert model.weight.grad is None - assert engine._accumulation_state is not None - assert engine._accumulation_state.status == "broken" - finally: - dist.destroy_process_group() - - -def _model_parallel_optimizer_fence_worker(rank: int, world_size: int, init_file: str) -> None: - dist.init_process_group( - "gloo", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - timeout=timedelta(seconds=20), - ) - real_finalize = engine_module.scale_grads_and_clip_grad_norm - try: - mesh_context = MeshContext.build( - MegatronFSDPConfig(), - ParallelismSizes(dp_size=1, tp_size=world_size), - world_size=world_size, - ) - steps = [] - - class RecordingSGD(torch.optim.SGD): - def step(self, closure=None): - steps.append("step") - return super().step(closure) - - model = ScaleModel() - engine = Engine( - model, - device="cpu", - mesh_context=mesh_context, - optimizers=RecordingSGD(model.parameters(), lr=0.1), - max_grad_norm=None, - ) - window = [_datum([1, 2])] - engine.begin_accumulation([window]) - engine.forward_backward(window, _identity_loss) - - finalize_calls = [] - - def recording_finalize(**_kwargs): - finalize_calls.append("finalize") - return torch.tensor(1.5) - - engine_module.scale_grads_and_clip_grad_norm = recording_finalize - - def rank_local_fence(): - if rank == 0: - raise ValueError("rank-local fence failure") - - expected = "rank-local fence failure" if rank == 0 else "another model-parallel rank" - with pytest.raises((ValueError, RuntimeError), match=expected): - engine.optim_step(before_optimizer_step=rank_local_fence) - - assert finalize_calls == ["finalize"] - assert steps == [] - torch.testing.assert_close(model.weight, torch.tensor(1.0)) - assert model.weight.grad is not None - - result = engine.optim_step() - assert finalize_calls == ["finalize"] - assert steps == ["step"] - torch.testing.assert_close(result.grad_norm, torch.tensor(1.5)) - torch.testing.assert_close(model.weight, torch.tensor(0.85)) - finally: - engine_module.scale_grads_and_clip_grad_norm = real_finalize - dist.destroy_process_group() - - -def _planned_cp_preflight_consensus_worker(rank: int, world_size: int, init_file: str) -> None: - dist.init_process_group( - "gloo", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - timeout=timedelta(seconds=20), - ) - try: - mesh_context = MeshContext.build( - MegatronFSDPConfig(), - ParallelismSizes(dp_size=2, cp_size=2), - world_size=world_size, - ) - model = ScaleModel() - engine = Engine( - model, - device="cpu", - mesh_context=mesh_context, - optimizers=torch.optim.SGD(model.parameters(), lr=0.1), - ) - # Only one CP subgroup disagrees. The full-control error reduction must - # stop all four ranks before any rank enters the later DP all-reduce. - weight = 2.0 if rank == 0 else 1.0 - window = [_datum([rank + 1], [weight])] - - with pytest.raises((ValueError, RuntimeError), match="context-parallel"): - engine.begin_accumulation([window]) - - assert engine._accumulation_state is None - assert model.forward_calls == 0 - assert model.weight.grad is None - dist.barrier() - finally: - dist.destroy_process_group() - - def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): mp.spawn( _distributed_worker, @@ -3470,6 +3531,15 @@ def test_data_parallel_window_uses_global_numerator_and_denominator(tmp_path): ) +def test_batch_context_receives_real_context_parallel_route_shards(tmp_path): + mp.spawn( + _batch_context_cp_worker, + args=(2, str(tmp_path / "engine_batch_context_cp_init")), + nprocs=2, + join=True, + ) + + def test_context_parallel_window_uses_dp_denominator_and_dp_cp_gradient_sum(tmp_path): mp.spawn( _context_parallel_worker, @@ -3513,39 +3583,3 @@ def test_data_parallel_output_errors_propagate_after_backward_without_hanging(tm nprocs=2, join=True, ) - - -def test_planned_accumulation_validation_errors_reach_every_data_rank(tmp_path): - mp.spawn( - _planned_accumulation_validation_consensus_worker, - args=(2, str(tmp_path / "engine_planned_validation_init")), - nprocs=2, - join=True, - ) - - -def test_planned_output_error_breaks_every_model_parallel_rank_after_backward(tmp_path): - mp.spawn( - _model_parallel_planned_output_error_worker, - args=(2, str(tmp_path / "engine_model_parallel_output_init")), - nprocs=2, - join=True, - ) - - -def test_optimizer_fence_failure_reaches_every_model_parallel_rank_and_can_retry(tmp_path): - mp.spawn( - _model_parallel_optimizer_fence_worker, - args=(2, str(tmp_path / "engine_model_parallel_fence_init")), - nprocs=2, - join=True, - ) - - -def test_planned_cp_preflight_failure_reaches_all_dp_cp_ranks(tmp_path): - mp.spawn( - _planned_cp_preflight_consensus_worker, - args=(4, str(tmp_path / "engine_planned_cp_preflight_init")), - nprocs=4, - join=True, - ) From ae130d0bd3d01a3c54a2b89b3913e7554e7aa669 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 01:18:24 -0700 Subject: [PATCH 23/34] feat(engine): support batch contexts with pipeline parallelism Signed-off-by: HuiyingLi --- .../distributed/pipelining/autopipeline.py | 99 ++++++++++++++- nemo_automodel/engine/__init__.py | 24 ++-- .../pipelining/test_autopipeline.py | 85 +++++++++++++ tests/unit_tests/test_engine.py | 114 +++++++++++++----- 4 files changed, 278 insertions(+), 44 deletions(-) diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index 00a3802307..77b3eb74fa 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -14,8 +14,10 @@ import inspect import logging +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass -from typing import Any, Callable, Literal +from typing import Any, Literal import torch import torch.nn as nn @@ -34,6 +36,64 @@ logger = logging.getLogger(__name__) +BatchContextFactory = Callable[[], AbstractContextManager[Any]] + + +@contextmanager +def _stage_batch_contexts( + stages: list[PipelineStage], + batch_context_fns: list[BatchContextFactory], +) -> Iterator[None]: + """Apply a repeatable batch context to each pipeline-stage phase. + + PyTorch schedules identify the logical microbatch with the first argument + to every stage ``forward_one_chunk``/``backward_*_one_chunk`` call. Wrapping + those instance methods keeps the context aligned when schedules interleave + microbatches or activation checkpointing recomputes a forward during + backward. Original instance state is restored even when the schedule fails. + + Args: + stages: Rank-local pipeline stages owned by one AutoPipeline. + batch_context_fns: One zero-argument context factory per logical + pipeline microbatch. + + Yields: + ``None`` while the pipeline schedule executes. + """ + missing = object() + originals: list[tuple[PipelineStage, str, Any]] = [] + method_names = ("forward_one_chunk", "backward_one_chunk", "backward_weight_one_chunk") + try: + for stage in stages: + for method_name in method_names: + method = getattr(stage, method_name, None) + if not callable(method): + continue + previous = getattr(stage, "__dict__", {}).get(method_name, missing) + + def wrapped( + chunk_id: int, + *args: Any, + _method: Callable[..., Any] = method, + **kwargs: Any, + ) -> Any: + if isinstance(chunk_id, bool) or not isinstance(chunk_id, int): + raise TypeError(f"pipeline microbatch id must be an integer, got {chunk_id!r}") + if chunk_id < 0 or chunk_id >= len(batch_context_fns): + raise IndexError(f"pipeline microbatch id {chunk_id} is outside [0, {len(batch_context_fns)})") + with batch_context_fns[chunk_id](): + return _method(chunk_id, *args, **kwargs) + + originals.append((stage, method_name, previous)) + setattr(stage, method_name, wrapped) + yield + finally: + for stage, method_name, previous in reversed(originals): + if previous is missing: + delattr(stage, method_name) + else: + setattr(stage, method_name, previous) + @dataclass class PipelineInfo: @@ -259,6 +319,7 @@ def step_microbatches( loss_fn: Callable[[Any, int], Any], losses: list[torch.Tensor] | None = None, return_outputs: bool = False, + batch_context_fns: list[BatchContextFactory] | None = None, ) -> Any: """Run a schedule step over already prepared model microbatches. @@ -276,6 +337,10 @@ def step_microbatches( losses: Mutable list populated by the schedule on the last stage. return_outputs: Whether the last stage returns merged model outputs when supported by the installed PyTorch version. + batch_context_fns: Optional repeatable context factory for each + logical microbatch. Each factory separately covers that + microbatch's stage forward and backward phases. The caller + wraps its loss callback with the same factory. Returns: The value returned by the underlying PyTorch pipeline schedule. @@ -286,6 +351,7 @@ def step_microbatches( losses=losses, return_outputs=return_outputs, schedule_method="step", + batch_context_fns=batch_context_fns, ) def eval_microbatches( @@ -295,6 +361,7 @@ def eval_microbatches( loss_fn: Callable[[Any, int], Any], losses: list[torch.Tensor] | None = None, return_outputs: bool = True, + batch_context_fns: list[BatchContextFactory] | None = None, ) -> Any: """Run forward-only evaluation over already prepared model microbatches. @@ -313,6 +380,10 @@ def eval_microbatches( losses: Mutable list populated by the schedule on the last stage. return_outputs: Whether the last stage returns merged model outputs when supported by the installed PyTorch version. + batch_context_fns: Optional repeatable context factory for each + logical microbatch. Each factory separately covers that + microbatch's stage forward phase. The caller wraps its loss + callback with the same factory. Returns: The value returned by the underlying PyTorch pipeline schedule. @@ -327,6 +398,7 @@ def eval_microbatches( losses=losses, return_outputs=return_outputs, schedule_method="eval", + batch_context_fns=batch_context_fns, ) def _run_prepared_microbatches( @@ -337,6 +409,7 @@ def _run_prepared_microbatches( losses: list[torch.Tensor] | None, return_outputs: bool, schedule_method: Literal["step", "eval"], + batch_context_fns: list[BatchContextFactory] | None, ) -> Any: """Run one schedule method with an exact prepared-microbatch split.""" schedule = self._info.schedule @@ -344,6 +417,10 @@ def _run_prepared_microbatches( raise RuntimeError("AutoPipeline.build() must be called before running a prepared PP schedule") if len(model_inputs) != self.num_microbatches: raise ValueError(f"Expected {self.num_microbatches} model input microbatches, got {len(model_inputs)}") + if batch_context_fns is not None and len(batch_context_fns) != self.num_microbatches: + raise ValueError( + f"Expected {self.num_microbatches} pipeline batch context factories, got {len(batch_context_fns)}" + ) model_args_chunks: list[tuple[Any, ...]] = [] model_kwargs_chunks: list[dict[str, Any]] = [] @@ -385,11 +462,21 @@ def indexed_loss(output: Any, microbatch_id: torch.Tensor) -> Any: schedule._split_inputs = lambda _args, _kwargs=None: (model_args_chunks, model_kwargs_chunks) schedule._loss_fn = indexed_loss try: - return run_schedule( - target=torch.arange(self.num_microbatches, device=self.device), - losses=losses, - **schedule_options, - ) + if batch_context_fns is None: + return run_schedule( + target=torch.arange(self.num_microbatches, device=self.device), + losses=losses, + **schedule_options, + ) + stages = self._info.stages + if stages is None: + raise RuntimeError("AutoPipeline.build() must create pipeline stages before using batch contexts") + with _stage_batch_contexts(stages, batch_context_fns): + return run_schedule( + target=torch.arange(self.num_microbatches, device=self.device), + losses=losses, + **schedule_options, + ) finally: schedule._loss_fn = previous_loss_fn schedule._split_inputs = previous_split_inputs diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py index bb6a979f73..f65d3c7427 100644 --- a/nemo_automodel/engine/__init__.py +++ b/nemo_automodel/engine/__init__.py @@ -21,6 +21,7 @@ from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass +from functools import partial from math import prod from typing import Any, TypeVar @@ -361,8 +362,9 @@ class Engine: batch_context_fn: Creates an optional context from the CP/packing- prepared model-input and loss/side-channel mappings. It covers eager model forward, loss, backward, and activation-checkpoint - recomputation. AutoPipeline is intentionally unsupported until it - can select one context payload per inner pipeline microbatch. + recomputation. With AutoPipeline, the factory is selected by the + logical inner-microbatch id and entered separately for every stage + forward, loss, and backward phase; it must therefore be repeatable. Rank-local callback failures are process-fatal in distributed execution, like failures from ``context_fn`` or model forward. defer_fsdp_grad_sync: Defer FSDP/DDP gradient synchronization until the @@ -417,11 +419,6 @@ def __init__( if isinstance(mtp_ignore_index, bool) or not isinstance(mtp_ignore_index, int): raise ValueError(f"mtp_ignore_index must be an integer, got {mtp_ignore_index!r}") self.pipeline = model if isinstance(model, AutoPipeline) else None - if self.pipeline is not None and batch_context_fn is not None: - raise NotImplementedError( - "batch_context_fn currently supports eager PP=1 execution only; " - "AutoPipeline needs per-inner-microbatch context routing" - ) self.model_parts = model.parts if self.pipeline is not None else [model] self.model = self.model_parts[0] self._summed_gradient_reduction = _resolve_summed_gradient_reduction(self.model_parts) @@ -1396,12 +1393,17 @@ def _pipeline_execute( batch_size=batch_size, ) ) + batch_context_fns: list[Callable[[], AbstractContextManager[Any]]] = [ + partial(self.batch_context_fn, model_inputs_mb, loss_inputs_mb) + for model_inputs_mb, loss_inputs_mb in zip(model_microbatches, loss_microbatches) + ] def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: loss_inputs_mb = loss_microbatches[microbatch_index] - numerator, batch_outputs, output_parse_error = _parse_loss_result( - loss_fn(output, loss_inputs_mb), loss_inputs_mb["weights"] - ) + with batch_context_fns[microbatch_index](): + numerator, batch_outputs, output_parse_error = _parse_loss_result( + loss_fn(output, loss_inputs_mb), loss_inputs_mb["weights"] + ) if loss_called_by_microbatch[microbatch_index]: raise RuntimeError(f"pipeline evaluated loss_fn twice for microbatch {microbatch_index}") loss_called_by_microbatch[microbatch_index] = True @@ -1419,6 +1421,7 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: loss_fn=pipeline_loss, losses=losses, return_outputs=False, + batch_context_fns=batch_context_fns, ) else: self.pipeline.step_microbatches( @@ -1426,6 +1429,7 @@ def pipeline_loss(output: Any, microbatch_index: int) -> torch.Tensor: loss_fn=pipeline_loss, losses=losses, return_outputs=False, + batch_context_fns=batch_context_fns, ) outputs: list[dict[str, Any]] = [] diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index deae8ac24d..02105e808a 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -13,6 +13,7 @@ # limitations under the License. import types +from contextlib import contextmanager from unittest.mock import Mock import pytest @@ -342,6 +343,39 @@ class _NoEvalSchedule(_LegacyStepSchedule): eval = None +class _ContextAwareStage: + def __init__(self, active_microbatch, events): + self.active_microbatch = active_microbatch + self.events = events + + def _record(self, phase, chunk_id): + assert self.active_microbatch == [chunk_id] + self.events.append((phase, chunk_id)) + + def forward_one_chunk(self, chunk_id, *_args, **_kwargs): + self._record("forward", chunk_id) + + def backward_one_chunk(self, chunk_id, *_args, **_kwargs): + self._record("backward", chunk_id) + + def backward_weight_one_chunk(self, chunk_id, *_args, **_kwargs): + self._record("weight_backward", chunk_id) + + +class _ContextAwareSchedule(_PreparedMicrobatchSchedule): + def __init__(self, stage): + super().__init__() + self.stage = stage + + def _run_schedule(self, *args, target=None, losses=None, return_outputs=True, **kwargs): + self.args_split, self.kwargs_split = self._split_inputs(args, kwargs) + self.stage.forward_one_chunk(1, (), {}) + self.stage.forward_one_chunk(0, (), {}) + self.stage.backward_one_chunk(0, None) + self.stage.backward_weight_one_chunk(1) + return "schedule-result" + + class TestAutoPipelinePreparedMicrobatches: def _pipeline_with_parts( self, @@ -404,6 +438,57 @@ def loss_fn(output, index): assert model_inputs[0]["input_ids"] is input_ids[0] assert model_inputs[1]["input_ids"] is input_ids[1] + def test_batch_context_follows_interleaved_stage_chunk_ids_and_restores_methods(self): + active_microbatch = [] + events = [] + stage = _ContextAwareStage(active_microbatch, events) + schedule = _ContextAwareSchedule(stage) + ap = self._pipeline_with_parts(nn.Module(), schedule=schedule) + ap._info.stages = [stage] + original_methods = { + name: getattr(stage, name) + for name in ("forward_one_chunk", "backward_one_chunk", "backward_weight_one_chunk") + } + + def make_context(index): + @contextmanager + def batch_context(): + assert not active_microbatch + active_microbatch.append(index) + events.append(("enter", index)) + try: + yield + finally: + events.append(("exit", index)) + active_microbatch.clear() + + return batch_context + + result = ap.step_microbatches( + [{"input_ids": torch.zeros(1, 8)} for _ in range(2)], + loss_fn=Mock(), + batch_context_fns=[make_context(0), make_context(1)], + ) + + assert result == "schedule-result" + assert events == [ + ("enter", 1), + ("forward", 1), + ("exit", 1), + ("enter", 0), + ("forward", 0), + ("exit", 0), + ("enter", 0), + ("backward", 0), + ("exit", 0), + ("enter", 1), + ("weight_backward", 1), + ("exit", 1), + ] + assert not active_microbatch + for name, original in original_methods.items(): + assert getattr(stage, name) == original + def test_step_microbatches_omits_primary_args_on_nonfirst_stage(self): schedule = _PreparedMicrobatchSchedule() ap = self._pipeline_with_parts(nn.Module(), schedule=schedule, has_first_stage=False) diff --git a/tests/unit_tests/test_engine.py b/tests/unit_tests/test_engine.py index 2db4af4b58..c77d3a618f 100644 --- a/tests/unit_tests/test_engine.py +++ b/tests/unit_tests/test_engine.py @@ -16,7 +16,7 @@ import pickle import sys -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from datetime import timedelta from functools import partial from types import SimpleNamespace @@ -157,7 +157,15 @@ def update_seq_len(self, seq_len, *, microbatch_size=None, input_tensor=None): self.updated_microbatch_sizes.append(microbatch_size) self.updated_input_shapes.append(tuple(input_tensor.shape) if input_tensor is not None else None) - def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): + def step_microbatches( + self, + model_inputs, + *, + loss_fn, + losses, + return_outputs, + batch_context_fns=None, + ): assert return_outputs is False assert len(model_inputs) == self.num_microbatches self.prepared_inputs.append(model_inputs) @@ -169,13 +177,24 @@ def step_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): inputs = dict(model_inputs[index]) primary_name = "inputs_embeds" if "inputs_embeds" in inputs else "input_ids" primary = inputs.pop(primary_name) - output = self.compute_model(primary, **inputs) + context_fn = nullcontext if batch_context_fns is None else batch_context_fns[index] + with context_fn(): + output = self.compute_model(primary, **inputs) scaled_loss = loss_fn(output, index) self.callback_losses.append(scaled_loss.detach()) - scaled_loss.backward() + with context_fn(): + scaled_loss.backward() self.backward_calls += 1 - def eval_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): + def eval_microbatches( + self, + model_inputs, + *, + loss_fn, + losses, + return_outputs, + batch_context_fns=None, + ): assert return_outputs is False assert len(model_inputs) == self.num_microbatches self.prepared_inputs.append(model_inputs) @@ -187,7 +206,9 @@ def eval_microbatches(self, model_inputs, *, loss_fn, losses, return_outputs): inputs = dict(model_inputs[index]) primary_name = "inputs_embeds" if "inputs_embeds" in inputs else "input_ids" primary = inputs.pop(primary_name) - output = self.compute_model(primary, **inputs) + context_fn = nullcontext if batch_context_fns is None else batch_context_fns[index] + with context_fn(): + output = self.compute_model(primary, **inputs) loss = loss_fn(output, index) self.callback_losses.append(loss.detach()) @@ -2797,33 +2818,70 @@ def loss_fn(output, _loss_fn_inputs): assert not active -def test_batch_context_rejects_auto_pipeline_at_construction(): - model = ScaleModel() - pipeline = _FakeAutoPipeline(model) +def test_pipeline_batch_context_tracks_each_microbatch_through_checkpoint_backward(): + active_route = [] + block_routes = [] + context_routes = [] - def unused_batch_context(_model_inputs, _loss_fn_inputs): - """Reject any attempted construction of a pipeline batch context. + class CheckpointedPipelineModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.tensor(1.0)) - Args: - _model_inputs: Pipeline model mapping; this callback must not receive it. - _loss_fn_inputs: Pipeline loss mapping; this callback must not receive it. + def block(self, values, weight): + assert len(active_route) == 1 + block_routes.append(active_route[0]) + return values * weight - Returns: - No context because AutoPipeline must fail during Engine construction. - """ - pytest.fail("pipeline batch context must not be constructed") + def forward(self, input_ids, **_kwargs): + expected_route = int((input_ids[0, 0] - 1) // 2) + assert active_route == [expected_route] + return checkpoint(self.block, input_ids.float(), self.weight, use_reentrant=False) - with pytest.raises(NotImplementedError, match="batch_context_fn.*eager PP=1"): - Engine( - pipeline, - device="cpu", - mesh_context=_pipeline_mesh_context(), - batch_context_fn=unused_batch_context, - ) + model = CheckpointedPipelineModel() + pipeline = _FakeAutoPipeline(model, num_microbatches=2, callback_order=[1, 0]) - assert pipeline.step_calls == 0 - assert pipeline.eval_calls == 0 - assert model.forward_calls == 0 + @contextmanager + def batch_context(_model_inputs, loss_fn_inputs): + route = int(loss_fn_inputs["route_id"].flatten()[0]) + assert not active_route + active_route.append(route) + context_routes.append(route) + try: + yield + finally: + active_route.clear() + + datum = Datum( + model_inputs={"input_ids": torch.tensor([[1, 2], [3, 4]])}, + loss_fn_inputs={ + "weights": torch.ones(2, 2), + "route_id": torch.tensor([[0, 0], [1, 1]]), + }, + loss_fn_input_layouts={ + "weights": LossInputLayout.PER_TOKEN, + "route_id": LossInputLayout.PER_TOKEN, + }, + ) + + def loss_fn(output, loss_fn_inputs): + route = int(loss_fn_inputs["route_id"].flatten()[0]) + assert active_route == [route] + return output + + result = Engine( + pipeline, + device="cpu", + mesh_context=_pipeline_mesh_context(), + collate_fn=collate_prebatched, + batch_context_fn=batch_context, + ).forward_backward([datum], loss_fn) + + assert result.loss.item() == pytest.approx(2.5) + assert model.weight.grad.item() == pytest.approx(2.5) + assert block_routes == [1, 1, 0, 0] + assert context_routes == [1, 1, 1, 0, 0, 0] + assert not active_route def test_window_sets_the_same_moe_aux_scale_as_the_recipes(monkeypatch): From 596876101f71fecb9d93d8d0c3d67ff9b48cae8e Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 01:24:01 -0700 Subject: [PATCH 24/34] feat(moe): bind routing replay across pipeline model parts Signed-off-by: HuiyingLi --- .../components/moe/router_replay.py | 176 ++++++++++-------- tests/unit_tests/moe/test_router_replay.py | 28 +++ 2 files changed, 122 insertions(+), 82 deletions(-) diff --git a/nemo_automodel/components/moe/router_replay.py b/nemo_automodel/components/moe/router_replay.py index 445290fb14..c636bf0584 100644 --- a/nemo_automodel/components/moe/router_replay.py +++ b/nemo_automodel/components/moe/router_replay.py @@ -51,13 +51,13 @@ model in the same process. For rollout-provided routing, :class:`RouterReplayAdapter` is the preferred -eager Engine interface. It maps global decoder-layer ids without using the -registry and consumes ``routed_experts`` only after the Engine has applied its -packing and context-parallel token transform. This adapter intentionally -supports PP=1 only. +Engine interface. It maps global decoder-layer ids without using the registry +and consumes ``routed_experts`` only after the Engine has applied its packing +and context-parallel token transform. It can bind either one eager model or the +rank-local model parts of an AutoPipeline. """ -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from enum import Enum @@ -262,7 +262,7 @@ class _RouterReplayBinding: class RouterReplayAdapter: - """Bind rollout routes to one model's MoE gates for eager Engine execution. + """Bind rollout routes to model-scoped MoE gates for Engine execution. The adapter is both the model-aware route formatter and the callable passed as ``Engine(batch_context_fn=...)``. It deliberately ignores the legacy @@ -270,102 +270,112 @@ class RouterReplayAdapter: hybrid MoE stacks and multiple models in one process remain unambiguous. Do not nest the legacy process-global ``record``/``replay`` contexts around an active adapter context when the same gate handles are registered. - AutoPipeline is not supported by this adapter; the Engine rejects that - combination before execution. Args: - model: Complete PP=1 model. Its primary decoder blocks must expose - numeric child ids or a consistent integer ``layer_idx``. A block - may contain at most one module with a ``router_replay`` slot. + model: One complete eager model, or the rank-local model parts of an + AutoPipeline. Primary decoder blocks must expose numeric child ids + or a consistent integer ``layer_idx``. A block may contain at most + one module with a ``router_replay`` slot. An explicit model-part + sequence may contain no local routed layers (for example an + embedding-only or LM-head-only pipeline rank). """ field_name = "routed_experts" - def __init__(self, model: nn.Module) -> None: - block_root = model - visited_roots: set[int] = set() - blocks: tuple[tuple[nn.Module, str, nn.Module], ...] = () - while id(block_root) not in visited_roots: - visited_roots.add(id(block_root)) - blocks = tuple(iter_transformer_blocks(block_root)) - if blocks: - break - wrapped = getattr(block_root, "module", None) - if not isinstance(wrapped, nn.Module): - break - block_root = wrapped + def __init__(self, model: nn.Module | Sequence[nn.Module]) -> None: + is_model_parts = not isinstance(model, nn.Module) + model_parts = tuple(model) if is_model_parts else (model,) + if not model_parts: + raise ValueError("RouterReplayAdapter requires at least one model part") + if not all(isinstance(part, nn.Module) for part in model_parts): + raise TypeError("RouterReplayAdapter model parts must all be nn.Module instances") bindings: list[_RouterReplayBinding] = [] seen_replays: set[int] = set() - for _parent, child_name, block in blocks: - slots = [module for module in block.modules() if hasattr(module, "router_replay")] - if not slots: - continue - if len(slots) != 1: - raise ValueError( - f"decoder block {child_name!r} has {len(slots)} router_replay slots; " - "RouterReplayAdapter requires one gate per routed layer" - ) + for model_part in model_parts: + block_root = model_part + visited_roots: set[int] = set() + blocks: tuple[tuple[nn.Module, str, nn.Module], ...] = () + while id(block_root) not in visited_roots: + visited_roots.add(id(block_root)) + blocks = tuple(iter_transformer_blocks(block_root)) + if blocks: + break + wrapped = getattr(block_root, "module", None) + if not isinstance(wrapped, nn.Module): + break + block_root = wrapped + + for _parent, child_name, block in blocks: + slots = [module for module in block.modules() if hasattr(module, "router_replay")] + if not slots: + continue + if len(slots) != 1: + raise ValueError( + f"decoder block {child_name!r} has {len(slots)} router_replay slots; " + "RouterReplayAdapter requires one gate per routed layer" + ) - declared_ids = { - layer_idx - for module in block.modules() - if isinstance((layer_idx := getattr(module, "layer_idx", None)), int) - and not isinstance(layer_idx, bool) - } - if len(declared_ids) > 1: - raise ValueError( - f"decoder block {child_name!r} contains conflicting layer_idx values {sorted(declared_ids)}" - ) - child_idx = int(child_name) if child_name.isdecimal() else None - declared_idx = next(iter(declared_ids), None) - if child_idx is not None and declared_idx is not None and child_idx != declared_idx: - raise ValueError(f"decoder block key {child_idx} disagrees with its layer_idx {declared_idx}") - layer_idx = declared_idx if declared_idx is not None else child_idx - if layer_idx is None: - raise ValueError(f"cannot resolve the global layer id for routed decoder block {child_name!r}") - if layer_idx < 0: - raise ValueError(f"routed decoder block {child_name!r} has negative layer_idx {layer_idx}") - - gate = slots[0] - topk = getattr(gate, "topk", None) - if not isinstance(topk, int) or isinstance(topk, bool) or topk <= 0: - raise ValueError(f"decoder block {layer_idx} replay gate must expose a positive integer topk") - num_experts = getattr(gate, "n_experts", getattr(gate, "num_experts", None)) - if num_experts is not None and ( - not isinstance(num_experts, int) or isinstance(num_experts, bool) or num_experts <= 0 - ): - raise ValueError(f"decoder block {layer_idx} replay gate has invalid expert count {num_experts!r}") - if getattr(gate, "use_routing_core", False): - raise RuntimeError( - "RouterReplayAdapter is incompatible with partial MoE router CUDA graphs; " - "disable the 'moe_router' graph module before enabling routing replay" - ) - replay = gate.router_replay - if replay is None: - replay = RouterReplay(register=False) - gate.router_replay = replay - if not isinstance(replay, RouterReplay): - raise TypeError( - f"decoder block {layer_idx} router_replay must be RouterReplay or None, got {type(replay).__name__}" - ) - if id(replay) in seen_replays: - raise ValueError("one RouterReplay handle is attached to more than one decoder block") - seen_replays.add(id(replay)) - bindings.append(_RouterReplayBinding(layer_idx, replay, topk, num_experts)) + declared_ids = { + layer_idx + for module in block.modules() + if isinstance((layer_idx := getattr(module, "layer_idx", None)), int) + and not isinstance(layer_idx, bool) + } + if len(declared_ids) > 1: + raise ValueError( + f"decoder block {child_name!r} contains conflicting layer_idx values {sorted(declared_ids)}" + ) + child_idx = int(child_name) if child_name.isdecimal() else None + declared_idx = next(iter(declared_ids), None) + if child_idx is not None and declared_idx is not None and child_idx != declared_idx: + raise ValueError(f"decoder block key {child_idx} disagrees with its layer_idx {declared_idx}") + layer_idx = declared_idx if declared_idx is not None else child_idx + if layer_idx is None: + raise ValueError(f"cannot resolve the global layer id for routed decoder block {child_name!r}") + if layer_idx < 0: + raise ValueError(f"routed decoder block {child_name!r} has negative layer_idx {layer_idx}") + + gate = slots[0] + topk = getattr(gate, "topk", None) + if not isinstance(topk, int) or isinstance(topk, bool) or topk <= 0: + raise ValueError(f"decoder block {layer_idx} replay gate must expose a positive integer topk") + num_experts = getattr(gate, "n_experts", getattr(gate, "num_experts", None)) + if num_experts is not None and ( + not isinstance(num_experts, int) or isinstance(num_experts, bool) or num_experts <= 0 + ): + raise ValueError(f"decoder block {layer_idx} replay gate has invalid expert count {num_experts!r}") + if getattr(gate, "use_routing_core", False): + raise RuntimeError( + "RouterReplayAdapter is incompatible with partial MoE router CUDA graphs; " + "disable the 'moe_router' graph module before enabling routing replay" + ) + replay = gate.router_replay + if replay is None: + replay = RouterReplay(register=False) + gate.router_replay = replay + if not isinstance(replay, RouterReplay): + raise TypeError( + f"decoder block {layer_idx} router_replay must be RouterReplay or None, " + f"got {type(replay).__name__}" + ) + if id(replay) in seen_replays: + raise ValueError("one RouterReplay handle is attached to more than one decoder block") + seen_replays.add(id(replay)) + bindings.append(_RouterReplayBinding(layer_idx, replay, topk, num_experts)) - if not bindings: + if not bindings and not is_model_parts: raise ValueError("RouterReplayAdapter found no MoE gate with a router_replay slot in the primary decoder") bindings.sort(key=lambda binding: binding.layer_idx) layer_ids = [binding.layer_idx for binding in bindings] if len(set(layer_ids)) != len(layer_ids): raise ValueError(f"multiple replay gates map to the same global decoder layer: {layer_ids}") topks = {binding.topk for binding in bindings} - if len(topks) != 1: + if bindings and len(topks) != 1: raise ValueError(f"all replay gates must use one topk, got {sorted(topks)}") self._bindings = tuple(bindings) self._layer_ids = tuple(layer_ids) - self._topk = next(iter(topks)) + self._topk = next(iter(topks), None) self._topology_tensors: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} @property @@ -424,6 +434,8 @@ def __call__( self._validate_integral_routes(routed_experts) if routed_experts.ndim < 3: raise ValueError("prepared routed_experts must have token axes followed by [global_layers, topk]") + if not self._bindings: + return nullcontext() weights = loss_fn_inputs.get("weights") if isinstance(weights, torch.Tensor) and weights.ndim > 0: diff --git a/tests/unit_tests/moe/test_router_replay.py b/tests/unit_tests/moe/test_router_replay.py index 035c763622..51bbce3e11 100644 --- a/tests/unit_tests/moe/test_router_replay.py +++ b/tests/unit_tests/moe/test_router_replay.py @@ -391,6 +391,34 @@ def __init__(self, module): assert isinstance(_adapter_gate(model, 1).router_replay, RouterReplay) +def test_adapter_binds_sparse_layers_across_pipeline_model_parts(): + first_part = _AdapterModel(num_layers=2, routed_layers=(1,)) + last_part = _AdapterModel(num_layers=5, routed_layers=(3,)) + adapter = RouterReplayAdapter([first_part, last_part]) + routes = torch.full((2, 5, 2), -1, dtype=torch.int16) + routes[:, 1] = torch.tensor([[1, 2], [3, 4]], dtype=torch.int16) + routes[:, 3] = torch.tensor([[5, 6], [7, 0]], dtype=torch.int16) + + assert adapter.layer_ids == (1, 3) + with adapter( + {"input_ids": torch.zeros(2, dtype=torch.long)}, + {"routed_experts": routes}, + ): + torch.testing.assert_close(_adapter_gate(first_part, 1).router_replay.target_indices, routes[:, 1]) + torch.testing.assert_close(_adapter_gate(last_part, 3).router_replay.target_indices, routes[:, 3]) + + +def test_adapter_allows_pipeline_rank_without_local_moe_layer(): + adapter = RouterReplayAdapter([_AdapterModel(num_layers=2, routed_layers=())]) + + assert adapter.layer_ids == () + with adapter( + {"input_ids": torch.zeros(2, dtype=torch.long)}, + {"routed_experts": torch.zeros(2, 3, 2, dtype=torch.int16)}, + ): + pass + + def test_adapter_rejects_partial_moe_router_cuda_graph_before_installing_handle(): model = _AdapterModel(num_layers=3, routed_layers=(1,)) gate = _adapter_gate(model, 1) From 35d9475559fa21fc431648d48f89948f5930e065 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 01:25:58 -0700 Subject: [PATCH 25/34] feat(data): support pinning Datum inputs Signed-off-by: HuiyingLi --- nemo_automodel/components/datasets/datum.py | 35 +++++ tests/unit_tests/datasets/test_datum.py | 142 ++++++++++++++++++++ 2 files changed, 177 insertions(+) diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index 7fb8bff6ca..cc08972c6f 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -36,6 +36,21 @@ __all__ = ["CollatedLossInputs", "Datum", "LossInputLayout", "collate_datums"] +def _pin_memory_tree(value: Any) -> Any: + """Pin tensors in the common container shapes accepted by Engine inputs.""" + if isinstance(value, torch.Tensor): + return value.pin_memory() + if isinstance(value, dict): + return {name: _pin_memory_tree(item) for name, item in value.items()} + if isinstance(value, list): + return [_pin_memory_tree(item) for item in value] + if isinstance(value, tuple): + pinned = tuple(_pin_memory_tree(item) for item in value) + return type(value)(*pinned) if hasattr(value, "_fields") else pinned + pin_memory = getattr(value, "pin_memory", None) + return pin_memory() if callable(pin_memory) else value + + class LossInputLayout(str, Enum): """How one loss input relates to the Datums being collated. @@ -255,6 +270,26 @@ def seq_len(self) -> int: return int(value.shape[0]) raise ValueError("cannot infer token length; use a model-specific collate_fn for this Datum") + def pin_memory(self) -> Datum: + """Pin tensor inputs when a DataLoader pins this custom batch object. + + PyTorch delegates custom batch pinning to ``pin_memory()``. Recursively + pin common model-input containers, then update this Datum only after + both input mappings succeed. Non-tensor metadata and loss layout + metadata remain unchanged. + + Returns: + This Datum with model and loss tensors replaced by their + pinned-memory counterparts. + """ + model_inputs = _pin_memory_tree(self.model_inputs) + loss_fn_inputs = _pin_memory_tree(self.loss_fn_inputs) + self.model_inputs.clear() + self.model_inputs.update(model_inputs) + self.loss_fn_inputs.clear() + self.loss_fn_inputs.update(loss_fn_inputs) + return self + def to_features(self, *, ignore_index: int = CROSS_ENTROPY_IGNORE_IDX) -> dict[str, Any]: """Convert one text Datum for the repository's canonical collaters. diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index 8d757e6a2a..7b029a8134 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -14,6 +14,7 @@ import pickle from copy import copy, deepcopy +from unittest.mock import Mock import pytest import torch @@ -76,6 +77,20 @@ def _routed_datums(): # ── Datum ───────────────────────────────────────────────────────────────── +def _clone_instead_of_pinning(tensor: torch.Tensor, _device: str | None = None) -> torch.Tensor: + """Return an observable CPU-safe stand-in for ``Tensor.pin_memory``. + + Args: + tensor: Tensor of arbitrary shape and dtype. + _device: Optional accelerator identifier accepted by PyTorch's pinning + dispatcher. It does not affect this CPU-safe test double. + + Returns: + Tensor of the same shape and dtype with independent storage. + """ + return tensor.clone() + + def test_datum_keeps_old_input_ids_convenience(): d = Datum(input_ids=[1, 2, 3], loss_fn_inputs={"weights": [1, 1, 0]}) assert isinstance(d.input_ids, torch.Tensor) @@ -107,6 +122,133 @@ def test_datum_accepts_model_specific_inputs(): assert datum.seq_len == 2 +def test_datum_pin_memory_handles_molt_prebatched_inputs(monkeypatch): + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[10, 11, 12], [20, 21, 0]]), + "attention_mask": torch.tensor([[1, 1, 1], [1, 1, 0]]), + "position_ids": torch.tensor([[0, 1, 2], [0, 1, 1]]), + }, + loss_fn_inputs={ + "labels": torch.tensor([[11, 12, -100], [21, -100, -100]]), + "weights": torch.tensor([[True, True, False], [True, False, False]]), + }, + loss_fn_input_layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"labels": -100}, + ) + original_model_tensors = dict(datum.model_inputs) + original_loss_tensors = dict(datum.loss_fn_inputs) + original_layouts = datum.loss_fn_input_layouts + original_pad_values = datum.loss_fn_input_pad_values + monkeypatch.setattr(torch.Tensor, "pin_memory", _clone_instead_of_pinning) + + result = datum.pin_memory() + + assert result is datum + for key, original in original_model_tensors.items(): + assert datum.model_inputs[key] is not original + torch.testing.assert_close(datum.model_inputs[key], original) + for key, original in original_loss_tensors.items(): + assert datum.loss_fn_inputs[key] is not original + torch.testing.assert_close(datum.loss_fn_inputs[key], original) + assert datum.loss_fn_input_layouts is original_layouts + assert datum.loss_fn_input_layouts == { + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } + assert datum.loss_fn_input_pad_values is original_pad_values + assert datum.loss_fn_input_pad_values == {"labels": -100} + + +def test_datum_pin_memory_uses_dataloader_recursion_for_model_inputs(monkeypatch): + custom_leaf = Mock() + pinned_custom_leaf = object() + custom_leaf.pin_memory.return_value = pinned_custom_leaf + image = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3) + grid = torch.tensor([[1, 2, 3]]) + auxiliary = torch.tensor([4.0, 5.0]) + datum = Datum( + model_inputs={ + "input_ids": torch.tensor([[1, 2, 3]]), + "media": { + "images": [image, None], + "details": (grid, {"auxiliary": auxiliary}), + "custom": custom_leaf, + }, + "batch_size": 1, + "qkv_format": "bshd", + }, + loss_fn_inputs={"weights": torch.ones(1, 3)}, + loss_fn_input_layouts={"weights": LossInputLayout.PER_TOKEN}, + ) + monkeypatch.setattr(torch.Tensor, "pin_memory", _clone_instead_of_pinning) + + result = datum.pin_memory() + + assert result is datum + assert isinstance(datum.model_inputs["media"], dict) + assert isinstance(datum.model_inputs["media"]["images"], list) + assert isinstance(datum.model_inputs["media"]["details"], tuple) + pinned_image = datum.model_inputs["media"]["images"][0] + pinned_grid = datum.model_inputs["media"]["details"][0] + pinned_auxiliary = datum.model_inputs["media"]["details"][1]["auxiliary"] + for pinned, original in ((pinned_image, image), (pinned_grid, grid), (pinned_auxiliary, auxiliary)): + assert pinned is not original + torch.testing.assert_close(pinned, original) + assert datum.model_inputs["media"]["images"][1] is None + assert datum.model_inputs["batch_size"] == 1 + assert datum.model_inputs["qkv_format"] == "bshd" + assert datum.model_inputs["media"]["custom"] is pinned_custom_leaf + custom_leaf.pin_memory.assert_called_once_with() + + +def test_datum_pin_memory_does_not_partially_commit_on_failure(monkeypatch): + model_tensor = torch.tensor([[1, 2, 3]]) + loss_tensor = torch.ones(1, 3) + datum = Datum( + model_inputs={"input_ids": model_tensor, "metadata": {"source": "molt"}}, + loss_fn_inputs={"weights": loss_tensor}, + loss_fn_input_layouts={"weights": LossInputLayout.PER_TOKEN}, + ) + original_model_inputs = datum.model_inputs + original_loss_fn_inputs = datum.loss_fn_inputs + original_layouts = datum.loss_fn_input_layouts + + def fail_for_loss_tensor(tensor: torch.Tensor, _device: str | None = None) -> torch.Tensor: + """Fail after model-input pinning reaches the loss-input mapping. + + Args: + tensor: Tensor of arbitrary shape and dtype. + _device: Optional accelerator identifier accepted by PyTorch's + pinning dispatcher. + + Returns: + An independent tensor when ``tensor`` is a model input. + + Raises: + RuntimeError: If ``tensor`` is the selected loss input. + """ + if tensor is loss_tensor: + raise RuntimeError("injected pin-memory failure") + return tensor.clone() + + monkeypatch.setattr(torch.Tensor, "pin_memory", fail_for_loss_tensor) + + with pytest.raises(RuntimeError, match="injected pin-memory failure"): + datum.pin_memory() + + assert datum.model_inputs is original_model_inputs + assert datum.loss_fn_inputs is original_loss_fn_inputs + assert datum.model_inputs["input_ids"] is model_tensor + assert datum.model_inputs["metadata"] == {"source": "molt"} + assert datum.loss_fn_inputs["weights"] is loss_tensor + assert datum.loss_fn_input_layouts is original_layouts + assert datum.loss_fn_input_layouts == {"weights": LossInputLayout.PER_TOKEN} + + def test_datum_accepts_optional_loss_input_layouts(): datum = Datum( input_ids=torch.tensor([1, 2]), From fd1077bd1e0ec6a7b409e12e0ae4d2814e9b4010 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 02:26:17 -0700 Subject: [PATCH 26/34] feat(model): add pre-FSDP structure hook Signed-off-by: HuiyingLi --- nemo_automodel/_transformers/auto_model.py | 24 +++ .../_transformers/infrastructure.py | 74 +++++++ .../components/checkpoint/checkpointing.py | 30 +++ .../_transformers/test_auto_model.py | 50 ++++- .../_transformers/test_infrastructure.py | 197 ++++++++++++++++++ .../checkpoint/test_checkpointing.py | 90 ++++++++ 6 files changed, 464 insertions(+), 1 deletion(-) diff --git a/nemo_automodel/_transformers/auto_model.py b/nemo_automodel/_transformers/auto_model.py index c07f35712c..1e2cf3de99 100644 --- a/nemo_automodel/_transformers/auto_model.py +++ b/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, ) diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index a854fc4cf5..4e46917182 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/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) diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py index 41e1c3b831..2e4f65a781 100644 --- a/nemo_automodel/components/checkpoint/checkpointing.py +++ b/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) ) diff --git a/tests/unit_tests/_transformers/test_auto_model.py b/tests/unit_tests/_transformers/test_auto_model.py index f75934dc45..e3acb09100 100644 --- a/tests/unit_tests/_transformers/test_auto_model.py +++ b/tests/unit_tests/_transformers/test_auto_model.py @@ -143,6 +143,44 @@ def test_infrastructure_forwards_frozen_multimodal_sharding_to_moe_parallelizer( class TestFromPretrainedDeviceMesh: + def test_forwards_pre_fsdp_hook_configuration_to_build_model(self): + hook = MagicMock() + task_head_prefixes = ("value_head.",) + + with ( + patch("torch.cuda.current_device", return_value=0), + patch("nemo_automodel._transformers.auto_model.instantiate_infrastructure") as mock_infra, + patch("nemo_automodel._transformers.auto_model.get_hf_config", return_value=MagicMock()), + patch("nemo_automodel._transformers.auto_model.get_is_hf_model", return_value=True), + patch("nemo_automodel._transformers.auto_model.resolve_sdpa_method", return_value=None), + patch.object(NeMoAutoModelForCausalLM, "_build_model", return_value=MagicMock()) as mock_build, + ): + mock_infra.return_value = (None, None, None, None) + + NeMoAutoModelForCausalLM.from_pretrained( + "test-model", + pre_fsdp_hook=hook, + skip_task_head_prefixes_for_base_model=task_head_prefixes, + ) + + assert mock_build.call_args.kwargs["pre_fsdp_hook"] is hook + assert mock_build.call_args.kwargs["skip_task_head_prefixes_for_base_model"] is task_head_prefixes + + def test_from_config_forwards_pre_fsdp_hook_configuration_to_build_model(self): + hook = MagicMock() + task_head_prefixes = ("value_head.",) + + with patch.object(_BaseNeMoAutoModelClass, "_build_model", return_value=MagicMock()) as mock_build: + _BaseNeMoAutoModelClass.from_config( + config=MagicMock(name_or_path="test"), + trust_remote_code=False, + pre_fsdp_hook=hook, + skip_task_head_prefixes_for_base_model=task_head_prefixes, + ) + + assert mock_build.call_args.kwargs["pre_fsdp_hook"] is hook + assert mock_build.call_args.kwargs["skip_task_head_prefixes_for_base_model"] is task_head_prefixes + def test_from_pretrained_accepts_device_mesh_as_topology_shortcut(self): device_mesh = _FakeMesh({MeshAxisName.DP_SHARD: 1, MeshAxisName.CP: 1, MeshAxisName.TP: 1}) sentinel_model = object() @@ -1538,6 +1576,12 @@ def test_retry_succeeds_within_limit(self): """When the retried call succeeds, the model is returned normally.""" build_kwargs, mock_config = self._make_build_kwargs() sentinel_model = MagicMock() + hook = MagicMock() + task_head_prefixes = ("value_head.",) + build_kwargs.update( + pre_fsdp_hook=hook, + skip_task_head_prefixes_for_base_model=task_head_prefixes, + ) with ( patch("nemo_automodel._transformers.auto_model._apply_preload_overrides", return_value=("eager", False)), patch("nemo_automodel._transformers.auto_model._init_model") as mock_init, @@ -1547,7 +1591,9 @@ def test_retry_succeeds_within_limit(self): "nemo_automodel._transformers.capabilities.attach_capabilities_and_validate", return_value=sentinel_model, ), - patch("nemo_automodel._transformers.auto_model.apply_model_infrastructure", return_value=sentinel_model), + patch( + "nemo_automodel._transformers.auto_model.apply_model_infrastructure", return_value=sentinel_model + ) as mock_apply, patch("torch.cuda.current_device", return_value=0), ): mock_init.side_effect = [ @@ -1557,6 +1603,8 @@ def test_retry_succeeds_within_limit(self): result = _BaseNeMoAutoModelClass._build_model(mock_config, **build_kwargs) assert result is sentinel_model assert mock_init.call_count == 2 + assert mock_apply.call_args.kwargs["pre_fsdp_hook"] is hook + assert mock_apply.call_args.kwargs["skip_task_head_prefixes_for_base_model"] is task_head_prefixes def test_build_model_applies_runtime_patches_before_infrastructure(self): """Model runtime hooks run after construction and before sharding/checkpoint infra.""" diff --git a/tests/unit_tests/_transformers/test_infrastructure.py b/tests/unit_tests/_transformers/test_infrastructure.py index 9468c88376..1ec8ca97e3 100644 --- a/tests/unit_tests/_transformers/test_infrastructure.py +++ b/tests/unit_tests/_transformers/test_infrastructure.py @@ -214,6 +214,203 @@ def _run_apply_model_infrastructure(*, is_meta_device, load_base_model, model_wr return result, mock_ckpt +class TestPreFSDPHook: + def test_load_before_shard_runs_hook_after_base_load_and_before_key_snapshot(self): + from nemo_automodel._transformers import infrastructure as infra + + with infra.init_empty_weights(): + model = _DummyModel() + + timeline = MagicMock() + + def materialize_model(model_to_materialize, device, **_kwargs): + model_to_materialize.to_empty(device=device) + + def mark_base_checkpoint_loaded(model_to_load, *_args, **_kwargs): + model_to_load.base_checkpoint_loaded = True + + def add_value_head(model_to_update): + assert model_to_update.base_checkpoint_loaded is True + assert all(parameter.device.type == "cpu" for parameter in model_to_update.parameters()) + model_to_update.value_head = torch.nn.Linear(4, 1) + + hook = MagicMock(side_effect=add_value_head) + snapshot = MagicMock(side_effect=lambda _model, state_dict, **_kwargs: state_dict) + shard = MagicMock(return_value=model) + timeline.attach_mock(hook, "hook") + timeline.attach_mock(snapshot, "snapshot") + timeline.attach_mock(shard, "shard") + + with ( + patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=1), + patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), + patch(f"{_INFRA_MODULE}.print_trainable_parameters"), + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=True), + patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf", snapshot), + patch(f"{_INFRA_MODULE}._shard_ep_fsdp", shard), + patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + ): + mock_ckpt = MockCheckpointer.return_value + mock_ckpt.config.dequantize_base_checkpoint = False + mock_ckpt.initialize_model_weights.side_effect = materialize_model + mock_ckpt.load_base_model.side_effect = mark_base_checkpoint_loaded + timeline.attach_mock(mock_ckpt.initialize_model_weights, "initialize") + timeline.attach_mock(mock_ckpt.load_base_model, "load") + + result = infra.apply_model_infrastructure( + model=model, + is_meta_device=True, + device=torch.device("cpu"), + load_base_model=True, + pretrained_model_name_or_path="test/model", + pre_fsdp_hook=hook, + skip_task_head_prefixes_for_base_model=("value_head.", "value_head."), + ) + + assert result is model + hook.assert_called_once_with(model) + assert [mock_call[0] for mock_call in timeline.mock_calls] == [ + "initialize", + "load", + "hook", + "snapshot", + "shard", + ] + assert snapshot.call_args.args[1]["value_head.weight"].device.type == "cpu" + assert model._pre_shard_hf_state_dict_keys == list(snapshot.call_args.args[1]) + checkpoint_config = MockCheckpointer.call_args.args[0] + assert checkpoint_config.skip_task_head_prefixes_for_base_model == ["value_head."] + + def test_meta_model_runs_hook_before_key_snapshot_and_creates_meta_head(self): + from nemo_automodel._transformers import infrastructure as infra + + with infra.init_empty_weights(): + model = _DummyModel() + + timeline = MagicMock() + + def add_value_head(model_to_update): + model_to_update.value_head = torch.nn.Linear(4, 1) + assert model_to_update.value_head.weight.device.type == "meta" + + hook = MagicMock(side_effect=add_value_head) + snapshot = MagicMock(side_effect=lambda _model, state_dict, **_kwargs: state_dict) + shard = MagicMock(return_value=model) + timeline.attach_mock(hook, "hook") + timeline.attach_mock(snapshot, "snapshot") + timeline.attach_mock(shard, "shard") + + with ( + patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=2), + patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), + patch(f"{_INFRA_MODULE}.print_trainable_parameters"), + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), + patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf", snapshot), + patch(f"{_INFRA_MODULE}._shard_ep_fsdp", shard), + patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + ): + mock_ckpt = MockCheckpointer.return_value + mock_ckpt.config.dequantize_base_checkpoint = False + + result = infra.apply_model_infrastructure( + model=model, + is_meta_device=True, + device=torch.device("cpu"), + load_base_model=False, + pre_fsdp_hook=hook, + ) + + assert result is model + hook.assert_called_once_with(model) + assert [mock_call[0] for mock_call in timeline.mock_calls] == ["hook", "snapshot", "shard"] + assert "value_head.weight" in snapshot.call_args.args[1] + assert "value_head.weight" in model._pre_shard_hf_state_dict_keys + + def test_rejects_hook_return_value_before_key_snapshot_or_sharding(self): + from nemo_automodel._transformers import infrastructure as infra + + model = _DummyModel() + hook = MagicMock(return_value=model) + + with ( + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), + patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf") as snapshot, + patch(f"{_INFRA_MODULE}._shard_ep_fsdp") as shard, + patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + ): + MockCheckpointer.return_value.config.dequantize_base_checkpoint = False + + with pytest.raises(TypeError, match="mutate the existing model in place and return None"): + infra.apply_model_infrastructure( + model=model, + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + pre_fsdp_hook=hook, + ) + + hook.assert_called_once_with(model) + snapshot.assert_not_called() + shard.assert_not_called() + + @pytest.mark.parametrize( + "mesh_overrides,infrastructure_overrides,unsupported_name", + [ + ({"tp_size": 2}, {}, "tensor parallelism"), + ({"cp_size": 2}, {}, "context parallelism"), + ({"ep_size": 2}, {}, "expert parallelism"), + ({"pp_size": 2}, {}, "pipeline parallelism"), + ({}, {"peft_config": object()}, "PEFT"), + ({}, {"quantization_config": object()}, "quantization"), + ({}, {"fp8_config": object()}, "FP8"), + ({}, {"qat_quantizer": object()}, "QAT"), + ], + ids=["tp", "cp", "ep", "pp", "peft", "quantization", "fp8", "qat"], + ) + def test_rejects_unsupported_parallelism_and_model_transforms( + self, + mesh_overrides, + infrastructure_overrides, + unsupported_name, + ): + from nemo_automodel._transformers import infrastructure as infra + + mesh_sizes = {"tp_size": 1, "cp_size": 1, "ep_size": 1, "pp_size": 1} + mesh_sizes.update(mesh_overrides) + hook = MagicMock() + + with pytest.raises(NotImplementedError, match=unsupported_name): + infra.apply_model_infrastructure( + model=_DummyModel(), + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + mesh=SimpleNamespace(**mesh_sizes), + pre_fsdp_hook=hook, + **infrastructure_overrides, + ) + + hook.assert_not_called() + + def test_rejects_checkpoint_native_quantization_before_running_hook(self): + from nemo_automodel._transformers import infrastructure as infra + + model = _DummyModel() + model.config.quantization_config = {"quant_method": "bitsandbytes"} + hook = MagicMock() + + with pytest.raises(NotImplementedError, match="quantization"): + infra.apply_model_infrastructure( + model=model, + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + pre_fsdp_hook=hook, + ) + + hook.assert_not_called() + + def test_apply_model_infrastructure_handles_unwrapped_single_rank_ddp_model(): """Single-rank DDP skips wrapping, so the returned model may not have ``.module``.""" from nemo_automodel._transformers.infrastructure import apply_model_infrastructure diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py index 8fad7924b6..4d60c83663 100644 --- a/tests/unit_tests/checkpoint/test_checkpointing.py +++ b/tests/unit_tests/checkpoint/test_checkpointing.py @@ -1528,6 +1528,96 @@ def test_bin_checkpoint_uses_fast_path(self, mock_load_full, mock_load_hf, mock_ mock_load_full.assert_called_once() mock_dcp_load.assert_not_called() + @pytest.mark.parametrize("checkpoint_format", ["safetensors", "bin"]) + def test_base_checkpoint_fast_path_preserves_skipped_task_head(self, checkpoint_format): + """Base checkpoint loading keeps task-head initialization for both HF fast paths.""" + checkpointer = self._make_checkpointer() + checkpointer.config.skip_task_head_prefixes_for_base_model = ["task_head."] + + model = torch.nn.Module() + model.backbone = torch.nn.Linear(4, 4) + model.task_head = torch.nn.Linear(4, 1) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.fill_(1.0) + + checkpoint_state = { + name: torch.full_like(tensor, 2.0 if name.startswith("backbone.") else 3.0) + for name, tensor in model.state_dict().items() + } + + with ( + patch("os.path.exists", return_value=True), + patch( + "nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", + return_value=checkpoint_format == "safetensors", + ), + patch( + "nemo_automodel.components.checkpoint.checkpointing._is_bin_checkpoint", + return_value=checkpoint_format == "bin", + ), + patch( + "nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype", + return_value=checkpoint_state, + ), + patch( + "nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model", + side_effect=lambda model_parts, state_dict: model_parts[0].load_state_dict(state_dict, strict=False), + ) as mock_load_full, + patch.object(checkpointer, "_do_load") as mock_dcp_load, + ): + checkpointer.load_model(model, model_path="/fake/path", is_init_step=True) + + loaded_state = mock_load_full.call_args.args[1] + assert set(loaded_state) == {"backbone.weight", "backbone.bias"} + assert all(torch.equal(parameter, torch.full_like(parameter, 2.0)) for parameter in model.backbone.parameters()) + assert all( + torch.equal(parameter, torch.full_like(parameter, 1.0)) for parameter in model.task_head.parameters() + ) + mock_dcp_load.assert_not_called() + + def test_training_checkpoint_restore_does_not_skip_task_head(self): + """Training checkpoint restore loads task heads even when base-load filtering is configured.""" + checkpointer = self._make_checkpointer() + checkpointer.config.skip_task_head_prefixes_for_base_model = ["task_head."] + + model = torch.nn.Module() + model.backbone = torch.nn.Linear(4, 4) + model.task_head = torch.nn.Linear(4, 1) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.fill_(1.0) + + checkpoint_state = { + name: torch.full_like(tensor, 2.0 if name.startswith("backbone.") else 3.0) + for name, tensor in model.state_dict().items() + } + + with ( + patch("os.path.exists", return_value=True), + patch( + "nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", + return_value=True, + ), + patch( + "nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype" + ) as mock_load_hf, + patch( + "nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model" + ) as mock_load_full, + patch.object(checkpointer, "_get_storage_reader", return_value=None), + patch.object(checkpointer, "_do_load", return_value=checkpoint_state) as mock_dcp_load, + ): + checkpointer.load_model(model, model_path="/fake/path", is_init_step=False) + + assert all(torch.equal(parameter, torch.full_like(parameter, 2.0)) for parameter in model.backbone.parameters()) + assert all( + torch.equal(parameter, torch.full_like(parameter, 3.0)) for parameter in model.task_head.parameters() + ) + mock_dcp_load.assert_called_once() + mock_load_hf.assert_not_called() + mock_load_full.assert_not_called() + @patch("nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", return_value=True) @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype") @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model") From 5420b30fd7372b015900adad07f5bb8ae522e9c4 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 02:34:06 -0700 Subject: [PATCH 27/34] feat(data): collate padded and packed VLM Datums Signed-off-by: HuiyingLi --- nemo_automodel/components/datasets/datum.py | 151 +++++++++++++++++++- tests/unit_tests/datasets/test_datum.py | 132 +++++++++++++++++ 2 files changed, 280 insertions(+), 3 deletions(-) diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index cc08972c6f..8f378b62d4 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -16,11 +16,11 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn.functional as F @@ -31,9 +31,12 @@ packed_sequence_thd_collater, ) +if TYPE_CHECKING: + from transformers import ProcessorMixin + CROSS_ENTROPY_IGNORE_IDX = -100 -__all__ = ["CollatedLossInputs", "Datum", "LossInputLayout", "collate_datums"] +__all__ = ["CollatedLossInputs", "Datum", "LossInputLayout", "collate_datums", "collate_vlm_datums"] def _pin_memory_tree(value: Any) -> Any: @@ -486,3 +489,145 @@ def collate_datums( item_to_datum=tuple(range(len(datums))), pad_values=loss_pad_values, ) + + +def collate_vlm_datums( + datums: list[Datum], + *, + processor: "ProcessorMixin", + packed: bool = False, + get_rope_index: Callable[..., object] | None = None, + sequence_alignment: int = 1, +) -> tuple[dict[str, Any], CollatedLossInputs]: + """Collate pre-tokenized VLM SFT Datums with processor-specific media inputs. + + The input Datums retain their unshifted token stream because + :func:`~nemo_automodel.components.datasets.vlm.collate_fns.pad_collate_fn` + owns the autoregressive shift. ``labels`` and ``weights`` therefore also + use the unshifted token axis, with zero weight at target position zero. + In packed mode, every Datum becomes one THD document and the collater + preserves its real and aligned sequence lengths. Common VLM fields are + handled by the canonical VLM collaters; additional processor tensor fields + retain their leading media axis and are concatenated. + + Args: + datums: Non-empty processor-ready VLM items. Each item has + ``input_ids``, ``labels``, and ``weights`` tensors of shape + ``[sequence]`` on matching unshifted token axes. Optional processor + fields carry their processor-defined token or leading media axes. + processor: Hugging Face processor (or compatible object) that supplies + the tokenizer padding token. + packed: Pack the Datums as THD documents instead of padding a batch. + get_rope_index: Optional bound model callable used to materialize + multi-axis VLM position IDs before packing. + sequence_alignment: Per-document THD alignment. Context-parallel + callers use ``2 * cp_size``. Multi-axis mRoPE with alignment above + one is rejected by the canonical packed-VLM implementation. + + Returns: + A pair of shifted/padded model inputs and layout-aware loss inputs. + In padded mode, token model fields, labels, and weights have shape + ``[batch, padded_sequence - 1]``. In packed mode they have shape + ``[1, aligned_tokens]`` and model inputs include THD sequence metadata. + Media tensors retain arbitrary trailing dimensions and are concatenated + on their leading media axis. + """ + if not datums: + raise ValueError("collate_vlm_datums requires at least one Datum") + if isinstance(sequence_alignment, bool) or not isinstance(sequence_alignment, int) or sequence_alignment < 1: + raise ValueError(f"sequence_alignment must be a positive integer, got {sequence_alignment!r}") + + from nemo_automodel.components.datasets.vlm.collate_fns import pad_collate_fn + + examples = [] + for datum in datums: + labels = datum.loss_fn_inputs.get("labels") + weights = datum.loss_fn_inputs.get("weights") + input_ids = datum.model_inputs.get("input_ids") + if not all(isinstance(value, torch.Tensor) and value.ndim == 1 for value in (input_ids, labels, weights)): + raise ValueError("VLM Datums require 1-D input_ids, labels, and weights") + if datum.seq_len < 2: + raise ValueError("VLM Datums require at least two tokens for autoregressive shifting") + if labels.shape != input_ids.shape or weights.shape != input_ids.shape: + raise ValueError("VLM labels and weights must match the unshifted input_ids shape") + if bool(weights[0] != 0): + raise ValueError("VLM weight at target position zero must be zero before autoregressive shift") + examples.append({**datum.model_inputs, "labels": labels}) + + if packed: + from nemo_automodel.components.datasets.vlm.collate_fns import packed_sequence_thd_vlm_collater + from nemo_automodel.components.datasets.vlm.neat_packing_vlm import ( + _aligned_length, + _build_packed_vlm_sample, + _compute_mrope_position_ids, + _shift_sample, + ) + + position_ids = ( + [_compute_mrope_position_ids(example, get_rope_index) for example in examples] + if get_rope_index is not None + else [None] * len(examples) + ) + if any(position is not None for position in position_ids) and not all( + position is not None for position in position_ids + ): + raise ValueError("get_rope_index must return position IDs for every VLM Datum or none of them") + has_mrope = bool(position_ids and position_ids[0] is not None) + shifted_examples = [] + shifted_weights = [] + for example, position, datum in zip(examples, position_ids, datums): + if position is not None: + example["position_ids"] = position + shifted_examples.append(_shift_sample(example, has_mrope=has_mrope)) + shifted_weights.append(datum.loss_fn_inputs["weights"][1:]) + + tokenizer = getattr(processor, "tokenizer", processor) + padding_idx = getattr(tokenizer, "pad_token_id", 0) or 0 + pack_size = sum(_aligned_length(item["input_ids"].shape[0], sequence_alignment) for item in shifted_examples) + packed_sample = _build_packed_vlm_sample( + shifted_examples, + pack_size=pack_size, + padding_idx=padding_idx, + has_mrope=has_mrope, + sequence_alignment=sequence_alignment, + ) + model_inputs = packed_sequence_thd_vlm_collater([packed_sample], padding_idx=padding_idx) + labels = model_inputs.pop("labels") + weights = torch.cat( + [ + F.pad(weight, (0, _aligned_length(weight.shape[0], sequence_alignment) - weight.shape[0])) + for weight in shifted_weights + ] + ).unsqueeze(0) + else: + model_inputs = pad_collate_fn(examples, processor) + labels = model_inputs.pop("labels") + + target_width = labels.shape[-1] + 1 + shifted_weights = [ + F.pad(datum.loss_fn_inputs["weights"], (0, target_width - datum.seq_len))[1:] for datum in datums + ] + weights = torch.stack(shifted_weights) + + unhandled_keys = set().union(*(datum.model_inputs.keys() for datum in datums)) - set(model_inputs) + unhandled_keys -= {"input_ids", "attention_mask"} + for key in sorted(unhandled_keys): + values = [datum.model_inputs[key] for datum in datums if key in datum.model_inputs] + if not values: + continue + if not all(isinstance(value, torch.Tensor) and value.ndim > 0 for value in values): + raise TypeError(f"VLM processor field {key!r} must contain tensors with a leading media axis") + try: + model_inputs[key] = torch.cat(values, dim=0) + except RuntimeError as exc: + raise ValueError(f"VLM processor field {key!r} cannot be concatenated across samples") from exc + + return model_inputs, CollatedLossInputs( + {"labels": labels, "weights": weights}, + layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + }, + item_to_datum=tuple(range(len(datums))), + pad_values={"labels": CROSS_ENTROPY_IGNORE_IDX}, + ) diff --git a/tests/unit_tests/datasets/test_datum.py b/tests/unit_tests/datasets/test_datum.py index 7b029a8134..a0401778e2 100644 --- a/tests/unit_tests/datasets/test_datum.py +++ b/tests/unit_tests/datasets/test_datum.py @@ -14,6 +14,7 @@ import pickle from copy import copy, deepcopy +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -25,6 +26,7 @@ Datum, LossInputLayout, collate_datums, + collate_vlm_datums, ) @@ -74,6 +76,36 @@ def _routed_datums(): ] +def _vlm_datum(input_ids, weights, **media_inputs): + """Build one unshifted VLM item for collater tests. + + Args: + input_ids: Token IDs with shape ``[sequence]``. + weights: Boolean supervision values with shape ``[sequence]``. + **media_inputs: Processor tensors with shape ``[media, ...]``. + + Returns: + A Datum whose token fields have shape ``[sequence]`` and whose media + fields preserve their input shapes. + """ + input_ids = torch.tensor(input_ids, dtype=torch.long) + weights = torch.tensor(weights, dtype=torch.bool) + labels = input_ids.clone().masked_fill(~weights, CROSS_ENTROPY_IGNORE_IDX) + return Datum( + model_inputs={ + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + **media_inputs, + }, + loss_fn_inputs={"labels": labels, "weights": weights}, + loss_fn_input_layouts={ + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + }, + loss_fn_input_pad_values={"labels": CROSS_ENTROPY_IGNORE_IDX}, + ) + + # ── Datum ───────────────────────────────────────────────────────────────── @@ -310,6 +342,106 @@ def test_to_features_native_python_ints(): # ── collate_datums delegates to the canonical collaters ───────────────────── +def test_collate_vlm_datums_delegates_padding_and_preserves_processor_fields(): + processor = SimpleNamespace(image_processor=object(), tokenizer=SimpleNamespace(pad_token_id=9)) + datums = [ + _vlm_datum( + [1, 2, 3], + [0, 1, 1], + pixel_values=torch.tensor([[1.0, 2.0]]), + image_grid_thw=torch.tensor([[1, 2, 2]]), + image_flags=torch.tensor([[1]]), + imgs_sizes=torch.tensor([[32, 64]]), + ), + _vlm_datum([4, 5], [0, 1]), + ] + + model_inputs, loss_inputs = collate_vlm_datums(datums, processor=processor) + + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([[1, 2], [4, 5]])) + torch.testing.assert_close(model_inputs["attention_mask"], torch.tensor([[1, 1], [1, 1]])) + torch.testing.assert_close(loss_inputs["labels"], torch.tensor([[2, 3], [5, -100]])) + torch.testing.assert_close(loss_inputs["weights"], torch.tensor([[True, True], [True, False]])) + torch.testing.assert_close(model_inputs["image_flags"], torch.tensor([[1]])) + torch.testing.assert_close(model_inputs["imgs_sizes"], torch.tensor([[32, 64]])) + assert model_inputs["pixel_values"].dtype == torch.bfloat16 + assert loss_inputs.item_to_datum == (0, 1) + assert loss_inputs.layouts == { + "labels": LossInputLayout.PER_TOKEN, + "weights": LossInputLayout.PER_TOKEN, + } + assert loss_inputs.pad_values == {"labels": CROSS_ENTROPY_IGNORE_IDX} + + +def test_collate_vlm_datums_rejects_invalid_unshifted_contract(): + processor = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=0)) + with pytest.raises(ValueError, match="at least one"): + collate_vlm_datums([], processor=processor) + with pytest.raises(ValueError, match="target position zero"): + collate_vlm_datums([_vlm_datum([1, 2], [1, 1])], processor=processor) + + +def test_collate_vlm_datums_packs_aligned_thd_documents(): + processor = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=9)) + datums = [ + _vlm_datum( + [1, 2, 3], + [0, 1, 1], + pixel_values=torch.tensor([[1.0, 2.0]]), + image_grid_thw=torch.tensor([[1, 2, 2]]), + image_flags=torch.tensor([[1]]), + ), + _vlm_datum([4, 5], [0, 1]), + ] + + model_inputs, loss_inputs = collate_vlm_datums( + datums, + processor=processor, + packed=True, + sequence_alignment=4, + ) + + torch.testing.assert_close(model_inputs["input_ids"], torch.tensor([[1, 2, 9, 9, 4, 9, 9, 9]])) + torch.testing.assert_close(model_inputs["position_ids"], torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]])) + torch.testing.assert_close(model_inputs["seq_lens"], torch.tensor([[2, 1]])) + torch.testing.assert_close(model_inputs["seq_lens_padded"], torch.tensor([[4, 4]])) + torch.testing.assert_close(loss_inputs["labels"], torch.tensor([[2, 3, -100, -100, 5, -100, -100, -100]])) + torch.testing.assert_close( + loss_inputs["weights"], + torch.tensor([[True, True, False, False, True, False, False, False]]), + ) + torch.testing.assert_close(model_inputs["image_flags"], torch.tensor([[1]])) + assert model_inputs["qkv_format"] == "thd" + assert loss_inputs.item_to_datum == (0, 1) + + +def test_collate_vlm_datums_packs_mrope_only_without_cp_alignment(): + processor = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=0)) + datums = [_vlm_datum([1, 2, 3], [0, 1, 1]), _vlm_datum([4, 5], [0, 1])] + + def get_rope_index(input_ids, **_kwargs): + sequence = input_ids.shape[-1] + positions = torch.arange(sequence).expand(3, 1, sequence) + return positions, torch.zeros(1) + + model_inputs, _ = collate_vlm_datums( + datums, + processor=processor, + packed=True, + get_rope_index=get_rope_index, + ) + + assert model_inputs["position_ids"].shape == (3, 1, 3) + with pytest.raises(NotImplementedError, match="multi-axis mRoPE"): + collate_vlm_datums( + datums, + processor=processor, + packed=True, + get_rope_index=get_rope_index, + sequence_alignment=2, + ) + + def test_collate_padded_uses_default_collater_schema(): batch, loss_inputs = collate_datums(_toy_datums()) assert isinstance(loss_inputs, CollatedLossInputs) From 174d7c175f55d5ad29af96892ee420f346a33112 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 12:10:42 -0700 Subject: [PATCH 28/34] feat(model): manage pre-FSDP task heads Signed-off-by: HuiyingLi --- nemo_automodel/__init__.py | 1 + nemo_automodel/_transformers/__init__.py | 2 + nemo_automodel/_transformers/auto_model.py | 20 +- .../_transformers/infrastructure.py | 207 ++++++--- .../checkpoint/stateful_wrappers.py | 7 +- .../components/distributed/parallelizer.py | 13 +- nemo_automodel/components/moe/fsdp_mixin.py | 8 + nemo_automodel/components/moe/parallelizer.py | 10 +- nemo_automodel/shared/task_heads.py | 184 ++++++++ .../_transformers/test_infrastructure.py | 433 +++++++++++++++--- .../checkpoint/test_checkpointing.py | 31 ++ tests/unit_tests/moe/test_fsdp_mixin.py | 33 ++ tests/unit_tests/moe/test_parallelizer.py | 42 +- tests/unit_tests/shared/test_task_heads.py | 87 ++++ 14 files changed, 933 insertions(+), 145 deletions(-) create mode 100644 nemo_automodel/shared/task_heads.py create mode 100644 tests/unit_tests/shared/test_task_heads.py diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 489c6c1624..378108dd51 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -62,6 +62,7 @@ "NeMoAutoModelCrossEncoder": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelCrossEncoder"), "NeMoAutoTokenizer": ("nemo_automodel._transformers.auto_tokenizer", "NeMoAutoTokenizer"), "PerTokenOutput": ("nemo_automodel.engine.outputs", "PerTokenOutput"), + "PreFSDPHookResult": ("nemo_automodel.shared.task_heads", "PreFSDPHookResult"), "NeMoAutoDiffusionPipeline": ("nemo_automodel._diffusers.auto_diffusion_pipeline", "NeMoAutoDiffusionPipeline"), "ModelCapabilities": ("nemo_automodel._transformers.model_capabilities", "ModelCapabilities"), "query_capabilities": ("nemo_automodel._transformers.model_capabilities", "query_capabilities"), diff --git a/nemo_automodel/_transformers/__init__.py b/nemo_automodel/_transformers/__init__.py index 991204e303..bf3792eff8 100644 --- a/nemo_automodel/_transformers/__init__.py +++ b/nemo_automodel/_transformers/__init__.py @@ -35,6 +35,7 @@ "NeMoAutoModelCrossEncoder": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelCrossEncoder"), "NeMoAutoTokenizer": ("nemo_automodel._transformers.auto_tokenizer", "NeMoAutoTokenizer"), "AutoMFU": ("nemo_automodel._transformers.mfu", "AutoMFU"), + "PreFSDPHookResult": ("nemo_automodel.shared.task_heads", "PreFSDPHookResult"), "RetrieverStudentWithProjection": ( "nemo_automodel._transformers.retrieval", "RetrieverStudentWithProjection", @@ -57,6 +58,7 @@ "NeMoAutoModelCrossEncoder", "NeMoAutoTokenizer", "AutoMFU", + "PreFSDPHookResult", "RetrieverStudentWithProjection", "RetrieverTeacherEmbeddingEncoder", ] diff --git a/nemo_automodel/_transformers/auto_model.py b/nemo_automodel/_transformers/auto_model.py index 1e2cf3de99..1af9b92433 100644 --- a/nemo_automodel/_transformers/auto_model.py +++ b/nemo_automodel/_transformers/auto_model.py @@ -64,6 +64,7 @@ init_empty_weights, resolve_trust_remote_code, ) +from nemo_automodel.shared.task_heads import PreFSDPHookResult # noqa: E402 from nemo_automodel.shared.utils import dtype_from_str # noqa: E402 if TYPE_CHECKING: @@ -400,7 +401,7 @@ def _build_model( fp8_config, compile_config, load_base_model, - pre_fsdp_hook: Callable[[torch.nn.Module], None] | None = None, + pre_fsdp_hook: Callable[[torch.nn.Module], PreFSDPHookResult | None] | None = None, skip_task_head_prefixes_for_base_model: Sequence[str] | None = None, _retry_depth=0, **kwargs, @@ -665,7 +666,7 @@ 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, + pre_fsdp_hook: Callable[[torch.nn.Module], PreFSDPHookResult | None] | None = None, skip_task_head_prefixes_for_base_model: Sequence[str] | None = None, **kwargs, ) -> PreTrainedModel: @@ -721,11 +722,14 @@ def from_pretrained( 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. + kernel/lower-precision transforms and any load-before-shard base-model + restore, but before state-key capture and FSDP wrapping. Returning + ``None`` keeps the legacy pure-data-parallel contract. Returning + :class:`PreFSDPHookResult` declares a fresh task module that AutoModel + excludes from TP and lower-precision transforms, keeps trainable with + PEFT, and owns in FP32 FSDP units. The hook must be deterministic on + every rank, avoid collectives, and keep weight-tying configuration + consistent with its structural changes. 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. @@ -820,7 +824,7 @@ 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, + pre_fsdp_hook: Callable[[torch.nn.Module], PreFSDPHookResult | None] | None = None, skip_task_head_prefixes_for_base_model: Sequence[str] | None = None, **kwargs, ) -> PreTrainedModel: diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index 4e46917182..49873cc42d 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -31,6 +31,8 @@ from typing import TYPE_CHECKING, Union import torch +import torch.distributed as dist +from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard from nemo_automodel._transformers.utils import _should_load_before_shard from nemo_automodel._transformers.v4_patches.kv_sharing import ( @@ -60,6 +62,7 @@ snapshot_distributed_param_attrs, ) from nemo_automodel.components.distributed.mesh import MeshContext +from nemo_automodel.components.distributed.mesh_utils import get_fsdp_dp_mesh from nemo_automodel.components.distributed.pipelining.autopipeline import AutoPipeline from nemo_automodel.components.distributed.pipelining.config import PipelineConfig from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy @@ -78,7 +81,13 @@ init_empty_weights, print_trainable_parameters, ) -from nemo_automodel.shared.tied_weights import ensure_tied_lm_head +from nemo_automodel.shared.task_heads import ( + PreFSDPHookResult, + is_task_head_parameter, + register_task_head_module, + task_head_module, +) +from nemo_automodel.shared.tied_weights import ensure_tied_lm_head, is_tied_word_embeddings if TYPE_CHECKING: from torchao.quantization.qat.linear import Int4WeightOnlyQATQuantizer, Int8DynActInt4WeightQATQuantizer @@ -206,6 +215,30 @@ def _shard_pp(autopipeline, model, loss_fn, parallelize_fn): def _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh: MeshContext): """Apply EP + FSDP sharding (non-PP path).""" + managed_task_head = task_head_module(model) + if managed_task_head is not None and get_world_size_safe() > 1: + if not isinstance(model_wrapper, FSDP2Manager): + raise NotImplementedError("The managed pre-FSDP task module currently requires FSDP2") + if mesh.device_mesh is None: + raise ValueError("The managed pre-FSDP task module requires a device mesh") + fsdp_mesh = get_fsdp_dp_mesh(mesh.device_mesh, "dp_replicate", "dp_shard_cp") + task_head_mp_policy = MixedPrecisionPolicy( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + output_dtype=torch.float32, + cast_forward_inputs=True, + ) + if isinstance(managed_task_head, FSDPModule): + raise RuntimeError("The managed task module must not already be FSDP-wrapped") + # A size-one data mesh still needs its own FSDP unit under pure TP: + # otherwise the later root unit inherits the backbone's BF16 policy. + fully_shard( + managed_task_head, + mesh=fsdp_mesh, + mp_policy=task_head_mp_policy, + offload_policy=model_wrapper.offload_policy, + reshard_after_forward=False, + ) if parallelize_fn is not None and get_world_size_safe() > 1: parallelize_fn( model, @@ -221,6 +254,36 @@ def _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh: MeshContext): return model +@torch.no_grad() +def _sync_task_head_replicas(model: torch.nn.Module, mesh: MeshContext, device: torch.device) -> None: + """Broadcast the local task-head shard across the TP replica axis.""" + managed_task_head = task_head_module(model) + if ( + managed_task_head is None + or mesh.device_mesh is None + or "tp" not in (mesh.device_mesh.mesh_dim_names or ()) + or not dist.is_available() + or not dist.is_initialized() + ): + return + + replica_mesh = mesh.device_mesh["tp"] + if replica_mesh.size() <= 1: + return + group = replica_mesh.get_group() + source_rank = dist.get_process_group_ranks(group)[0] + backend = str(dist.get_backend(group)).lower() + for parameter in managed_task_head.parameters(): + to_local = getattr(parameter, "to_local", None) + local = to_local() if callable(to_local) else parameter + if "nccl" in backend and local.device.type != "cuda": + communication_tensor = local.to(device=device) + dist.broadcast(communication_tensor, src=source_rank, group=group) + local.copy_(communication_tensor.to(device=local.device)) + else: + dist.broadcast(local, src=source_rank, group=group) + + # Infrastructure instantiation (config -> runtime objects) def _instantiate_distributed( config: DistributedStrategyConfig | None, @@ -484,7 +547,7 @@ 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, + pre_fsdp_hook: Callable[[torch.nn.Module], PreFSDPHookResult | None] | None = None, skip_task_head_prefixes_for_base_model: Sequence[str] | None = None, **_kwargs, ): @@ -526,13 +589,14 @@ def apply_model_infrastructure( 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. + restore, but before state-key capture and FSDP wrapping. Returning + ``None`` retains the legacy pure-data-parallel structural-hook contract. + Returning :class:`PreFSDPHookResult` declares a fresh task module that + AutoModel excludes from TP/PEFT/quantization transforms and manages as + trainable FP32 FSDP units. The hook must make the same deterministic + change on every rank and avoid collectives. It 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. @@ -547,30 +611,8 @@ def apply_model_infrastructure( 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 pre_fsdp_hook is not None and (autopipeline is not None or mesh.pp_size != 1): + raise NotImplementedError("pre_fsdp_hook does not support pipeline parallelism") 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.") @@ -604,17 +646,12 @@ def apply_model_infrastructure( getattr(model_wrapper, "moe_mesh", None), process_group=getattr(mesh, "process_group", None), ) + if checkpointer.config.dequantize_base_checkpoint is None: + checkpointer.config.dequantize_base_checkpoint = hasattr(getattr(model, "config", None), "quantization_config") - # Handle checkpointer config updates if checkpointer is provided - if checkpointer is not None: - if checkpointer.config.dequantize_base_checkpoint is None: - checkpointer.config.dequantize_base_checkpoint = hasattr( - getattr(model, "config", None), "quantization_config" - ) - - # Apply PEFT and lower precision if configured - # When on meta device, wrap in init_empty_weights() so new LoRA modules are also on meta device - # This allows copy operations between meta tensors to succeed (they're no-ops) + # Apply PEFT and lower precision if configured. The task hook runs after + # these transforms so a fresh task head is not silently LoRA-wrapped, + # quantized, or converted to FP8/QAT. peft_ctx = init_empty_weights() if is_meta_device else nullcontext() with peft_ctx: model = _apply_peft_and_lower_precision( @@ -666,23 +703,71 @@ 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 + task_head_needing_reset: str | None = None 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. + pre_hook_module_ids = {id(module) for module in model.modules()} + pre_hook_parameter_ids = {id(parameter) for parameter in model.parameters()} + # A load-before-shard model has already been materialized above. A + # post-shard-load model is still meta, so construct its fresh modules on + # meta too and initialize them only after FSDP materialization. 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() + # Native/BnB quantized models may already reside on the training device + # and cannot be moved afterwards. Make implicit task-module construction + # follow that device; meta models must stay meta until FSDP materializes + # them below. + hook_context = init_empty_weights() if has_meta_tensors else torch.device(device) 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__}." + + if hook_result is 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), + ("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( + "A pre_fsdp_hook that returns None supports only the legacy unquantized, non-PEFT, " + f"tp_size=cp_size=ep_size=1 contract; unsupported: {', '.join(unsupported)}. " + "Return PreFSDPHookResult to declare a managed task module." + ) + elif isinstance(hook_result, PreFSDPHookResult): + name = register_task_head_module( + model, + hook_result, + pre_hook_module_ids=pre_hook_module_ids, + pre_hook_parameter_ids=pre_hook_parameter_ids, ) + if is_tied_word_embeddings(model): + raise ValueError( + "A managed task module requires an untied output head; the hook must disable supported " + "embedding tying and tied-only architectures are unsupported" + ) + if any(parameter.device.type == "meta" for parameter in model.get_submodule(name).parameters()): + task_head_needing_reset = name + reset_parameters = getattr(model.get_submodule(name), "reset_parameters", None) + if not callable(reset_parameters): + raise TypeError( + f"Managed task module {name!r} was created on meta and must implement reset_parameters()" + ) + skip_task_head_prefixes = list(dict.fromkeys([*(skip_task_head_prefixes or ()), f"{name}."])) + checkpointer.config.skip_task_head_prefixes_for_base_model = skip_task_head_prefixes + else: + raise TypeError(f"pre_fsdp_hook must return None or PreFSDPHookResult; 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) @@ -796,13 +881,25 @@ def apply_model_infrastructure( checkpoint_loaded=bool(checkpoint_already_loaded or weights_already_loaded or should_load_checkpoint), ) + # Base checkpoints intentionally omit the fresh task module. Initialize it + # only when the hook created it on meta, after materialization and base load, + # so architecture-specific model initializers cannot skip or overwrite them. + if task_head_needing_reset is not None: + managed_task_head = task_head_module(model) + if managed_task_head is None: + raise RuntimeError("The managed task module disappeared during model parallelization") + managed_task_head.to(dtype=torch.float32) + if any(parameter.dtype != torch.float32 for parameter in managed_task_head.parameters()): + raise RuntimeError("Failed to restore the managed task module to float32 after model initialization") + managed_task_head.reset_parameters() + # Freeze parameters after checkpoint loading and parallelization # This catches params created during parallelization (e.g., GroupedExpertsTE in init_token_dispatcher) if peft_config is not None: models_to_freeze = model.parts if hasattr(model, "parts") else [model] for mp in models_to_freeze: for name, param in mp.named_parameters(): - if "lora_" not in name and param.requires_grad: + if "lora_" not in name and not is_task_head_parameter(mp, name) and param.requires_grad: param.requires_grad_(False) if autopipeline is None: @@ -835,6 +932,8 @@ def apply_model_infrastructure( else: raise + _sync_task_head_replicas(model, mesh, device) + # Configure dense attention parallelism. Transformer Engine must know its # TP head partition even without CP; for CP, TE owns THD communication while # SDPA uses DTensor context-parallel hooks. diff --git a/nemo_automodel/components/checkpoint/stateful_wrappers.py b/nemo_automodel/components/checkpoint/stateful_wrappers.py index a8dc70ca70..a70082841b 100644 --- a/nemo_automodel/components/checkpoint/stateful_wrappers.py +++ b/nemo_automodel/components/checkpoint/stateful_wrappers.py @@ -65,6 +65,7 @@ def _safe_op_set_extra_state(self, state): materialize_missing_tied_lm_head, ) from nemo_automodel.shared.parameter_names import canonical_parameter_fqn +from nemo_automodel.shared.task_heads import is_task_head_parameter _PREFIX = "model." _OPTIMIZER_PARTS_KEY = "optimizer_parts" @@ -440,7 +441,11 @@ def state_dict(self) -> dict[str, Any]: # this filtering removes them. # TODO: this is a hack and we should find a better way to do this. if self.is_peft: - model_state_dict = {k: v for k, v in model_state_dict.items() if "lora_" in k} + model_state_dict = { + k: v + for k, v in model_state_dict.items() + if "lora_" in k or any(is_task_head_parameter(part, k) for part in self.model) + } # Pipeline parallelism partitions layers across PP ranks, so each rank's # local adapter (collected above) only covers its own stages. Gather the diff --git a/nemo_automodel/components/distributed/parallelizer.py b/nemo_automodel/components/distributed/parallelizer.py index 4db6da3b0d..850d3be301 100644 --- a/nemo_automodel/components/distributed/parallelizer.py +++ b/nemo_automodel/components/distributed/parallelizer.py @@ -83,6 +83,7 @@ class Gemma4ForConditionalGeneration: # type: ignore[no-redef] module_parameters, normalize_frozen_multimodal_sharding, ) +from nemo_automodel.shared.task_heads import exclude_task_heads_from_tp_plan from nemo_automodel.shared.tied_weights import ensure_tied_lm_head from nemo_automodel.shared.torch_patches import ( patch_fsdp_accumulated_grad_guard as _patch_fsdp_accumulated_grad_guard, @@ -337,11 +338,14 @@ def parallelize( # Generate or use tensor parallel plan model_parallel_plan = { k: translate_to_lora(v) - for k, v in _get_parallel_plan( + for k, v in exclude_task_heads_from_tp_plan( model, - sequence_parallel, - tp_shard_plan, - tp_size=tp_mesh.size(), + _get_parallel_plan( + model, + sequence_parallel, + tp_shard_plan, + tp_size=tp_mesh.size(), + ), ).items() } @@ -520,6 +524,7 @@ def parallelize( model_tp_plan: dict[str, ParallelStyle] = { "lm_head": translate_to_lora(ColwiseParallel(output_layouts=Shard(-1), use_local_output=False)), } + model_tp_plan = exclude_task_heads_from_tp_plan(model, model_tp_plan) mlp_tp_plan: dict[str, ParallelStyle] = { "mixer.up_proj": translate_to_lora(ColwiseParallel()), diff --git a/nemo_automodel/components/moe/fsdp_mixin.py b/nemo_automodel/components/moe/fsdp_mixin.py index 78d507554f..6e708a070b 100644 --- a/nemo_automodel/components/moe/fsdp_mixin.py +++ b/nemo_automodel/components/moe/fsdp_mixin.py @@ -22,6 +22,7 @@ from nemo_automodel.components.models.common.utils import get_is_optim_step from nemo_automodel.shared.multimodal_fsdp import iter_multimodal_modules +from nemo_automodel.shared.task_heads import task_head_module def _iter_fsdp_modules(module: torch.nn.Module) -> Iterator[FSDPModule]: @@ -57,6 +58,13 @@ def _iter_fsdp_module_candidates(module: torch.nn.Module) -> Iterator[FSDPModule if hasattr(module, "lm_head") and isinstance(module.lm_head, FSDPModule): yield module.lm_head + # A managed task head may live below an architecture-specific wrapper (for + # example ``thinker.lm_head``). Its independent FSDP unit must participate + # in the same deferred-sync lifecycle as the MoE backbone. + managed_task_head = task_head_module(module) + if isinstance(managed_task_head, FSDPModule): + yield managed_task_head + # Multimodal towers/projectors. ``moe/parallelizer.apply_fsdp`` gives every # recognized *trainable* multimodal module its own FSDP unit, and under the # ``per_layer`` frozen policy it shards their layer containers instead of the diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index e95ed49221..58d4748794 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -26,7 +26,7 @@ checkpoint_wrapper as ptd_checkpoint_wrapper, ) from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.fsdp import fully_shard +from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.fsdp._fully_shard import MixedPrecisionPolicy, OffloadPolicy from torch.distributed.tensor import Replicate, Shard, distribute_module, distribute_tensor from torch.distributed.tensor.parallel import ParallelStyle, parallelize_module @@ -52,6 +52,7 @@ normalize_frozen_multimodal_sharding, shard_multimodal_module, ) +from nemo_automodel.shared.task_heads import exclude_task_heads_from_tp_plan from nemo_automodel.shared.tied_weights import ensure_tied_lm_head from nemo_automodel.shared.torch_patches import ( patch_fsdp_accumulated_grad_guard as _patch_fsdp_accumulated_grad_guard, @@ -821,7 +822,11 @@ def apply_fsdp( if embed_norm is not None: fully_shard_default(embed_norm) - if lm_head is not None and not tied_input_output_embeddings: + if isinstance(lm_head, FSDPModule): + # Managed task heads are wrapped centrally before this model-specific + # path so dense and MoE models share the same FP32 FSDP policy. + pass + elif lm_head is not None and not tied_input_output_embeddings: # Use custom mixed precision policy for lm_head if lm_head_precision is specified if lm_head_precision == torch.float32: lm_head_mp_policy = MixedPrecisionPolicy( @@ -1014,6 +1019,7 @@ def parallelize_model( tp_shard_plan=tp_shard_plan, tp_size=tp_mesh.size(), ) + model_parallel_plan = exclude_task_heads_from_tp_plan(model, model_parallel_plan) # Every custom-MoE TP plan keeps the token path (attention, router) # replicated across TP ranks, so each expert gradient accumulates # tp_size identical contributions through the EP all-gather. diff --git a/nemo_automodel/shared/task_heads.py b/nemo_automodel/shared/task_heads.py new file mode 100644 index 0000000000..e6f792d72a --- /dev/null +++ b/nemo_automodel/shared/task_heads.py @@ -0,0 +1,184 @@ +# 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. + +"""Model-scoped metadata for task heads installed before FSDP.""" + +from __future__ import annotations + +from dataclasses import dataclass +from fnmatch import fnmatchcase +from typing import Mapping + +import torch +from torch import nn + +from nemo_automodel.shared.parameter_names import canonical_parameter_fqn + +__all__ = ["PreFSDPHookResult"] + +_TASK_HEAD_MODULE_NAME = "_nemo_task_head_module_name" + + +@dataclass(frozen=True) +class PreFSDPHookResult: + """Declare a fresh task module installed by ``pre_fsdp_hook``. + + The declared module is excluded from tensor- and expert-specific sharding, + replicated over tensor parallelism, and placed in its own FP32 FSDP unit + over the data/context-parallel mesh. It stays trainable under PEFT and is + included in resumable training checkpoints. The module must be newly + created by the hook, attached exactly once below the model root, and expose + ``reset_parameters()`` when it is created on the meta device. Under tensor + or sequence parallelism, the module remains responsible for accepting the + distributed activation layout supplied by its owning model. Consolidated + Hugging Face export remains the owning model's responsibility. + + Args: + task_module: Fresh parameter-owning module attached to the model by the + hook. A compound task head can be represented by one container + module. Every parameter must be trainable, float32, and owned only + within that module's subtree. + """ + + task_module: nn.Module + + +def register_task_head_module( + model: nn.Module, + result: PreFSDPHookResult, + *, + pre_hook_module_ids: set[int], + pre_hook_parameter_ids: set[int], +) -> str: + """Validate and record the task module returned by a pre-FSDP hook.""" + module = result.task_module + if not isinstance(module, nn.Module): + raise TypeError("PreFSDPHookResult.task_module must be a torch.nn.Module") + + paths_by_id: dict[int, list[str]] = {} + for name, candidate in model.named_modules(remove_duplicate=False): + paths_by_id.setdefault(id(candidate), []).append(name) + + paths = [path for path in paths_by_id.get(id(module), ()) if path] + if not paths: + raise ValueError("The declared task module must be attached below the model root") + if len(paths) != 1: + raise ValueError( + f"The declared task module must have exactly one model FQN; found aliases: {', '.join(sorted(paths))}" + ) + name = paths[0] + parameters = tuple(module.parameters()) + if not parameters: + raise ValueError(f"Declared task module {name!r} owns no parameters") + if any(id(parameter) in pre_hook_parameter_ids for parameter in parameters): + raise ValueError(f"Declared task module {name!r} reuses pre-hook parameters; the task module must be fresh") + if id(module) in pre_hook_module_ids: + raise ValueError(f"Declared task module {name!r} existed before the hook; the task module must be fresh") + parameter_paths: dict[int, list[str]] = {} + for parameter_name, parameter in model.named_parameters(remove_duplicate=False): + parameter_paths.setdefault(id(parameter), []).append(parameter_name) + task_prefix = f"{name}." + for parameter in parameters: + outside_aliases = [ + parameter_name + for parameter_name in parameter_paths.get(id(parameter), ()) + if not parameter_name.startswith(task_prefix) + ] + if outside_aliases: + raise ValueError( + f"Declared task module {name!r} shares parameters outside its subtree: " + f"{', '.join(sorted(outside_aliases))}" + ) + if any(parameter.dtype != torch.float32 for parameter in parameters): + raise ValueError(f"Declared task module {name!r} must own only float32 parameters") + if any(not parameter.requires_grad for parameter in parameters): + raise ValueError(f"Declared task module {name!r} must own only trainable parameters") + + setattr(model, _TASK_HEAD_MODULE_NAME, name) + return name + + +def _task_head_owner_and_name(model: nn.Module) -> tuple[nn.Module, str | None]: + """Resolve the module that owns task-head metadata and its local FQN.""" + pending = [model] + seen: set[int] = set() + while pending: + owner = pending.pop() + if id(owner) in seen: + continue + seen.add(id(owner)) + name = vars(owner).get(_TASK_HEAD_MODULE_NAME) + if name is not None: + return owner, name + children = getattr(owner, "_modules", {}) or {} + for child_name in ("module", "_orig_mod"): + child = children.get(child_name) + if isinstance(child, nn.Module): + pending.append(child) + return model, None + + +def task_head_module_name(model: nn.Module) -> str | None: + """Return the task-module FQN recorded on ``model`` or its transparent wrapper.""" + return _task_head_owner_and_name(model)[1] + + +def task_head_module(model: nn.Module) -> nn.Module | None: + """Resolve the declared task module on ``model``.""" + owner, name = _task_head_owner_and_name(model) + return owner.get_submodule(name) if name is not None else None + + +def is_task_head_parameter(model: nn.Module, name: str) -> bool: + """Return whether ``name`` belongs to a declared task module.""" + normalized = canonical_parameter_fqn(name) + while normalized.startswith("_orig_mod."): + normalized = normalized.removeprefix("_orig_mod.") + task_name = task_head_module_name(model) + return task_name is not None and normalized.startswith(f"{task_name}.") + + +def exclude_task_heads_from_tp_plan( + model: nn.Module, + plan: Mapping[str, object], +) -> dict[str, object]: + """Remove TP rules that resolve exclusively inside the declared task module. + + A wildcard that also resolves outside a task module is rejected because + removing it would silently disable tensor parallelism for unrelated model + modules. + """ + task_root = task_head_module_name(model) + if task_root is None: + return dict(plan) + + module_names = [name for name, _ in model.named_modules() if name] + + def _is_task_subtree(name: str) -> bool: + return name == task_root or name.startswith(task_root + ".") + + filtered: dict[str, object] = {} + for pattern, style in plan.items(): + matches = [name for name in module_names if fnmatchcase(name, pattern)] + task_matches = [name for name in matches if _is_task_subtree(name)] + if not task_matches: + filtered[pattern] = style + continue + non_task_matches = [name for name in matches if not _is_task_subtree(name)] + if non_task_matches: + raise ValueError( + f"Tensor-parallel rule {pattern!r} matches both the managed task module and backbone modules; " + "use a narrower rule so the task module remains replicated" + ) + return filtered diff --git a/tests/unit_tests/_transformers/test_infrastructure.py b/tests/unit_tests/_transformers/test_infrastructure.py index 1ec8ca97e3..af38bc63e7 100644 --- a/tests/unit_tests/_transformers/test_infrastructure.py +++ b/tests/unit_tests/_transformers/test_infrastructure.py @@ -215,8 +215,9 @@ def _run_apply_model_infrastructure(*, is_meta_device, load_base_model, model_wr class TestPreFSDPHook: - def test_load_before_shard_runs_hook_after_base_load_and_before_key_snapshot(self): + def test_typed_hook_resets_meta_head_after_post_shard_base_load(self): from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.shared.task_heads import PreFSDPHookResult with infra.init_empty_weights(): model = _DummyModel() @@ -225,14 +226,34 @@ def test_load_before_shard_runs_hook_after_base_load_and_before_key_snapshot(sel def materialize_model(model_to_materialize, device, **_kwargs): model_to_materialize.to_empty(device=device) + # Architecture initializers may cast the whole sharded model to the + # backbone dtype. Infrastructure must restore the isolated task head + # to FP32 before sampling its fresh initialization. + model_to_materialize.to(dtype=torch.bfloat16) def mark_base_checkpoint_loaded(model_to_load, *_args, **_kwargs): + assert all(parameter.device.type == "cpu" for parameter in model_to_load.parameters()) + assert mock_ckpt.config.skip_task_head_prefixes_for_base_model == ["manual_head.", "value_head."] model_to_load.base_checkpoint_loaded = True def add_value_head(model_to_update): - assert model_to_update.base_checkpoint_loaded is True - assert all(parameter.device.type == "cpu" for parameter in model_to_update.parameters()) - model_to_update.value_head = torch.nn.Linear(4, 1) + assert not hasattr(model_to_update, "base_checkpoint_loaded") + assert all(parameter.device.type == "meta" for parameter in model_to_update.parameters()) + + value_head = torch.nn.Linear(4, 1) + assert value_head.weight.device.type == "meta" + original_reset = value_head.reset_parameters + + def reset_after_base_load(): + assert model_to_update.base_checkpoint_loaded is True + assert all(parameter.device.type == "cpu" for parameter in value_head.parameters()) + assert all(parameter.dtype == torch.float32 for parameter in value_head.parameters()) + original_reset() + + value_head.reset_parameters = MagicMock(side_effect=reset_after_base_load) + timeline.attach_mock(value_head.reset_parameters, "reset") + model_to_update.value_head = value_head + return PreFSDPHookResult(task_module=value_head) hook = MagicMock(side_effect=add_value_head) snapshot = MagicMock(side_effect=lambda _model, state_dict, **_kwargs: state_dict) @@ -245,7 +266,7 @@ def add_value_head(model_to_update): patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=1), patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), patch(f"{_INFRA_MODULE}.print_trainable_parameters"), - patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=True), + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf", snapshot), patch(f"{_INFRA_MODULE}._shard_ep_fsdp", shard), patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, @@ -264,69 +285,91 @@ def add_value_head(model_to_update): load_base_model=True, pretrained_model_name_or_path="test/model", pre_fsdp_hook=hook, - skip_task_head_prefixes_for_base_model=("value_head.", "value_head."), + skip_task_head_prefixes_for_base_model=("manual_head.",), ) assert result is model hook.assert_called_once_with(model) assert [mock_call[0] for mock_call in timeline.mock_calls] == [ - "initialize", - "load", "hook", "snapshot", "shard", + "initialize", + "load", + "reset", ] - assert snapshot.call_args.args[1]["value_head.weight"].device.type == "cpu" + assert snapshot.call_args.args[1]["value_head.weight"].device.type == "meta" assert model._pre_shard_hf_state_dict_keys == list(snapshot.call_args.args[1]) - checkpoint_config = MockCheckpointer.call_args.args[0] - assert checkpoint_config.skip_task_head_prefixes_for_base_model == ["value_head."] + assert mock_ckpt.config.skip_task_head_prefixes_for_base_model == ["manual_head.", "value_head."] + model.value_head.reset_parameters.assert_called_once_with() + assert all(parameter.dtype == torch.float32 for parameter in model.value_head.parameters()) - def test_meta_model_runs_hook_before_key_snapshot_and_creates_meta_head(self): + @pytest.mark.parametrize( + "invalid_declaration,error_match", + [ + ("reused", "reuses pre-hook parameters"), + ("existing_module", "existed before the hook"), + ("unattached", "attached below the model root"), + ("external_alias", "shares parameters outside its subtree"), + ("bf16", "only float32 parameters"), + ("frozen", "only trainable parameters"), + ], + ids=["reused_parameters", "reused_module", "not_attached", "external_alias", "bf16", "frozen"], + ) + def test_typed_result_validates_fresh_attached_module( + self, + invalid_declaration, + error_match, + ): from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.shared.task_heads import PreFSDPHookResult - with infra.init_empty_weights(): - model = _DummyModel() - - timeline = MagicMock() - - def add_value_head(model_to_update): - model_to_update.value_head = torch.nn.Linear(4, 1) - assert model_to_update.value_head.weight.device.type == "meta" - - hook = MagicMock(side_effect=add_value_head) - snapshot = MagicMock(side_effect=lambda _model, state_dict, **_kwargs: state_dict) - shard = MagicMock(return_value=model) - timeline.attach_mock(hook, "hook") - timeline.attach_mock(snapshot, "snapshot") - timeline.attach_mock(shard, "shard") + model = _DummyModel() + if invalid_declaration == "existing_module": + model.value_head = torch.nn.Identity() + + def declare_invalid_modules(model_to_update): + if invalid_declaration == "reused": + return PreFSDPHookResult(task_module=model_to_update.linear) + if invalid_declaration == "existing_module": + model_to_update.value_head.projection = torch.nn.Linear(4, 1) + return PreFSDPHookResult(task_module=model_to_update.value_head) + if invalid_declaration == "unattached": + return PreFSDPHookResult(task_module=torch.nn.Linear(4, 1)) + if invalid_declaration == "external_alias": + model_to_update.value_head = torch.nn.Linear(4, 1) + model_to_update.register_parameter("aliased_weight", model_to_update.value_head.weight) + return PreFSDPHookResult(task_module=model_to_update.value_head) + if invalid_declaration == "bf16": + model_to_update.value_head = torch.nn.Linear(4, 1, dtype=torch.bfloat16) + return PreFSDPHookResult(task_module=model_to_update.value_head) + if invalid_declaration == "frozen": + model_to_update.value_head = torch.nn.Linear(4, 1) + model_to_update.value_head.requires_grad_(False) + return PreFSDPHookResult(task_module=model_to_update.value_head) + raise AssertionError(f"unexpected declaration: {invalid_declaration}") + + hook = MagicMock(side_effect=declare_invalid_modules) with ( - patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=2), - patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), - patch(f"{_INFRA_MODULE}.print_trainable_parameters"), - patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), - patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf", snapshot), - patch(f"{_INFRA_MODULE}._shard_ep_fsdp", shard), - patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf") as snapshot, + patch(f"{_INFRA_MODULE}._shard_ep_fsdp") as shard, + patch(f"{_INFRA_MODULE}.Checkpointer"), ): - mock_ckpt = MockCheckpointer.return_value - mock_ckpt.config.dequantize_base_checkpoint = False - - result = infra.apply_model_infrastructure( - model=model, - is_meta_device=True, - device=torch.device("cpu"), - load_base_model=False, - pre_fsdp_hook=hook, - ) + with pytest.raises(ValueError, match=error_match): + infra.apply_model_infrastructure( + model=model, + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + pre_fsdp_hook=hook, + ) - assert result is model hook.assert_called_once_with(model) - assert [mock_call[0] for mock_call in timeline.mock_calls] == ["hook", "snapshot", "shard"] - assert "value_head.weight" in snapshot.call_args.args[1] - assert "value_head.weight" in model._pre_shard_hf_state_dict_keys + snapshot.assert_not_called() + shard.assert_not_called() - def test_rejects_hook_return_value_before_key_snapshot_or_sharding(self): + def test_rejects_invalid_hook_result_before_key_snapshot_or_sharding(self): from nemo_automodel._transformers import infrastructure as infra model = _DummyModel() @@ -340,7 +383,7 @@ def test_rejects_hook_return_value_before_key_snapshot_or_sharding(self): ): MockCheckpointer.return_value.config.dequantize_base_checkpoint = False - with pytest.raises(TypeError, match="mutate the existing model in place and return None"): + with pytest.raises(TypeError, match="must return None or PreFSDPHookResult"): infra.apply_model_infrastructure( model=model, is_meta_device=False, @@ -354,62 +397,304 @@ def test_rejects_hook_return_value_before_key_snapshot_or_sharding(self): shard.assert_not_called() @pytest.mark.parametrize( - "mesh_overrides,infrastructure_overrides,unsupported_name", + "infrastructure_overrides,native_quantization", [ - ({"tp_size": 2}, {}, "tensor parallelism"), - ({"cp_size": 2}, {}, "context parallelism"), - ({"ep_size": 2}, {}, "expert parallelism"), - ({"pp_size": 2}, {}, "pipeline parallelism"), - ({}, {"peft_config": object()}, "PEFT"), - ({}, {"quantization_config": object()}, "quantization"), - ({}, {"fp8_config": object()}, "FP8"), - ({}, {"qat_quantizer": object()}, "QAT"), + ({"peft_config": object()}, False), + ({"quantization_config": object()}, False), + ({"fp8_config": object()}, False), + ({"qat_quantizer": object()}, False), + ({}, True), ], - ids=["tp", "cp", "ep", "pp", "peft", "quantization", "fp8", "qat"], + ids=["peft", "quantization", "fp8", "qat", "native_quantization"], ) - def test_rejects_unsupported_parallelism_and_model_transforms( + def test_managed_result_runs_after_and_is_allowed_with_model_transforms( self, - mesh_overrides, infrastructure_overrides, - unsupported_name, + native_quantization, ): from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.shared.task_heads import PreFSDPHookResult, task_head_module_name - mesh_sizes = {"tp_size": 1, "cp_size": 1, "ep_size": 1, "pp_size": 1} - mesh_sizes.update(mesh_overrides) - hook = MagicMock() + model = _DummyModel() + if native_quantization: + model.config.quantization_config = {"quant_method": "bitsandbytes"} - with pytest.raises(NotImplementedError, match=unsupported_name): - infra.apply_model_infrastructure( - model=_DummyModel(), + timeline = MagicMock() + + def apply_transforms(model_to_update, *_args): + model_to_update.transforms_complete = True + return model_to_update + + def add_value_head(model_to_update): + assert model_to_update.transforms_complete is True + model_to_update.value_head = torch.nn.Linear(4, 1) + return PreFSDPHookResult(task_module=model_to_update.value_head) + + transform = MagicMock(side_effect=apply_transforms) + hook = MagicMock(side_effect=add_value_head) + timeline.attach_mock(transform, "transform") + timeline.attach_mock(hook, "hook") + + with ( + patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=1), + patch(f"{_INFRA_MODULE}._apply_peft_and_lower_precision", transform), + patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), + patch(f"{_INFRA_MODULE}.print_trainable_parameters"), + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), + patch(f"{_INFRA_MODULE}._shard_ep_fsdp", return_value=model), + patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + ): + MockCheckpointer.return_value.config.dequantize_base_checkpoint = False + result = infra.apply_model_infrastructure( + model=model, is_meta_device=False, device=torch.device("cpu"), load_base_model=False, - mesh=SimpleNamespace(**mesh_sizes), pre_fsdp_hook=hook, **infrastructure_overrides, ) - hook.assert_not_called() + assert result is model + assert [mock_call[0] for mock_call in timeline.mock_calls] == ["transform", "hook"] + assert task_head_module_name(model) == "value_head" + assert model.value_head.weight.dtype == torch.float32 + assert all(parameter.requires_grad for parameter in model.value_head.parameters()) - def test_rejects_checkpoint_native_quantization_before_running_hook(self): + def test_pipeline_parallelism_rejects_hook_before_it_runs(self): from nemo_automodel._transformers import infrastructure as infra model = _DummyModel() - model.config.quantization_config = {"quant_method": "bitsandbytes"} hook = MagicMock() - with pytest.raises(NotImplementedError, match="quantization"): + with pytest.raises(NotImplementedError, match="pipeline parallelism"): infra.apply_model_infrastructure( model=model, is_meta_device=False, device=torch.device("cpu"), load_base_model=False, + mesh=SimpleNamespace(pp_size=2), pre_fsdp_hook=hook, ) hook.assert_not_called() + def test_managed_result_rejects_tied_output_head(self): + from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.shared.task_heads import PreFSDPHookResult + + model = _DummyModel() + model.config.tie_word_embeddings = True + + def add_value_head(model_to_update): + model_to_update.value_head = torch.nn.Linear(4, 1) + return PreFSDPHookResult(task_module=model_to_update.value_head) + + hook = MagicMock(side_effect=add_value_head) + + with ( + patch(f"{_INFRA_MODULE}._maybe_adapt_state_dict_to_hf") as snapshot, + patch(f"{_INFRA_MODULE}._shard_ep_fsdp") as shard, + patch(f"{_INFRA_MODULE}.Checkpointer"), + ): + with pytest.raises(ValueError, match="untied output head"): + infra.apply_model_infrastructure( + model=model, + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + pre_fsdp_hook=hook, + ) + + hook.assert_called_once_with(model) + snapshot.assert_not_called() + shard.assert_not_called() + + def test_peft_freeze_keeps_lora_and_managed_task_head_trainable(self): + from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.shared.task_heads import PreFSDPHookResult + + model = _DummyModel() + + def apply_peft(model_to_update, *_args): + model_to_update.linear.register_parameter("lora_A", torch.nn.Parameter(torch.ones(1))) + return model_to_update + + def add_value_head(model_to_update): + assert hasattr(model_to_update.linear, "lora_A") + model_to_update.value_head = torch.nn.Linear(4, 1) + return PreFSDPHookResult(task_module=model_to_update.value_head) + + with ( + patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=1), + patch(f"{_INFRA_MODULE}._apply_peft_and_lower_precision", side_effect=apply_peft), + patch(f"{_INFRA_MODULE}._supports_logits_to_keep", return_value=True), + patch(f"{_INFRA_MODULE}.print_trainable_parameters"), + patch(f"{_INFRA_MODULE}._should_load_before_shard", return_value=False), + patch(f"{_INFRA_MODULE}._shard_ep_fsdp", return_value=model), + patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer, + ): + MockCheckpointer.return_value.config.dequantize_base_checkpoint = False + result = infra.apply_model_infrastructure( + model=model, + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=False, + peft_config=SimpleNamespace(), + pre_fsdp_hook=add_value_head, + ) + + assert result is model + assert model.linear.weight.requires_grad is False + assert model.linear.bias.requires_grad is False + assert model.linear.lora_A.requires_grad is True + assert all(parameter.requires_grad for parameter in model.value_head.parameters()) + + +def _register_test_task_head(model: torch.nn.Module) -> torch.nn.Module: + from nemo_automodel.shared.task_heads import PreFSDPHookResult, register_task_head_module + + pre_hook_module_ids = {id(module) for module in model.modules()} + pre_hook_parameter_ids = {id(parameter) for parameter in model.parameters()} + model.value_head = torch.nn.Linear(4, 1) + register_task_head_module( + model, + PreFSDPHookResult(task_module=model.value_head), + pre_hook_module_ids=pre_hook_module_ids, + pre_hook_parameter_ids=pre_hook_parameter_ids, + ) + return model.value_head + + +def test_shard_ep_fsdp_wraps_fp32_task_head_on_size_one_data_mesh_for_pure_tp(): + from nemo_automodel._transformers import infrastructure as infra + from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager + + model = _DummyModel() + task_head = _register_test_task_head(model) + manager = object.__new__(FSDP2Manager) + manager.offload_policy = object() + data_mesh = SimpleNamespace(size=lambda: 1) + world_mesh = object() + parallelize_fn = MagicMock() + mesh = SimpleNamespace( + device_mesh=world_mesh, + moe_mesh=None, + parallelize_axis_kwargs=MagicMock(return_value={}), + ) + + with ( + patch(f"{_INFRA_MODULE}.get_world_size_safe", return_value=4), + patch(f"{_INFRA_MODULE}.get_fsdp_dp_mesh", return_value=data_mesh) as get_data_mesh, + patch(f"{_INFRA_MODULE}.MixedPrecisionPolicy", return_value="task-head-fp32") as policy, + patch(f"{_INFRA_MODULE}.fully_shard") as fully_shard, + ): + result = infra._shard_ep_fsdp(model, manager, parallelize_fn, mesh) + + assert result is model + assert data_mesh.size() == 1 + assert all(parameter.dtype == torch.float32 for parameter in task_head.parameters()) + get_data_mesh.assert_called_once_with(world_mesh, "dp_replicate", "dp_shard_cp") + policy.assert_called_once_with( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + output_dtype=torch.float32, + cast_forward_inputs=True, + ) + fully_shard.assert_called_once_with( + task_head, + mesh=data_mesh, + mp_policy="task-head-fp32", + offload_policy=manager.offload_policy, + reshard_after_forward=False, + ) + parallelize_fn.assert_called_once_with(model, world_mesh=world_mesh, moe_mesh=None) + + +def test_sync_task_head_replicas_broadcasts_over_tp_but_not_ep(): + from nemo_automodel._transformers import infrastructure as infra + + model = _DummyModel() + task_head = _register_test_task_head(model) + tp_group = object() + ep_group = object() + + class _Submesh: + def __init__(self, group): + self.group = group + + def size(self): + return 2 + + def get_group(self): + return self.group + + class _RootMesh: + mesh_dim_names = ("dp_shard", "ep", "tp") + + def __init__(self): + self.requested_axes = [] + + def __getitem__(self, axis): + self.requested_axes.append(axis) + return {"tp": _Submesh(tp_group), "ep": _Submesh(ep_group)}[axis] + + root_mesh = _RootMesh() + mesh = SimpleNamespace(device_mesh=root_mesh) + + with ( + patch(f"{_INFRA_MODULE}.dist.is_available", return_value=True), + patch(f"{_INFRA_MODULE}.dist.is_initialized", return_value=True), + patch(f"{_INFRA_MODULE}.dist.get_process_group_ranks", return_value=[7, 9]), + patch(f"{_INFRA_MODULE}.dist.get_backend", return_value="gloo"), + patch(f"{_INFRA_MODULE}.dist.broadcast") as broadcast, + ): + infra._sync_task_head_replicas(model, mesh, torch.device("cpu")) + + assert root_mesh.requested_axes == ["tp"] + assert broadcast.call_count == len(tuple(task_head.parameters())) + assert all(mock_call.kwargs == {"src": 7, "group": tp_group} for mock_call in broadcast.call_args_list) + + +def _run_two_rank_task_head_tp_sync(rank, init_file): + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + + from nemo_automodel._transformers.infrastructure import _sync_task_head_replicas + + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=2, + ) + try: + model = _DummyModel() + task_head = _register_test_task_head(model) + with torch.no_grad(): + for parameter in task_head.parameters(): + parameter.fill_(10.0 + rank) + + world_mesh = init_device_mesh("cpu", mesh_shape=(1, 2), mesh_dim_names=("ep", "tp")) + _sync_task_head_replicas( + model, + SimpleNamespace(device_mesh=world_mesh), + torch.device("cpu"), + ) + + for parameter in task_head.parameters(): + torch.testing.assert_close(parameter, torch.full_like(parameter, 10.0), rtol=0.0, atol=0.0) + finally: + dist.destroy_process_group() + + +def test_sync_task_head_replicas_two_rank_gloo(tmp_path): + torch.multiprocessing.spawn( + _run_two_rank_task_head_tp_sync, + args=(str(tmp_path / "task_head_tp_sync"),), + nprocs=2, + join=True, + ) + def test_apply_model_infrastructure_handles_unwrapped_single_rank_ddp_model(): """Single-rank DDP skips wrapping, so the returned model may not have ``.module``.""" diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py index 4d60c83663..07319b3d5c 100644 --- a/tests/unit_tests/checkpoint/test_checkpointing.py +++ b/tests/unit_tests/checkpoint/test_checkpointing.py @@ -74,6 +74,7 @@ materialize_missing_tied_lm_head, ) from nemo_automodel.components.training.rng import RNGState, StatefulRNG, init_all_rng +from nemo_automodel.shared.task_heads import PreFSDPHookResult, register_task_head_module CLOUD_PATH_MODEL = "msc://bucket/step-100/model" CLOUD_PATH_OPTIM = "msc://bucket/step-100/optim" @@ -906,6 +907,36 @@ def fake_get_model_state_dict(model_part, options=None): assert "model.embed_tokens.weight" in saved_state_dict +def test_peft_model_state_saves_lora_and_managed_task_head(): + model = torch.nn.Module() + model.backbone = torch.nn.Linear(2, 2) + model.backbone.weight.requires_grad_(False) + model.backbone.bias.requires_grad_(False) + model.backbone.register_parameter("lora_A", torch.nn.Parameter(torch.ones(1))) + + pre_hook_module_ids = {id(module) for module in model.modules()} + pre_hook_parameter_ids = {id(parameter) for parameter in model.parameters()} + model.task_head = torch.nn.Linear(2, 1) + register_task_head_module( + model, + PreFSDPHookResult(task_module=model.task_head), + pre_hook_module_ids=pre_hook_module_ids, + pre_hook_parameter_ids=pre_hook_parameter_ids, + ) + + with patch( + "nemo_automodel.components.checkpoint.stateful_wrappers.get_model_state_dict", + return_value=model.state_dict(), + ): + saved_state_dict = ModelState(model, is_peft=True).state_dict() + + assert set(saved_state_dict) == { + "base_model.model.backbone.lora_A", + "base_model.model.task_head.bias", + "base_model.model.task_head.weight", + } + + @pytest.mark.parametrize("cpu_offload", [False, True]) def test_model_state_passes_cpu_offload_to_dcp(cpu_offload): model = torch.nn.Linear(2, 2) diff --git a/tests/unit_tests/moe/test_fsdp_mixin.py b/tests/unit_tests/moe/test_fsdp_mixin.py index a0da9965ab..4748c0be70 100644 --- a/tests/unit_tests/moe/test_fsdp_mixin.py +++ b/tests/unit_tests/moe/test_fsdp_mixin.py @@ -77,6 +77,12 @@ def __init__(self, backend, model, has_lm_head=False, has_embed_tokens=False): if has_embed_tokens: model.embed_tokens = MockFSDPModule() + def get_submodule(self, target): + module = self + for name in target.split("."): + module = getattr(module, name) + return module + class MockOuterFSDPMoEModel(MockFSDPModule, MoEFSDPSyncMixin): """Mock MoE model whose FSDP root is the outer wrapper.""" @@ -260,6 +266,33 @@ def isinstance_side_effect(obj, cls): assert model in modules assert moe_model.lm_head in modules + @patch("nemo_automodel.components.moe.fsdp_mixin.isinstance") + def test_iterates_managed_task_head_and_drives_sync_state(self, mock_isinstance): + mock_isinstance.side_effect = _mock_fsdp_isinstance + + model = MockFSDPModule() + task_head = MockFSDPModule() + moe_model = MockMoEModel(MockBackend(enable_fsdp_optimizations=True), model) + moe_model.value_head = task_head + moe_model._nemo_task_head_module_name = "value_head" + + modules = list(_iter_fsdp_modules(moe_model)) + + assert task_head in modules + + task_head._is_last_backward = True + task_head._reshard_after_backward = True + task_head._requires_gradient_sync = True + moe_model.prepare_for_grad_accumulation() + assert task_head._is_last_backward is False + assert task_head._reshard_after_backward is False + assert task_head._requires_gradient_sync is False + + moe_model.prepare_for_final_backward() + assert task_head._is_last_backward is True + assert task_head._reshard_after_backward is True + assert task_head._requires_gradient_sync is True + @patch("nemo_automodel.components.moe.fsdp_mixin.isinstance") def test_iterates_model_embeddings_lm_head(self, mock_isinstance): def isinstance_side_effect(obj, cls): diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index d4770b5793..009647b7b7 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -154,9 +154,13 @@ def __init__(self, *args, **kwargs): # fsdp fsdp_stub = types.ModuleType("torch.distributed.fsdp") + class FSDPModule: + pass + def fully_shard(*args, **kwargs): return None + fsdp_stub.FSDPModule = FSDPModule fsdp_stub.fully_shard = fully_shard fsdp_fully_stub = types.ModuleType("torch.distributed.fsdp._fully_shard") @@ -225,6 +229,7 @@ class CheckpointImpl: REENTRANT = "reentrant" cpw_stub.checkpoint_wrapper = checkpoint_wrapper + cpw_stub._CHECKPOINT_PREFIX = "_checkpoint_wrapped_module." # components/distributed/activation_checkpointing.py imports this at module # scope; without it that module only imports when an earlier test happened to # cache it under real torch, making this file order-dependent. @@ -341,6 +346,14 @@ class GroupedExpertsTE: experts_stub.GroupedExpertsTE = GroupedExpertsTE monkeypatch.setitem(sys.modules, "nemo_automodel.components.moe.experts", experts_stub) + mok_experts_stub = types.ModuleType("nemo_automodel.components.moe.mok_experts") + + class GroupedExpertsMoK: + pass + + mok_experts_stub.GroupedExpertsMoK = GroupedExpertsMoK + monkeypatch.setitem(sys.modules, "nemo_automodel.components.moe.mok_experts", mok_experts_stub) + def _import_parallelizer_with_stubs(monkeypatch): import importlib @@ -463,8 +476,8 @@ def reject_unsupported_mtp_cp_pp(model): activation_checkpointing_stub = types.ModuleType("nemo_automodel.components.distributed.activation_checkpointing") activation_checkpointing_stub.ensure_fsdp_ops_sac_ignored = lambda: None activation_checkpointing_stub.ensure_profiler_ops_sac_ignored = lambda: None - activation_checkpointing_stub.transformer_engine_attention_backend_snapshot_context_fn = ( - lambda context_fn=None: context_fn() if context_fn is not None else (nullcontext(), nullcontext()) + activation_checkpointing_stub.transformer_engine_attention_backend_snapshot_context_fn = lambda context_fn=None: ( + context_fn() if context_fn is not None else (nullcontext(), nullcontext()) ) monkeypatch.setitem( sys.modules, @@ -966,6 +979,31 @@ def test_apply_fsdp_skips_separate_wrapping_for_tied_embeddings(monkeypatch): assert outer_call is not None and outer_call[1]["mesh"] is fsdp_mesh +def test_apply_fsdp_does_not_double_wrap_pre_wrapped_task_head(monkeypatch): + P = _import_parallelizer_with_stubs(monkeypatch) + monkeypatch.setattr(P, "MoE", DummyMoE) + + fully_shard_mock = MagicMock() + monkeypatch.setattr(P, "fully_shard", fully_shard_mock) + monkeypatch.setattr(P, "MixedPrecisionPolicy", MagicMock(return_value="MP_POLICY")) + + block = DummyBlock(mlp=DummyMoE()) + task_head = P.FSDPModule() + model = DummyModel([block], lm_head=task_head) + fsdp_mesh = object() + + P.apply_fsdp( + model=model, + fsdp_mesh=fsdp_mesh, + ep_enabled=True, + ep_shard_enabled=False, + ) + + assert _find_call_by_first_arg(fully_shard_mock, task_head) is None + assert _find_call_by_first_arg(fully_shard_mock, block) is not None + assert _find_call_by_first_arg(fully_shard_mock, model) is not None + + def test_apply_fsdp_rejects_cross_root_tied_embeddings_without_outer_wrap(monkeypatch): P = _import_parallelizer_with_stubs(monkeypatch) monkeypatch.setattr(P, "MoE", DummyMoE) diff --git a/tests/unit_tests/shared/test_task_heads.py b/tests/unit_tests/shared/test_task_heads.py new file mode 100644 index 0000000000..d71c3b6433 --- /dev/null +++ b/tests/unit_tests/shared/test_task_heads.py @@ -0,0 +1,87 @@ +# 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. + +import pytest +from torch import nn + +from nemo_automodel.shared.task_heads import ( + PreFSDPHookResult, + exclude_task_heads_from_tp_plan, + is_task_head_parameter, + register_task_head_module, + task_head_module, +) + + +class _ProjectionBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.projection = nn.Linear(4, 4) + + +class _ModelWithTaskHead(nn.Module): + def __init__(self) -> None: + super().__init__() + self.backbone = _ProjectionBlock() + + +def _model_with_registered_task_head() -> _ModelWithTaskHead: + model = _ModelWithTaskHead() + pre_hook_module_ids = {id(module) for module in model.modules()} + pre_hook_parameter_ids = {id(parameter) for parameter in model.parameters()} + model.task_head = _ProjectionBlock() + register_task_head_module( + model, + PreFSDPHookResult(task_module=model.task_head), + pre_hook_module_ids=pre_hook_module_ids, + pre_hook_parameter_ids=pre_hook_parameter_ids, + ) + return model + + +def test_exclude_task_heads_from_tp_plan_keeps_only_backbone_rules() -> None: + model = _model_with_registered_task_head() + backbone_style = object() + + filtered = exclude_task_heads_from_tp_plan( + model, + { + "backbone.projection": backbone_style, + "task_head.projection": object(), + "task_head.*": object(), + }, + ) + + assert filtered == {"backbone.projection": backbone_style} + + +def test_exclude_task_heads_from_tp_plan_rejects_mixed_wildcard() -> None: + model = _model_with_registered_task_head() + + with pytest.raises(ValueError, match="matches both the managed task module and backbone modules"): + exclude_task_heads_from_tp_plan(model, {"*.projection": object()}) + + +def test_task_head_module_resolves_below_transparent_wrapper() -> None: + model = _model_with_registered_task_head() + wrapper = nn.Module() + wrapper.module = model + + assert task_head_module(wrapper) is model.task_head + + compiled_wrapper = nn.Module() + compiled_wrapper._orig_mod = model + + assert task_head_module(compiled_wrapper) is model.task_head + assert is_task_head_parameter(compiled_wrapper, "_orig_mod.task_head.projection.weight") From c4856fe280b9139cdcecc26d4d21d3de119950b5 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sat, 22 Aug 2026 14:11:14 -0700 Subject: [PATCH 29/34] feat(data): align VLM Datum side channels Signed-off-by: HuiyingLi --- nemo_automodel/components/datasets/datum.py | 307 +++++++++++++--- .../components/datasets/vlm/collate_fns.py | 17 +- .../components/datasets/vlm/datasets.py | 75 +--- .../datasets/vlm/neat_packing_vlm.py | 34 +- .../components/datasets/vlm/utils.py | 87 +++++ tests/unit_tests/datasets/test_datum.py | 337 +++++++++++++++++- tests/unit_tests/test_engine.py | 62 ++++ 7 files changed, 772 insertions(+), 147 deletions(-) diff --git a/nemo_automodel/components/datasets/datum.py b/nemo_automodel/components/datasets/datum.py index 8f378b62d4..f63a97dd90 100644 --- a/nemo_automodel/components/datasets/datum.py +++ b/nemo_automodel/components/datasets/datum.py @@ -491,6 +491,177 @@ def collate_datums( ) +def _pad_leading_axis(value: torch.Tensor, length: int, pad_value: float | int | bool) -> torch.Tensor: + """Right-pad a token tensor while preserving arbitrary trailing axes. + + Args: + value: Tensor with shape ``[tokens, ...]``. + length: Requested leading-axis length. + pad_value: Scalar used for the appended token rows. + + Returns: + A tensor with shape ``[length, ...]``. + """ + padding = length - value.shape[0] + if padding < 0: + raise ValueError(f"cannot pad a leading axis of length {value.shape[0]} to {length}") + if padding == 0: + return value + tail = torch.full( + (padding, *value.shape[1:]), + pad_value, + dtype=value.dtype, + device=value.device, + ) + return torch.cat((value, tail), dim=0) + + +def _resolve_vlm_loss_contract( + datums: list[Datum], +) -> tuple[ + dict[str, LossInputLayout], + dict[str, float | int | bool], + dict[str, str], +]: + """Resolve VLM loss layouts and their pre-shift token conventions. + + Args: + datums: Processor-ready VLM items with unshifted token streams. + + Returns: + The resolved layouts, explicit padding values, and ``S`` versus + ``S-1`` convention for each ``PER_TOKEN`` field. + """ + loss_keys = set(datums[0].loss_fn_inputs) + missing_required = {"labels", "weights"} - loss_keys + if missing_required: + raise ValueError(f"VLM Datums require loss inputs {sorted(missing_required)}") + for index, datum in enumerate(datums[1:], start=1): + if set(datum.loss_fn_inputs) != loss_keys: + raise ValueError( + "every VLM Datum must contain the same loss_fn_inputs keys; " + f"Datum 0 has {sorted(loss_keys)} and Datum {index} has {sorted(datum.loss_fn_inputs)}" + ) + + layouts: dict[str, LossInputLayout] = {} + pad_values: dict[str, float | int | bool] = {} + token_conventions: dict[str, str] = {} + for key in sorted(loss_keys): + declared_layouts = [datum.loss_fn_input_layouts[key] for datum in datums if key in datum.loss_fn_input_layouts] + if declared_layouts and len(declared_layouts) != len(datums): + raise ValueError(f"every Datum must declare the layout for VLM loss input {key!r}, or none may declare it") + if len(set(declared_layouts)) > 1: + raise ValueError(f"every VLM Datum must use the same explicit layout for loss input {key!r}") + if declared_layouts: + layout = declared_layouts[0] + elif key in {"labels", "weights"}: + layout = LossInputLayout.PER_TOKEN + else: + raise ValueError(f"VLM loss input {key!r} must declare an explicit LossInputLayout") + if key in {"labels", "weights"} and layout is not LossInputLayout.PER_TOKEN: + raise ValueError(f"VLM {key} must use the PER_TOKEN layout") + layouts[key] = layout + + declared_pad_values = [ + datum.loss_fn_input_pad_values[key] for datum in datums if key in datum.loss_fn_input_pad_values + ] + if declared_pad_values and len(declared_pad_values) != len(datums): + raise ValueError( + f"every Datum must declare the pad value for VLM loss input {key!r}, or none may declare it" + ) + if len(set(declared_pad_values)) > 1: + raise ValueError(f"every VLM Datum must use the same pad value for loss input {key!r}") + if declared_pad_values: + pad_values[key] = declared_pad_values[0] + elif key == "labels": + pad_values[key] = CROSS_ENTROPY_IGNORE_IDX + + if layout is not LossInputLayout.PER_TOKEN: + continue + + values = [datum.loss_fn_inputs[key] for datum in datums] + if key in {"labels", "weights"} and any(value.ndim != 1 for value in values): + raise ValueError(f"VLM {key} must be one-dimensional") + conventions = [] + for value, datum in zip(values, datums): + if value.ndim < 1: + raise ValueError(f"PER_TOKEN VLM loss input {key!r} must have a leading token axis") + if value.shape[0] == datum.seq_len: + conventions.append("S") + elif value.shape[0] == datum.seq_len - 1: + conventions.append("S-1") + else: + raise ValueError( + f"PER_TOKEN VLM loss input {key!r} must use leading length S or S-1; " + f"got shape {tuple(value.shape)} for S={datum.seq_len}" + ) + if len(set(conventions)) != 1: + raise ValueError(f"PER_TOKEN VLM loss input {key!r} cannot mix S and S-1 conventions across Datums") + token_conventions[key] = conventions[0] + + if pad_values.get("labels", CROSS_ENTROPY_IGNORE_IDX) != CROSS_ENTROPY_IGNORE_IDX: + raise ValueError(f"VLM labels must use pad value {CROSS_ENTROPY_IGNORE_IDX}") + if pad_values.get("weights", 0) != 0: + raise ValueError("VLM weights must use pad value 0") + if token_conventions["weights"] == "S" and any(bool(datum.loss_fn_inputs["weights"][0] != 0) for datum in datums): + raise ValueError("VLM weight at target position zero must be zero before autoregressive shift") + return layouts, pad_values, token_conventions + + +def _collate_vlm_loss_inputs( + datums: list[Datum], + *, + layouts: dict[str, LossInputLayout], + pad_values: dict[str, float | int | bool], + token_conventions: dict[str, str], + packed: bool, + sequence_alignment: int, + padding_idx: int, +) -> CollatedLossInputs: + """Apply the VLM autoregressive/packing plan to every loss side channel. + + Args: + datums: Processor-ready VLM items with token length ``S``. + layouts: Resolved layout for every loss field. + pad_values: Explicit per-field token padding values. + token_conventions: ``S`` or ``S-1`` convention for token fields. + packed: Whether to concatenate aligned THD documents. + sequence_alignment: Per-document alignment in packed mode. + padding_idx: Token id used in the throwaway alignment template. + + Returns: + Layout-aware loss inputs on the shifted prediction-token axis. + """ + if packed: + from nemo_automodel.components.datasets.vlm.neat_packing_vlm import _aligned_length + + normalized = [] + for datum in datums: + prediction_length = datum.seq_len - 1 + output_length = _aligned_length(prediction_length, sequence_alignment) if packed else prediction_length + loss_inputs = {} + for key, value in datum.loss_fn_inputs.items(): + if layouts[key] is LossInputLayout.PER_TOKEN: + value = value[1:] if token_conventions[key] == "S" else value + if packed: + value = _pad_leading_axis(value, output_length, pad_values.get(key, 0)) + loss_inputs[key] = value + + input_ids = datum.input_ids[:-1] + if packed: + input_ids = _pad_leading_axis(input_ids, output_length, padding_idx) + normalized.append( + Datum( + model_inputs={"input_ids": input_ids}, + loss_fn_inputs=loss_inputs, + loss_fn_input_layouts=layouts, + loss_fn_input_pad_values=pad_values, + ) + ) + + return collate_datums(normalized, packed=packed)[1] + + def collate_vlm_datums( datums: list[Datum], *, @@ -499,22 +670,31 @@ def collate_vlm_datums( get_rope_index: Callable[..., object] | None = None, sequence_alignment: int = 1, ) -> tuple[dict[str, Any], CollatedLossInputs]: - """Collate pre-tokenized VLM SFT Datums with processor-specific media inputs. + """Collate processor-ready VLM Datums with aligned loss side channels. The input Datums retain their unshifted token stream because :func:`~nemo_automodel.components.datasets.vlm.collate_fns.pad_collate_fn` - owns the autoregressive shift. ``labels`` and ``weights`` therefore also - use the unshifted token axis, with zero weight at target position zero. - In packed mode, every Datum becomes one THD document and the collater - preserves its real and aligned sequence lengths. Common VLM fields are - handled by the canonical VLM collaters; additional processor tensor fields - retain their leading media axis and are concatenated. + owns the autoregressive shift. A ``PER_TOKEN`` loss field may use the + unshifted target-token axis ``[S, ...]`` (the collater takes ``[1:]``) or + the already shifted prediction axis ``[S-1, ...]``. One field must use the + same convention in every Datum. Source-aligned metadata such as routing + replay must be converted by its owner to ``[:-1]`` before being passed as + an ``S-1`` field; its direction cannot be inferred from shape alone. + + In packed mode, every Datum becomes one THD document. Each loss field is + padded inside that document before concatenation, preserving nonzero pad + sentinels and arbitrary trailing feature axes. ``token_type_ids`` and + ``mm_token_type_ids`` follow the same model-token shift/alignment. Other + processor media tensors retain their leading media axis and are + concatenated. Args: datums: Non-empty processor-ready VLM items. Each item has - ``input_ids``, ``labels``, and ``weights`` tensors of shape - ``[sequence]`` on matching unshifted token axes. Optional processor - fields carry their processor-defined token or leading media axes. + 1-D ``input_ids`` and ``attention_mask`` plus mandatory ``labels`` + and ``weights`` loss inputs. Additional loss inputs must declare a + :class:`LossInputLayout`; ``PER_TOKEN`` values use a leading ``S`` + or ``S-1`` axis. Optional processor fields carry their + processor-defined token or leading media axes. processor: Hugging Face processor (or compatible object) that supplies the tokenizer padding token. packed: Pack the Datums as THD documents instead of padding a batch. @@ -526,11 +706,12 @@ def collate_vlm_datums( Returns: A pair of shifted/padded model inputs and layout-aware loss inputs. - In padded mode, token model fields, labels, and weights have shape - ``[batch, padded_sequence - 1]``. In packed mode they have shape - ``[1, aligned_tokens]`` and model inputs include THD sequence metadata. - Media tensors retain arbitrary trailing dimensions and are concatenated - on their leading media axis. + In padded mode, token model fields and ``PER_TOKEN`` loss fields have + shape ``[batch, padded_sequence - 1, ...]``. In packed mode they have + shape ``[1, aligned_tokens, ...]`` and model inputs include THD sequence + metadata. ``PER_DATUM`` and ``REPLICATED`` values retain their generic + Datum semantics. Media tensors retain arbitrary trailing dimensions and + are concatenated on their leading media axis. """ if not datums: raise ValueError("collate_vlm_datums requires at least one Datum") @@ -538,21 +719,48 @@ def collate_vlm_datums( raise ValueError(f"sequence_alignment must be a positive integer, got {sequence_alignment!r}") from nemo_automodel.components.datasets.vlm.collate_fns import pad_collate_fn + from nemo_automodel.components.datasets.vlm.utils import _media_token_mismatch + + tokenizer = getattr(processor, "tokenizer", processor) + padding_idx = getattr(tokenizer, "pad_token_id", 0) or 0 examples = [] - for datum in datums: - labels = datum.loss_fn_inputs.get("labels") - weights = datum.loss_fn_inputs.get("weights") + for index, datum in enumerate(datums): input_ids = datum.model_inputs.get("input_ids") - if not all(isinstance(value, torch.Tensor) and value.ndim == 1 for value in (input_ids, labels, weights)): - raise ValueError("VLM Datums require 1-D input_ids, labels, and weights") + attention_mask = datum.model_inputs.get("attention_mask") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 1: + raise ValueError("VLM Datums require 1-D input_ids") + if ( + not isinstance(attention_mask, torch.Tensor) + or attention_mask.ndim != 1 + or attention_mask.shape != input_ids.shape + ): + raise ValueError("VLM Datums require a 1-D attention_mask matching input_ids") if datum.seq_len < 2: raise ValueError("VLM Datums require at least two tokens for autoregressive shifting") - if labels.shape != input_ids.shape or weights.shape != input_ids.shape: - raise ValueError("VLM labels and weights must match the unshifted input_ids shape") - if bool(weights[0] != 0): - raise ValueError("VLM weight at target position zero must be zero before autoregressive shift") - examples.append({**datum.model_inputs, "labels": labels}) + if packed and bool((attention_mask == 0).any()): + raise ValueError("packed VLM Datums cannot contain pre-existing attention_mask padding") + + mismatch = _media_token_mismatch(input_ids, datum.model_inputs, processor) + if mismatch is not None: + raise ValueError(f"VLM media token mismatch for Datum {index}: {mismatch}") + shifted_mismatch = _media_token_mismatch(input_ids[:-1], datum.model_inputs, processor) + if shifted_mismatch is not None: + raise ValueError( + f"VLM media token mismatch after autoregressive shift for Datum {index}: {shifted_mismatch}" + ) + + examples.append( + { + **datum.model_inputs, + # Canonical VLM collaters own the model-input shift and require + # a labels field. Real labels are collated below with every + # other loss side channel. + "labels": torch.full_like(input_ids, CROSS_ENTROPY_IGNORE_IDX), + } + ) + + layouts, loss_pad_values, token_conventions = _resolve_vlm_loss_contract(datums) if packed: from nemo_automodel.components.datasets.vlm.collate_fns import packed_sequence_thd_vlm_collater @@ -574,15 +782,11 @@ def collate_vlm_datums( raise ValueError("get_rope_index must return position IDs for every VLM Datum or none of them") has_mrope = bool(position_ids and position_ids[0] is not None) shifted_examples = [] - shifted_weights = [] - for example, position, datum in zip(examples, position_ids, datums): + for example, position in zip(examples, position_ids): if position is not None: example["position_ids"] = position shifted_examples.append(_shift_sample(example, has_mrope=has_mrope)) - shifted_weights.append(datum.loss_fn_inputs["weights"][1:]) - tokenizer = getattr(processor, "tokenizer", processor) - padding_idx = getattr(tokenizer, "pad_token_id", 0) or 0 pack_size = sum(_aligned_length(item["input_ids"].shape[0], sequence_alignment) for item in shifted_examples) packed_sample = _build_packed_vlm_sample( shifted_examples, @@ -592,22 +796,27 @@ def collate_vlm_datums( sequence_alignment=sequence_alignment, ) model_inputs = packed_sequence_thd_vlm_collater([packed_sample], padding_idx=padding_idx) - labels = model_inputs.pop("labels") - weights = torch.cat( - [ - F.pad(weight, (0, _aligned_length(weight.shape[0], sequence_alignment) - weight.shape[0])) - for weight in shifted_weights - ] - ).unsqueeze(0) + model_inputs.pop("labels") else: model_inputs = pad_collate_fn(examples, processor) - labels = model_inputs.pop("labels") - - target_width = labels.shape[-1] + 1 - shifted_weights = [ - F.pad(datum.loss_fn_inputs["weights"], (0, target_width - datum.seq_len))[1:] for datum in datums - ] - weights = torch.stack(shifted_weights) + model_inputs.pop("labels") + attention_mask = model_inputs.get("attention_mask") + if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2: + # pad_collate_fn pads before shifting, so a shorter row otherwise + # retains its final no-target source token as an apparent real + # prediction position. Keep model and loss/output axes identical. + for row, datum in enumerate(datums): + attention_mask[row, datum.seq_len - 1 :] = 0 + + loss_inputs = _collate_vlm_loss_inputs( + datums, + layouts=layouts, + pad_values=loss_pad_values, + token_conventions=token_conventions, + packed=packed, + sequence_alignment=sequence_alignment, + padding_idx=padding_idx, + ) unhandled_keys = set().union(*(datum.model_inputs.keys() for datum in datums)) - set(model_inputs) unhandled_keys -= {"input_ids", "attention_mask"} @@ -622,12 +831,4 @@ def collate_vlm_datums( except RuntimeError as exc: raise ValueError(f"VLM processor field {key!r} cannot be concatenated across samples") from exc - return model_inputs, CollatedLossInputs( - {"labels": labels, "weights": weights}, - layouts={ - "labels": LossInputLayout.PER_TOKEN, - "weights": LossInputLayout.PER_TOKEN, - }, - item_to_datum=tuple(range(len(datums))), - pad_values={"labels": CROSS_ENTROPY_IGNORE_IDX}, - ) + return model_inputs, loss_inputs diff --git a/nemo_automodel/components/datasets/vlm/collate_fns.py b/nemo_automodel/components/datasets/vlm/collate_fns.py index 59241901f2..7ba2e446cb 100644 --- a/nemo_automodel/components/datasets/vlm/collate_fns.py +++ b/nemo_automodel/components/datasets/vlm/collate_fns.py @@ -1434,8 +1434,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: @@ -1670,8 +1670,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,6 +1758,14 @@ 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: diff --git a/nemo_automodel/components/datasets/vlm/datasets.py b/nemo_automodel/components/datasets/vlm/datasets.py index 7bc761b25d..7071fbd34f 100644 --- a/nemo_automodel/components/datasets/vlm/datasets.py +++ b/nemo_automodel/components/datasets/vlm/datasets.py @@ -39,6 +39,7 @@ from nemo_automodel.components.datasets.vlm.utils import ( _build_video_metadata, _lmdb_env_cache, + _media_token_mismatch, _preload_media, json2token, ) @@ -1308,80 +1309,6 @@ def _load_one_dataset(ds_name, ds_config): return result -def _resolve_processor_token_id(processor, attr_names, token_names): - """Resolve a model-specific media token id from processor/config/tokenizer.""" - tokenizer = getattr(processor, "tokenizer", processor) - unk_id = getattr(tokenizer, "unk_token_id", None) - - config = getattr(processor, "config", None) - for source in (processor, config, tokenizer, getattr(tokenizer, "config", None)): - if source is None: - continue - for attr in attr_names: - value = getattr(source, attr, None) - if isinstance(value, int): - return value - if isinstance(value, str): - try: - token_id = tokenizer.convert_tokens_to_ids(value) - except Exception: - continue - if isinstance(token_id, int) and token_id != unk_id: - return token_id - for token in token_names: - try: - token_id = tokenizer.convert_tokens_to_ids(token) - except Exception: - continue - if isinstance(token_id, int) and token_id != unk_id: - return token_id - return None - - -def _grid_media_token_count(grid, merge_size: int) -> int: - if grid is None: - return 0 - grid_t = torch.as_tensor(grid) - if grid_t.numel() == 0: - return 0 - if grid_t.ndim == 1: - grid_t = grid_t.view(1, -1) - merge_len = int(merge_size) ** 2 - return int((grid_t.to(torch.long).prod(dim=-1) // merge_len).sum().item()) - - -def _media_token_mismatch(input_ids, result, processor) -> str | None: - """Return a mismatch description if media grids survived without tokens.""" - image_token_id = _resolve_processor_token_id( - processor, - ("image_token_id", "image_token_index", "image_token"), - ("<|image_pad|>", "", "<|image|>"), - ) - video_token_id = _resolve_processor_token_id( - processor, - ("video_token_id", "video_token_index", "video_token"), - ("<|video_pad|>", "