diff --git a/README.md b/README.md index ec4bbcd4a9..50ac075dd2 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,7 @@ High-throughput scalable training - ✅ **DCP** - Supports PyTorch DCP and SafeTensors, sharded and consolidated layouts, merge/reshard utilities, and Hugging Face-compatible outputs. - ✅ **Async checkpointing** - Can write checkpoints in the background to reduce training stalls caused by I/O. - ✅ **Dion and Muon optimizers** - Distributed optimizer integrations with typed recipe configuration. +- ✅ **Training Engine** - DeepSpeed-style eager `forward`, `backward`, and `step` over already-distributed models. - ✅ **Environment Support** - SLURM, interactive, SkyPilot, and Kubernetes (via SkyPilot) launchers. SOTA algorithms @@ -190,7 +191,7 @@ Agentic Development and UX - ✅ **Agent-friendly skills** - Curated [`skills/`](https://github.com/NVIDIA-NeMo/Automodel/tree/main/skills) for common dev tasks (recipe runs, model onboarding, CI). Planned for 26.08 -- 🔜 **Unified Engine API and recipes** - Introduce a common engine and consolidate the LLM and VLM recipe paths. +- 🔜 **Pipeline and recipe consolidation** - Consolidate the LLM and VLM pipeline-training lifecycle behind one clear recipe boundary. - 🔜 **Composable component configuration** - Complete the typed config and `.build()` refactor across data and remaining components. - 🔜 **Packed long-context training with CP** - Combine THD sequence packing with context parallelism, including DeepSeek V4 coverage. - 🔜 **Kernel and runtime upgrades** - Add partial CUDA graphs, evaluate FlashAttention 3/4, and upgrade to DeepEP v2. diff --git a/docs/api-reference/index.mdx b/docs/api-reference/index.mdx index 4cac36488e..3999e4462c 100644 --- a/docs/api-reference/index.mdx +++ b/docs/api-reference/index.mdx @@ -19,5 +19,6 @@ This reference is built from docstrings in the [source code](https://github.com/ | `nemo_automodel.components.checkpoint` | Async DCP and SafeTensors checkpointing | | `nemo_automodel.components.quantization` | FP8, QAT, calibration | | `nemo_automodel.components._peft` | LoRA, QLoRA adapters | +| `nemo_automodel.engine` | DeepSpeed-style eager model wrapper for backward, gradient accumulation, clipping, optimizer updates, and scheduler steps | | `nemo_automodel.components.launcher` | Interactive, SkyPilot, and NeMo-Run job launchers; Slurm jobs use the repository's `slurm.sub` script | | `nemo_automodel.cli` | `uv run automodel [--nproc-per-node N] [--key.subkey=override ...]` entry point | diff --git a/engine_integration_before_after.md b/engine_integration_before_after.md new file mode 100644 index 0000000000..e1b76557f9 --- /dev/null +++ b/engine_integration_before_after.md @@ -0,0 +1,194 @@ +# Molt ↔ AutoModel Engine Integration: Responsibility Changes + +Current status of the two branches (2026-08-30): + +| Repo | Branch / PR | Diff vs main | Status | +| --- | --- | --- | --- | +| NVIDIA-NeMo/Automodel | `huiyingl/feat/datum-forward-backward` / [#3614](https://github.com/NVIDIA-NeMo/Automodel/pull/3614) | 53 files, +3762/−935 | draft, CI green, full unit suite 13339 passed | +| NVIDIA-NeMo/labs-molt | `huiyingl/feat/sft-automodel-engine-integration` / [#92](https://github.com/NVIDIA-NeMo/labs-molt/pull/92) | 49 files, +2599/−2939 | draft, CI green, 234 unit tests passed, requirements pin the AM branch HEAD | +| verl-project/verl | `huiyingl/automodel-engine-v2` | 2 files, +1505/−336 | branch only, 48 CPU adapter tests passed against the AM branch | + +Related PRs: HybridEP token equalization [#3641](https://github.com/NVIDIA-NeMo/Automodel/pull/3641) (merged into AM main); VLM PP validation [#3759](https://github.com/NVIDIA-NeMo/Automodel/pull/3759) and the NemotronParse loss fix [#3760](https://github.com/NVIDIA-NeMo/Automodel/pull/3760) (independent drafts split out of #3614). + +--- + +## Part 1 — RL (molt) side: before → after + +Only responsibilities that changed hands are listed. Optimizer CPU offload, the critic value head, RL losses, RL input packing, and checkpoint/refit/Ray orchestration belong to molt both before and after, so they are not in this list. + +### 1. Training execution loop + +**Before**: `molt/trainer/fsdp/strategy.py` implemented the whole execution path in its `backward()` / `optimizer_step()` — gradient-accumulation bookkeeping, deferred FSDP sync via `no_sync`, distributed gradient finalization across EP/TP (calling AM's low-level `scale_grads_and_clip_grad_norm`), clipping, `optimizer.step()`, `zero_grad`, and scheduler advancement. Trainers had to interact with this state machine on every step. + +**After**: the whole execution path belongs to AM's **`nemo_automodel/engine/_engine.py` (Engine)**. Molt trainers write the algorithm in its natural order: + +```python +out = actor(...) # Engine.forward (ordinary nn.Module call) +loss = policy_loss(...) # RL loss stays in the molt trainer +actor.model.backward(loss) # Engine.backward +actor.model.step() # Engine: finalize + clip + step + scheduler at the accumulation boundary +``` + +`FsdpStrategy` no longer implements training execution; it keeps only runtime glue: optimizer CPU offload (`CpuOptimizerOffloader`) and the value head's `sync_replicated_grads`. + +### 2. Log-probs / entropy under TP (vocabulary-parallel statistics) + +**Before**: `molt/trainer/fsdp/packing.py::log_probs_from_vocab_parallel_logits` — when TP shards logits into a vocab-sharded DTensor, molt computed selected-token log-probs itself without gathering the vocabulary; entropy went through molt's `compute_entropy` + `unshard_dtensor` (which first materializes the vocab axis). + +**After**: `molt/trainer/fsdp/packing.py` is deleted. The computation belongs to AM's **`nemo_automodel/components/loss/vocab_parallel.py`** — `token_log_probs` / `token_entropy`, one entry point for both dense tensors and vocab-sharded DTensors, and entropy never gathers the vocabulary either. The molt actor is down to two imports and two calls. + +### 3. Routing Replay (R3) lifecycle management + +**Before**: `molt/models/base.py` (~lines 452-476) managed AM's per-gate `RouterReplay` handles itself: clearing the global registry, walking the module tree to install a handle on every MoE gate, subclassing a `_SentinelRouterReplay` to implement the `-1` (keep-live-selection) sentinel, and relying on construction order for layer alignment. Molt had to understand the internal structure of AM gates. + +**After**: this belongs to AM's **`nemo_automodel/components/moe/router_replay.py::RouterReplayAdapter`** — gate discovery by decoder-layer id, handle binding, `-1`-row fallback to live routing, trailing-padded-token handling, and state restoration on exceptions are all inside AM. Molt only does `adapter = RouterReplayAdapter(model)` and wraps the forward in `with adapter.replay(routes):`, with routes prepared in the model input's token order. + +### 4. VLM media merging / on-the-fly packing primitives + +**Before**: `molt/utils/vlm_utils.py` carried its own `_pad_to_common_hw` and merging logic to pad variable-resolution image patches to a common shape and concatenate them into a batch tensor; there was no reusable component for packed-VLM sample boxing. + +**After**: merging/padding belongs to AM's **`components/datasets/vlm/utils.py::merge_media_values`** (including 4-D variable-resolution padding); sample boxing belongs to AM's **`components/datasets/vlm/neat_packing_vlm.py::pack_vlm_samples`** plus the two collaters. Molt's `pack_vlm_batch` (`molt/models/packing.py`) keeps only the RL-semantic part (valid-span extraction, restore indices) and calls AM for the physical packing. + +### 5. HybridEP token equalization for dynamic batches + +**Before**: nobody owned this — rollout-produced variable-length/packed batches give each rank a different token count, and the hybridep backend deadlocks or SIGABRTs (a 4-token alignment constraint), so RL dynamic batches could not train on hybridep at all. + +**After**: this belongs to AM's **`components/moe/megatron/token_dispatcher.py::_HybridEPManager.dispatch()`** (#3641, already in AM main): all-reduce the EP-group max token count, round up to the 4-token alignment, pad rows route to no expert, and `combine()` slices the padding back off. Callers (molt and AM recipes alike) are completely unaware of the constraint. + +--- + +## Part 2 — veRL side: before → after + +Branch `huiyingl/automodel-engine-v2` updates veRL's AutoModel backend (upstream PR #5407 lineage) to the same boundary. The diff vs verl main is deliberately narrow: only `verl/workers/engine/automodel/transformer_impl.py` (±930 lines) plus a new CPU contract-test file (+911 lines). + +### 1. Training execution loop + +**Before**: verl main's AutoModel adapter hand-rolled the execution loop out of AM internals — manual `prepare_for_grad_accumulation` / `prepare_for_final_backward` calls, a manually set `MoEAuxLossAutoScaler.main_loss_backward_scale` (and only when ep>1), and an `optimizer_step` that invoked `scale_grads_and_clip_grad_norm` itself. The same shape molt main had: the RL framework reimplements the execution layer. + +**After**: the loop is the standard calls — `engine(**inputs)` → veRL's own loss → `engine.backward(loss, scale_wrt_gas=False)`, with `set_gradient_accumulation_steps(n)` making the last microbatch the boundary forward (deferred FSDP gradient sync re-enabled there so its backward reduces). The boundary microstep stays pending until veRL's `optimizer_step()`, which closes the window — Engine finalizes, clips, and updates there — preserving the split contract Tinker-style callers rely on (optimizer adjustments between backward and step). Gradients are numerically identical to main's raw `loss.backward()` under FSDP2's averaging reducer, while MoE aux-loss scaling and MegatronFSDP summed-gradient compensation move into the Engine. A failed microstep resets the window (`Engine.reset_accumulation`) so an OOM cannot poison later windows; a second `forward_backward_batch` before the step fails closed. The scheduler stays veRL-owned. + +### 2. Log-probs / entropy under TP + +**Before**: `prepare_model_outputs` called `full_tensor()` on TP-sharded logits — gathering the whole vocabulary before computing log-probs, which blows up memory on large-vocab models. + +**After**: vocab-sharded DTensor logits go through AM's `token_log_probs` / `token_entropy` (the same primitives molt uses), never gathering the vocabulary. + +### 3. Optimizer construction + +**Before**: imported `build_optimizer` from `nemo_automodel.recipes.llm.train_ft` and went through `ConfigNode` — reaching into AM recipe internals. + +**After**: `components.optim.build_optimizer` public API with explicit kwargs, an `override_optimizer_config` escape hatch, and fail-fast rejection of fp16 optimizer states. + +### 4. Adapter hardening (from the intermediate commit this migration completes) + +- Systematic fail-closed validation: PP, CP, LoRA, activation offload, router replay, and fp16 precisions raise explicitly instead of running wrong. +- Packed inputs distinguish TE THD from FlashAttention indexed-mask layouts. +- vLLM refit's `get_per_tensor_param` converts per tensor via the AM adapter's `convert_single_tensor_to_hf` (main's `convert_weight_keys` mis-maps custom-model weight layouts) and fails fast on non-DTensor EP expert weights. +- Checkpointing goes through AM's `Checkpointer` (DCP/safetensors/consolidated); veRL keeps ownership of the scheduler/RNG/step extra payload. + +### 5. Tests + +**Before**: no unit coverage for the adapter. **After**: 48 CPU tests — mock-Engine tests pinning the call-sequence contract (one GAS window per mini-batch, `scale_wrt_gas=False`, forward-only never touches the training state machine), plus one end-to-end test driving a real nemo-automodel Engine through two accumulating microbatches with exactly one optimizer update at the boundary. + + +--- + +## Part 3 — AM side: the Engine and the modules serving RL + +### Class hierarchy — inputs, outputs, ownership + +``` +RL framework (molt / veRL) ── owns: batch prep, losses, window timing +│ +├── Engine (nn.Module) engine/_engine.py — STATEFUL +│ construct(module, optimizer, lr_scheduler?, mesh_context?, max_grad_norm, gas, defer_fsdp_grad_sync) +│ forward(**model_inputs) → whatever the model returns (logits may be a vocab-sharded DTensor) +│ backward(loss, scale_wrt_gas) → None (loss: scalar; scale_wrt_gas=False when caller pre-normalized) +│ step() → None (non-boundary: count; boundary: finalize+clip+update+zero+sched) +│ zero_grad() / get_global_grad_norm() / set_gradient_accumulation_steps(n) +│ is_gradient_accumulation_boundary() / reset_accumulation() +│ owns: accumulation window state machine, deferred-FSDP sync, MoE aux-loss scaling, +│ summed-reducer compensation, gradient finalize (EP/TP factors), clip, +│ non-finite-norm update skip, optimizer + scheduler advance +│ does NOT own: loss math, collation/packing/CP prep, pipeline scheduling +│ +├── token_log_probs / token_entropy components/loss/vocab_parallel.py — STATELESS fns +│ in : logits Tensor|DTensor [.., vocab], targets int64 [..] (global ids), temperature +│ out: fp32 [..] selected-token log-probs / entropy, replicated on every TP rank +│ owns: the TP vocab-shard layout contract (Shard(-1), even chunks) and the +│ no-gather reduction; caller owns target construction and temperature semantics +│ AM-internal consumers: none (RL-only; SFT losses consume log-probs inside fused CE) +│ +├── RouterReplayAdapter components/moe/router_replay.py — per-model instance +│ construct(model) → binds one handle per MoE gate, keyed by decoder layer_idx +│ replay(routes int [tokens, global_layers, topk] | None) → context manager +│ -1 row = keep live routing; context must span forward AND backward (AC recompute) +│ owns: gate discovery, layer-id → route-slice mapping, sentinel/trailing fallback, +│ handle state restore on exit/exception +│ caller owns: recording routes, storing them with old_log_probs, packed token-order alignment +│ AM-internal consumers of the adapter: none; the per-gate RouterReplay handle +│ underneath IS AM-internal (gates call replay_selection in their forward) +│ +├── pack_vlm_samples / merge_media_values / collaters components/datasets/vlm/ — STATELESS fns +│ in : per-sample dicts (input_ids, labels, media) + get_rope_index +│ out: one THD-packed physical batch (cu_seqlens, positions, media side channels) +│ owns: physical boxing/media merge; caller owns sample selection and restore semantics +│ +└── state_dict_adapter.convert_single_tensor_to_hf per model family — STATELESS method + in : (fqn, full_tensor, exclude_key_regex, quantization) + out: [(hf_name, tensor), ...] for vLLM refit streaming + owns: custom-layout → HF key/shape mapping; caller owns the gather and the refit protocol +``` + +Not in the tree on purpose: `_HybridEPManager.dispatch/combine` (HybridEP token +equalization, #3641). It is reached only through the MoE layer's own forward — no RL +framework ever sees or calls it, which is exactly the fix: the kernel constraint moved +from caller-visible padding code into invisible dispatcher plumbing. + +The ownership rule behind every node: a component owns exactly the knowledge that is +private to AM (model layout, gate structure, kernel constraints, wrapper conventions); +everything expressible in the RL framework's own terms (losses, advantages, when a window +opens, what a sample means) stays with the caller. + +### `nemo_automodel/engine/_engine.py` — Engine (core, new) + +Wraps an **already-distributed** eager model. Public surface: `forward` / `backward(loss)` / `step()` / `zero_grad()` / `get_global_grad_norm()` / `set_gradient_accumulation_steps()`. + +Responsibilities: +- gradient-accumulation window management; non-boundary microsteps defer FSDP gradient sync through the wrapper's `no_sync`; +- backward scaling (the caller can disable GAS normalization with `scale_wrt_gas=False` when it has already normalized the whole window — the RL case); per-microstep scaling of the MoE auxiliary loss; +- at the accumulation boundary: distributed gradient finalization (EP/TP expert replication factors), clipping, `optimizer.step`, `zero_grad`, MoE gate-bias update, FP8 scale precompute, and scheduler advancement; +- compensation for MegatronFSDP's summed-gradient semantics (detected via `calculate_per_token_loss`). + +Explicitly **not** its job: loss computation, RL semantics, collation/packing/CP sharding (the caller's job), and pipeline scheduling (PP runs through the AutoPipeline schedule; the Engine raises explicitly). + +### `components/loss/vocab_parallel.py` (new) + +`token_log_probs(logits, targets)` / `token_entropy(logits)`: selected-token log-probabilities and categorical entropy, one entry point for dense and vocab-sharded DTensor logits; the sharded path uses distributed reductions and never gathers the vocabulary; the dense path upcasts to fp32 in 256-row chunks to bound peak memory. RL actor/reference/critic scoring and the training forward all depend on it. + +### `components/moe/router_replay.py` (extended) + +`RouterReplayAdapter`: replays rollout-recorded expert selections through the training forward (forward / activation recomputation / backward), with layer-id mapping, sentinel fallback, trailing-token handling, and state restoration built in. The importance-sampling correctness of GRPO/GSPO on MoE models depends on it. + +### `components/moe/megatron/token_dispatcher.py` (in main via #3641) + +Per-rank token-count equalization plus 4-token alignment inside HybridEP dispatch, making dynamic batches (RL rollouts) usable on the hybridep backend. + +### `components/datasets/vlm/` (extended) + +`pack_vlm_samples` (sample → THD physical boxing, including mRoPE positions and media side channels), `neat_packed_vlm_collater` / `packed_sequence_thd_vlm_collater`, and `merge_media_values` (variable-resolution media merge/pad). Molt's on-the-fly VLM packing calls these directly at runtime. + +### Supporting model/infrastructure changes (small) + +- **state_dict adapters** (llama/qwen2/qwen3 gain `convert_single_tensor_to_hf`): per-tensor HF-name conversion for vLLM refit (called by molt's `policy_actor`; the interface itself predates this PR — this PR fills in the pure-passthrough models). +- **`_transformers/model_init.py`**: an in-memory config passed to `from_pretrained(config=...)` is no longer forwarded twice — molt's `load_automodel` loads with a supplied config. +- **`models/muse_glimmer`**: CP preparation for packed TE THD (molt's Muse path). +- **`utils/model_utils.py::squeeze_input_for_thd`**: skips media keys, fixing the item axis of a single-media batch being squeezed away. +- **`optim/scheduler.py`**: `step(increment=1)` default (the Engine calls the scheduler with no arguments). +- **`quantization/fp8.py`**: capability-assignment ordering so the Engine can read the FP8 precompute flag. +- **`distributed/pipelining/autopipeline.py`**: new `eval()` (forward-only schedule entry sharing `step()`'s implementation) — used by AM's own LLM recipe for PP validation, and the hook for future RL PP support. +- **`context_parallel/magi.py`**: fixes including `pad_value=-100` on the labels dispatch, explicit CP token-index exposure, and cp_group refresh for Engine callers. + +### AM's own beneficiaries + +The LLM/VLM SFT recipes (`train_ft.py` / `finetune.py` / `benchmark.py`) drive their eager paths through the same Engine and drop their duplicated backward/clip/step implementations — the Engine is not an RL-specific component; it is AM's general training-execution layer. diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 740728f570..3ce35cfe04 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -37,9 +37,12 @@ # Heavy dependencies (e.g., torch/transformers) are intentionally imported lazily # via __getattr__ so importing tokenizers doesn't pull in the full training stack. -_SUBMODULES = {"recipes", "shared", "components", "models"} +_SUBMODULES = {"recipes", "shared", "components", "models", "engine"} _LAZY_ATTRS: dict[str, tuple[str, str]] = { + "AutoMFU": ("nemo_automodel._transformers.mfu", "AutoMFU"), + "Engine": ("nemo_automodel.engine", "Engine"), + "get_is_hf_model": ("nemo_automodel._transformers.model_init", "get_is_hf_model"), "NeMoAutoModelForCausalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForCausalLM"), "NeMoAutoModelForImageTextToText": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForImageTextToText"), "NeMoAutoModelForMultimodalLM": ("nemo_automodel._transformers.auto_model", "NeMoAutoModelForMultimodalLM"), diff --git a/nemo_automodel/_transformers/model_init.py b/nemo_automodel/_transformers/model_init.py index 03fac687ca..ced0379b8d 100644 --- a/nemo_automodel/_transformers/model_init.py +++ b/nemo_automodel/_transformers/model_init.py @@ -1252,6 +1252,11 @@ def __init_model( _download_model_weights(hf_config, pretrained_model_name_or_path, process_group=process_group) logger.info(f"Using custom model implementation for {architectures[0]}") kwargs.pop("trust_remote_code", None) + # ``hf_config`` is passed positionally below. Keep an in-memory + # config supplied to ``from_pretrained`` out of ``**kwargs`` so + # custom constructors that accept arbitrary keywords do not see + # the same argument twice. + kwargs.pop("config", None) # Treat config-related kwargs as config overrides (HF behavior) and # avoid forwarding them into model __init__. init_param_names = _get_init_param_names(model_cls) diff --git a/nemo_automodel/components/datasets/vlm/collate_fns.py b/nemo_automodel/components/datasets/vlm/collate_fns.py index 59241901f2..0020780eb3 100644 --- a/nemo_automodel/components/datasets/vlm/collate_fns.py +++ b/nemo_automodel/components/datasets/vlm/collate_fns.py @@ -56,7 +56,7 @@ mask_fake_vision_tokens_batch, ) from nemo_automodel.components.datasets.vlm.samplers import _smart_resize_image -from nemo_automodel.components.datasets.vlm.utils import default_stop_tokens +from nemo_automodel.components.datasets.vlm.utils import default_stop_tokens, merge_media_values # --------------------------------------------------------------------------- # Patch BaseVideoProcessor.fetch_videos to use decord (decord2) instead of @@ -1362,22 +1362,6 @@ def default_collate_fn( return batch -def _merge_media_values(values: list[Any]) -> torch.Tensor | list[Any]: - """Merge fixed-shape patch tensors or preserve variable-resolution media lists.""" - if not values: - raise ValueError("Media merge requires at least one value.") - if all(isinstance(value, torch.Tensor) for value in values): - return torch.cat(values, dim=0).to(torch.bfloat16) - if all(isinstance(value, (list, tuple)) for value in values): - return [ - item.to(torch.bfloat16) if isinstance(item, torch.Tensor) else item for value in values for item in value - ] - raise TypeError( - "VLM media values must be consistently tensors or variable-resolution lists, " - f"got {[type(value).__name__ for value in values]}." - ) - - def pad_collate_fn( examples: Sequence[Dict[str, Any]], processor, @@ -1434,8 +1418,8 @@ def pad_collate_fn( "attention_mask": torch.stack(padded_attention_mask), } - # Pad sequence-length tensors that mirror input_ids (e.g. mm_token_type_ids) - for seq_key in ("mm_token_type_ids",): + # Pad sequence-length tensors that mirror input_ids. + for seq_key in ("mm_token_type_ids", "token_type_ids"): if any(seq_key in ex for ex in examples): padded = [] for ex in examples: @@ -1465,7 +1449,7 @@ def pad_collate_fn( for key in ("pixel_values", "pixel_values_videos"): tensors = [ex[key] for ex in examples if key in ex and ex[key] is not None] if tensors: - batch[key] = _merge_media_values(tensors) + batch[key] = merge_media_values(tensors, field_name=key) # Per-sample image counts from image_grid_thw shapes (before concat) image_grid_per_sample = [ @@ -1621,7 +1605,7 @@ def _pad_mrope(pos, target_len): for key in ("pixel_values", "pixel_values_videos"): tensors = [x[key] for x in batch if key in x and x[key] is not None] if tensors: - result[key] = _merge_media_values(tensors) + result[key] = merge_media_values(tensors, field_name=key) for key in ("image_grid_thw", "image_position_ids", "video_grid_thw", "second_per_grid_ts"): tensors = [x[key] for x in batch if key in x and x[key] is not None] @@ -1670,8 +1654,9 @@ def packed_sequence_thd_vlm_collater( (default -1000); filtered downstream in ``process_input_for_thd``. Returns: - Dict with ``input_ids``/``labels`` ``[batch, seq]``, ``position_ids`` - ``[batch, seq]`` or ``[3, batch, seq]``, ``seq_lens``/``seq_lens_padded`` + Dict with ``input_ids``/``labels`` ``[batch, seq]``, optional token-type + fields ``[batch, seq]``, ``position_ids`` ``[batch, seq]`` or + ``[3, batch, seq]``, ``seq_lens``/``seq_lens_padded`` ``[batch, max_packs]``, ``qkv_format='thd'``, and concatenated media tensors. """ if not batch: @@ -1757,10 +1742,18 @@ def _pad_seq(tensor, pad_value, target_len, seq_dim=-1): "qkv_format": "thd", } + for key in ("mm_token_type_ids", "token_type_ids"): + if any(key in item and item[key] is not None for item in batch): + values = [ + item[key] if item.get(key) is not None else torch.zeros_like(torch.as_tensor(item["input_ids"])) + for item in batch + ] + result[key] = torch.stack([_pad_seq(value, 0, max_len) for value in values]) + for key in ("pixel_values", "pixel_values_videos"): tensors = [x[key] for x in batch if key in x and x[key] is not None] if tensors: - result[key] = _merge_media_values(tensors) + result[key] = merge_media_values(tensors, field_name=key) for key in ("image_grid_thw", "image_position_ids", "video_grid_thw", "second_per_grid_ts"): tensors = [x[key] for x in batch if key in x and x[key] is not None] diff --git a/nemo_automodel/components/datasets/vlm/neat_packing_vlm.py b/nemo_automodel/components/datasets/vlm/neat_packing_vlm.py index 1ed8885912..10eabfc1d0 100644 --- a/nemo_automodel/components/datasets/vlm/neat_packing_vlm.py +++ b/nemo_automodel/components/datasets/vlm/neat_packing_vlm.py @@ -53,6 +53,7 @@ _smart_resize_image, _smart_resize_video, ) +from nemo_automodel.components.datasets.vlm.utils import merge_media_values logger = logging.getLogger(__name__) @@ -65,6 +66,8 @@ "second_per_grid_ts", ) +_TOKEN_TYPE_KEYS = ("mm_token_type_ids", "token_type_ids") + # --------------------------------------------------------------------------- # Visual-token-balanced greedy knapsack # --------------------------------------------------------------------------- @@ -345,9 +348,10 @@ def _shift_sample(sample: dict, has_mrope: bool = False) -> dict: out["labels"] = sample["labels"][1:] out["attention_mask"] = sample["attention_mask"][:-1] - if (mm_ttids := sample.get("mm_token_type_ids")) is not None: - mm_ttids = torch.as_tensor(mm_ttids) - out["mm_token_type_ids"] = mm_ttids[0, :-1] if mm_ttids.ndim == 2 else mm_ttids[:-1] + for key in _TOKEN_TYPE_KEYS: + if (token_types := sample.get(key)) is not None: + token_types = torch.as_tensor(token_types) + out[key] = token_types[0, :-1] if token_types.ndim == 2 else token_types[:-1] if has_mrope and "position_ids" in sample and sample["position_ids"] is not None: out["position_ids"] = sample["position_ids"][:, :-1] @@ -367,7 +371,6 @@ def _aligned_length(length: int, alignment: int) -> int: def _build_packed_vlm_sample( samples: list[dict], - pack_size: int, padding_idx: int, has_mrope: bool = False, sequence_alignment: int = 1, @@ -379,7 +382,10 @@ def _build_packed_vlm_sample( all_input_ids: list[int] = [] all_labels: list[int] = [] all_attention_mask: list[int] = [] - all_mm_token_type_ids: list[int] = [] + all_token_type_ids: dict[str, list[int]] = {key: [] for key in _TOKEN_TYPE_KEYS} + present_token_type_keys = tuple( + key for key in _TOKEN_TYPE_KEYS if any(sample.get(key) is not None for sample in samples) + ) all_position_ids_1d: list[int] = [] mrope_position_ids_list: list[torch.Tensor] = [] seq_lens: list[int] = [] @@ -411,12 +417,13 @@ def _build_packed_vlm_sample( all_labels.extend(labs + [-100] * pad) all_attention_mask.extend([seq_idx] * padded_seq_len) - mm_ttids = sample.get("mm_token_type_ids") - if mm_ttids is not None: - mm_ttids = mm_ttids.tolist() if isinstance(mm_ttids, torch.Tensor) else list(mm_ttids) - all_mm_token_type_ids.extend(mm_ttids + [0] * pad) - else: - all_mm_token_type_ids.extend([0] * padded_seq_len) + for key in present_token_type_keys: + token_types = sample.get(key) + if token_types is None: + all_token_type_ids[key].extend([0] * padded_seq_len) + continue + token_types = torch.as_tensor(token_types).reshape(-1).tolist() + all_token_type_ids[key].extend(token_types + [0] * pad) if has_mrope and "position_ids" in sample: mrope_position_ids_list.append(sample["position_ids"]) @@ -443,26 +450,22 @@ def _build_packed_vlm_sample( "input_ids": torch.tensor(all_input_ids, dtype=torch.long), "labels": torch.tensor(all_labels, dtype=torch.long), "attention_mask": torch.tensor(all_attention_mask, dtype=torch.long), - "mm_token_type_ids": torch.tensor(all_mm_token_type_ids, dtype=torch.long), "seq_lens": seq_lens, "seq_lens_padded": seq_lens_padded, "n_images": n_images, "n_videos": n_videos, } + for key in present_token_type_keys: + packed[key] = torch.tensor(all_token_type_ids[key], dtype=torch.long) if has_mrope and mrope_position_ids_list: packed["position_ids"] = torch.cat(mrope_position_ids_list, dim=1) else: packed["position_ids"] = torch.tensor(all_position_ids_1d, dtype=torch.long) - if pixel_values_list and all(isinstance(value, torch.Tensor) for value in pixel_values_list): - packed["pixel_values"] = torch.cat(pixel_values_list, dim=0) - elif pixel_values_list and all(isinstance(value, (list, tuple)) for value in pixel_values_list): - packed["pixel_values"] = [item for value in pixel_values_list for item in value] - elif pixel_values_list: - raise TypeError("Packed VLM pixel_values must be consistently tensors or variable-resolution lists.") - else: - packed["pixel_values"] = None + packed["pixel_values"] = ( + merge_media_values(pixel_values_list, field_name="pixel_values") if pixel_values_list else None + ) packed["image_grid_thw"] = torch.cat(image_grid_thw_list, dim=0) if image_grid_thw_list else None packed["image_position_ids"] = torch.cat(image_position_ids_list, dim=0) if image_position_ids_list else None packed["pixel_values_videos"] = torch.cat(pixel_values_videos_list, dim=0) if pixel_values_videos_list else None @@ -472,6 +475,43 @@ def _build_packed_vlm_sample( return packed +def pack_vlm_samples( + samples: Sequence[dict[str, object]], + *, + padding_idx: int, + get_rope_index: Callable[..., object] | None = None, + sequence_alignment: int = 1, +) -> dict[str, object]: + """Shift and concatenate pretokenized VLM samples into one packed sample. + + Args: + samples: Processor outputs before the autoregressive input/label shift. + padding_idx: Token ID used for per-document alignment padding. + get_rope_index: Optional model callback that builds multi-axis position IDs. + sequence_alignment: Alignment applied independently to every shifted sample. + + Returns: + One packed sample containing token, position, sequence, and media fields. + """ + has_mrope = get_rope_index is not None + shifted_samples = [] + for sample in samples: + prepared = dict(sample) + if get_rope_index is not None: + position_ids = _compute_mrope_position_ids(prepared, get_rope_index) + if position_ids is None: + raise ValueError("get_rope_index must accept input_ids and return VLM position_ids") + prepared["position_ids"] = position_ids + shifted_samples.append(_shift_sample(prepared, has_mrope=has_mrope)) + + return _build_packed_vlm_sample( + shifted_samples, + padding_idx, + has_mrope=has_mrope, + sequence_alignment=sequence_alignment, + ) + + # --------------------------------------------------------------------------- # PackedDatasetWrapper — lazy packing via __getitem__ # --------------------------------------------------------------------------- @@ -555,8 +595,7 @@ def __init__( raise ValueError(f"sequence_alignment must be at least 1, got {sequence_alignment}.") self.sequence_alignment = sequence_alignment self.get_rope_index = get_rope_index - self.has_mrope = get_rope_index is not None - if self.has_mrope and self.sequence_alignment > 1: + if get_rope_index is not None and self.sequence_alignment > 1: raise NotImplementedError("Context-parallel THD packing for multi-axis mRoPE VLMs is not yet implemented.") self.max_retries = max_retries @@ -566,18 +605,11 @@ def __len__(self): def __getitem__(self, pack_idx: int) -> dict: """Materialize one pack: tokenize + shift + concat all samples in the bin.""" bin_indices = self.bins[pack_idx] - shifted_samples: list[dict] = [] + samples: list[dict] = [] for sample_idx in bin_indices: sample = self.inner[sample_idx] # tokenize + load media - - if self.has_mrope and self.get_rope_index is not None: - mrope_pos = _compute_mrope_position_ids(sample, self.get_rope_index) - if mrope_pos is not None: - sample["position_ids"] = mrope_pos - - shifted = _shift_sample(sample, has_mrope=self.has_mrope) - seq_len = shifted["input_ids"].shape[0] + seq_len = len(sample["input_ids"]) - 1 aligned_seq_len = _aligned_length(seq_len, self.sequence_alignment) # The aligned length is the actual capacity consumed by THD CP. @@ -592,16 +624,16 @@ def __getitem__(self, pack_idx: int) -> dict: ) continue - shifted_samples.append(shifted) + samples.append(sample) # Truncate if total exceeds pack_size (estimation was wrong) total = 0 kept: list[dict] = [] - for s in shifted_samples: - slen = s["input_ids"].shape[0] + for sample in samples: + slen = len(sample["input_ids"]) - 1 aligned_slen = _aligned_length(slen, self.sequence_alignment) if total + aligned_slen <= self.pack_size: - kept.append(s) + kept.append(sample) total += aligned_slen else: logger.debug( @@ -615,13 +647,18 @@ def __getitem__(self, pack_idx: int) -> dict: if not kept: # Fallback: return a padding-only pack - kept = [{"input_ids": torch.tensor([], dtype=torch.long), "labels": torch.tensor([], dtype=torch.long)}] - - return _build_packed_vlm_sample( + kept = [ + { + "input_ids": torch.tensor([self.padding_idx], dtype=torch.long), + "labels": torch.tensor([self.padding_idx], dtype=torch.long), + "attention_mask": torch.ones(1, dtype=torch.long), + } + ] + + return pack_vlm_samples( kept, - self.pack_size, - self.padding_idx, - has_mrope=self.has_mrope, + padding_idx=self.padding_idx, + get_rope_index=self.get_rope_index, sequence_alignment=self.sequence_alignment, ) diff --git a/nemo_automodel/components/datasets/vlm/pp_media.py b/nemo_automodel/components/datasets/vlm/pp_media.py index 9d20f89a4c..cbf0b18c73 100644 --- a/nemo_automodel/components/datasets/vlm/pp_media.py +++ b/nemo_automodel/components/datasets/vlm/pp_media.py @@ -39,21 +39,76 @@ 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 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: + media_counts = n_images_per_sample.to(dtype=torch.long, device="cpu") + + 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 and 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 +285,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 +317,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 +330,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 +342,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 +354,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/datasets/vlm/utils.py b/nemo_automodel/components/datasets/vlm/utils.py index 1e54bafb02..44ae238fdc 100644 --- a/nemo_automodel/components/datasets/vlm/utils.py +++ b/nemo_automodel/components/datasets/vlm/utils.py @@ -14,9 +14,10 @@ import io import logging -from typing import Iterable +from typing import Any, Iterable import torch +import torch.nn.functional as F from PIL import Image logger = logging.getLogger(__name__) @@ -32,6 +33,51 @@ _lmdb_env_cache: dict[str, "lmdb.Environment"] = {} +def merge_media_values(values: list[Any], *, field_name: str) -> torch.Tensor | list[Any]: + """Merge media values while preserving the explicit ragged-list contract. + + Four-dimensional ``pixel_values`` use ``[media, channels, height, width]``. + Some processors emit a different image resolution for each sample, so pad + those tensors on the right and bottom to the largest spatial extent before + concatenating their media axes. Other tensor fields retain their processor- + defined layout and must already be concatenable. + + Args: + values: Non-empty per-sample media values. Tensor ``pixel_values`` + use ``[media, channels, height, width]``; other tensor fields use + ``[media, ...]``. List or tuple items retain their processor-defined + layout. + field_name: Processor field being merged. + + Returns: + One BF16 tensor on the input device, concatenated on its media axis. + Four-dimensional ``pixel_values`` have shape + ``[sum_media, channels, max_height, max_width]``; other tensor fields + have shape ``[sum_media, ...]``. Processors that explicitly represent + variable-resolution media as lists produce one flattened list whose + tensor items are converted to BF16 on their original devices. + """ + if not values: + raise ValueError("Media merge requires at least one value.") + if all(isinstance(value, torch.Tensor) for value in values): + tensors = values + if field_name == "pixel_values" and all(value.ndim == 4 for value in tensors): + max_height = max(value.shape[-2] for value in tensors) + max_width = max(value.shape[-1] for value in tensors) + tensors = [ + F.pad(value, (0, max_width - value.shape[-1], 0, max_height - value.shape[-2])) for value in tensors + ] + return torch.cat(tensors, dim=0).to(torch.bfloat16) + if all(isinstance(value, (list, tuple)) for value in values): + return [ + item.to(torch.bfloat16) if isinstance(item, torch.Tensor) else item for value in values for item in value + ] + raise TypeError( + "VLM media values must be consistently tensors or variable-resolution lists, " + f"got {[type(value).__name__ for value in values]}." + ) + + def _resolve_lmdb_image(path): """Read an image from an LMDB database. diff --git a/nemo_automodel/components/distributed/context_parallel/magi.py b/nemo_automodel/components/distributed/context_parallel/magi.py index a6e41d505f..7aa7e1f6d4 100644 --- a/nemo_automodel/components/distributed/context_parallel/magi.py +++ b/nemo_automodel/components/distributed/context_parallel/magi.py @@ -528,8 +528,8 @@ def magi_prepare_batch( # pragma: no cover - requires GPU + magi_attention chunk_size: dispatch solver chunk size. 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, local_indices)`` where ``local_indices`` maps 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 +567,9 @@ 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 = 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 +582,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 def _packed_cp_doc_seqlens(batch: dict, total_len: int) -> list: @@ -625,8 +628,8 @@ def magi_prepare_packed_cp(model, batch: dict, cp_group): # pragma: no cover - into the global loss (like TE-CP). 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, local_indices)`` where ``local_indices`` maps the + local token stream to the global input. """ from magi_attention.api import dispatch, get_position_ids @@ -643,12 +646,16 @@ 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 = 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 +666,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 def _iter_language_model_attention(model): @@ -750,10 +757,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 @@ -761,6 +768,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 @@ -779,17 +791,24 @@ 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, _ = 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) 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) + 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) 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. @@ -851,6 +870,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/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index b88a509123..3afaa2e5fc 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/nemo_automodel/components/distributed/pipelining/autopipeline.py @@ -14,7 +14,7 @@ import logging from dataclasses import dataclass -from typing import Any, Callable, Literal +from typing import Callable, Literal import torch import torch.nn as nn @@ -218,24 +218,8 @@ def update_seq_len(self, seq_len: int) -> None: self._pp_current_seq_len = seq_len logger.debug(f"PP stage shapes updated for seq_len={seq_len}") - 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. - """ + def _get_schedule_kwargs_chunk_spec(self, kwargs: dict[str, object]) -> dict[str, object] | None: + """Build schedule chunking metadata for model-owned keyword layouts.""" model_parts = self._info.model_parts if not model_parts: raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step") @@ -248,7 +232,6 @@ def _get_schedule_kwargs_chunk_spec(self, kwargs: dict[str, Any]) -> dict[str, A 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 @@ -262,46 +245,57 @@ def default_spec(value): kwargs_chunk_spec[key] = TensorChunkSpec(split_dim) return kwargs_chunk_spec - def step( + def _run_schedule( self, + method_name: Literal["step", "eval"], 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. - """ + target: torch.Tensor | None, + losses: list[torch.Tensor] | None, + kwargs: dict[str, object], + ) -> object: + """Run one schedule method with first-stage and model-owned chunking rules.""" 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 PP schedule") + method = getattr(schedule, method_name, None) + if not callable(method): + raise NotImplementedError(f"The configured pipeline schedule does not support {method_name}()") 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) + return method(*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) + return method(*schedule_args, target=target, losses=losses, **kwargs) finally: schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec + def step( + self, + model_input: torch.Tensor, + *, + target: torch.Tensor | None = None, + losses: list[torch.Tensor] | None = None, + **kwargs: object, + ) -> object: + """Run a training schedule step with model-owned input chunking.""" + return self._run_schedule("step", model_input, target=target, losses=losses, kwargs=kwargs) + + def eval( + self, + model_input: torch.Tensor, + *, + target: torch.Tensor | None = None, + losses: list[torch.Tensor] | None = None, + **kwargs: object, + ) -> object: + """Run a forward-only schedule step with model-owned input chunking.""" + return self._run_schedule("eval", model_input, target=target, losses=losses, kwargs=kwargs) + @property def parts(self) -> list[nn.Module]: if self._info.model_parts is None: diff --git a/nemo_automodel/components/distributed/utils.py b/nemo_automodel/components/distributed/utils.py index e526933e41..ac7709d91e 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/components/loss/__init__.py b/nemo_automodel/components/loss/__init__.py index 1327315fae..8cd7862b9d 100644 --- a/nemo_automodel/components/loss/__init__.py +++ b/nemo_automodel/components/loss/__init__.py @@ -23,6 +23,10 @@ build_loss_config, build_loss_module, ) +from nemo_automodel.components.loss.vocab_parallel import ( + token_entropy, + token_log_probs, +) __all__ = [ "LOSS_CONFIG_REGISTRY", @@ -34,4 +38,6 @@ "TEParallelCEConfig", "build_loss_config", "build_loss_module", + "token_entropy", + "token_log_probs", ] diff --git a/nemo_automodel/components/loss/vocab_parallel.py b/nemo_automodel/components/loss/vocab_parallel.py new file mode 100644 index 0000000000..5f9a634837 --- /dev/null +++ b/nemo_automodel/components/loss/vocab_parallel.py @@ -0,0 +1,282 @@ +# 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. + +"""Token statistics for dense and vocabulary-sharded logits.""" + +import math + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor, Partial, Replicate, Shard + +_DENSE_TOKEN_CHUNK_SIZE = 256 + + +def _sum_vocab_shards(values: torch.Tensor, logits: DTensor, vocab_mesh_dim: int) -> torch.Tensor: + """Sum rank-local statistics over the vocabulary-shard mesh dimension. + + Wraps ``values`` as a ``Partial`` DTensor on ``vocab_mesh_dim`` so that + redistributing to ``Replicate`` performs the all-reduce. + """ + partial_placements = [Replicate() for _ in logits.placements] + partial_placements[vocab_mesh_dim] = Partial() + return ( + DTensor.from_local(values, logits.device_mesh, partial_placements, run_check=False) + .redistribute(placements=[Replicate() for _ in logits.placements]) + .to_local() + ) + + +def _shifted_local_logits( + logits: DTensor, + temperature: float, +) -> tuple[torch.Tensor, int, int]: + """Scale local vocabulary-sharded logits and subtract the global max. + + Returns a rank-local fp32 tensor of shape [..., local_vocab], this rank's + global vocabulary offset, and the vocabulary-shard mesh dimension. + """ + shard_error = ( + f"logits must have exactly one Shard placement on the last vocabulary axis; got placements={logits.placements}" + ) + vocab_mesh_dim = None + for mesh_dim, placement in enumerate(logits.placements): + if isinstance(placement, Shard): + shard_dim = placement.dim if placement.dim >= 0 else placement.dim + logits.ndim + if shard_dim != logits.ndim - 1 or vocab_mesh_dim is not None: + raise ValueError(shard_error) + vocab_mesh_dim = mesh_dim + elif not isinstance(placement, Replicate): + raise ValueError( + f"logits must be replicated on every non-vocabulary mesh dimension; got placements={logits.placements}" + ) + if vocab_mesh_dim is None: + raise ValueError(shard_error) + + global_vocab_size = int(logits.shape[-1]) + mesh = logits.device_mesh + shard_count = mesh.size(vocab_mesh_dim) + shard_rank = mesh.get_local_rank(vocab_mesh_dim) + chunk_size = (global_vocab_size + shard_count - 1) // shard_count + shard_offset = min(shard_rank * chunk_size, global_vocab_size) + shard_size = min(chunk_size, global_vocab_size - shard_offset) + local_logits = logits.to_local() + # The shard offset above assumes DTensor's even-chunk layout; this guards + # that assumption against a differently laid-out local shard. + if local_logits.shape[-1] != shard_size: + raise ValueError( + "logits local vocabulary size does not match its DTensor metadata: " + f"expected {shard_size}, got {local_logits.shape[-1]}" + ) + + scaled_logits = local_logits.float() / temperature + if shard_size == 0: + global_max = torch.full( + scaled_logits.shape[:-1], + -torch.inf, + dtype=scaled_logits.dtype, + device=scaled_logits.device, + ) + else: + global_max = scaled_logits.detach().amax(dim=-1) + group = mesh.get_group(vocab_mesh_dim) + dist.all_reduce(global_max, op=dist.ReduceOp.MAX, group=group) + return scaled_logits - global_max.unsqueeze(-1), shard_offset, vocab_mesh_dim + + +def _vocab_parallel_log_probs( + logits: DTensor, + targets: torch.Tensor, + *, + temperature: float = 1.0, +) -> torch.Tensor: + """Selected-token log probabilities without gathering the vocabulary. + + Each rank contributes its shard's exp-sum, the selected logit when the + target falls in its shard, and a target-coverage flag; one all-reduce + combines the three. Uncovered (out-of-range) targets yield ``NaN``. + """ + shifted_logits, shard_offset, vocab_mesh_dim = _shifted_local_logits(logits, temperature) + local_vocab_size = shifted_logits.shape[-1] + local_denom = shifted_logits.exp().sum(dim=-1) + in_local_shard = (targets >= shard_offset) & (targets < shard_offset + local_vocab_size) + if local_vocab_size == 0: + local_selected = torch.zeros_like(local_denom) + else: + local_targets = (targets - shard_offset).clamp(min=0, max=local_vocab_size - 1) + local_selected = shifted_logits.gather(dim=-1, index=local_targets.unsqueeze(-1)).squeeze(-1) + local_selected = torch.where(in_local_shard, local_selected, torch.zeros_like(local_selected)) + + global_stats = _sum_vocab_shards( + torch.stack((local_denom, local_selected, in_local_shard.to(local_denom.dtype)), dim=-1), + logits, + vocab_mesh_dim, + ) + log_probs = global_stats[..., 1] - global_stats[..., 0].log() + return log_probs.masked_fill(global_stats[..., 2] != 1, torch.nan) + + +def _vocab_parallel_entropy( + logits: DTensor, + *, + temperature: float = 1.0, +) -> torch.Tensor: + """Categorical entropy without gathering the vocabulary: each rank + contributes its shard's exp-sum and probability-weighted logit sum.""" + shifted_logits, _, vocab_mesh_dim = _shifted_local_logits(logits, temperature) + local_weights = shifted_logits.exp() + local_denom = local_weights.sum(dim=-1) + local_weighted_logits = (local_weights * shifted_logits).sum(dim=-1) + global_stats = _sum_vocab_shards( + torch.stack((local_denom, local_weighted_logits), dim=-1), + logits, + vocab_mesh_dim, + ) + return global_stats[..., 0].log() - global_stats[..., 1] / global_stats[..., 0] + + +def _validate_logits(logits: torch.Tensor | DTensor, temperature: float) -> None: + if not isinstance(logits, torch.Tensor): + raise TypeError(f"logits must be a torch.Tensor or DTensor, got {type(logits).__name__}") + if not logits.is_floating_point(): + raise TypeError(f"logits must have a floating-point dtype, got {logits.dtype}") + if logits.ndim == 0: + raise ValueError("logits must have shape [..., vocab]") + if logits.shape[-1] <= 0: + raise ValueError(f"logits vocabulary size must be positive, got {logits.shape[-1]}") + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError(f"temperature must be positive and finite, got {temperature}") + + +def token_log_probs( + logits: torch.Tensor | DTensor, + targets: torch.Tensor, + *, + temperature: float = 1.0, +) -> torch.Tensor: + """Compute selected-token log probabilities for dense or vocabulary-sharded logits. + + Dense logits upcast to fp32 in chunks of at most 256 token rows, so the + forward peak never holds the full vocabulary tensor in fp32 (with autograd + on, saved activations still accumulate across chunks). Vocabulary-sharded + DTensors use distributed reductions without gathering the vocabulary. + + Args: + logits: Floating-point tensor with global shape [..., vocab], with + arbitrary leading dimensions. A DTensor must have exactly one + ``Shard`` placement on the vocabulary axis and ``Replicate`` on + every other mesh dimension; its per-rank local shape is + [..., local_vocab]. + targets: Rank-local int64 tensor of shape [...], matching the leading + dimensions of ``logits`` and containing global vocabulary indices. + For DTensor logits, targets must be replicated across the + vocabulary-shard mesh dimension. + temperature: Positive finite scale applied before normalization. + + Returns: + Differentiable fp32 tensor of shape [...] containing selected-token log + probabilities. Invalid target positions contain ``NaN``. The function + does not mutate or gather ``logits``. + + Raises: + TypeError: If an input has an invalid type or dtype. + ValueError: If a shape, device, DTensor placement, vocabulary size, or + temperature is invalid. + """ + _validate_logits(logits, temperature) + if isinstance(targets, DTensor): + raise TypeError("targets must be a rank-local torch.Tensor, not a DTensor") + if not isinstance(targets, torch.Tensor): + raise TypeError(f"targets must be a torch.Tensor, got {type(targets).__name__}") + if targets.dtype != torch.long: + raise TypeError(f"targets must have dtype torch.int64, got {targets.dtype}") + if targets.device != logits.device: + raise ValueError(f"targets and logits must be on the same device, got {targets.device} and {logits.device}") + if tuple(targets.shape) != tuple(logits.shape[:-1]): + raise ValueError( + f"targets shape must match logits leading shape {tuple(logits.shape[:-1])}, got {tuple(targets.shape)}" + ) + if isinstance(logits, DTensor): + return _vocab_parallel_log_probs(logits, targets, temperature=temperature) + + leading_shape = logits.shape[:-1] + vocab_size = logits.shape[-1] + flat_logits = logits.reshape(-1, vocab_size) + flat_targets = targets.reshape(-1) + if flat_targets.numel() == 0: + # Route through the input so an empty batch still produces a gradient. + return flat_logits.float().sum(dim=-1).reshape(leading_shape) + log_probs = torch.empty(flat_targets.shape, dtype=torch.float32, device=logits.device) + for start in range(0, flat_targets.numel(), _DENSE_TOKEN_CHUNK_SIZE): + end = min(start + _DENSE_TOKEN_CHUNK_SIZE, flat_targets.numel()) + chunk = flat_logits[start:end].float() + if temperature != 1.0: + chunk = chunk / temperature + chunk_targets = flat_targets[start:end] + valid_targets = (chunk_targets >= 0) & (chunk_targets < vocab_size) + safe_targets = chunk_targets.clamp(min=0, max=vocab_size - 1) + selected_logits = chunk.gather(dim=-1, index=safe_targets.unsqueeze(-1)).squeeze(-1) + chunk_log_probs = selected_logits - torch.logsumexp(chunk, dim=-1) + log_probs[start:end] = chunk_log_probs.masked_fill(~valid_targets, torch.nan) + return log_probs.reshape(leading_shape) + + +def token_entropy( + logits: torch.Tensor | DTensor, + *, + temperature: float = 1.0, +) -> torch.Tensor: + """Compute categorical entropy for dense or vocabulary-sharded logits. + + Dense logits upcast to fp32 in chunks of at most 256 token rows, so the + forward peak never holds the full vocabulary tensor in fp32 (with autograd + on, saved activations still accumulate across chunks). Vocabulary-sharded + DTensors use distributed reductions without gathering the vocabulary. + + Args: + logits: Floating-point tensor with global shape [..., vocab], with + arbitrary leading dimensions. A DTensor must have exactly one + ``Shard`` placement on the vocabulary axis and ``Replicate`` on + every other mesh dimension; its per-rank local shape is + [..., local_vocab]. + temperature: Positive finite scale applied before normalization. + + Returns: + Differentiable fp32 tensor of shape [...] containing categorical entropy. + The function does not mutate or gather ``logits``. + + Raises: + TypeError: If ``logits`` has an invalid type or dtype. + ValueError: If a shape, DTensor placement, vocabulary size, or + temperature is invalid. + """ + _validate_logits(logits, temperature) + if isinstance(logits, DTensor): + return _vocab_parallel_entropy(logits, temperature=temperature) + + leading_shape = logits.shape[:-1] + vocab_size = logits.shape[-1] + flat_logits = logits.reshape(-1, vocab_size) + if flat_logits.shape[0] == 0: + # Route through the input so an empty batch still produces a gradient. + return flat_logits.float().sum(dim=-1).reshape(leading_shape) + entropy = torch.empty(flat_logits.shape[0], dtype=torch.float32, device=logits.device) + for start in range(0, flat_logits.shape[0], _DENSE_TOKEN_CHUNK_SIZE): + end = min(start + _DENSE_TOKEN_CHUNK_SIZE, flat_logits.shape[0]) + chunk = flat_logits[start:end].float() + if temperature != 1.0: + chunk = chunk / temperature + log_distribution = torch.log_softmax(chunk, dim=-1) + entropy[start:end] = -(log_distribution.exp() * log_distribution).sum(dim=-1) + return entropy.reshape(leading_shape) diff --git a/nemo_automodel/components/models/deepseek_v4/cp.py b/nemo_automodel/components/models/deepseek_v4/cp.py index 44baf63fc4..dd0f561784 100644 --- a/nemo_automodel/components/models/deepseek_v4/cp.py +++ b/nemo_automodel/components/models/deepseek_v4/cp.py @@ -252,7 +252,6 @@ def _repad_dsv4_packed_batch( cp_size: int, pad_multiple: int, padding_token_id: int, - sync_packed_length: bool = False, loss_mask: torch.Tensor | None = None, ) -> tuple[dict, torch.Tensor | None, torch.Tensor]: """Insert DSV4 compression-safe padding into packed BSHD rows before CP slicing. @@ -261,9 +260,7 @@ def _repad_dsv4_packed_batch( compression additionally needs document boundaries to align to compressor windows; CSA then uses ``packed_seq_ids`` to reset the previous-window overlap. This routine rebuilds each row from real sequence spans, pads every span to - ``pad_multiple``, and appends row-level pack padding with sequence ID 0. When - requested for HybridEP, the final physical length is max-reduced across ranks - before CP slicing so every rank in a flattened DP x CP expert group is uniform. + ``pad_multiple``, and appends row-level pack padding with sequence ID 0. """ if "seq_lens" not in batch: raise KeyError("DSV4 packed context parallelism requires `seq_lens` in the batch.") @@ -377,13 +374,6 @@ def _repad_dsv4_packed_batch( row_lengths.append(rebuilt_primary[-1].shape[0]) total_seq_len = _pad_length(max(row_lengths), cp_size * pad_multiple) - if sync_packed_length and dist.is_available() and dist.is_initialized(): - # HybridEP flattens DP x CP into one EP group and requires every rank to - # contribute the same number of tokens. Different DP packs can acquire - # different amounts of per-document compression padding. - length = torch.tensor(total_seq_len, dtype=torch.int64, device=primary.device) - dist.all_reduce(length, op=dist.ReduceOp.MAX) - total_seq_len = int(length.item()) def _right_pad_rows(rows: list[torch.Tensor], fill_value, *, dtype=None) -> torch.Tensor: padded = [] @@ -429,21 +419,36 @@ def make_dsv4_contiguous_shard_cp_batch_and_ctx( loss_mask=None, padding_token_id: int = 0, pad_multiple: int | None = None, - sync_packed_length: bool = False, ): """Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism. Exposed as ``ContextParallelSharder.shard_batch`` (via ``functools.partial`` to bind - ``pad_multiple``) and invoked by the CP dispatch. HybridEP can - first max-reduce packed lengths so every rank contributes a uniform token count. - Each CP rank then keeps one ``seq_start:seq_end`` slice; DSV4 attention all-gathers - K/V across CP ranks during forward. Returns ``(nullcontext, batch)``. + ``pad_multiple``) and invoked by the CP dispatch. Each CP rank then keeps one + ``seq_start:seq_end`` slice; DSV4 attention all-gathers K/V across CP ranks + during forward. Returns ``(nullcontext, batch, layout)``; + at CP size one, ``layout`` records the unchanged sequence length so token + outputs can still use ``gather_token_tensor(trim=True)``. ``pad_multiple`` is the required *per-CP-rank* shard multiple (from ``dsv4_cp_local_seq_multiple``); the global sequence is padded so it is divisible by ``cp_size`` and each local shard is divisible by ``pad_multiple`` (>= 2). At CP size one, the native THD route only marks packed input as THD and leaves its tensors and packing metadata unchanged. + + Args: + cp_mesh: Context-parallel mesh that owns contiguous sequence shards. + tp_mesh: Tensor-parallel mesh, forwarded to the shared contiguous sharder. + batch: Mapping with ``input_ids`` of shape [batch, sequence] or + ``inputs_embeds`` of shape [batch, sequence, hidden], ``labels`` of + shape [batch, sequence], and optional packed-length metadata. + loss_mask: Optional tensor of shape [batch, sequence] sharded with the + model inputs. + padding_token_id: Token value used for physical sequence padding. + pad_multiple: Required local sequence-length multiple. + + Returns: + A null context factory, the local model-input mapping, and a + :class:`ShardLayout` describing the global padded sequence. """ import contextlib @@ -462,7 +467,14 @@ def make_dsv4_contiguous_shard_cp_batch_and_ctx( batch["labels"] = loss_mask elif loss_mask is not None: batch["loss_mask"] = loss_mask - return contextlib.nullcontext, batch, None + primary_name = "inputs_embeds" if isinstance(batch.get("inputs_embeds"), torch.Tensor) else "input_ids" + primary = batch.get(primary_name) + layout = None + if isinstance(primary, torch.Tensor) and primary.ndim > 0: + is_batched = primary.ndim > (2 if primary_name == "inputs_embeds" else 1) + seq_len = primary.shape[1] if is_batched else primary.shape[0] + layout = ShardLayout(original_seq_len=seq_len, padded_seq_len=seq_len) + return contextlib.nullcontext, batch, layout local_multiple = max(int(pad_multiple or 2), 2) @@ -482,7 +494,6 @@ def make_dsv4_contiguous_shard_cp_batch_and_ctx( cp_size=cp_size, pad_multiple=local_multiple, padding_token_id=padding_token_id, - sync_packed_length=sync_packed_length, loss_mask=loss_mask, ) diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index ba08a17e8d..5022b68647 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -893,7 +893,6 @@ def prepare_model_inputs_for_cp( shard_batch=partial( make_dsv4_contiguous_shard_cp_batch_and_ctx, pad_multiple=dsv4_cp_local_seq_multiple(self.config), - sync_packed_length=self.backend.dispatcher == "hybridep", ), local_token_global_indices=contiguous_local_indices, ) diff --git a/nemo_automodel/components/models/llama/state_dict_adapter.py b/nemo_automodel/components/models/llama/state_dict_adapter.py index 9e69ac5eb0..60b3cc6f42 100644 --- a/nemo_automodel/components/models/llama/state_dict_adapter.py +++ b/nemo_automodel/components/models/llama/state_dict_adapter.py @@ -75,3 +75,17 @@ def to_hf( if exclude_key_regex is not None: return {k: v for k, v in state_dict.items() if not re.search(exclude_key_regex, k)} return dict(state_dict) + + def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]: + """Return one already-HuggingFace-format tensor. + + Args: + fqn: Fully qualified HuggingFace parameter or buffer name. + tensor: Tensor with arbitrary parameter shape, dtype, and device. + **kwargs: Conversion options accepted by :meth:`to_hf`. + + Returns: + The unchanged name and tensor, or an empty list when the name is + excluded by ``exclude_key_regex``. + """ + return list(self.to_hf({fqn: tensor}, **kwargs).items()) diff --git a/nemo_automodel/components/models/muse_glimmer/model.py b/nemo_automodel/components/models/muse_glimmer/model.py index 1eeee7c951..acf4bfe7d3 100644 --- a/nemo_automodel/components/models/muse_glimmer/model.py +++ b/nemo_automodel/components/models/muse_glimmer/model.py @@ -803,7 +803,17 @@ def prepare_model_inputs_for_cp( *, num_chunks: int = 1, ) -> dict[str, Any]: - """Select native MuseGlimmer CP preparation for BSHD or packed TE THD.""" + """Select native MuseGlimmer CP preparation for BSHD or packed TE THD. + + Args: + batch: Mapping whose packed VLM ``input_ids`` and derived vision + mask have shape [batch, sequence]. + num_chunks: Number of pipeline microbatch token streams. + + Returns: + Model-input updates. For packed VLM, the global vision mask remains + [batch, sequence]. + """ del num_chunks if batch.get("qkv_format") == "thd": max_real_seqlen = self._validate_thd_documents(batch) diff --git a/nemo_automodel/components/models/qwen2/state_dict_adapter.py b/nemo_automodel/components/models/qwen2/state_dict_adapter.py index 036447f437..a48f70aa6c 100644 --- a/nemo_automodel/components/models/qwen2/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen2/state_dict_adapter.py @@ -74,3 +74,17 @@ def to_hf( if exclude_key_regex is not None: return {k: v for k, v in state_dict.items() if not re.search(exclude_key_regex, k)} return dict(state_dict) + + def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]: + """Return one already-HuggingFace-format tensor. + + Args: + fqn: Fully qualified HuggingFace parameter or buffer name. + tensor: Tensor with arbitrary parameter shape, dtype, and device. + **kwargs: Conversion options accepted by :meth:`to_hf`. + + Returns: + The unchanged name and tensor, or an empty list when the name is + excluded by ``exclude_key_regex``. + """ + return list(self.to_hf({fqn: tensor}, **kwargs).items()) diff --git a/nemo_automodel/components/models/qwen3/state_dict_adapter.py b/nemo_automodel/components/models/qwen3/state_dict_adapter.py index 04f0f3fc30..84a24b0c5d 100644 --- a/nemo_automodel/components/models/qwen3/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3/state_dict_adapter.py @@ -17,6 +17,7 @@ from __future__ import annotations import re +from typing import Any import torch from transformers import Qwen3Config @@ -71,3 +72,17 @@ def to_hf( if exclude_key_regex is None: return dict(state_dict) return {key: value for key, value in state_dict.items() if not re.search(exclude_key_regex, key)} + + def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]: + """Return one already-HuggingFace-format tensor. + + Args: + fqn: Fully qualified HuggingFace parameter or buffer name. + tensor: Tensor with arbitrary parameter shape, dtype, and device. + **kwargs: Conversion options accepted by :meth:`to_hf`. + + Returns: + The unchanged name and tensor, or an empty list when the name is + excluded by ``exclude_key_regex``. + """ + return list(self.to_hf({fqn: tensor}, **kwargs).items()) diff --git a/nemo_automodel/components/moe/router_replay.py b/nemo_automodel/components/moe/router_replay.py index 1f3492fb33..527553ff44 100644 --- a/nemo_automodel/components/moe/router_replay.py +++ b/nemo_automodel/components/moe/router_replay.py @@ -49,15 +49,25 @@ 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` maps global +decoder-layer ids without using the registry. Callers prepare routes in the +same token order as the model input, then keep the adapter's replay context +active through forward and backward. """ -from contextlib import contextmanager +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass from enum import Enum -from typing import Iterator, List +from math import prod 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 +86,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 +107,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 +121,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 +159,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 +170,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 +191,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 +233,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 +249,180 @@ 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 + + +class RouterReplayAdapter: + """Bind rollout routes to model-scoped MoE gates. + + The adapter 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. + + Args: + model: One complete model. 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. + """ + + def __init__(self, model: nn.Module) -> None: + # Descend through .module wrappers (e.g. an Engine) until decoder + # blocks are visible. + block_root = model + while True: + 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") + 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)) + # Replayed routes cover only the live tokens the caller prepared; + # trailing padded tokens keep live routing. + replay._allow_trailing_live_tokens = True + bindings.append(_RouterReplayBinding(layer_idx, replay, topk)) + + 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)) + + @property + def layer_ids(self) -> tuple[int, ...]: + """Global decoder-layer ids, in the model's replay order.""" + return self._layer_ids + + def replay(self, prepared_routes: torch.Tensor | None) -> AbstractContextManager[None]: + """Replay routes prepared in the model input's token order. + + Args: + prepared_routes: Signed integer expert ids with arbitrary token + axes followed by ``[global_layers, topk]``. Its flattened token + order must match the model input after any padding, packing, or + context-parallel sharding. ``-1`` keeps a token's complete live + top-k selection. ``None`` selects live routing. + + Returns: + A context that replays this model's layer targets through forward, + backward, and activation-checkpoint recomputation. + """ + if prepared_routes is None: + return nullcontext() + if not isinstance(prepared_routes, torch.Tensor): + raise TypeError("prepared_routes must be a Tensor or None") + if prepared_routes.dtype not in {torch.int8, torch.int16, torch.int32, torch.int64}: + raise TypeError(f"prepared_routes must use a signed integer dtype, got {prepared_routes.dtype}") + if prepared_routes.ndim < 3: + raise ValueError("prepared_routes must have token axes followed by [global_layers, topk]") + + route_tokens = prod(prepared_routes.shape[:-2]) + num_layers, route_topk = prepared_routes.shape[-2:] + max_layer_idx = self._bindings[-1].layer_idx + if num_layers <= max_layer_idx: + raise ValueError( + f"prepared_routes has {num_layers} global layers but this model requires layer {max_layer_idx}" + ) + per_token = prepared_routes.reshape(route_tokens, num_layers, route_topk) + if route_topk != self._topk: + raise ValueError(f"prepared_routes topk {route_topk} does not match model topk {self._topk}") + layer_indices = torch.tensor(self.layer_ids, device=per_token.device, dtype=torch.long) + selected = per_token.index_select(1, layer_indices) + targets = list(selected.unbind(dim=1)) + return self._activate(targets) + + @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) for binding in self._bindings] + try: + for binding, target in zip(self._bindings, targets): + binding.replay.target_indices = target + binding.replay.mode = RouterReplayMode.REPLAY + yield + finally: + for binding, (mode, target) in zip(self._bindings, previous): + binding.replay.mode = mode + binding.replay.target_indices = target + + 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/components/optim/scheduler.py b/nemo_automodel/components/optim/scheduler.py index 3031857cbd..c14fde10fa 100644 --- a/nemo_automodel/components/optim/scheduler.py +++ b/nemo_automodel/components/optim/scheduler.py @@ -246,7 +246,7 @@ def get_lr(self, param_group: dict[str, Any]) -> float: return min_lr + coeff * delta_lr - def step(self, increment: int) -> None: + def step(self, increment: int = 1) -> None: """ Set lr for all parameters groups. diff --git a/nemo_automodel/components/quantization/fp8.py b/nemo_automodel/components/quantization/fp8.py index 8e3e30fe22..9604bbf274 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/components/utils/model_utils.py b/nemo_automodel/components/utils/model_utils.py index 0102930d76..4acbb57640 100644 --- a/nemo_automodel/components/utils/model_utils.py +++ b/nemo_automodel/components/utils/model_utils.py @@ -577,12 +577,24 @@ def squeeze_input_for_thd(input_ids, position_ids, padding_mask, attn_kwargs, se This function modifies attn_kwargs in-place. If you need to preserve the original dictionary, pass a copy. """ + # Media tensors are indexed by media item, not by the placeholder THD batch + # dimension, so squeezing dim 0 would drop a single item. + media_keys = { + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "image_position_ids", + "second_per_grid_ts", + } if input_ids is not None: input_ids = input_ids.squeeze(0) position_ids = position_ids.squeeze(0) if isinstance(padding_mask, torch.Tensor): padding_mask = padding_mask.squeeze(0) for key, value in attn_kwargs.items(): + if key in media_keys: + continue if isinstance(value, torch.Tensor): attn_kwargs[key] = value.squeeze(0) if key in ["cu_seqlens", "cu_seqlens_padded"]: diff --git a/nemo_automodel/engine/__init__.py b/nemo_automodel/engine/__init__.py new file mode 100644 index 0000000000..7dddc8afcb --- /dev/null +++ b/nemo_automodel/engine/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Public training Engine API.""" + +from nemo_automodel.engine._engine import Engine + +__all__ = ["Engine"] diff --git a/nemo_automodel/engine/_engine.py b/nemo_automodel/engine/_engine.py new file mode 100644 index 0000000000..2596d02bf5 --- /dev/null +++ b/nemo_automodel/engine/_engine.py @@ -0,0 +1,310 @@ +# 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 training engine for an already-distributed eager model.""" + +from __future__ import annotations + +import logging +import math +import sys +from collections.abc import Callable +from contextlib import AbstractContextManager +from typing import Any + +import torch +import torch.distributed as dist +from torch import nn + +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.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.shared.import_utils import MISSING_TORCHAO_MSG, safe_import_from + +logger = logging.getLogger(__name__) + + +def _resolve_fp8_scale_precompute(module: nn.Module) -> Callable[[nn.Module], None] | None: + """Resolve torchao's post-step helper when the module opts into it.""" + if not getattr(module, "precompute_float8_dynamic_scale_for_fsdp", False): + 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 precompute + + +def _uses_summed_gradient_reduction(module: nn.Module) -> bool: + """Detect MegatronFSDP's summed-gradient reduction contract. It publishes + ``calculate_per_token_loss`` on its wrapper, which may sit below other + wrapper layers, so walk the whole module tree.""" + return any(getattr(child, "calculate_per_token_loss", False) for child in module.modules()) + + +class Engine(nn.Module): + """Train an already-distributed eager model. + + The interface intentionally mirrors a regular PyTorch module plus the + training operations exposed by DeepSpeed: callers prepare batches and + losses, while the engine owns backward, gradient accumulation, distributed + gradient finalization, optimizer updates, and scheduler advancement. The + engine preserves the usual averaged-gradient backward semantics, while + compensating when a backend instead sums distributed gradients. + + Args: + module: Model whose forward method is exposed by this engine. It must + already be parallelized and wrapped for distributed training. + optimizer: Optimizer for ``module``. A forward-only engine may omit it. + lr_scheduler: Optional scheduler advanced after each optimizer update. + mesh_context: Runtime topology used to finalize distributed gradients. + max_grad_norm: Maximum global gradient norm. ``None`` disables clipping + while preserving distributed expert-gradient finalization. + gradient_accumulation_steps: Number of forward/backward microsteps per + optimizer update. + defer_fsdp_grad_sync: Defer FSDP gradient synchronization until the + final microstep in an accumulation window. Wrappers exposing a + ``no_sync`` context use it on non-final microsteps. + + Note: + Batch collation, device transfer, context-parallel input sharding, + packing, loss normalization, and output restoration belong to the + caller. Pipeline execution has a different scheduling contract and is + not supported by this eager engine. + """ + + def __init__( + self, + module: nn.Module, + *, + optimizer: torch.optim.Optimizer | None = None, + lr_scheduler: OptimizerParamScheduler | torch.optim.lr_scheduler.LRScheduler | None = None, + mesh_context: MeshContext | None = None, + max_grad_norm: float | None = 1.0, + gradient_accumulation_steps: int = 1, + defer_fsdp_grad_sync: bool = True, + ) -> None: + super().__init__() + if isinstance(module, AutoPipeline): + raise NotImplementedError( + "Engine supports eager modules only; execute AutoPipeline with its pipeline schedule" + ) + if mesh_context is not None and mesh_context.pp_enabled: + raise NotImplementedError( + "Engine supports eager modules only; execute pipeline stages with their pipeline schedule" + ) + + self.module = module + self.optimizer = optimizer + self.lr_scheduler = lr_scheduler + self.mesh_context = mesh_context + self.max_grad_norm = max_grad_norm + self.defer_fsdp_grad_sync = defer_fsdp_grad_sync + + self._micro_step = 0 + self._backward_context: AbstractContextManager[Any] | None = None + self._global_grad_norm: torch.Tensor | float | None = None + self._precompute_fp8_scale_fn = _resolve_fp8_scale_precompute(module) + self._summed_gradient_reduction = _uses_summed_gradient_reduction(module) + self.set_gradient_accumulation_steps(gradient_accumulation_steps) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Delegate a raw forward call to :attr:`module`.""" + if self.optimizer is None or not self.module.training or not torch.is_grad_enabled(): + return self.module(*args, **kwargs) + if self._backward_context is not None: + raise RuntimeError("call Engine.backward() before starting another training forward") + + if self._micro_step == 0: + prepare_for_grad_accumulation([self.module], pp_enabled=False) + if self.is_gradient_accumulation_boundary(): + prepare_for_final_backward([self.module], pp_enabled=False) + + context = get_sync_ctx( + self.module, + self.is_gradient_accumulation_boundary(), + defer_fsdp_grad_sync=self.defer_fsdp_grad_sync, + ) + context.__enter__() + try: + output = self.module(*args, **kwargs) + except BaseException: + context.__exit__(*sys.exc_info()) + raise + self._backward_context = context + return output + + def backward( + self, + loss: torch.Tensor, + retain_graph: bool = False, + scale_wrt_gas: bool = True, + ) -> None: + """Backpropagate a scalar loss. + + Args: + loss: Scalar loss to backpropagate. Its scale follows the usual + averaged-gradient distributed-training convention. + retain_graph: Preserve the autograd graph after backward. + scale_wrt_gas: Divide ``loss`` by the configured gradient + accumulation steps. Disable this when the caller has already + normalized the complete accumulation window. + """ + if self.optimizer is None: + self._close_backward_context() + raise RuntimeError("Engine.backward requires an optimizer") + if not isinstance(loss, torch.Tensor) or loss.numel() != 1: + self._close_backward_context() + raise ValueError(f"loss must be a scalar Tensor, got {loss!r}") + + # Compensate backends that sum distributed gradients instead of averaging. + reduction_compensation = 1.0 / self._gradient_group_size() if self._summed_gradient_reduction else 1.0 + scale = reduction_compensation + if scale_wrt_gas: + scale /= self._gradient_accumulation_steps + # The auxiliary loss is injected by its own autograd node on every + # microstep, so it always needs accumulation scaling even when the + # caller has already normalized the main loss for the full window. Its + # local mean also needs the same reducer compensation. + MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor( + self._context_parallel_size() * reduction_compensation / self._gradient_accumulation_steps, + device=loss.device, + ) + try: + (loss * scale).backward(retain_graph=retain_graph) + finally: + self._close_backward_context(*sys.exc_info()) + + if self._micro_step == 0: + prepare_after_first_microbatch() + + @torch.no_grad() + def step(self) -> None: + """Advance one microstep and update parameters at its GAS boundary.""" + if self.optimizer is None: + raise RuntimeError("Engine.step requires an optimizer") + if self._backward_context is not None: + raise RuntimeError("call Engine.backward() before Engine.step()") + + if not self.is_gradient_accumulation_boundary(): + self._micro_step += 1 + return + + 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 + self._global_grad_norm = scale_grads_and_clip_grad_norm( + max_grad_norm=self.max_grad_norm, + model_parts=[self.module], + norm_type=2.0, + pp_enabled=False, + 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=None, + foreach=True, + num_label_tokens=None, + dp_group_size=self._gradient_group_size(), + expert_tp_replication_factor=get_expert_tp_replication_factor([self.module], device_mesh), + ) + + grad_norm = self._global_grad_norm + grad_is_finite = ( + bool(torch.isfinite(grad_norm).all()) if isinstance(grad_norm, torch.Tensor) else math.isfinite(grad_norm) + ) + if grad_is_finite: + self.optimizer.step() + else: + # The gradients themselves are inf/NaN; stepping would write NaN + # into the weights, and recovering then requires a checkpoint. + logger.warning("skipping the optimizer update: non-finite gradient norm %s", grad_norm) + self.zero_grad() + + update_moe_gate_bias = getattr(self.module, "update_moe_gate_bias", None) + if callable(update_moe_gate_bias): + update_moe_gate_bias() + + if self._precompute_fp8_scale_fn is not None: + self._precompute_fp8_scale_fn(self.module) + + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + self._micro_step = 0 + + def zero_grad(self) -> None: + """Clear gradients through the optimizer.""" + if self.optimizer is None: + raise RuntimeError("Engine.zero_grad requires an optimizer") + self.optimizer.zero_grad(set_to_none=True) + + def set_gradient_accumulation_steps(self, steps: int) -> None: + """Set the number of microsteps in the next optimizer window.""" + if steps < 1: + raise ValueError(f"gradient accumulation steps must be a positive integer, got {steps!r}") + if self._micro_step != 0: + raise RuntimeError("gradient accumulation steps cannot change during an active window") + self._gradient_accumulation_steps = steps + + def is_gradient_accumulation_boundary(self) -> bool: + """Return whether the current microstep completes the optimizer window.""" + return self._micro_step + 1 == self._gradient_accumulation_steps + + def reset_accumulation(self) -> None: + """Abandon the current accumulation window after a failed microstep. + + A microstep that raises (e.g. an OOM between forward and backward) can + leave an open gradient-sync context and a nonzero microstep counter, + which would poison every later window. Callers recovering from such a + failure reset here and clear gradients before retrying. + """ + self._close_backward_context() + self._micro_step = 0 + if self.optimizer is not None: + self.optimizer.zero_grad(set_to_none=True) + + def get_global_grad_norm(self) -> torch.Tensor | float | None: + """Return the global gradient norm measured at the latest optimizer step.""" + return self._global_grad_norm + + def _close_backward_context(self, *exc_info: Any) -> None: + context, self._backward_context = self._backward_context, None + if context is not None: + context.__exit__(*(exc_info or (None, None, None))) + + def _context_parallel_size(self) -> int: + return self.mesh_context.cp_size if self.mesh_context is not None else 1 + + def _gradient_group_size(self) -> int: + if self.mesh_context is not None and self.mesh_context.device_mesh is not None: + axis = "dp_cp" if self._context_parallel_size() > 1 else "dp" + return int(get_flat_mesh(self.mesh_context.device_mesh, axis).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 dist.get_world_size(group=group) + return 1 diff --git a/nemo_automodel/recipes/llm/benchmark.py b/nemo_automodel/recipes/llm/benchmark.py index ec9633864d..4b37ed792b 100644 --- a/nemo_automodel/recipes/llm/benchmark.py +++ b/nemo_automodel/recipes/llm/benchmark.py @@ -308,8 +308,13 @@ def run_benchmark(self): # Calculate gradient accumulation steps dp_size = self._get_dp_group_size() - ga_steps = global_batch_size // (local_batch_size * dp_size) - assert ga_steps > 0, "Global batch size must be divisible by local batch size * dp_size" + optimizer_microbatch_size = local_batch_size * dp_size + if global_batch_size < optimizer_microbatch_size or global_batch_size % optimizer_microbatch_size != 0: + raise ValueError( + f"global_batch_size ({global_batch_size}) must be a positive multiple of " + f"local_batch_size * dp_size ({optimizer_microbatch_size})" + ) + ga_steps = global_batch_size // optimizer_microbatch_size if rank == 0: logger.info(f"Running {steps} iterations with {warmup_steps} warmup steps") @@ -334,49 +339,59 @@ def run_benchmark(self): 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 + # Materialize the optimizer window: num_label_tokens must be + # all-reduced before the first forward. + batches = [next(dataloader_iter) for _ in range(ga_steps)] + num_label_tokens = sum((batch["labels"] != -100).sum().item() for batch in batches) + num_label_tokens_tensor = torch.tensor(num_label_tokens, dtype=torch.long, device=device) + num_label_tokens = self._dp_allreduce(num_label_tokens_tensor).item() loss_buffer = [] - prepare_for_grad_accumulation(self.model_parts, pp_enabled=self.pp_enabled) + if self.pp_enabled: + self._set_moe_aux_loss_backward_scale( + num_batches=ga_steps, + num_label_tokens=num_label_tokens, + ) + prepare_for_grad_accumulation(self.model_parts, pp_enabled=True) + else: + self.engine.set_gradient_accumulation_steps(ga_steps) - 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) + for ga_step_idx, batch in enumerate(batches): + if self.pp_enabled and ga_step_idx == ga_steps - 1: + prepare_for_final_backward(self.model_parts, pp_enabled=True) - # 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_label_tokens=num_label_tokens, num_batches=ga_steps, is_train=True, ) torch.cuda.nvtx.range_pop() - if ga_step_idx == 0: + if self.pp_enabled and ga_step_idx == 0: prepare_after_first_microbatch() - # Optimizer step - with self.timers("optimizer", log_level=2): - for opt in self.optimizer: - opt.step() - logger.debug("Optimizer step") + if not self.pp_enabled: + if self.engine.is_gradient_accumulation_boundary(): + self.checkpointer.maybe_wait_for_staging() + with self.timers("optimizer", log_level=2): + self.engine.step() + + if self.pp_enabled: + with self.timers("optimizer", log_level=2): + self._step_pipeline_optimizer( + num_label_tokens=num_label_tokens, + max_grad_norm=self.max_grad_norm, + ) + logger.debug("Optimizer step") # Match the training-loop lifecycle: record one complete eager # optimizer step, then capture outside the measured iteration. @@ -384,24 +399,18 @@ 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(torch.float32) / num_label_tokens + if num_label_tokens > 0 + else reporting_loss.to(torch.float32) * 0.0 + ) 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 = self._broadcast_from_last_pp_stage(reporting_loss) reporting_loss = reporting_loss.cpu().item() diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index ff396c7f8d..a597d380bf 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/nemo_automodel/recipes/llm/train_ft.py @@ -62,7 +62,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 ( @@ -95,6 +95,7 @@ filter_forward_kwargs, resolve_trust_remote_code, ) +from nemo_automodel.engine import Engine 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 @@ -148,6 +149,23 @@ def _should_pack_validation( ) +def _validate_pipeline_thd_model(model: nn.Module) -> None: + """Require a model-owned THD path for pipeline training.""" + 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.""" # TODO: Replace model-type exceptions with a shared mask-ownership capability. @@ -506,6 +524,12 @@ def setup(self): if not self._should_setup_training_components(): return + if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): + raise ValueError( + "Pipeline finetuning applies external global-token normalization and 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) @@ -544,9 +568,17 @@ 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" + ) + elif 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 ( @@ -561,9 +593,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 @@ -622,6 +653,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,6 +688,8 @@ def setup(self): # Loss-function capability check self.loss_fn = _maybe_downgrade_loss_fn(self.loss_fn, self.model_parts[0], self.pp is not None) + if getattr(self.loss_fn, "reduction", None) != "sum": + raise ValueError("Global-token-normalized finetuning requires a loss with reduction='sum'") # 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 @@ -728,7 +764,9 @@ def materialize_loader(config): else self.cfg.get("distributed.cp_size", 1) ), attn_implementation=attn_implementation, - collate_wrapper=collate_wrapper, + # THD already encodes document boundaries in cu_seqlens; + # a dense PP causal mask is redundant and prohibitively large. + collate_wrapper=None if getattr(config, "emits_thd", False) else collate_wrapper, ) self.dataloader = materialize_loader(self.cfg.dataloader) @@ -770,6 +808,18 @@ def materialize_loader(config): else None ) + self.engine = None + if not self.pp_enabled: + self.engine = Engine( + self.model_parts[0], + optimizer=self.optimizer[0], + lr_scheduler=self.lr_scheduler[0] if self.lr_scheduler is not None else None, + mesh_context=self.mesh_context, + max_grad_norm=self.max_grad_norm, + gradient_accumulation_steps=self.step_scheduler.grad_acc_steps, + 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) @@ -1098,17 +1148,9 @@ def _forward_backward_step( 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) + self.pp.step(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.eval(input_ids, target=targets, losses=losses, **batch_filtered) if self.pp.info.has_last_stage: local_loss = torch.sum(torch.stack(losses)) @@ -1118,26 +1160,17 @@ def _forward_backward_step( 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: + with train_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) + out = self.engine(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) + out = self.engine(**batch) # Gather the LM head once and share it across the main loss and # all MTP depths (FusedLinearCrossEntropy path) to avoid redundant @@ -1190,7 +1223,10 @@ def _forward_backward_step( ) loss_buffer.append(local_loss.clone().detach()) if is_train: - (local_loss * self._get_dp_group_size(include_cp=True)).backward() + self.engine.backward( + local_loss * self._get_dp_group_size(include_cp=True), + scale_wrt_gas=False, + ) 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.""" @@ -1199,6 +1235,48 @@ 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 _step_pipeline_optimizer(self, *, num_label_tokens: int, max_grad_norm: float | None) -> torch.Tensor | float: + """Finalize pipeline gradients and perform one complete optimizer update.""" + 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=True, + 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", + foreach=True, + num_label_tokens=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), + ) + + self.checkpointer.maybe_wait_for_staging() + for optimizer in self.optimizer: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + if hasattr(self.model_parts[0], "update_moe_gate_bias"): + for model_part in self.model_parts: + model_part.update_moe_gate_bias() + + if self.lr_scheduler is not None: + for scheduler in self.lr_scheduler: + scheduler.step(1) + + 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]) + + return grad_norm + def _run_train_optim_step(self, batches, max_grad_norm: float | None = None): """Execute a single training step. @@ -1213,8 +1291,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = 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. @@ -1224,61 +1300,34 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = 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) + if self.pp_enabled: + 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=True) + else: + self.engine.set_gradient_accumulation_steps(num_batches) for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) + if self.pp_enabled and i == num_batches - 1: + prepare_for_final_backward(self.model_parts, pp_enabled=True) self._forward_backward_step( i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches ) - if i == 0: + if self.pp_enabled and i == 0: prepare_after_first_microbatch() + if not self.pp_enabled: + if i == num_batches - 1: + self.checkpointer.maybe_wait_for_staging() + self.engine.step() - 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=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), - ) - - # 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) - - # 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]) + if self.pp_enabled: + grad_norm = self._step_pipeline_optimizer( + num_label_tokens=num_label_tokens, + max_grad_norm=max_grad_norm, + ) + else: + grad_norm = self.engine.get_global_grad_norm() # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api # self.model_parts[0].install_optimized_model_weights() @@ -1319,7 +1368,7 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = None): 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 / num_label_tokens if num_label_tokens > 0 else reporting_loss * 0.0 reporting_loss = reporting_loss.to(self.dist_env.device) reporting_loss = self._broadcast_from_last_pp_stage(reporting_loss) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index ac0dd0ce79..17ade7348a 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -28,7 +28,7 @@ import logging import pathlib import time -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from typing import TYPE_CHECKING, Any, Protocol import mlflow @@ -57,7 +57,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 ( @@ -82,6 +82,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 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 @@ -459,6 +460,12 @@ def setup(self): if not self._should_setup_training_components(): return + if self.pp_enabled and getattr(self.pipeline_config, "scale_grads_in_schedule", False): + raise ValueError( + "Pipeline VLM finetuning applies external global-token normalization and 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). @@ -487,13 +494,20 @@ 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" + ) + elif 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 @@ -554,6 +568,8 @@ def setup(self): if not _supports_logits_to_keep(model) and not isinstance(self.loss_fn, MaskedCrossEntropy): logger.warning("logits_to_keep not found in model.forward. Using MaskedCrossEntropy instead.") self.loss_fn = MaskedCrossEntropy() + if getattr(self.loss_fn, "reduction", None) != "sum": + raise ValueError("Global-token-normalized VLM finetuning requires a loss with reduction='sum'") if isinstance(model, AutoPipeline): self.model_parts = model.parts @@ -659,6 +675,18 @@ def setup(self): else None ) + self.engine = None + if not self.pp_enabled: + self.engine = Engine( + self.model_parts[0], + optimizer=self.optimizer[0], + lr_scheduler=self.lr_scheduler[0] if self.lr_scheduler is not None else None, + mesh_context=self.mesh_context, + max_grad_norm=self.max_grad_norm, + gradient_accumulation_steps=self.step_scheduler.grad_acc_steps, + 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) @@ -943,27 +971,18 @@ def _forward_backward_step( 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(): + with 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) + out = self.engine(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) + out = self.engine(**batch) grad_reduce_group = self._get_dp_group(include_cp=True) if is_train else None shared_lm_weight = ( @@ -1030,7 +1049,10 @@ def _forward_backward_step( loss_buffer.append(local_loss.clone().detach()) if is_train: - (local_loss * self._get_dp_group_size(include_cp=True)).backward() + self.engine.backward( + local_loss * self._get_dp_group_size(include_cp=True), + scale_wrt_gas=False, + ) def _configure_pipeline_loss_fn(self): if self.pp is None or not self.pp.info.has_last_stage: @@ -1050,6 +1072,48 @@ def _configure_pipeline_loss_fn(self): grad_reduce_group=self._get_dp_group(include_cp=True), ) + def _step_pipeline_optimizer(self, *, num_label_tokens: int, max_grad_norm: float | None) -> torch.Tensor | float: + """Finalize pipeline gradients and perform one complete optimizer update.""" + 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=True, + 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", + foreach=True, + num_label_tokens=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), + ) + + self.checkpointer.maybe_wait_for_staging() + for optimizer in self.optimizer: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + if hasattr(self.model_parts[0], "update_moe_gate_bias"): + for model_part in self.model_parts: + model_part.update_moe_gate_bias() + + if self.lr_scheduler is not None: + for scheduler in self.lr_scheduler: + scheduler.step(1) + + 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]) + + return grad_norm + def _run_train_optim_step(self, batches, max_grad_norm: float | None = None): """Execute a single training step. @@ -1063,8 +1127,6 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = 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. @@ -1074,60 +1136,34 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = 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) + if self.pp_enabled: + 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=True) + else: + self.engine.set_gradient_accumulation_steps(num_batches) for i, batch in enumerate(batches): - if i == num_batches - 1: - prepare_for_final_backward(self.model_parts, pp_enabled=self.pp_enabled) + if self.pp_enabled and i == num_batches - 1: + prepare_for_final_backward(self.model_parts, pp_enabled=True) self._forward_backward_step( i, batch, loss_buffer=loss_buffer, num_label_tokens=num_label_tokens, num_batches=num_batches ) - if i == 0: + if self.pp_enabled and i == 0: prepare_after_first_microbatch() + if not self.pp_enabled: + if i == num_batches - 1: + self.checkpointer.maybe_wait_for_staging() + self.engine.step() - 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=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), - ) - - # 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 ( - 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]) + if self.pp_enabled: + grad_norm = self._step_pipeline_optimizer( + num_label_tokens=num_label_tokens, + max_grad_norm=max_grad_norm, + ) + else: + grad_norm = self.engine.get_global_grad_norm() # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api # self.model.install_optimized_model_weights() 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/moe/test_experts_ep_tp_grad_parity.py b/tests/functional_tests/moe/test_experts_ep_tp_grad_parity.py index f41c5c9f19..2e4a6f80ff 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 +an ordinary Engine forward, backward, and step with tp=2 replicated tokens and +a 2-rank EP mesh. It asserts that the gradient finalizer restores the +single-process fp32 loss, global gradient norm, and one-step parameter update. """ from __future__ import annotations @@ -37,12 +36,10 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import Shard, distribute_tensor +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 _TP_SIZE = 2 _WORLD_SIZE = _TP_SIZE @@ -51,6 +48,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 +119,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 +133,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 +154,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 +176,29 @@ 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 task owns its inputs and scalar loss; Engine only executes the + # module, backward, distributed finalization, and optimizer update. 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) - - y_local.sum().backward() + loss_weights = _loss_weights() + model = _ExpertModel(experts) + model._nemo_moe_tp_requires_replica_sync = True + optimizer = torch.optim.SGD(model.parameters(), lr=_LEARNING_RATE) + engine = Engine( + model, + mesh_context=MeshContext.from_meshes(world_mesh, ep_mesh), + optimizer=optimizer, + max_grad_norm=1e6, + ) + output = engine( + input_ids=x, + token_mask=token_mask, + router_weights=weights, + router_indices=indices, + ) + loss = (output.square() * loss_weights).sum() / loss_weights.sum() + engine.backward(loss, scale_wrt_gas=False) + torch.testing.assert_close(output.detach(), y_ref, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(loss.detach(), loss_ref, rtol=1e-5, atol=1e-7) n_local_experts = _N_EXPERTS // world_size start = rank * n_local_experts @@ -163,28 +217,30 @@ def _ep_tp_grad_parity_worker(rank: int, world_size: int, port: int) -> None: 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 - scale_grads_and_clip_grad_norm( - max_grad_norm=None, - model_parts=[experts], - moe_mesh=ep_mesh, - ep_axis_name="ep", - dp_group_size=1, - expert_tp_replication_factor=replication_factor, + # Engine.step removes exactly that factor before parameter mutation. + engine.step() + torch.testing.assert_close(engine.get_global_grad_norm(), reference_grad_norm, rtol=1e-5, atol=1e-7) + 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.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.gate_and_up_projs.grad.to_local(), gate_up_grad_ref_local, rtol=1e-4, atol=1e-5 + experts.down_projs.to_local(), + reference_experts.down_projs[start:end], + 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) 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/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..c034831969 --- /dev/null +++ b/tests/functional_tests/parallelism/run_megatron_fsdp_per_token_loss.py @@ -0,0 +1,209 @@ +# 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.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) -> torch.Tensor: + """Map numeric token rows ``[batch, 4]`` to outputs of the same shape.""" + return self.block(input_ids.to(torch.float32)) + + +def _windows(rank: int, device: torch.device) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """Build two rank-local windows with unequal global token denominators. + + Args: + rank: Data-parallel rank, zero or one. + device: CUDA device on which to construct each microbatch. + + Returns: + Two ``(input_ids, weights)`` microbatches with shape ``[1, 4]``. + Across DP, their weight sums are 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 microbatch(values: list[int], weights: list[float]) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.tensor([values], device=device), + torch.tensor([weights], device=device), + ) + + return microbatch(values_a, weights_a), microbatch(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 _weighted_identity_loss(output: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """Return the weighted model-output numerator. + + Args: + output: Model values shaped ``[batch=1, sequence=4]``. + weights: Per-token weights with the same ``[batch=1, sequence=4]`` layout. + + Returns: + Scalar rank-local weighted numerator. + """ + assert output.shape == weights.shape + return (output * weights.to(output)).sum() + + +def _run_mode( + mesh_context: MeshContext, + *, + summed_gradients: bool, +) -> 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. + summed_gradients: Value passed as ``calculate_per_token_loss``. True + selects SUM gradient collectives; false selects averaged gradients. + + Returns: + 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() + device = torch.device("cuda", torch.cuda.current_device()) + 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) + windows = _windows(dist.get_rank(), device) + engine = Engine( + model, + optimizer=optimizer, + mesh_context=mesh_context, + max_grad_norm=1e9, + gradient_accumulation_steps=len(windows), + ) + + global_weight_sum = sum(weights.sum() for _, weights in windows) + dist.all_reduce(global_weight_sum) + reduction_scale = dist.get_world_size() / global_weight_sum + + local_loss_sum = torch.zeros((), device=device) + for input_ids, weights in windows: + loss_sum = _weighted_identity_loss(engine(input_ids=input_ids), weights) + local_loss_sum += loss_sum.detach() + engine.backward(loss_sum * reduction_scale, scale_wrt_gas=False) + engine.step() + + global_loss_sum = local_loss_sum.clone() + dist.all_reduce(global_loss_sum) + + statistics = ( + global_loss_sum.item(), + global_weight_sum.item(), + (global_loss_sum / global_weight_sum).item(), + ) + return statistics, float(engine.get_global_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] == 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) + + 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 01acc138c2..4a886a8692 100644 --- a/tests/functional_tests/parallelism/test_parallelism.py +++ b/tests/functional_tests/parallelism/test_parallelism.py @@ -32,6 +32,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" DEEPSEEK_V4_PP2_PARITY_FILENAME = "L2_Parallelism_DeepSeekV4_PP2_Parity.sh" DEEPSEEK_V4_EP2_PARITY_FILENAME = "L2_Parallelism_DeepSeekV4_EP2_Parity.sh" QWEN3_5_MOE_PP2_PARITY_FILENAME = "L2_Parallelism_Qwen3_5MoE_PP2_Parity.sh" @@ -48,6 +49,9 @@ 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) + def test_deepseek_v4_pp2_parity(self): run_test_script(TEST_FOLDER, DEEPSEEK_V4_PP2_PARITY_FILENAME) diff --git a/tests/unit_tests/_transformers/test_model_init.py b/tests/unit_tests/_transformers/test_model_init.py index f6e1045bdf..4b8dbaedf4 100644 --- a/tests/unit_tests/_transformers/test_model_init.py +++ b/tests/unit_tests/_transformers/test_model_init.py @@ -664,6 +664,33 @@ def fake_model_cls(config, **kwargs): # otherwise collide with the positional config and/or raise TypeError). assert "config" not in captured_kwargs + @patch("nemo_automodel._transformers.model_init._download_model_weights") + @patch("nemo_automodel._transformers.model_init._resolve_custom_model_cls_for_config") + def test_config_object_is_passed_only_positionally(self, mock_resolve_cls, _mock_download): + hf_config = self._make_config() + captured_kwargs = {} + + def fake_model_cls(config, **kwargs): + assert config is hf_config + captured_kwargs.update(kwargs) + return MagicMock() + + fake_model_cls.__module__ = "nemo_automodel.components.models.fake" + mock_resolve_cls.return_value = fake_model_cls + + is_custom, _ = _init_model( + cls=MagicMock(), + pretrained_model_name_or_path_or_config="fake/model", + attn_implementation="flash_attention_2", + torch_dtype="auto", + quantization_config=None, + force_hf=False, + config=hf_config, + ) + + assert is_custom is True + assert "config" not in captured_kwargs + class TestSetupBnbLoadingKwargs: """_setup_bnb_loading_kwargs sets a per-GPU device_map and disables HF async weight loading.""" diff --git a/tests/unit_tests/datasets/vlm/test_neat_packing_vlm.py b/tests/unit_tests/datasets/vlm/test_neat_packing_vlm.py index c0019df1d0..6349294dc2 100644 --- a/tests/unit_tests/datasets/vlm/test_neat_packing_vlm.py +++ b/tests/unit_tests/datasets/vlm/test_neat_packing_vlm.py @@ -22,6 +22,7 @@ _compute_mrope_position_ids, _shift_sample, neat_pack_dataset_vlm, + pack_vlm_samples, ) @@ -105,7 +106,7 @@ def test_basic(self): "labels": torch.tensor([105, 106]), }, ] - result = _build_packed_vlm_sample(samples, pack_size=8, padding_idx=0) + result = _build_packed_vlm_sample(samples, padding_idx=0) # No pre-padding: length = sum of samples (3 + 2 = 5) assert result["input_ids"].shape == (5,) @@ -128,7 +129,7 @@ def test_media_concat(self): "image_grid_thw": torch.tensor([[1, 224, 224]]), }, ] - result = _build_packed_vlm_sample(samples, pack_size=4, padding_idx=0) + result = _build_packed_vlm_sample(samples, padding_idx=0) assert result["pixel_values"].shape[0] == 3 # 2 + 1 assert result["image_grid_thw"].shape[0] == 3 @@ -149,7 +150,7 @@ def test_variable_resolution_media_lists_are_preserved(self): }, ] - result = _build_packed_vlm_sample(samples, pack_size=4, padding_idx=0) + result = _build_packed_vlm_sample(samples, padding_idx=0) assert isinstance(result["pixel_values"], list) assert [tuple(value.shape) for value in result["pixel_values"]] == [(3, 8, 12), (3, 16, 8)] @@ -167,7 +168,7 @@ def test_mm_token_type_ids_propagated(self): "mm_token_type_ids": torch.tensor([1, 1]), }, ] - result = _build_packed_vlm_sample(samples, pack_size=8, padding_idx=0) + result = _build_packed_vlm_sample(samples, padding_idx=0) assert result["mm_token_type_ids"].tolist() == [0, 1, 0, 1, 1] def test_sequence_alignment_pads_each_document_and_preserves_real_lengths(self): @@ -178,7 +179,6 @@ def test_sequence_alignment_pads_each_document_and_preserves_real_lengths(self): result = _build_packed_vlm_sample( samples, - pack_size=8, padding_idx=0, sequence_alignment=4, ) @@ -190,6 +190,45 @@ def test_sequence_alignment_pads_each_document_and_preserves_real_lengths(self): assert result["position_ids"].tolist() == [0, 1, 2, 3, 0, 1, 2, 3] +def test_pack_vlm_samples_applies_shift_alignment_and_media_merge(): + samples = [ + _make_vlm_sample(4, has_image=True), + _make_vlm_sample(3), + ] + + result = pack_vlm_samples(samples, padding_idx=0, sequence_alignment=4) + + assert result["input_ids"].tolist() == [1, 2, 3, 0, 1, 2, 0, 0] + assert result["labels"].tolist() == [102, 103, 104, -100, 102, 103, -100, -100] + assert result["seq_lens"] == [3, 2] + assert result["seq_lens_padded"] == [4, 4] + assert result["n_images"] == 2 + + +def test_pack_vlm_samples_builds_and_shifts_mrope_positions(): + def get_rope_index(input_ids, attention_mask=None): + seq_len = input_ids.shape[1] + positions = torch.arange(seq_len).expand(3, 1, seq_len) + return positions, torch.zeros(1) + + result = pack_vlm_samples( + [_make_vlm_sample(4)], + padding_idx=0, + get_rope_index=get_rope_index, + ) + + assert result["position_ids"].shape == (3, 3) + assert result["position_ids"][0].tolist() == [0, 1, 2] + + with pytest.raises(NotImplementedError, match="multi-axis mRoPE"): + pack_vlm_samples( + [_make_vlm_sample(4)], + padding_idx=0, + get_rope_index=get_rope_index, + sequence_alignment=2, + ) + + class TestNeatPackDatasetVlm: def test_end_to_end(self): samples = [ @@ -436,7 +475,7 @@ def test_build_packed_with_mrope(self): ), }, ] - result = _build_packed_vlm_sample(samples, pack_size=7, padding_idx=0, has_mrope=True) + result = _build_packed_vlm_sample(samples, padding_idx=0, has_mrope=True) # No pre-padding: [3, 5] (3+2 tokens) assert result["position_ids"].shape == (3, 5) diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 9bc08cce91..97ec2d0376 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -18,7 +18,7 @@ import pytest import torch import torch.nn as nn -from torch.distributed.pipelining.microbatch import TensorChunkSpec, split_args_kwargs_into_chunks +from torch.distributed.pipelining.microbatch import split_args_kwargs_into_chunks from nemo_automodel.components.distributed.pipelining.autopipeline import AutoPipeline from nemo_automodel.components.distributed.pipelining.functional import ( @@ -118,6 +118,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): @@ -222,153 +228,76 @@ 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): +class _MropeChunkingPart(nn.Module): def get_pipeline_kwargs_chunk_dims(self, kwargs): - return {"unknown": 0} + position_ids = kwargs.get("position_ids") + return {"position_ids": 1} if isinstance(position_ids, torch.Tensor) and position_ids.ndim == 3 else {} -class _KwargsChunkSchedule: - def __init__(self, *, fail_on_step: bool = False): - self._kwargs_chunk_spec = None - self.fail_on_step = fail_on_step - self.args_during_step = None - self.kwargs_chunk_spec_during_step = None +class _ChunkingSchedule: + def __init__(self): + self._kwargs_chunk_spec = {"original": object()} self.kwargs_split = None - def step(self, *args, target=None, losses=None, **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. - losses: Optional mutable list populated with scalar loss tensors. - **kwargs: Keyword schedule inputs. Tensor values have arbitrary - model-defined layouts. - - Returns: - A sentinel string identifying the schedule result. - """ - del target, losses - self.args_during_step = args - self.kwargs_chunk_spec_during_step = self._kwargs_chunk_spec - if self.fail_on_step: - raise RuntimeError("schedule failed") + def _run(self, *args, **kwargs): + kwargs.pop("target", None) + kwargs.pop("losses", None) _, self.kwargs_split = split_args_kwargs_into_chunks( args, kwargs, 2, kwargs_chunk_spec=self._kwargs_chunk_spec, ) - return "schedule-result" - -class TestAutoPipelineKwargsChunkSpec: - def _pipeline_with_parts(self, *parts: nn.Module, schedule=None, has_first_stage: bool = True): - ap = AutoPipeline( - world_mesh=FakeDeviceMesh(), - pp_axis_name="pp", - pp_schedule="1f1b", - pp_microbatch_size=1, - pp_batch_size=2, - device=torch.device("cpu"), - ) - ap._info.schedule = schedule or _KwargsChunkSchedule() - 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_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)) + def step(self, *args, **kwargs): + self._run(*args, **kwargs) + + def eval(self, *args, **kwargs): + self._run(*args, **kwargs) + + +@pytest.mark.parametrize("method_name", ["step", "eval"]) +def test_schedule_methods_split_mrope_position_ids_on_batch_axis(method_name): + """The model hook keeps every mRoPE axis in each PP microbatch.""" + pipeline = AutoPipeline( + world_mesh=FakeDeviceMesh(), + pp_axis_name="pp", + pp_schedule="1f1b", + pp_microbatch_size=1, + pp_batch_size=2, + device=torch.device("cpu"), + ) + schedule = _ChunkingSchedule() + original_chunk_spec = schedule._kwargs_chunk_spec + pipeline._info.schedule = schedule + pipeline._info.model_parts = [_MropeChunkingPart()] + pipeline._info.has_first_stage = True + input_ids = torch.zeros(2, 8, dtype=torch.long) + position_ids = torch.arange(8).view(1, 1, 8).expand(3, 2, -1).clone() + + getattr(pipeline, method_name)( + input_ids, + position_ids=position_ids, + attention_mask=torch.ones(2, 8, dtype=torch.bool), + qkv_format="thd", + ) + + assert schedule.kwargs_split[0]["position_ids"].shape == (3, 1, 8) + assert schedule.kwargs_split[1]["position_ids"].shape == (3, 1, 8) + torch.testing.assert_close(schedule.kwargs_split[0]["position_ids"], position_ids[:, :1]) + torch.testing.assert_close(schedule.kwargs_split[1]["position_ids"], position_ids[:, 1:]) + assert schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) + assert schedule.kwargs_split[0]["qkv_format"] == "thd" + assert schedule._kwargs_chunk_spec is original_chunk_spec # ----------------------------- -# 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.""" @@ -388,8 +317,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") @@ -504,11 +433,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_magi_attn_utils.py b/tests/unit_tests/distributed/test_magi_attn_utils.py index c7206377ff..ccf600446a 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 @@ -79,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) # --------------------------------------------------------------------------- # @@ -165,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): + calls.append((model, batch, cp_group)) + 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)] + + 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 @@ -252,6 +337,53 @@ 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), + ) + + 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.""" @@ -325,6 +457,41 @@ 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]), + } + out, returned_key, local_indices = mu.magi_prepare_packed_cp(model, batch, _FakeGroup(2)) + + 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/loss/test_vocab_parallel.py b/tests/unit_tests/loss/test_vocab_parallel.py new file mode 100644 index 0000000000..0204268c4d --- /dev/null +++ b/tests/unit_tests/loss/test_vocab_parallel.py @@ -0,0 +1,297 @@ +# 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 math + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard + +from nemo_automodel.components.loss import ( + token_entropy, + token_log_probs, +) + + +def _run_vocab_parallel_parity(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 = DeviceMesh("cpu", list(range(world_size)), mesh_dim_names=("tp",)) + temperature = 0.7 + cases = (((6,), 6), ((2, 3), 5)) + + for case_index, (leading_shape, vocab_size) in enumerate(cases): + torch.manual_seed(1234 + case_index) + full_logits = torch.randn(*leading_shape, vocab_size, dtype=torch.float32) + targets = (torch.arange(math.prod(leading_shape), dtype=torch.long) % vocab_size).reshape(leading_shape) + targets.reshape(-1)[-1] = vocab_size - 1 + chunk_size = (vocab_size + world_size - 1) // world_size + shard_offset = min(rank * chunk_size, vocab_size) + shard_size = min(chunk_size, vocab_size - shard_offset) + full_stride = full_logits.stride() + + local_log_prob_logits = ( + full_logits[..., shard_offset : shard_offset + shard_size].detach().clone().requires_grad_() + ) + distributed_log_prob_logits = DTensor.from_local( + local_log_prob_logits, + mesh, + [Shard(-1)], + run_check=False, + shape=full_logits.shape, + stride=full_stride, + ) + actual_log_probs = token_log_probs( + distributed_log_prob_logits, + targets, + temperature=temperature, + ) + + reference_log_prob_logits = full_logits.detach().clone().requires_grad_() + reference_log_probs = torch.log_softmax(reference_log_prob_logits / temperature, dim=-1) + reference_log_probs = reference_log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1) + assert actual_log_probs.dtype == torch.float32 + torch.testing.assert_close(actual_log_probs, reference_log_probs, rtol=1e-6, atol=1e-6) + + upstream = torch.linspace(0.25, 1.25, targets.numel(), dtype=torch.float32).reshape(leading_shape) + (actual_log_probs * upstream).sum().backward() + (reference_log_probs * upstream).sum().backward() + assert local_log_prob_logits.grad is not None + torch.testing.assert_close( + local_log_prob_logits.grad, + reference_log_prob_logits.grad[..., shard_offset : shard_offset + shard_size], + rtol=2e-6, + atol=2e-6, + ) + + local_entropy_logits = ( + full_logits[..., shard_offset : shard_offset + shard_size].detach().clone().requires_grad_() + ) + distributed_entropy_logits = DTensor.from_local( + local_entropy_logits, + mesh, + [Shard(-1)], + run_check=False, + shape=full_logits.shape, + stride=full_stride, + ) + actual_entropy = token_entropy(distributed_entropy_logits, temperature=temperature) + + reference_entropy_logits = full_logits.detach().clone().requires_grad_() + reference_log_distribution = torch.log_softmax(reference_entropy_logits / temperature, dim=-1) + reference_distribution = reference_log_distribution.exp() + reference_entropy = -(reference_distribution * reference_log_distribution).sum(dim=-1) + assert actual_entropy.dtype == torch.float32 + torch.testing.assert_close(actual_entropy, reference_entropy, rtol=2e-6, atol=2e-6) + + (actual_entropy * upstream).sum().backward() + (reference_entropy * upstream).sum().backward() + assert local_entropy_logits.grad is not None + torch.testing.assert_close( + local_entropy_logits.grad, + reference_entropy_logits.grad[..., shard_offset : shard_offset + shard_size], + rtol=3e-6, + atol=3e-6, + ) + + with torch.no_grad(): + no_grad_log_probs = token_log_probs( + distributed_log_prob_logits, + targets, + temperature=temperature, + ) + no_grad_entropy = token_entropy(distributed_entropy_logits, temperature=temperature) + assert not no_grad_log_probs.requires_grad + assert not no_grad_entropy.requires_grad + finally: + dist.destroy_process_group() + + +def test_vocab_parallel_forward_and_backward_match_dense_reference(tmp_path) -> None: + mp.spawn( + _run_vocab_parallel_parity, + args=(2, str(tmp_path / "vocab_parallel_pg")), + nprocs=2, + join=True, + ) + + +@pytest.fixture +def one_rank_mesh(): + dist.init_process_group("gloo", rank=0, world_size=1, store=dist.HashStore()) + try: + yield DeviceMesh("cpu", [0], mesh_dim_names=("tp",)) + finally: + dist.destroy_process_group() + + +def test_vocab_parallel_rejects_invalid_placements(one_rank_mesh) -> None: + replicated = DTensor.from_local(torch.randn(2, 3), one_rank_mesh, [Replicate()], run_check=False) + token_sharded = DTensor.from_local(torch.randn(2, 3), one_rank_mesh, [Shard(0)], run_check=False) + targets = torch.tensor([0, 1]) + + with pytest.raises(ValueError, match="exactly one Shard placement"): + token_log_probs(replicated, targets) + with pytest.raises(ValueError, match="last vocabulary axis"): + token_entropy(token_sharded) + + +@pytest.mark.parametrize("temperature", [0.0, -1.0, math.inf, math.nan]) +def test_vocab_parallel_rejects_invalid_temperature(one_rank_mesh, temperature) -> None: + logits = DTensor.from_local(torch.randn(2, 3), one_rank_mesh, [Shard(-1)], run_check=False) + + with pytest.raises(ValueError, match="positive and finite"): + token_entropy(logits, temperature=temperature) + + +@pytest.mark.parametrize( + ("targets", "error_type", "message"), + [ + (torch.tensor([[0, 1]]), ValueError, "targets shape"), + (torch.tensor([0.0, 1.0]), TypeError, "torch.int64"), + ], +) +def test_vocab_parallel_rejects_invalid_targets(one_rank_mesh, targets, error_type, message) -> None: + logits = DTensor.from_local(torch.randn(2, 3), one_rank_mesh, [Shard(-1)], run_check=False) + + with pytest.raises(error_type, match=message): + token_log_probs(logits, targets) + + +@pytest.mark.parametrize( + ("targets", "invalid_index"), + [(torch.tensor([-1, 1]), 0), (torch.tensor([0, 3]), 1)], +) +def test_vocab_parallel_marks_out_of_range_targets_nan(one_rank_mesh, targets, invalid_index) -> None: + logits = DTensor.from_local(torch.randn(2, 3), one_rank_mesh, [Shard(-1)], run_check=False) + + result = token_log_probs(logits, targets) + + assert torch.isnan(result[invalid_index]) + assert torch.isfinite(result[1 - invalid_index]) + + +def test_dense_forward_and_backward_match_fp32_reference_across_chunks() -> None: + torch.manual_seed(1234) + temperature = 0.7 + logits = torch.randn(3, 97, 19, dtype=torch.float32, requires_grad=True) + targets = torch.randint(0, logits.shape[-1], logits.shape[:-1]) + upstream_log_probs = torch.randn(logits.shape[:-1]) + upstream_entropy = torch.randn(logits.shape[:-1]) + + actual_log_probs = token_log_probs(logits, targets, temperature=temperature) + actual_entropy = token_entropy(logits, temperature=temperature) + + reference_logits = logits.detach().clone().requires_grad_() + reference_log_distribution = torch.log_softmax(reference_logits.float() / temperature, dim=-1) + reference_log_probs = reference_log_distribution.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1) + reference_entropy = -(reference_log_distribution.exp() * reference_log_distribution).sum(dim=-1) + + assert actual_log_probs.dtype == torch.float32 + assert actual_entropy.dtype == torch.float32 + torch.testing.assert_close(actual_log_probs, reference_log_probs) + torch.testing.assert_close(actual_entropy, reference_entropy) + + (actual_log_probs * upstream_log_probs + actual_entropy * upstream_entropy).sum().backward() + (reference_log_probs * upstream_log_probs + reference_entropy * upstream_entropy).sum().backward() + torch.testing.assert_close(logits.grad, reference_logits.grad) + + +def test_dense_bfloat16_token_stats_return_fp32() -> None: + torch.manual_seed(1234) + logits = torch.randn(257, 19, dtype=torch.bfloat16) + targets = torch.randint(0, logits.shape[-1], logits.shape[:-1]) + + log_probs = token_log_probs(logits, targets) + entropy = token_entropy(logits) + reference_log_distribution = torch.log_softmax(logits.float(), dim=-1) + + assert log_probs.dtype == torch.float32 + assert entropy.dtype == torch.float32 + torch.testing.assert_close( + log_probs, + reference_log_distribution.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1), + ) + torch.testing.assert_close(entropy, -(reference_log_distribution.exp() * reference_log_distribution).sum(dim=-1)) + + +def test_dense_token_stats_support_arbitrary_leading_dimensions() -> None: + logits = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) + target = torch.tensor(2) + + log_prob = token_log_probs(logits, target) + entropy = token_entropy(logits) + + assert log_prob.shape == torch.Size([]) + assert entropy.shape == torch.Size([]) + torch.testing.assert_close(log_prob, torch.log_softmax(logits, dim=-1)[target]) + expected_entropy = -(torch.softmax(logits, dim=-1) * torch.log_softmax(logits, dim=-1)).sum() + torch.testing.assert_close(entropy, expected_entropy) + + +def test_dense_token_stats_preserve_empty_token_gradient() -> None: + logits = torch.empty(0, 5, requires_grad=True) + targets = torch.empty(0, dtype=torch.long) + + log_probs = token_log_probs(logits, targets) + entropy = token_entropy(logits) + (log_probs.sum() + entropy.sum()).backward() + + assert log_probs.shape == torch.Size([0]) + assert entropy.shape == torch.Size([0]) + assert logits.grad is not None + assert logits.grad.shape == logits.shape + + +def test_dense_log_probs_marks_out_of_range_targets_nan() -> None: + logits = torch.randn(3, 5) + targets = torch.tensor([-1, 4, 5]) + + result = token_log_probs(logits, targets) + + assert torch.isnan(result[0]) + assert torch.isfinite(result[1]) + assert torch.isnan(result[2]) + + +@pytest.mark.parametrize("temperature", [0.0, -1.0, math.inf, math.nan]) +def test_dense_token_stats_reject_invalid_temperature(temperature) -> None: + logits = torch.randn(2, 3) + + with pytest.raises(ValueError, match="positive and finite"): + token_log_probs(logits, torch.tensor([0, 1]), temperature=temperature) + with pytest.raises(ValueError, match="positive and finite"): + token_entropy(logits, temperature=temperature) + + +@pytest.mark.parametrize( + ("operation", "message"), + [ + (lambda: token_log_probs(torch.ones(2, 3, dtype=torch.long), torch.tensor([0, 1])), "floating-point"), + (lambda: token_log_probs(torch.randn(2, 3), torch.tensor([0.0, 1.0])), "torch.int64"), + (lambda: token_log_probs(torch.randn(2, 3), torch.tensor([[0, 1]])), "targets shape"), + (lambda: token_entropy(torch.empty(2, 0)), "vocabulary size"), + ], +) +def test_dense_token_stats_reject_invalid_inputs(operation, message) -> None: + with pytest.raises((TypeError, ValueError), match=message): + operation() diff --git a/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py index 6d8ae0aa1c..0f5ee23694 100644 --- a/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py @@ -23,11 +23,16 @@ from __future__ import annotations import contextlib +from functools import partial from types import SimpleNamespace import pytest import torch +from nemo_automodel.components.distributed.context_parallel.sharder import ( + ContextParallelSharder, + contiguous_local_indices, +) from nemo_automodel.components.models.deepseek_v4 import cp as cpmod from nemo_automodel.components.models.deepseek_v4.cp import ( dsv4_cp_enabled, @@ -70,6 +75,15 @@ def get_group(self) -> object: return self._group +class _FakeDeviceMesh(dict): + """Minimal named device mesh containing one CP submesh.""" + + mesh_dim_names = ("cp",) + + def __init__(self, cp_mesh: _FakeMesh): + super().__init__(cp=cp_mesh) + + def _shard(batch, *, cp_size, local_rank, **kwargs): mesh = _FakeMesh(cp_size, local_rank) return make_dsv4_contiguous_shard_cp_batch_and_ctx(mesh, None, batch, **kwargs) @@ -182,7 +196,14 @@ def test_cp_size_one_native_thd_preserves_packed_batch_without_padding(): } expected = {key: value.clone() for key, value in batch.items()} - ctx, out, _ = _shard(batch, cp_size=1, local_rank=0, pad_multiple=8, padding_token_id=99) + cp_mesh = _FakeMesh(1) + sharder = ContextParallelSharder( + device_mesh=_FakeDeviceMesh(cp_mesh), + shard_batch=partial(make_dsv4_contiguous_shard_cp_batch_and_ctx, pad_multiple=8), + local_token_global_indices=contiguous_local_indices, + padding_token_id=99, + ) + ctx, out = sharder.shard(batch) assert ctx is contextlib.nullcontext assert out["qkv_format"] == "thd" @@ -191,6 +212,10 @@ def test_cp_size_one_native_thd_preserves_packed_batch_without_padding(): assert "padding_mask" not in out for key, value in expected.items(): torch.testing.assert_close(out[key], value) + assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (8, 8) + token_values = torch.arange(8, dtype=torch.float32).view(1, 8) + torch.testing.assert_close(sharder.gather_token_tensor(token_values, trim=True), token_values) + def test_contiguous_shard_pads_to_divisor(): @@ -420,58 +445,6 @@ def test_repad_packed_batch_validates_labels_and_metadata_extent(): ) -def test_contiguous_shard_syncs_packed_length_for_hybridep(monkeypatch): - monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - - def _all_reduce_max(length, op): - assert op == torch.distributed.ReduceOp.MAX - length.fill_(16) - - monkeypatch.setattr(torch.distributed, "all_reduce", _all_reduce_max) - batch = { - "input_ids": torch.arange(8).view(1, 8), - "labels": torch.arange(8).view(1, 8), - "qkv_format": "thd", - "seq_lens": torch.tensor([[3, 2]]), - "seq_lens_padded": torch.tensor([[3, 5]]), - } - - _, out, _ = _shard( - batch, - cp_size=2, - local_rank=0, - pad_multiple=4, - padding_token_id=99, - sync_packed_length=True, - ) - - assert out["input_ids"].shape == (1, 8) - torch.testing.assert_close(out["input_ids"], torch.tensor([[0, 1, 2, 99, 3, 4, 99, 99]])) - torch.testing.assert_close(out["labels"], torch.tensor([[0, 1, 2, -100, 3, 4, -100, -100]])) - - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 1) - batch = { - "input_ids": torch.arange(8).view(1, 8), - "labels": torch.arange(8).view(1, 8), - "qkv_format": "thd", - "seq_lens": torch.tensor([[3, 2]]), - "seq_lens_padded": torch.tensor([[3, 5]]), - } - _, out, _ = _shard( - batch, - cp_size=2, - local_rank=1, - pad_multiple=4, - padding_token_id=99, - sync_packed_length=True, - ) - - torch.testing.assert_close(out["input_ids"], torch.full((1, 8), 99)) - torch.testing.assert_close(out["labels"], torch.full((1, 8), -100)) - torch.testing.assert_close(out["padding_mask"], torch.ones((1, 8), dtype=torch.bool)) - torch.testing.assert_close(out["packed_seq_ids"], torch.zeros((1, 8), dtype=torch.long)) - def test_contiguous_shard_requires_exactly_one_primary_key(): # both input_ids and inputs_embeds -> assertion @@ -511,7 +484,6 @@ def test_prepare_model_inputs_for_cp_returns_sharder(): fn = sharder.shard_batch # the partial binds the config-derived per-rank multiple (lcm(8,128) == 128) assert fn.keywords["pad_multiple"] == 128 - assert fn.keywords["sync_packed_length"] is True assert fn.func is make_dsv4_contiguous_shard_cp_batch_and_ctx # the bound fn shards a batch end-to-end with a real (fake-mesh) divisor. @@ -527,7 +499,6 @@ def test_prepare_model_inputs_for_cp_binds_shard_multiple(): fake_self = SimpleNamespace(config=cfg, backend=SimpleNamespace(dispatcher="deepep")) out = DeepseekV4ForCausalLM.prepare_model_inputs_for_cp(fake_self, {"input_ids": torch.arange(8).view(1, 8)}) assert out["cp_sharder"].shard_batch.keywords["pad_multiple"] == 8 - assert out["cp_sharder"].shard_batch.keywords["sync_packed_length"] is False def test_setup_cp_attention_stores_group(): 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 9fb0dad7e4..f7950cae62 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 @@ -491,7 +491,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/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py index 9b0034f7cd..34d9f98780 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py @@ -13,6 +13,7 @@ from __future__ import annotations import types +from unittest.mock import MagicMock import pytest import torch @@ -68,6 +69,41 @@ def test_declares_cp_vision_frame_sharding_support(): assert capabilities.supports_cp_vision_frame_sharding is True +@pytest.mark.parametrize( + ("pixel_key", "grid_key"), + (("pixel_values", "image_grid_thw"), ("pixel_values_videos", "video_grid_thw")), +) +def test_thd_forward_preserves_single_media_axes(pixel_key, grid_key): + model = Qwen3_5MoeForConditionalGeneration.__new__(Qwen3_5MoeForConditionalGeneration) + nn.Module.__init__(model) + model.config = types.SimpleNamespace( + text_config=types.SimpleNamespace(output_hidden_states=False), + image_token_id=99, + vision_start_token_id=98, + ) + model.model = MagicMock(return_value=types.SimpleNamespace(last_hidden_state=torch.randn(3, 4))) + model.cp_mesh = None + model.mtp = None + model.lm_head = None + + pixel_values = torch.randn(1, 12) + grid_thw = torch.tensor([[1, 1, 1]]) + result = model( + input_ids=torch.tensor([[99, 1, 2]]), + position_ids=torch.arange(3).unsqueeze(0), + padding_mask=torch.zeros(1, 3, dtype=torch.bool), + qkv_format="thd", + **{pixel_key: pixel_values, grid_key: grid_thw}, + ) + + assert result.logits.shape == (3, 4) + model_kwargs = model.model.call_args.kwargs + assert model_kwargs[pixel_key] is pixel_values + assert model_kwargs[pixel_key].shape == (1, 12) + assert model_kwargs[grid_key] is grid_thw + assert model_kwargs[grid_key].shape == (1, 3) + + class TestPrepareModelInputsForCP: def test_requires_input_ids(self): model = _build_model() diff --git a/tests/unit_tests/models/qwen3_vl_moe/test_qwen3_vl_moe_model.py b/tests/unit_tests/models/qwen3_vl_moe/test_qwen3_vl_moe_model.py index ed1ba6fb73..cd82663445 100644 --- a/tests/unit_tests/models/qwen3_vl_moe/test_qwen3_vl_moe_model.py +++ b/tests/unit_tests/models/qwen3_vl_moe/test_qwen3_vl_moe_model.py @@ -505,11 +505,17 @@ def test_forward_handles_thd_format(self, vl_config, backend_config, moe_config, position_ids = torch.arange(seq_len, device=device).unsqueeze(0) attention_mask = torch.ones(batch, seq_len, device=device) padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool, device=device) + pixel_values = torch.randn(1, 4, vl_config.vision_config.in_channels, device=device) + image_grid_thw = torch.tensor([[1, 2, 2]], device=device) squeezed_ids = torch.randint(0, vl_config.text_config.vocab_size, (batch, seq_len), device=device) squeezed_position_ids = torch.arange(seq_len, device=device).unsqueeze(0) squeezed_padding_mask = torch.ones(batch, seq_len, dtype=torch.bool, device=device) - squeezed_kwargs = {"foo": "bar"} + squeeze_input_kwargs = {} + + def fake_squeeze(input_ids, position_ids, padding_mask, kwargs): + squeeze_input_kwargs.update(kwargs) + return squeezed_ids, squeezed_position_ids, squeezed_padding_mask, kwargs # Mock the model.model.forward to avoid internal tensor operations mock_hidden = torch.randn(batch, seq_len, vl_config.text_config.hidden_size, device=device, dtype=model_dtype) @@ -517,7 +523,7 @@ def test_forward_handles_thd_format(self, vl_config, backend_config, moe_config, with ( patch( "nemo_automodel.components.models.qwen3_vl_moe.model.squeeze_input_for_thd", - return_value=(squeezed_ids, squeezed_position_ids, squeezed_padding_mask, squeezed_kwargs), + side_effect=fake_squeeze, ) as mock_squeeze, patch.object(model.model, "forward") as mock_model_forward, ): @@ -531,6 +537,8 @@ def test_forward_handles_thd_format(self, vl_config, backend_config, moe_config, attention_mask=attention_mask, padding_mask=padding_mask, qkv_format="thd", + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, ) # Result should be logits from lm_head @@ -542,9 +550,15 @@ def test_forward_handles_thd_format(self, vl_config, backend_config, moe_config, assert squeeze_args[1] is position_ids assert squeeze_args[2] is padding_mask assert squeeze_args[3]["qkv_format"] == "thd" + # Media tensors flow through squeeze_input_for_thd, which skips them. + assert squeeze_input_kwargs["pixel_values"] is pixel_values + assert squeeze_input_kwargs["image_grid_thw"] is image_grid_thw # Verify model.model.forward was called mock_model_forward.assert_called_once() + model_kwargs = mock_model_forward.call_args.kwargs + assert model_kwargs["pixel_values"] is pixel_values + assert model_kwargs["image_grid_thw"] is image_grid_thw def test_initialize_weights_invokes_language_model(self, vl_config, backend_config, moe_config): model = Qwen3VLMoeForConditionalGeneration(vl_config, backend=backend_config, moe_config=moe_config) diff --git a/tests/unit_tests/models/test_passthrough_state_dict_adapters.py b/tests/unit_tests/models/test_passthrough_state_dict_adapters.py new file mode 100644 index 0000000000..354efea405 --- /dev/null +++ b/tests/unit_tests/models/test_passthrough_state_dict_adapters.py @@ -0,0 +1,47 @@ +# 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 +import torch +from transformers import LlamaConfig, Qwen2Config, Qwen3Config + +from nemo_automodel.components.models.llama.state_dict_adapter import LlamaStateDictAdapter +from nemo_automodel.components.models.qwen2.state_dict_adapter import Qwen2StateDictAdapter +from nemo_automodel.components.models.qwen3.state_dict_adapter import Qwen3StateDictAdapter + + +@pytest.mark.parametrize( + "adapter", + [ + LlamaStateDictAdapter(LlamaConfig()), + Qwen2StateDictAdapter(Qwen2Config()), + Qwen3StateDictAdapter(Qwen3Config()), + ], +) +def test_passthrough_adapter_streams_one_hf_tensor(adapter): + tensor = torch.randn(2, 3) + + converted = adapter.convert_single_tensor_to_hf("model.weight", tensor) + + assert len(converted) == 1 + assert converted[0][0] == "model.weight" + assert converted[0][1] is tensor + assert ( + adapter.convert_single_tensor_to_hf( + "model.weight", + tensor, + exclude_key_regex=r"model\..*", + ) + == [] + ) diff --git a/tests/unit_tests/moe/test_router_replay.py b/tests/unit_tests/moe/test_router_replay.py index 27fc1c4f8a..f41cc389c0 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,235 @@ def test_multilayer_distribute_and_collect(): assert torch.equal(rep, rec) +# --------------------------------------------------------------------------- # +# Model-scoped adapter: global mapping and scoped lifecycle +# --------------------------------------------------------------------------- # + + +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, + ) + prepared = torch.full((batch, sequence, num_layers, topk), -1, dtype=torch.int16) + prepared[:, :, 1] = layer_one + prepared[:, :, 3] = layer_three + + 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.replay(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_replays_real_gate_and_preserves_router_gradient(): + 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 one selected-router weight for a padded token batch. + + Args: + input_ids: Token ids with shape ``[batch, sequence]``. + + Returns: + First selected probability 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[..., 0].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)) + with adapter.replay(target): + output = model(input_ids) + output.sum().backward() + + 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.replay(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.replay(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, 4, dtype=torch.int16), ValueError, "token axes"), + (torch.zeros(1, 3, 2), TypeError, "signed integer dtype"), + (torch.zeros(1, 3, 2, dtype=torch.uint8), TypeError, "signed integer dtype"), + ([], TypeError, "Tensor or None"), + ], +) +def test_adapter_replay_rejects_invalid_routes(routes, error, match): + adapter = RouterReplayAdapter(_AdapterModel()) + with pytest.raises(error, match=match): + adapter.replay(routes) + + +@pytest.mark.parametrize( + ("routes", "match"), + [ + (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.replay(routes) + + # --------------------------------------------------------------------------- # # Error handling + context-manager cleanup # --------------------------------------------------------------------------- # 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/engine_stub.py b/tests/unit_tests/recipes/engine_stub.py new file mode 100644 index 0000000000..02cac83abd --- /dev/null +++ b/tests/unit_tests/recipes/engine_stub.py @@ -0,0 +1,41 @@ +# 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. + + +class RecipeEngineStub: + """Minimal Engine stand-in for recipe unit tests, matching Engine's signatures.""" + + def __init__(self, module, *, optimizer=None, grad_norm=0.0): + self.module = module + self.optimizer = optimizer + self.grad_norm = grad_norm + self.gradient_accumulation_steps = None + + def __call__(self, *args, **kwargs): + return self.module(*args, **kwargs) + + def backward(self, loss, retain_graph=False, scale_wrt_gas=True): + del scale_wrt_gas + loss.backward(retain_graph=retain_graph) + + def step(self): + if self.optimizer is not None: + self.optimizer.step() + self.optimizer.zero_grad(set_to_none=True) + + def get_global_grad_norm(self): + return self.grad_norm + + def set_gradient_accumulation_steps(self, steps): + self.gradient_accumulation_steps = steps diff --git a/tests/unit_tests/recipes/llm/test_benchmark.py b/tests/unit_tests/recipes/llm/test_benchmark.py index 81fae306fb..3bd3af39f8 100644 --- a/tests/unit_tests/recipes/llm/test_benchmark.py +++ b/tests/unit_tests/recipes/llm/test_benchmark.py @@ -134,6 +134,24 @@ def mock_recipe(mock_config, monkeypatch): intermediate_size=3072, ) recipe.optimizer = [MagicMock()] + engine_state = {"micro_step": 0} + + def is_gradient_accumulation_boundary(): + return engine_state["micro_step"] + 1 == 8 + + def engine_step(): + engine_state["micro_step"] += 1 + if engine_state["micro_step"] == 8: + recipe.optimizer[0].step() + recipe.optimizer[0].zero_grad(set_to_none=True) + engine_state["micro_step"] = 0 + + recipe.engine = SimpleNamespace( + is_gradient_accumulation_boundary=MagicMock(side_effect=is_gradient_accumulation_boundary), + step=MagicMock(side_effect=engine_step), + set_gradient_accumulation_steps=MagicMock(), + ) + recipe.checkpointer = SimpleNamespace(maybe_wait_for_staging=MagicMock()) recipe.dataloader = MagicMock() recipe.val_dataloader = None recipe.pp_enabled = False @@ -157,6 +175,26 @@ def mock_recipe(mock_config, monkeypatch): return recipe +def _configure_one_iteration(recipe, batches): + """Trim the benchmark fixture to one two-microbatch optimizer window.""" + recipe.cfg.step_scheduler.local_batch_size = 2 + recipe.cfg.step_scheduler.global_batch_size = 4 + recipe._get_dp_group_size = MagicMock(return_value=1) + recipe._bench_steps = 1 + recipe._bench_warmup_steps = 0 + recipe.dataloader.__iter__ = MagicMock(return_value=iter(batches)) + recipe.timers._get_global_min_max_time = MagicMock( + return_value={"iteration_warmup": (0.0, 1.0), "iteration": (0.0, 1.0)} + ) + timer = MagicMock() + timer.active_time.return_value = 1.0 + recipe.timers._timers = { + "setup": timer, + "iteration": timer, + "iteration_warmup": timer, + } + + @pytest.mark.usefixtures("patch_torch_distributed_for_benchmark") class TestBenchmarkingRecipeInitialization: """Test initialization of BenchmarkingRecipeForNextTokenPrediction.""" @@ -422,6 +460,72 @@ def mock_forward_backward_step(ga_step_idx, batch, loss_buffer=None, **kwargs): 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 + assert ( + mock_recipe.engine.set_gradient_accumulation_steps.call_args_list == [((expected_ga_steps,), {})] * 30 + ) + + def test_run_benchmark_normalizes_each_microbatch_by_global_window_tokens(self, mock_recipe, capsys): + 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, -100, -100]])}, + ] + _configure_one_iteration(mock_recipe, batches) + denominators = [] + numerators = [3.0, 6.0] + + def forward_backward_step(index, batch, *, loss_buffer, num_label_tokens, **kwargs): + del batch, kwargs + denominators.append(num_label_tokens) + loss_buffer.append(torch.tensor(numerators[index] / num_label_tokens)) + + mock_recipe._forward_backward_step = MagicMock(side_effect=forward_backward_step) + + with patch("torch.distributed.barrier"): + mock_recipe.run_benchmark() + + assert denominators == [3, 3] + mock_recipe.engine.set_gradient_accumulation_steps.assert_called_once_with(2) + assert "num_label_tokens=3 | loss=3.0000" in capsys.readouterr().out + + def test_run_benchmark_pipeline_uses_training_optimizer_lifecycle(self, mock_recipe, monkeypatch, capsys): + 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, -100, -100]])}, + ] + _configure_one_iteration(mock_recipe, batches) + mock_recipe.pp_enabled = True + mock_recipe.pp = SimpleNamespace(pp_batch_size=2, pp_microbatch_size=1) + mock_recipe.max_grad_norm = 0.75 + mock_recipe._broadcast_from_last_pp_stage = MagicMock(side_effect=lambda tensor: tensor) + mock_recipe._step_pipeline_optimizer = MagicMock(return_value=torch.tensor(1.0)) + prepare = MagicMock() + prepare_final = MagicMock() + prepare_after_first = MagicMock() + monkeypatch.setattr("nemo_automodel.recipes.llm.benchmark.prepare_for_grad_accumulation", prepare) + monkeypatch.setattr("nemo_automodel.recipes.llm.benchmark.prepare_for_final_backward", prepare_final) + monkeypatch.setattr("nemo_automodel.recipes.llm.benchmark.prepare_after_first_microbatch", prepare_after_first) + denominators = [] + + def forward_backward_step(index, batch, *, loss_buffer, num_label_tokens, **kwargs): + del batch, kwargs + denominators.append(num_label_tokens) + loss_buffer.append(torch.tensor([3.0, 6.0][index])) + + mock_recipe._forward_backward_step = MagicMock(side_effect=forward_backward_step) + + with patch("torch.distributed.barrier"): + mock_recipe.run_benchmark() + + assert denominators == [3, 3] + prepare.assert_called_once_with(mock_recipe.model_parts, pp_enabled=True) + prepare_final.assert_called_once_with(mock_recipe.model_parts, pp_enabled=True) + prepare_after_first.assert_called_once_with() + mock_recipe._step_pipeline_optimizer.assert_called_once_with( + num_label_tokens=3, + max_grad_norm=0.75, + ) + mock_recipe._broadcast_from_last_pp_stage.assert_called_once() + assert "num_label_tokens=3 | loss=3.0000" in capsys.readouterr().out def test_run_benchmark_zero_grads_per_iteration(self, mock_recipe): """Test that gradients are zeroed at the start of each iteration.""" @@ -555,7 +659,7 @@ class TestBenchmarkingRecipeHelpers: """Test helper methods and edge cases.""" def test_benchmark_with_invalid_ga_config(self, mock_recipe): - """Test that invalid gradient accumulation config raises assertion.""" + """Test that invalid gradient accumulation config raises a clear error.""" mock_recipe._get_dp_group_size = MagicMock(return_value=8) # Set invalid batch sizes that will cause assertion error # global_batch_size=16, local_batch_size=4, dp_size=8 @@ -563,7 +667,7 @@ def test_benchmark_with_invalid_ga_config(self, mock_recipe): mock_recipe.cfg.step_scheduler.local_batch_size = 4 mock_recipe.cfg.step_scheduler.global_batch_size = 16 - with pytest.raises(AssertionError, match="Global batch size must be divisible"): + with pytest.raises(ValueError, match="global_batch_size.*positive multiple"): mock_recipe.run_benchmark() def test_init_requires_benchmark_section(self): 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..da5b8acee6 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -12,26 +12,18 @@ # 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``. + +The VLM recipe owns context-parallel sharding and executes pipeline schedules +directly. These tests cover pipeline media staging setup, vision-frame context +publication, and the validation handoff plus epoch-level DP aggregation. """ from __future__ import annotations from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -39,9 +31,23 @@ 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 +class _UnsupportedVisionModel: + supports_cp_vision_frame_sharding = False + + +class _SupportedVisionModel: + supports_cp_vision_frame_sharding = True + + +class _PackedCPModel: + def __init__(self, *, supported): + self.supports_cp_with_sequence_packing = supported + + def _identity_cp_shard(sharder, batch): """Bypass CP transport while preserving constructor-side strategy resolution. @@ -110,19 +116,6 @@ def __getitem__(self, key): return SimpleNamespace(size=lambda: 2, get_group=lambda: "cp-group") -class _UnsupportedVisionModel: - supports_cp_vision_frame_sharding = False - - -class _SupportedVisionModel: - supports_cp_vision_frame_sharding = True - - -class _PackedCPModel: - def __init__(self, *, supported): - self.supports_cp_with_sequence_packing = supported - - def test_cp_vision_frame_sharding_rejects_model_without_capability(): policy = CpVisionFrameShardingConfig(enabled=True) @@ -240,82 +233,6 @@ def step(self, model_input, *, target=None, losses=None, **kwargs): 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).""" @@ -410,10 +327,11 @@ 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: +class _StageWithCPPreembedInForward(torch.nn.Module): # Sunk VLM (minimax/qwen3_5/qwen3_5_moe/step3p7): embeds + shards per # microbatch inside forward and pulls media from the PP side channel, so # media MUST still be staged for PP under CP. @@ -421,7 +339,7 @@ def prepare_model_inputs_for_cp(self): return {} -class _StageWithoutCPPrepare: +class _StageWithoutCPPrepare(torch.nn.Module): pass @@ -437,7 +355,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 +370,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,74 +492,14 @@ 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 - - -# ----------------------------------------------------------------------------- -# val-side wiring (the bug-fix territory) -# ----------------------------------------------------------------------------- - - -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) + assert trainer.engine is None + assert isinstance(trainer.pp, _FakePPModel) 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. monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext()) monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _identity_cp_shard) @@ -688,54 +546,3 @@ def _fake_allreduce(tensor, include_cp=False): 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 *, out, base_loss, labels, model, num_label_tokens: 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 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 3f4130827b..0ad382b4b8 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -45,6 +45,7 @@ _get_model_name, build_model, ) +from tests.unit_tests.recipes.engine_stub import RecipeEngineStub as _RecipeEngineStub def build_optimizer(model, cfg_opt, distributed_config, device_mesh): @@ -379,6 +380,7 @@ class _TensorModel(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.nn.Parameter(torch.zeros(1)) + self.supports = SimpleNamespace(mtp_enabled=False) def forward(self, **batch): return torch.zeros((), requires_grad=True) @@ -395,6 +397,7 @@ def test_run_train_step_supports_tensor_outputs(monkeypatch): recipe.model_parts = [model] # Now uses model_parts instead of model recipe.pp_enabled = False # Pipeline parallelism disabled recipe.optimizer = [_DummyOptimizer()] # Now a list + recipe.engine = _RecipeEngineStub(model, optimizer=recipe.optimizer[0], grad_norm=2.5) # ``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. @@ -426,39 +429,18 @@ def fake_calculate_loss(*args, **kwargs): "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, - ) - 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) assert recipe.optimizer[0].step_called assert recipe.optimizer[0].zero_grad_called + assert recipe.engine.gradient_accumulation_steps == len(batches) @pytest.mark.cuda(False) @@ -469,6 +451,7 @@ def test_forward_backward_step_routes_thd_batch_through_te(monkeypatch): recipe.mesh_context = SimpleNamespace(cp_size=2) recipe.processor = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=7)) recipe.model_parts = [_TensorModel()] + recipe.engine = _RecipeEngineStub(recipe.model_parts[0]) recipe.pp_enabled = False recipe.magi = SimpleNamespace(enabled=False) recipe.distributed_config = None @@ -483,7 +466,6 @@ def make_thd_batch(model, device_mesh, batch, **kwargs): 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), @@ -1548,31 +1530,15 @@ class _MockAutoPipeline: 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 step(self, model_input, **kwargs): + args = (model_input,) if self.info.has_first_stage else () + return self.info.schedule.step(*args, **kwargs) def _create_pp_recipe(model=None): @@ -1695,7 +1661,9 @@ def step_side_effect(*args, **kwargs): # Verify schedule.step was called pp_recipe.pp.info.schedule.step.assert_called_once() - assert pp_recipe.pp.step_batches == [{}] + schedule_call = pp_recipe.pp.info.schedule.step.call_args + assert len(schedule_call.args) == 1 + assert set(schedule_call.kwargs) == {"target", "losses"} # Verify loss was computed assert len(loss_buffer) == 1 @@ -1725,9 +1693,10 @@ def test_pp_step_receives_remaining_kwargs(self, pp_recipe, monkeypatch): 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) + schedule_call = pp_recipe.pp.info.schedule.step.call_args + assert len(schedule_call.args) == 1 + assert set(schedule_call.kwargs) == {"target", "losses", "position_ids"} + assert torch.equal(schedule_call.kwargs["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.""" @@ -2267,6 +2236,7 @@ def _create_non_pp_recipe(model, device="cpu"): recipe.__dict__["distributed_config"] = None recipe.__dict__["cp_vision_frame_sharding"] = CpVisionFrameShardingConfig(enabled=True) recipe.__dict__["model_parts"] = [model] + recipe.__dict__["engine"] = _RecipeEngineStub(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 @@ -2412,10 +2382,6 @@ def get_output_embeddings(self): "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)), @@ -2459,10 +2425,6 @@ def forward(self, logits_to_keep=None, **kwargs): "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)), @@ -2502,10 +2464,6 @@ def forward(self, **kwargs): "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)), @@ -2862,7 +2820,7 @@ def _stub_build_checkpoint_config(*a, **k): monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ScopedRNG", lambda **kwargs: nullcontext()) monkeypatch.setattr( "nemo_automodel.components.training.step_scheduler.StepSchedulerConfig.build", - lambda self, *a, **k: SimpleNamespace(step=0, epoch=0, epochs=[]), + lambda self, *a, **k: SimpleNamespace(step=0, epoch=0, epochs=[], grad_acc_steps=1), ) monkeypatch.setattr("nemo_automodel.components.optim.optimizer.LRSchedulerConfig.build", lambda self, *a, **k: []) monkeypatch.setattr( @@ -2923,6 +2881,96 @@ def _minimal_vlm_cfg( return ConfigNode(cfg) +def _patch_vlm_distributed_setup( + monkeypatch, + *, + pp_enabled: bool, + 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=None, + pipeline_config=pipeline_config, + moe_parallel_config=None, + activation_checkpointing=False, + ), + ) + + +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.step_scheduler.global_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 + + class DummyAutoPipeline(SimpleNamespace): + pass + + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.AutoPipeline", DummyAutoPipeline) + model = DummyModel() + pipeline = DummyAutoPipeline( + parts=[model], + pp_batch_size=1, + pp_microbatch_size=1, + scale_grads_in_schedule=False, + info=SimpleNamespace( + has_first_stage=True, + has_last_stage=False, + stages=[SimpleNamespace(is_first=True, is_last=False)], + schedule=MagicMock(), + ), + ) + 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 is None + + 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}) diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py index 7bbfadf0e4..c5a6450566 100644 --- a/tests/unit_tests/recipes/test_train_ft.py +++ b/tests/unit_tests/recipes/test_train_ft.py @@ -48,6 +48,7 @@ build_model, compute_trust_remote_code_from_model, ) +from tests.unit_tests.recipes.engine_stub import RecipeEngineStub as _RecipeEngineStub def test_recipe_config_resolves_mfu_settings(): @@ -979,7 +980,7 @@ def _stub_build_checkpoint_config(*a, **k): monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.ScopedRNG", lambda **kwargs: nullcontext()) monkeypatch.setattr( "nemo_automodel.components.training.step_scheduler.StepSchedulerConfig.build", - lambda self, *a, **k: SimpleNamespace(step=0, epoch=0, epochs=[]), + lambda self, *a, **k: SimpleNamespace(step=0, epoch=0, epochs=[], grad_acc_steps=1), ) monkeypatch.setattr( "nemo_automodel.components.optim.optimizer.LRSchedulerConfig.build", @@ -1091,6 +1092,136 @@ def patch_fn(model, name=None, add_backward_hooks=True): assert patch_calls == [] +@pytest.mark.parametrize( + ( + "dataloader_emits_thd", + "pipeline_thd_kind", + "fused_loss", + "scale_grads_in_schedule", + "magi_enabled", + "local_batch_size", + "error_match", + ), + [ + (False, None, False, False, False, 2, None), + (True, "native", False, False, False, 2, None), + (True, "stock_hf", False, False, False, 2, "do not consume packed document boundaries"), + (False, None, True, False, False, 2, None), + (False, None, False, True, False, 2, "scale_grads_in_schedule=False"), + (False, None, False, False, True, 2, "Magi pipeline training requires"), + (False, None, False, False, True, 1, None), + ], +) +def test_setup_pipeline_matrix( + monkeypatch, + dataloader_emits_thd, + pipeline_thd_kind, + fused_loss, + scale_grads_in_schedule, + magi_enabled, + local_batch_size, + error_match, +): + """Pipeline setup keeps one schedule path and preserves its layout guards.""" + from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy + + cfg = _minimal_cfg_with_nvtx(nvtx_value=False) + cfg.step_scheduler.local_batch_size = local_batch_size + cfg.step_scheduler.global_batch_size = local_batch_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( + 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) + + 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=build_dataloader, + 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) + parts = [DummyModel()] + parts[0]._pp_return_hidden_states_supported = True + if pipeline_thd_kind == "native": + parts[0].supports_thd = True + pipeline = DummyAutoPipeline( + parts=parts, + pp_batch_size=local_batch_size, + 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=1, + pp_size=2, + ), + strategy_config=None, + pipeline_config=SimpleNamespace( + pp_seq_len=None, + scale_grads_in_schedule=scale_grads_in_schedule, + ), + moe_parallel_config=None, + activation_checkpointing=False, + ), + ) + + trainer = TrainFinetuneRecipeForNextTokenPrediction(cfg) + if error_match is not None: + with pytest.raises(ValueError, match=error_match): + trainer.setup() + return + + trainer.setup() + + if fused_loss_fn is not None: + assert trainer.loss_fn is fused_loss_fn + assert trainer.engine is None + assert trainer.pp is pipeline + assert dataloader_build_kwargs[0]["collate_wrapper"] is (None if dataloader_emits_thd else pp_collate_wrapper) + + 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") @@ -1335,16 +1466,19 @@ def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info): 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, - ), + pp = SimpleNamespace( + info=pp_info, + pp_batch_size=1, + pp_microbatch_size=1, + update_seq_len=lambda seq_len: None, + ) + pp.step = lambda model_input, **kwargs: pp_info.schedule.step( + *((model_input,) if pp_info.has_first_stage else ()), **kwargs ) + pp.eval = lambda model_input, **kwargs: pp_info.schedule.eval( + *((model_input,) if pp_info.has_first_stage else ()), **kwargs + ) + object.__setattr__(recipe, "pp", pp) object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) object.__setattr__(recipe, "te_fp8", None) @@ -2090,7 +2224,7 @@ def _make_recipe( object.__setattr__( recipe, "optimizer", - [SimpleNamespace(step=lambda: None, zero_grad=lambda: None, param_groups=[{"lr": 0.01}])], + [SimpleNamespace(step=lambda: None, zero_grad=lambda **_kwargs: None, param_groups=[{"lr": 0.01}])], ) object.__setattr__(recipe, "lr_schedulers", []) object.__setattr__(recipe, "step_scheduler", SimpleNamespace(step=1, epoch=0)) @@ -2114,6 +2248,16 @@ def _make_recipe( # 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) + else: + object.__setattr__( + recipe, + "engine", + SimpleNamespace( + set_gradient_accumulation_steps=MagicMock(), + step=MagicMock(), + get_global_grad_norm=MagicMock(return_value=torch.tensor(1.0)), + ), + ) object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) monkeypatch.setattr( @@ -2163,37 +2307,85 @@ 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) - @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) - + def test_pp_zero_label_window_reports_finite_zero(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=True) batches = [ - {"input_ids": torch.tensor([[1, 2, 3, 4]]), "labels": torch.tensor([[1, 2, 3, -100]])} for _ in range(4) + {"input_ids": torch.tensor([[1, 2]]), "labels": torch.tensor([[-100, -100]])}, + {"input_ids": torch.tensor([[3, 4]]), "labels": torch.tensor([[-100, -100]])}, ] - 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 + metrics = recipe._run_train_optim_step(batches) - recipe = self._make_recipe( - monkeypatch, - pp_enabled=False, - dp_group_size=8, - cp_group_size=2, + assert metrics.metrics["num_label_tokens"] == 0 + assert metrics.metrics["loss"] == 0.0 + + def test_pipeline_optimizer_helper_owns_complete_update_lifecycle(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=True) + + class GatePart(nn.Linear): + def __init__(self): + super().__init__(2, 2) + self.update_moe_gate_bias = MagicMock() + + model_parts = [GatePart(), GatePart()] + optimizers = [MagicMock(), MagicMock()] + schedulers = [MagicMock(), MagicMock()] + checkpointer = SimpleNamespace(maybe_wait_for_staging=MagicMock()) + dp_shard_mesh = SimpleNamespace(size=lambda: 2) + device_mesh = MagicMock() + device_mesh.__getitem__.side_effect = lambda name: dp_shard_mesh if name == "dp_shard" else MagicMock() + scale_and_clip = MagicMock(return_value=torch.tensor(2.5)) + precompute_fp8 = MagicMock() + + object.__setattr__(recipe, "model_parts", model_parts) + object.__setattr__(recipe, "optimizer", optimizers) + object.__setattr__(recipe, "lr_scheduler", schedulers) + object.__setattr__(recipe, "checkpointer", checkpointer) + object.__setattr__(recipe, "device_mesh", device_mesh) + object.__setattr__( + recipe, + "cfg", + SimpleNamespace( + get=lambda key, default=None: { + "fp8": { + "enabled": True, + "precompute_float8_dynamic_scale_for_fsdp": True, + } + }.get(key, default) + ), + ) + monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.scale_grads_and_clip_grad_norm", scale_and_clip) + monkeypatch.setattr( + "nemo_automodel.recipes.llm.train_ft.precompute_float8_dynamic_scale_for_fsdp", + precompute_fp8, ) + + grad_norm = recipe._step_pipeline_optimizer(num_label_tokens=17, max_grad_norm=0.75) + + assert grad_norm.item() == pytest.approx(2.5) + assert scale_and_clip.call_args.kwargs["num_label_tokens"] == 17 + assert scale_and_clip.call_args.kwargs["max_grad_norm"] == 0.75 + checkpointer.maybe_wait_for_staging.assert_called_once_with() + for optimizer in optimizers: + optimizer.step.assert_called_once_with() + optimizer.zero_grad.assert_called_once_with(set_to_none=True) + for scheduler in schedulers: + scheduler.step.assert_called_once_with(1) + for model_part in model_parts: + model_part.update_moe_gate_bias.assert_called_once_with() + precompute_fp8.assert_called_once_with(model_parts[0]) + + def test_eager_short_window_reconfigures_engine(self, monkeypatch): + recipe = self._make_recipe(monkeypatch, pp_enabled=False) batches = [ - {"input_ids": torch.tensor([[1, 2, 3, 4]]), "labels": torch.tensor([[1, 2, 3, -100]])} for _ in range(4) + {"input_ids": torch.tensor([[1, 2]]), "labels": torch.tensor([[1, -100]])}, + {"input_ids": torch.tensor([[3, 4]]), "labels": torch.tensor([[3, -100]])}, ] recipe._run_train_optim_step(batches) - assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(0.5) + recipe.engine.set_gradient_accumulation_steps.assert_called_once_with(2) + assert recipe.engine.step.call_count == 2 def test_tps_per_gpu_divides_global_tps_by_world_size(self, monkeypatch): """Under pp>1 the per-GPU divisor must be the full world size, not dp*cp. @@ -2594,6 +2786,7 @@ def get_local_rank(self): object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) object.__setattr__(recipe, "te_fp8", None) object.__setattr__(recipe, "model_parts", [model]) + object.__setattr__(recipe, "engine", _RecipeEngineStub(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_size", lambda include_cp=False: 1) @@ -2611,7 +2804,6 @@ def _fake_calc_loss(loss_fn, *, logits, labels, model, hidden_states, lm_weight, ) 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)} @@ -2712,6 +2904,7 @@ def shard_token_tensor(self, tensor, seq_dim=1, fill=None): object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0)) object.__setattr__(recipe, "te_fp8", None) object.__setattr__(recipe, "model_parts", [model]) + object.__setattr__(recipe, "engine", _RecipeEngineStub(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) @@ -2731,8 +2924,6 @@ def _fake_calculate_mtp_loss(loss_fn, *, mtp_per_depth_h, mtp_per_depth_targets, 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) - monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.get_sync_ctx", lambda *args, **kwargs: nullcontext()) - batch = { "input_ids": torch.tensor([[10, 11, 12, 20, 21, 22]]), "labels": torch.tensor([[11, 12, -100, 21, 22, -100]]), 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..008d02eb0a 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, max_grad_norm=None: ( 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 new file mode 100644 index 0000000000..a4b7d4a573 --- /dev/null +++ b/tests/unit_tests/test_engine.py @@ -0,0 +1,384 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +import nemo_automodel.engine._engine as engine_module +from nemo_automodel import Engine as PublicEngine +from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler +from nemo_automodel.engine import Engine + + +class _Scale(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.tensor(1.0)) + + def forward(self, values: torch.Tensor, *, offset: float = 0.0) -> torch.Tensor: + """Scale a tensor and add a scalar offset. + + Args: + values: Input values with arbitrary shape. + offset: Scalar added to every output value. + + Returns: + Tensor with the same shape as ``values``. + """ + return values * self.weight + offset + + +class _CountingSGD(torch.optim.SGD): + def __init__(self, parameters) -> None: + super().__init__(parameters, lr=0.1) + self.step_calls = 0 + self.zero_grad_calls = 0 + + def step(self, closure=None): + self.step_calls += 1 + return super().step(closure) + + def zero_grad(self, *args, **kwargs): + self.zero_grad_calls += 1 + return super().zero_grad(*args, **kwargs) + + +class _CountingScheduler: + def __init__(self) -> None: + self.calls = 0 + + def step(self) -> None: + self.calls += 1 + + +def _skip_gradient_finalization(monkeypatch, norm: float = 2.0) -> None: + monkeypatch.setattr( + engine_module, + "scale_grads_and_clip_grad_norm", + lambda **_kwargs: torch.tensor(norm), + ) + + +def test_engine_is_a_public_module_and_raw_forward_delegate() -> None: + module = _Scale() + engine = Engine(module) + + output = engine(torch.tensor([1.0, 2.0]), offset=3.0) + + assert PublicEngine is Engine + assert isinstance(engine, nn.Module) + assert engine.module is module + torch.testing.assert_close(output, torch.tensor([4.0, 5.0])) + + +def test_forward_only_engine_rejects_training_operations() -> None: + engine = Engine(_Scale()) + + with pytest.raises(RuntimeError, match="backward requires an optimizer"): + engine.backward(torch.tensor(1.0, requires_grad=True)) + with pytest.raises(RuntimeError, match="step requires an optimizer"): + engine.step() + + +def test_engine_rejects_pipeline_contract(monkeypatch) -> None: + class _Pipeline(nn.Module): + pass + + monkeypatch.setattr(engine_module, "AutoPipeline", _Pipeline) + + with pytest.raises(NotImplementedError, match="pipeline schedule"): + Engine(_Pipeline()) + + with pytest.raises(NotImplementedError, match="pipeline stages"): + Engine(nn.Linear(2, 2), mesh_context=SimpleNamespace(pp_enabled=True)) + + +def test_gradient_accumulation_updates_only_at_boundary(monkeypatch) -> None: + _skip_gradient_finalization(monkeypatch) + module = _Scale() + optimizer = _CountingSGD(module.parameters()) + scheduler = _CountingScheduler() + engine = Engine( + module, + optimizer=optimizer, + lr_scheduler=scheduler, + gradient_accumulation_steps=2, + ) + + engine.backward(engine(torch.ones(())) ** 2) + assert not engine.is_gradient_accumulation_boundary() + engine.step() + + assert module.weight.item() == 1.0 + assert optimizer.step_calls == 0 + assert optimizer.zero_grad_calls == 0 + assert scheduler.calls == 0 + + engine.backward(engine(torch.ones(())) ** 2) + assert engine.is_gradient_accumulation_boundary() + engine.step() + + torch.testing.assert_close(module.weight, torch.tensor(0.8)) + assert optimizer.step_calls == 1 + assert optimizer.zero_grad_calls == 1 + assert scheduler.calls == 1 + torch.testing.assert_close(engine.get_global_grad_norm(), torch.tensor(2.0)) + + +def test_short_accumulation_window_updates_after_runtime_reconfiguration(monkeypatch) -> None: + _skip_gradient_finalization(monkeypatch) + module = _Scale() + optimizer = _CountingSGD(module.parameters()) + engine = Engine(module, optimizer=optimizer, gradient_accumulation_steps=4) + engine.set_gradient_accumulation_steps(2) + + for _ in range(2): + engine.backward(engine(torch.ones(())) ** 2) + engine.step() + + assert optimizer.step_calls == 1 + torch.testing.assert_close(module.weight, torch.tensor(0.8)) + + +def test_backward_can_skip_main_loss_gas_scaling() -> None: + module = _Scale() + engine = Engine( + module, + optimizer=torch.optim.SGD(module.parameters(), lr=0.1), + gradient_accumulation_steps=4, + ) + + engine.backward(engine(torch.tensor(3.0)), scale_wrt_gas=False) + + torch.testing.assert_close(module.weight.grad, torch.tensor(3.0)) + torch.testing.assert_close(MoEAuxLossAutoScaler.main_loss_backward_scale, torch.tensor(0.25)) + + +@pytest.mark.parametrize( + ("declared_mode", "expected_gradient", "expected_aux_scale"), + [ + pytest.param(None, 2.0, 2.0 / 3.0, id="undeclared-averaged"), + pytest.param(False, 2.0, 2.0 / 3.0, id="explicit-averaged"), + pytest.param(True, 0.25, 1.0 / 12.0, id="summed"), + ], +) +def test_backward_compensates_backend_gradient_reduction( + monkeypatch, + declared_mode, + expected_gradient, + expected_aux_scale, +) -> None: + module = _Scale() + if declared_mode is not None: + module.calculate_per_token_loss = declared_mode + engine = Engine( + module, + optimizer=torch.optim.SGD(module.parameters(), lr=0.1), + gradient_accumulation_steps=3, + ) + monkeypatch.setattr(engine, "_gradient_group_size", lambda: 8) + monkeypatch.setattr(engine, "_context_parallel_size", lambda: 2) + + engine.backward(engine(torch.tensor(2.0)), scale_wrt_gas=False) + + torch.testing.assert_close(module.weight.grad, torch.tensor(expected_gradient)) + assert MoEAuxLossAutoScaler.main_loss_backward_scale.item() == pytest.approx(expected_aux_scale) + + +def test_summed_gradient_declaration_is_found_below_wrapper() -> None: + module = _Scale() + module.distributed_backend = nn.Identity() + module.distributed_backend.calculate_per_token_loss = True + + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1)) + + assert engine._summed_gradient_reduction + + +def test_forward_keeps_sync_context_open_through_backward(monkeypatch) -> None: + events: list[str] = [] + boundaries: list[tuple[bool, bool]] = [] + + @contextmanager + def sync_context(_module, boundary, defer_fsdp_grad_sync): + boundaries.append((boundary, defer_fsdp_grad_sync)) + events.append("enter") + try: + yield + finally: + events.append("exit") + + class _ObservedScale(_Scale): + def forward(self, values: torch.Tensor, *, offset: float = 0.0) -> torch.Tensor: + """Record forward before delegating to ``_Scale``.""" + events.append("forward") + return super().forward(values, offset=offset) + + module = _ObservedScale() + module.weight.register_hook(lambda gradient: events.append("backward") or gradient) + optimizer = torch.optim.SGD(module.parameters(), lr=0.1) + engine = Engine(module, optimizer=optimizer, gradient_accumulation_steps=2) + monkeypatch.setattr(engine_module, "get_sync_ctx", sync_context) + 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("first_done")) + _skip_gradient_finalization(monkeypatch) + + engine.backward(engine(torch.ones(()))) + engine.step() + engine.backward(engine(torch.ones(()))) + engine.step() + + assert boundaries == [(False, True), (True, True)] + assert events == [ + "prepare", + "enter", + "forward", + "backward", + "exit", + "first_done", + "final", + "enter", + "forward", + "backward", + "exit", + ] + + +@pytest.mark.parametrize("disable_training", ["no_grad", "eval"]) +def test_non_training_forward_does_not_open_training_context(monkeypatch, disable_training) -> None: + module = _Scale().eval() if disable_training == "eval" else _Scale() + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1)) + monkeypatch.setattr( + engine_module, + "get_sync_ctx", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("unexpected training context")), + ) + + if disable_training == "no_grad": + with torch.no_grad(): + output = engine(torch.tensor(2.0)) + else: + output = engine(torch.tensor(2.0)) + + torch.testing.assert_close(output, torch.tensor(2.0)) + + +def test_forward_failure_closes_sync_context(monkeypatch) -> None: + exits = 0 + + @contextmanager + def sync_context(*_args, **_kwargs): + nonlocal exits + try: + yield + finally: + exits += 1 + + class _FailOnce(_Scale): + def __init__(self) -> None: + super().__init__() + self.fail = True + + def forward(self, values: torch.Tensor, *, offset: float = 0.0) -> torch.Tensor: + """Fail once, then return a tensor shaped like ``values``.""" + if self.fail: + self.fail = False + raise RuntimeError("forward failed") + return super().forward(values, offset=offset) + + module = _FailOnce() + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1)) + monkeypatch.setattr(engine_module, "get_sync_ctx", sync_context) + + with pytest.raises(RuntimeError, match="forward failed"): + engine(torch.ones(())) + engine.backward(engine(torch.ones(()))) + + assert exits == 2 + + +def test_accumulation_steps_validate_and_cannot_change_mid_window() -> None: + module = _Scale() + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1), gradient_accumulation_steps=2) + + with pytest.raises(ValueError, match="positive integer"): + engine.set_gradient_accumulation_steps(0) + engine.backward(engine(torch.ones(()))) + engine.step() + with pytest.raises(RuntimeError, match="active window"): + engine.set_gradient_accumulation_steps(3) + + +def test_optimizer_boundary_runs_model_post_step_hooks(monkeypatch) -> None: + _skip_gradient_finalization(monkeypatch) + precomputed: list[nn.Module] = [] + + class _PostStepModel(_Scale): + precompute_float8_dynamic_scale_for_fsdp = True + + def __init__(self) -> None: + super().__init__() + self.gate_updates = 0 + + def update_moe_gate_bias(self) -> None: + self.gate_updates += 1 + + monkeypatch.setattr( + engine_module, + "safe_import_from", + lambda *_args, **_kwargs: (True, lambda module: precomputed.append(module)), + ) + module = _PostStepModel() + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1)) + + engine.backward(engine(torch.ones(()))) + engine.step() + + assert module.gate_updates == 1 + assert precomputed == [module] + + +def test_reset_accumulation_recovers_a_failed_window(): + module = _Scale() + engine = Engine(module, optimizer=torch.optim.SGD(module.parameters(), lr=0.1), gradient_accumulation_steps=2) + + engine(torch.tensor(2.0)) # opens the sync context; backward never runs + with pytest.raises(RuntimeError, match="before starting another training forward"): + engine(torch.tensor(2.0)) + + engine.reset_accumulation() + + # A fresh window runs cleanly end to end. + engine.set_gradient_accumulation_steps(1) + engine.backward(engine(torch.tensor(2.0))) + engine.step() + assert module.weight.grad is None + + +def test_nonfinite_gradient_norm_skips_the_update(monkeypatch) -> None: + _skip_gradient_finalization(monkeypatch, norm=float("nan")) + module = _Scale() + optimizer = _CountingSGD(module.parameters()) + engine = Engine(module, optimizer=optimizer) + + engine.backward(engine(torch.ones(())) ** 2) + engine.step() + + assert optimizer.step_calls == 0 + assert optimizer.zero_grad_calls == 1 + assert module.weight.item() == 1.0