> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.recipes.llm.train_ft

## Module Contents

### Classes

| Name                                                                                                                          | Description                                               |
| ----------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------- |
| [`TrainFinetuneRecipeForNextTokenPrediction`](#nemo_automodel-recipes-llm-train_ft-TrainFinetuneRecipeForNextTokenPrediction) | Recipe for fine-tuning a model for next-token prediction. |

### Functions

| Name                                                                                                                | Description                                                                             |
| ------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| [`_build_partial_cuda_graph_manager`](#nemo_automodel-recipes-llm-train_ft-_build_partial_cuda_graph_manager)       | Build partial CUDA-graph state only for explicitly enabled model backends.              |
| [`_build_pp_collate_wrapper`](#nemo_automodel-recipes-llm-train_ft-_build_pp_collate_wrapper)                       | Return a collate-fn wrapper that precomputes pipeline-parallel causal masks, or `None`. |
| [`_build_tokenizer`](#nemo_automodel-recipes-llm-train_ft-_build_tokenizer)                                         | -                                                                                       |
| [`_get_domain_mixture_blend`](#nemo_automodel-recipes-llm-train_ft-_get_domain_mixture_blend)                       | Extract the explicit Megatron training blend for domain-mixture construction.           |
| [`_get_model_name`](#nemo_automodel-recipes-llm-train_ft-_get_model_name)                                           | -                                                                                       |
| [`_maybe_downgrade_loss_fn`](#nemo_automodel-recipes-llm-train_ft-_maybe_downgrade_loss_fn)                         | Downgrade to MaskedCrossEntropy when the requested loss cannot run.                     |
| [`_should_pack_validation`](#nemo_automodel-recipes-llm-train_ft-_should_pack_validation)                           | Return whether validation must use the configured training packer.                      |
| [`_should_precompute_pp_causal_masks`](#nemo_automodel-recipes-llm-train_ft-_should_precompute_pp_causal_masks)     | Return whether the recipe should attach PP causal-mask precomputation.                  |
| [`_supports_loss_weights`](#nemo_automodel-recipes-llm-train_ft-_supports_loss_weights)                             | Return whether `loss_fn` accepts the per-token `loss_weights` contract.                 |
| [`build_model`](#nemo_automodel-recipes-llm-train_ft-build_model)                                                   | Build and initialize a model.                                                           |
| [`compute_trust_remote_code_from_model`](#nemo_automodel-recipes-llm-train_ft-compute_trust_remote_code_from_model) | Compute the value of trust\_remote\_code based on the model configuration.              |
| [`main`](#nemo_automodel-recipes-llm-train_ft-main)                                                                 | Main entry point for the fine-tuning recipe.                                            |

### Data

[`logger`](#nemo_automodel-recipes-llm-train_ft-logger)

### API

```python
class nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction(
    cfg
)
```

**Bases:** [BaseRecipe](/nemo-automodel/nemo_automodel/recipes/base_recipe#nemo_automodel-recipes-base_recipe-BaseRecipe)

Recipe for fine-tuning a model for next-token prediction.

This class orchestrates training, from setup to main training loop.

**`cfg`**

---

**`magi`** `= MagiState()`

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._broadcast_from_last_pp_stage(
    tensor: torch.Tensor
) -> torch.Tensor
```

Broadcast a PP last-stage scalar to the other ranks in its pipeline group.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._collect_moe_load_balance()
```

Collect MoE load balance metrics with DP all-reduce.

Must be called on ALL ranks (the all-reduce is collective).
Stores the result in `self._moe_layer_loads` for rank-0 logging.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._configure_packing() -> nemo_automodel.components.models.common.packing.PackingCapabilities
```

Configure every local model stage and return its NEAT data requirements.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._configure_pipeline_loss_fn()
```

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._create_distributed_setup() -> nemo_automodel.components.distributed.config.DistributedSetup
```

Create the distributed setup used by this recipe rank.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._enable_qat_if_delayed(
    step: int
)
```

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._forward_backward_step(
    idx,
    batch,
    loss_buffer,
    num_label_tokens,
    num_batches,
    is_train: bool = True
)
```

Run one local batch and accumulate its loss and optional gradients.

**Parameters:**

**`idx`**

Microbatch index in the accumulation window.

---

**`batch`**

Input mapping with token IDs, labels, and physical NEAT
document IDs of shape \[batch, sequence]. NEAT attention metadata
is batch-major; legacy THD inputs are flattened by the sharder.
THD MTP requires physical boundaries in cu\_seqlens\_padded of
shape \[num\_sequences + 1] or \[1, num\_sequences + 1], or
\_packed\_seq\_ids \[batch, sequence] supplied by the native model
sharder from physical seq\_lens\_padded \[batch, num\_sequences].

---

**`loss_buffer`**

List receiving the detached scalar loss.

---

**`num_label_tokens`**

Global supervised-token count for loss normalization.

---

**`num_batches`**

Number of microbatches in the accumulation window.

---

**`is_train`** `bool` — default: True

Whether to backpropagate the combined main and MTP loss.

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._log_moe_metrics(
    step: int,
    wandb_log_fn
) -> None
```

Log MoE load balance metrics to wandb.

Call after `_collect_moe_load_balance`.  Only logs when
`_moe_layer_loads` is populated and a wandb log function is provided.

**Parameters:**

**`step`** `int`

Current training/benchmark step for wandb x-axis.

---

**`wandb_log_fn`**

Callable like `wandb.log` or `wandb_run.log`.

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._run_train_optim_step(
    batches,
    max_grad_norm: float | None = None
)
```

Execute a single training step.

**Parameters:**

**`batches`**

List of batches of training data.

---

**`max_grad_norm`** `float | None` — default: None

Gradient clipping norm. Optional, if None will not clip gradients.

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._run_validation_epoch(
    val_dataloader
)
```

Run one pass over a single validation dataloader.

**Parameters:**

**`val_name`**

Name of the validation dataset.

---

**`val_dataloader`**

DataLoader for the validation dataset.

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._setup_qat(
    cfg,
    model_parts: list[torch.nn.Module]
)
```

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction._should_setup_training_components() -> bool
```

Whether this rank owns the trainable model and its components.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction.log_train_metrics(
    log_data
)
```

Log metrics to wandb and other loggers.

**Parameters:**

**`log_data`**

MetricsSample object, containing:
step: int, the current step.
epoch: int, the current epoch.
metrics: Dict\[str, float], containing:
"loss": Training loss.
"grad\_norm": Grad norm from the training step.
"lr": Learning rate.
"mem": Memory allocated.
"tps": Tokens per second.
"tps\_per\_gpu": Tokens per second per GPU.
"num\_label\_tokens": Number of label tokens.

---

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction.log_val_metrics(
    val_name,
    log_data,
    metric_logger = None
)
```

Log metrics to wandb, MLflow and other loggers
Args:
log\_data: MetricsSample object, containing:
step: int, the current step.
epoch: int, the current epoch.
metrics: Dict\[str, float], containing:
"val\_loss": Validation loss.
"lr": Learning rate.
"num\_label\_tokens": Number of label tokens.
"mem": Memory allocated.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction.run_train_validation_loop()
```

Run the training loop over all epochs and batches.

For each batch, perform a forward pass, compute loss, backpropagate,
and update model parameters when necessary. Also prints loss every gradient step.

```python
nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction.setup()
```

Builds all components needed for training/validation/logging/checkpointing/etc.

This is the last place where self.cfg should be referenced.

**Raises:**

* `NotImplemented`: Raises if it tries to restore a checkpoint; will be removed.

```python
nemo_automodel.recipes.llm.train_ft._build_partial_cuda_graph_manager(
    model_parts: list[torch.nn.Module],
    activation_checkpointing: bool,
    pipeline_parallel: bool
) -> nemo_automodel.components.cuda_graphs.PartialCudaGraphManager | None
```

Build partial CUDA-graph state only for explicitly enabled model backends.

**Parameters:**

**`model_parts`** `list[nn.Module]`

Fully initialized model roots or pipeline-local model parts.

---

**`activation_checkpointing`** `bool`

Whether PyTorch activation checkpointing is enabled.

---

**`pipeline_parallel`** `bool`

Whether pipeline parallelism is enabled.

---

**Returns:** `PartialCudaGraphManager | None`

An armed manager when any backend selects CUDA-graph scopes, otherwise `None`.

```python
nemo_automodel.recipes.llm.train_ft._build_pp_collate_wrapper(
    cfg_model,
    pp_enabled: bool
)
```

Return a collate-fn wrapper that precomputes pipeline-parallel causal masks, or `None`.

`None` when PP is disabled, the model config can't be loaded, or the model
computes masks internally (e.g. `deepseek_v4` or `glm_moe_dsa`).  Passed to
`DataloaderConfig.build` as `collate_wrapper`.

```python
nemo_automodel.recipes.llm.train_ft._build_tokenizer(
    cfg_model,
    cfg_ds
)
```

```python
nemo_automodel.recipes.llm.train_ft._get_domain_mixture_blend(
    dataloader_config: nemo_automodel.components.datasets.loader.DataloaderConfig
) -> tuple[list[str], list[float]]
```

Extract the explicit Megatron training blend for domain-mixture construction.

```python
nemo_automodel.recipes.llm.train_ft._get_model_name(
    cfg_model
)
```

```python
nemo_automodel.recipes.llm.train_ft._maybe_downgrade_loss_fn(
    loss_fn: torch.nn.Module,
    probe_module: torch.nn.Module,
    pp_enabled: bool
) -> torch.nn.Module
```

Downgrade to MaskedCrossEntropy when the requested loss cannot run.

```python
nemo_automodel.recipes.llm.train_ft._should_pack_validation(
    training_dataloader: nemo_automodel.components.datasets.loader.DataloaderConfig | None,
    validation_dataloader: nemo_automodel.components.datasets.loader.DataloaderConfig,
    model: torch.nn.Module
) -> bool
```

Return whether validation must use the configured training packer.

```python
nemo_automodel.recipes.llm.train_ft._should_precompute_pp_causal_masks(
    model_config: typing.Any
) -> bool
```

Return whether the recipe should attach PP causal-mask precomputation.

```python
nemo_automodel.recipes.llm.train_ft._supports_loss_weights(
    loss_fn: torch.nn.Module
) -> bool
```

Return whether `loss_fn` accepts the per-token `loss_weights` contract.

```python
nemo_automodel.recipes.llm.train_ft.build_model(
    cfg_model,
    cfg_peft,
    seed,
    has_packed_sequence = False,
    cfg_fp8 = None,
    cfg_compile = None,
    cfg_quantization = None,
    distributed_setup: nemo_automodel.components.distributed.config.DistributedSetup | None = None,
    cfg_qat = None,
    cfg_freeze: nemo_automodel.components.config.loader.ConfigNode | dict[str, typing.Any] | nemo_automodel.components.utils.model_utils.FreezeConfig | None = None,
    sdpa_method: list[str] | None = None,
    device_mesh = None
) -> tuple[torch.nn.Module | nemo_automodel.components.distributed.pipelining.AutoPipeline, list['Optimizer']]
```

Build and initialize a model.

**Parameters:**

**`cfg_model`**

Configuration for model instantiation.

---

**`cfg_peft`**

Configuration for PEFT.

---

**`seed`**

Random seed.

---

**`has_packed_sequence`** — default: False

Whether using packed sequences.

---

**`cfg_fp8`** — default: None

Configuration for FP8.

---

**`cfg_compile`** — default: None

Configuration for torch.compile.

---

**`cfg_quantization`** — default: None

Configuration for BitsAndBytes quantization.

---

**`distributed_setup`** `DistributedSetup | None` — default: None

Resolved distributed topology and policy object.

---

**`cfg_qat`** — default: None

Configuration for QAT (will be instantiated to QATConfig).

---

**`cfg_freeze`** `ConfigNode | dict[str, Any] | FreezeConfig | None` — default: None

Freeze configuration (`freeze_config` YAML section as a
mapping, or a typed FreezeConfig) controlling parameter trainability.

---

**`sdpa_method`** `list[str] | None` — default: None

Explicit list of SDPA backend name strings (e.g.
`["flash_attention", "efficient_attention"]`), or `None` to
auto-select based on CP / activation checkpointing.

---

**`device_mesh`** — default: None

Pre-created device mesh forwarded when `distributed_setup` is not provided.

---

```python
nemo_automodel.recipes.llm.train_ft.compute_trust_remote_code_from_model(
    cfg_model
)
```

Compute the value of trust\_remote\_code based on the model configuration.

**Parameters:**

**`cfg_model`** `ConfigNode`

Model configuration.

---

**Returns:**

Whether to trust remote code.

```python
nemo_automodel.recipes.llm.train_ft.main(
    config_path = None
)
```

Main entry point for the fine-tuning recipe.

Loads the configuration, sets up the trainer, and initiates the training loop.

```python
nemo_automodel.recipes.llm.train_ft.logger = logging.getLogger(__name__)
```