> 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.diffusion.train

## Module Contents

### Classes

| Name                                                                                   | Description                           |
| -------------------------------------------------------------------------------------- | ------------------------------------- |
| [`TrainDiffusionRecipe`](#nemo_automodel-recipes-diffusion-train-TrainDiffusionRecipe) | Training recipe for diffusion models. |

### Functions

| Name                                                                                                                   | Description                                                                   |
| ---------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_build_diffusion_mesh_context`](#nemo_automodel-recipes-diffusion-train-_build_diffusion_mesh_context)               | Build the diffusion transformer's sole parallelization input.                 |
| [`_build_transformer_engine_fp8_recipe`](#nemo_automodel-recipes-diffusion-train-_build_transformer_engine_fp8_recipe) | Build a Transformer Engine FP8 recipe from CLI-friendly config values.        |
| [`_calculate_throughput_metrics`](#nemo_automodel-recipes-diffusion-train-_calculate_throughput_metrics)               | Calculate directly measured training throughput metrics.                      |
| [`_count_local_batch_group_samples`](#nemo_automodel-recipes-diffusion-train-_count_local_batch_group_samples)         | Count local samples processed by one optimizer step.                          |
| [`_get_diffusion_microbatch_size`](#nemo_automodel-recipes-diffusion-train-_get_diffusion_microbatch_size)             | Return the number of samples in one local diffusion micro-batch.              |
| [`_reject_removed_diffusion_keys`](#nemo_automodel-recipes-diffusion-train-_reject_removed_diffusion_keys)             | Raise with old->new key mapping when a removed diffusion YAML key is present. |
| [`_resolve_model_dtypes`](#nemo_automodel-recipes-diffusion-train-_resolve_model_dtypes)                               | Resolve model storage and compute dtypes from the recipe config.              |
| [`_resolve_transformer_engine_autocast`](#nemo_automodel-recipes-diffusion-train-_resolve_transformer_engine_autocast) | Resolve Transformer Engine's quantization autocast context manager.           |
| [`_validate_precision_configuration`](#nemo_automodel-recipes-diffusion-train-_validate_precision_configuration)       | Reject split storage/compute dtypes on paths without FSDP param casting.      |
| [`build_diffusion_pipeline`](#nemo_automodel-recipes-diffusion-train-build_diffusion_pipeline)                         | Build the sharded diffusion pipeline (model + parallel scheme).               |

### Data

[`_REMOVED_KEY_MIGRATIONS`](#nemo_automodel-recipes-diffusion-train-_REMOVED_KEY_MIGRATIONS)

### API

```python
class nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe(
    cfg
)
```

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

Training recipe for diffusion models.

**`cfg`**

---

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._autocast_context() -> typing.Any
```

Return the per-forward autocast context used when FSDP2 does not cast parameters.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._count_global_samples(
    local_samples: int
) -> int
```

Count samples processed across the data-parallel group.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._elapsed_seconds_since(
    start_time: float
) -> tuple[float, float]
```

Return the max elapsed wall-clock seconds across ranks since start\_time.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._get_collective_device() -> torch.device
```

Return a tensor device compatible with the active distributed backend.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._get_dp_group_size(
    include_cp: bool = False
) -> int
```

Get data parallel world size, handling DDP mode where device\_mesh is None.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._get_dp_rank(
    include_cp: bool = False
) -> int
```

Get data parallel rank, handling DDP mode where device\_mesh is None.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._get_memory_metrics() -> typing.Dict[str, float]
```

Return PyTorch CUDA allocator memory counters, max-reduced across ranks.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._run_validation_epoch(
    global_step: int
) -> float
```

Score the held-out set with the training flow-matching objective.

Seeding by data rank mirrors training: every evaluation draws the same timesteps, noise,
and CFG dropout, so `val_loss` tracks the model rather than the sampling, while
data-parallel ranks stay decorrelated and context-parallel peers of one rank draw
identical values for their shared batch. `ScopedRNG` restores the training RNG state
afterwards, leaving the training trajectory unchanged.

The forward runs under the same compute-dtype autocast as training, so a single-rank
split-dtype config evaluates the way it trains. FP8 autocast is left out: delayed-scaling
amax history is training state and must not be updated from held-out batches.

**Parameters:**

**`global_step`** `int`

Current optimizer step, forwarded to the flow-matching pipeline for
its periodic diagnostics.

---

**Returns:** `float`

Mean loss over the validation batches of the data-parallel group, comparable to the

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._sync_device() -> None
```

Wait for queued CUDA work so timing reflects completed training work.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe._transformer_engine_fp8_context() -> typing.Any
```

Return the per-forward Transformer Engine FP8 context.

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe.run_train_validation_loop()
```

```python
nemo_automodel.recipes.diffusion.train.TrainDiffusionRecipe.setup()
```

```python
nemo_automodel.recipes.diffusion.train._build_diffusion_mesh_context(
    fsdp_cfg: typing.Dict[str, typing.Any] | None,
    ddp_cfg: typing.Dict[str, typing.Any] | None,
    world_size: int,
    dtype: torch.dtype,
    compute_dtype: torch.dtype | None = None,
    lora_enabled: bool
) -> nemo_automodel.components.distributed.MeshContext
```

Build the diffusion transformer's sole parallelization input.

```python
nemo_automodel.recipes.diffusion.train._build_transformer_engine_fp8_recipe(
    recipe_name: str,
    amax_history_len: int,
    amax_compute_algo: str
) -> typing.Any
```

Build a Transformer Engine FP8 recipe from CLI-friendly config values.

```python
nemo_automodel.recipes.diffusion.train._calculate_throughput_metrics(
    elapsed_seconds: float,
    optimizer_steps: int,
    global_samples: int,
    world_size: int
) -> typing.Dict[str, float]
```

Calculate directly measured training throughput metrics.

```python
nemo_automodel.recipes.diffusion.train._count_local_batch_group_samples(
    batch_group: list[typing.Dict[str, typing.Any]]
) -> int
```

Count local samples processed by one optimizer step.

```python
nemo_automodel.recipes.diffusion.train._get_diffusion_microbatch_size(
    batch: typing.Dict[str, typing.Any]
) -> int
```

Return the number of samples in one local diffusion micro-batch.

```python
nemo_automodel.recipes.diffusion.train._reject_removed_diffusion_keys(
    cfg: typing.Any
) -> None
```

Raise with old->new key mapping when a removed diffusion YAML key is present.

**Parameters:**

**`cfg`** `Any`

Raw recipe config (`ConfigNode` or `RecipeConfig`) as loaded from YAML.

---

**Raises:**

* `ValueError`: If any key from `_REMOVED_KEY_MIGRATIONS` is present.

```python
nemo_automodel.recipes.diffusion.train._resolve_model_dtypes(
    cfg: typing.Any
) -> tuple[torch.dtype, torch.dtype]
```

Resolve model storage and compute dtypes from the recipe config.

```python
nemo_automodel.recipes.diffusion.train._resolve_transformer_engine_autocast() -> typing.Any
```

Resolve Transformer Engine's quantization autocast context manager.

```python
nemo_automodel.recipes.diffusion.train._validate_precision_configuration(
    dtype: torch.dtype,
    compute_dtype: torch.dtype,
    ddp_cfg: typing.Dict[str, typing.Any] | None,
    peft_cfg: typing.Any
) -> None
```

Reject split storage/compute dtypes on paths without FSDP param casting.

```python
nemo_automodel.recipes.diffusion.train.build_diffusion_pipeline(
    model_id: str,
    finetune_mode: bool,
    device: torch.device,
    dtype: torch.dtype,
    compute_dtype: torch.dtype | None = None,
    cpu_offload: bool = False,
    fsdp_cfg: typing.Dict[str, typing.Any] | None = None,
    ddp_cfg: typing.Dict[str, typing.Any] | None = None,
    attention_backend: str | None = None,
    transformer_engine_linear: bool = False,
    transformer_engine_fp8_safe_only: bool = False,
    fuse_qkv_projections: bool = False,
    compact_fused_qkv_projections: bool = False,
    pipeline_spec: typing.Dict[str, typing.Any] | None = None,
    peft_cfg = None,
    model_type = None,
    active_transformer: str | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    config_overrides: dict[str, nemo_automodel._diffusers.auto_diffusion_pipeline.ConfigFieldValue] | None = None
) -> tuple[nemo_automodel._diffusers.auto_diffusion_pipeline.NeMoAutoDiffusionPipeline, nemo_automodel.components.distributed.MeshContext]
```

Build the sharded diffusion pipeline (model + parallel scheme).

The optimizer is built separately by the recipe via
`OptimizerConfig.build(...)` on the returned pipeline's transformer, so
that parameters are collected after FSDP2 wrapping.

**Parameters:**

**`model_id`** `str`

Pretrained model name or path.

---

**`finetune_mode`** `bool`

Whether to load for finetuning (True) or pretraining (False).

---

**`device`** `torch.device`

Target device.

---

**`dtype`** `torch.dtype`

Model parameter storage dtype.

---

**`compute_dtype`** `torch.dtype | None` — default: None

Forward/FSDP compute dtype. Defaults to dtype when unset.

---

**`cpu_offload`** `bool` — default: False

Whether to enable CPU offload (FSDP only).

---

**`fsdp_cfg`** `Dict[str, Any] | None` — default: None

FSDP configuration dict. Mutually exclusive with ddp\_cfg.

---

**`ddp_cfg`** `Dict[str, Any] | None` — default: None

DDP configuration dict. Mutually exclusive with fsdp\_cfg.

---

**`attention_backend`** `str | None` — default: None

Optional attention backend override.

---

**`transformer_engine_linear`** `bool` — default: False

Whether to replace transformer torch.nn.Linear modules with Transformer Engine Linear.

---

**`transformer_engine_fp8_safe_only`** `bool` — default: False

Whether to skip TE conversion for known FP8-incompatible modules.

---

**`fuse_qkv_projections`** `bool` — default: False

Whether to call Diffusers QKV projection fusion on the transformer before FSDP.

---

**`compact_fused_qkv_projections`** `bool` — default: False

Whether to remove original projection modules after QKV fusion.

---

**`pipeline_spec`** `Dict[str, Any] | None` — default: None

Pipeline specification for pretraining (from\_config).
Required when finetune\_mode is False. Should contain:

* transformer\_cls: str (e.g., "WanTransformer3DModel", "FluxTransformer2DModel")
* subfolder: str (e.g., "transformer")
* Optional: pipeline\_cls, load\_full\_pipeline

---

**`peft_cfg`** — default: None

PeftConfig instance or None. When provided, only LoRA params
are trained; base weights are frozen and sharded by FSDP2 for memory.

---

**`model_type`** — default: None

"flux" | "flux2" | "wan" | "hunyuan" | "ltx2". Required when peft\_cfg is provided.

---

**`active_transformer`** `str | None` — default: None

For two-transformer pipelines (Wan2.2), select which
transformer to finetune. `"transformer"` (default for Wan2.2 = high-noise)
or `"transformer_2"` (low-noise). The unused transformer is dropped
before device placement so only one transformer lives on GPU.

---

**`backend`** `BackendConfig | None` — default: None

`BackendConfig` for a transformer with a custom Automodel
implementation (for example MoE diffusion transformers that need expert parallelism).

---

**`config_overrides`** `dict[str, ConfigFieldValue] | None` — default: None

Config attributes overridden before building a custom-model transformer.

---

**Returns:** `NeMoAutoDiffusionPipeline`

Tuple of (pipeline, resolved MeshContext). The mesh context carries both the

**Raises:**

* `ValueError`: If both fsdp\_cfg and ddp\_cfg are provided.
* `ValueError`: If finetune\_mode is False and pipeline\_spec is not provided.

```python
nemo_automodel.recipes.diffusion.train._REMOVED_KEY_MIGRATIONS = {'optim.learning_rate': 'optimizer.lr', 'optim.optimizer': 'optimizer (with an e...
```