Merge branch 'main' into codex/refactor-component-builders

Author: akoumpa

examples/llm_finetune/agent/qwen2_5_3b_function_calling.yaml

@@ -114,17 +114,23 @@ validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   collate_fn: nemo_automodel.components.datasets.utils.default_collater
 
-# Generation-based tool-call accuracy eval. Runs at every val step alongside
-# val_loss. Catches the case where loss is decreasing but the model is
-# emitting wrong tool names or malformed argument JSON — a regression
-# loss alone cannot detect. Remove or comment out this block to disable.
+# Generation-based tool-call accuracy eval. Catches the case where loss is
+# decreasing but the model emits wrong tool names or malformed argument JSON —
+# something loss alone cannot detect. Remove or comment out this block to disable.
+#
+# NOTE: this recipe uses `distributed.strategy: fsdp2`, which SKIPS in-loop
+# generation eval by default — running generate() on the sharded training model
+# repeatedly unshards params and can OOM the next step. It runs every val step
+# under DDP. To force it under FSDP2, set `run_on_fsdp2: true` below (slower,
+# higher memory); or run the evaluator offline against a saved checkpoint.
 tool_call_eval:
   _target_: nemo_automodel.components.eval.tool_call_evaluator.ToolCallAccuracyEvaluator
   dataset_name: llamafactory/glaive_toolcall_en
   split: train[:128]
   max_eval_samples: 128
   max_new_tokens: 256
   max_prompt_tokens: 3584
+  # run_on_fsdp2: true   # opt in to in-loop eval under FSDP2 (otherwise skipped)
 
 optimizer:
   _target_: torch.optim.AdamW

nemo_automodel/components/datasets/llm/agent_chat.py

@@ -275,38 +275,6 @@ def _truncate_messages_to_fit(
     return system + history[boundaries[-1] :]
 
 
-def _mask_labels_to_last_turn(labels: List[int], ignore_index: int = -100) -> List[int]:
-    """Restrict the loss to the final assistant turn (``mask_history``).
-
-    ``labels`` come from :func:`format_chat_template` with every assistant
-    turn supervised; non-assistant tokens are already ``ignore_index``.
-    Because the chat template renders each assistant message as a single
-    contiguous span, supervised tokens form one maximal run per assistant
-    turn separated by ``ignore_index`` runs. This keeps only the last such
-    run and masks every earlier supervised token in place.
-
-    Args:
-        labels: per-token labels (``ignore_index`` marks unsupervised tokens).
-        ignore_index: the value marking unsupervised tokens.
-
-    Returns:
-        The same list, mutated so only the final supervised run is kept.
-    """
-    last = -1
-    for i in range(len(labels) - 1, -1, -1):
-        if labels[i] != ignore_index:
-            last = i
-            break
-    if last < 0:
-        return labels
-    start = last
-    while start - 1 >= 0 and labels[start - 1] != ignore_index:
-        start -= 1
-    for i in range(start):
-        labels[i] = ignore_index
-    return labels
-
-
 def _format_example(
     example: Dict[str, Any],
     tokenizer,
@@ -400,10 +368,8 @@ def _format_example_impl(
         truncation=truncation,
         answer_only_loss_mask=True,
         mask_reasoning_content=mask_reasoning_content,
+        train_on_last_turn_only=train_on_last_turn_only,
     )
-    if train_on_last_turn_only:
-        _mask_labels_to_last_turn(tokenized["labels"])
-
     # Truncation (or over-aggressive last-turn masking) can leave a sample with no
     # supervised tokens at all — every label is ``ignore_index`` (-100). A single
     # such sample is harmless: the loss normalizes by the batch's supervised-token

nemo_automodel/components/datasets/llm/formatting_utils.py

@@ -279,6 +279,41 @@ def _build_reasoning_mask(
     return reasoning_mask
 
 
+def _mask_labels_to_last_turn(mask: List[int], ignore_index: int = -100) -> List[int]:
+    """Restrict supervision to the final assistant turn (``mask_history``).
+
+    Operates on any per-token sequence where ``ignore_index`` marks
+    unsupervised positions: a label list (``ignore_index=-100``) or a 0/1
+    assistant mask (``ignore_index=0``). Each assistant turn renders as a
+    single contiguous supervised span, so this keeps only the last such run
+    and rewrites every earlier supervised position to ``ignore_index``.
+
+    Apply this to the assistant mask **before** any reasoning_content holes are
+    punched into it; running it on already-holed labels would treat the
+    reasoning gap as a turn boundary and drop in-turn content before the hole.
+
+    Args:
+        mask: per-token labels or 0/1 mask (``ignore_index`` marks unsupervised).
+        ignore_index: the value marking unsupervised positions.
+
+    Returns:
+        The same list, mutated so only the final supervised run is kept.
+    """
+    last = -1
+    for i in range(len(mask) - 1, -1, -1):
+        if mask[i] != ignore_index:
+            last = i
+            break
+    if last < 0:
+        return mask
+    start = last
+    while start - 1 >= 0 and mask[start - 1] != ignore_index:
+        start -= 1
+    for i in range(start):
+        mask[i] = ignore_index
+    return mask
+
+
 @torch.no_grad()
 def _get_right_trailing_pad_mask(
     sequence: torch.Tensor,
@@ -583,6 +618,7 @@ def format_chat_template(
     tools: Optional[List[Dict]] = None,
     answer_only_loss_mask: bool = True,
     mask_reasoning_content: bool = False,
+    train_on_last_turn_only: bool = False,
     unshifted: bool = False,
 ) -> Dict[str, List[int]]:
     """
@@ -597,6 +633,9 @@ def format_chat_template(
         tools: Optional list of tool definitions for function calling.
         answer_only_loss_mask: Whether to compute the loss mask only on the answer tokens.
         mask_reasoning_content: Whether to exclude rendered reasoning_content tokens from loss.
+        train_on_last_turn_only: Whether to supervise only the final assistant turn,
+            masking every earlier assistant turn (``mask_history``). Applied to the
+            assistant mask before reasoning_content is masked out.
 
     Returns:
         A dictionary with the formatted example.
@@ -667,6 +706,11 @@ def format_chat_template(
             if not tokenizer_attn_mask[i]:
                 mask[i] = 0
 
+    # Restrict to the last assistant turn before reasoning is masked, so the
+    # contiguous-run heuristic sees a hole-free mask (one run per turn).
+    if train_on_last_turn_only:
+        _mask_labels_to_last_turn(mask, ignore_index=0)
+
     if mask_reasoning_content and has_reasoning_content:
         reasoning_mask = _build_reasoning_mask(
             tokenizer,

nemo_automodel/components/distributed/utils.py

@@ -16,12 +16,14 @@
 import logging
 from contextlib import ContextDecorator, nullcontext
 from datetime import datetime, timedelta
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 import torch.distributed
 import torch.distributed as dist
 
+from nemo_automodel.components.distributed.config import DDPConfig
+
 logger = logging.getLogger(__name__)
 
 
@@ -245,3 +247,17 @@ def get_sync_ctx(model, is_optim_step, defer_fsdp_grad_sync: bool):
     elif isinstance(model, torch.nn.parallel.DistributedDataParallel) and not is_optim_step:
         sync_ctx = model.no_sync()
     return sync_ctx
+
+
+def dp_eval_sample_shard(distributed_config: Any, dp_rank: int, dp_size: int) -> tuple | None:
+    """Return ``(dp_rank, dp_size)`` to shard eval samples across DP ranks, else ``None``.
+
+    Only DDP replicates the full model on every rank, so only there can each rank
+    score a disjoint subset independently. Under FSDP2 / MegatronFSDP the model is
+    sharded and ``generate()`` issues per-layer all-gather collectives that must
+    stay in lockstep across the group — different samples per rank would desync
+    them and hang — so every rank scores the same samples (``None``).
+    """
+    if isinstance(distributed_config, DDPConfig) and dp_size > 1:
+        return (dp_rank, dp_size)
+    return None

nemo_automodel/components/eval/tool_call_evaluator.py

@@ -46,16 +46,6 @@
 logger = logging.getLogger(__name__)
 
 
-_METRIC_KEYS = (
-    "has_call",
-    "name_correct",
-    "args_json_valid",
-    "args_field_recall",
-    "args_field_precision",
-    "args_exact_match",
-)
-
-
 class ToolCallAccuracyEvaluator:
     """Generation-based tool-call accuracy evaluator for agent SFT.
 
@@ -86,6 +76,19 @@ class ToolCallAccuracyEvaluator:
             returned ``_count`` and weighted-summed metrics.
     """
 
+    #: Fixed metric names returned by :meth:`evaluate` (one mean each). Exposed
+    #: so a caller can all-reduce a stable key set across ranks instead of the
+    #: data-dependent ``_skip_<reason>`` diagnostics, which differ per rank and
+    #: would desync collectives.
+    METRIC_KEYS = (
+        "has_call",
+        "name_correct",
+        "args_json_valid",
+        "args_field_recall",
+        "args_field_precision",
+        "args_exact_match",
+    )
+
     def __init__(
         self,
         *,
@@ -120,6 +123,20 @@ def __init__(
 
         self._samples_cache: Optional[List[Dict[str, Any]]] = None
 
+    @property
+    def sample_shard(self) -> Optional[tuple]:
+        """``(rank, world_size)`` shard, or ``None`` to score every sample.
+
+        The training recipe sets this so each data-parallel rank scores a
+        disjoint subset, but only when the model is replicated per rank (DDP);
+        sharded strategies must keep every rank on the same samples.
+        """
+        return self._sample_shard
+
+    @sample_shard.setter
+    def sample_shard(self, value: Optional[tuple]) -> None:
+        self._sample_shard = value
+
     def _cleanup_cuda(self) -> None:
         gc.collect()
         if torch.cuda.is_available():
@@ -278,7 +295,7 @@ def evaluate(self, model, tokenizer) -> Dict[str, float]:
             samples actually scored on this rank.
         """
         samples = self._iter_my_samples()
-        sums = {k: 0.0 for k in _METRIC_KEYS}
+        sums = {k: 0.0 for k in self.METRIC_KEYS}
         n_scored = 0
         skip_reasons: Dict[str, int] = {}
 
@@ -369,7 +386,7 @@ def evaluate(self, model, tokenizer) -> Dict[str, float]:
             decoded = tokenizer.decode(new_tokens, skip_special_tokens=False)
             pred_calls = parse_tool_calls(decoded)
             metrics = compute_sample_metrics(pred_calls, sample["gt_tool_calls"])
-            for k in _METRIC_KEYS:
+            for k in self.METRIC_KEYS:
                 sums[k] += metrics[k]
             n_scored += 1
             del input_ids, attention_mask, output
@@ -389,10 +406,10 @@ def evaluate(self, model, tokenizer) -> Dict[str, float]:
 
         result: Dict[str, float] = {}
         if n_scored > 0:
-            for k in _METRIC_KEYS:
+            for k in self.METRIC_KEYS:
                 result[f"{self.metric_prefix}/{k}"] = sums[k] / n_scored
         else:
-            for k in _METRIC_KEYS:
+            for k in self.METRIC_KEYS:
                 result[f"{self.metric_prefix}/{k}"] = 0.0
         result[f"{self.metric_prefix}/_count"] = float(n_scored)
         result[f"{self.metric_prefix}/_skipped"] = float(n_skipped)

nemo_automodel/recipes/llm/train_eagle1.py

@@ -21,6 +21,7 @@
 import math
 import os
 import pathlib
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import torch
@@ -68,6 +69,35 @@ def _optim_steps_per_epoch(num_batches_per_epoch: int, grad_accumulation_steps:
     return -(-num_batches_per_epoch // grad_accumulation_steps)
 
 
+def _should_sync_grads(
+    *,
+    pending_micro_batches: int,
+    grad_accumulation_steps: int,
+    batch_idx: int,
+    batches_per_epoch: int | None,
+    is_ddp: bool,
+) -> bool:
+    """Return True when this micro-batch's backward should all-reduce gradients.
+
+    Under DDP with gradient accumulation only the micro-batch immediately
+    followed by an ``optimizer.step()`` needs to synchronize: at that point the
+    locally-accumulated ``.grad`` already holds the whole window's contribution,
+    so a single all-reduce averages the complete window and the intervening
+    micro-batches can run under ``no_sync()`` -- saving ``grad_accumulation_steps - 1``
+    all-reduces per window. That step is either the window closer
+    (``pending_micro_batches + 1 == grad_accumulation_steps``) or the epoch's
+    final batch (which the trailing-flush step consumes). When the dataloader
+    length is unknown we cannot identify the final batch, so we sync every step
+    (correct, just no speedup). With a single process (no DDP) there is nothing
+    to synchronize, so this is always True.
+    """
+    if not is_ddp or batches_per_epoch is None:
+        return True
+    closes_window = pending_micro_batches + 1 == grad_accumulation_steps
+    is_last_batch = batch_idx == batches_per_epoch - 1
+    return closes_window or is_last_batch
+
+
 def _all_reduce_mean(value: torch.Tensor) -> torch.Tensor:
     if dist.is_available() and dist.is_initialized():
         dist.all_reduce(value, op=dist.ReduceOp.SUM)
@@ -506,6 +536,11 @@ def run_train_validation_loop(self):
             if self.dist_env.is_main:
                 logger.info("All %d epochs already completed; nothing to do.", self.num_epochs)
             return
+        try:
+            batches_per_epoch = len(self.train_dataloader)
+        except TypeError:
+            batches_per_epoch = None
+        is_ddp = isinstance(self.trainer_module, DistributedDataParallel)
         for epoch_idx in range(start_epoch, self.num_epochs):
             if hasattr(self.train_dataloader, "sampler") and hasattr(self.train_dataloader.sampler, "set_epoch"):
                 self.train_dataloader.sampler.set_epoch(epoch_idx)
@@ -525,16 +560,28 @@ def run_train_validation_loop(self):
                     attention_mask=batch["attention_mask"],
                     loss_mask=batch["loss_mask"],
                 )
-                metrics = self.trainer_module(
-                    input_ids=target_batch.input_ids,
-                    attention_mask=target_batch.attention_mask,
-                    loss_mask=target_batch.loss_mask,
-                    input_hidden_states=target_batch.input_hidden_states,
-                    target_hidden_states=target_batch.target_hidden_states,
-                    target_logits=target_batch.target_logits,
+                # Skip DDP's per-micro-batch all-reduce on every micro-batch
+                # except the one an optimizer step immediately follows; that
+                # step's all-reduce covers the whole locally-accumulated window.
+                sync_grads = _should_sync_grads(
+                    pending_micro_batches=pending_micro_batches,
+                    grad_accumulation_steps=self.grad_accumulation_steps,
+                    batch_idx=batch_idx,
+                    batches_per_epoch=batches_per_epoch,
+                    is_ddp=is_ddp,
                 )
-                loss = metrics.loss / self.grad_accumulation_steps
-                loss.backward()
+                sync_ctx = nullcontext() if sync_grads else self.trainer_module.no_sync()
+                with sync_ctx:
+                    metrics = self.trainer_module(
+                        input_ids=target_batch.input_ids,
+                        attention_mask=target_batch.attention_mask,
+                        loss_mask=target_batch.loss_mask,
+                        input_hidden_states=target_batch.input_hidden_states,
+                        target_hidden_states=target_batch.target_hidden_states,
+                        target_logits=target_batch.target_logits,
+                    )
+                    loss = metrics.loss / self.grad_accumulation_steps
+                    loss.backward()
 
                 running_loss += metrics.loss.detach().item()
                 running_acc += metrics.accuracy.detach().item()