> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/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.components.distributed.pipelining

## Submodules

* **[`nemo_automodel.components.distributed.pipelining.autopipeline`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/autopipeline)**
* **[`nemo_automodel.components.distributed.pipelining.config`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/config)**
* **[`nemo_automodel.components.distributed.pipelining.functional`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/functional)**
* **[`nemo_automodel.components.distributed.pipelining.hf_utils`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/hf_utils)**
* **[`nemo_automodel.components.distributed.pipelining.recv_buffer_pool`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/recv_buffer_pool)**
* **[`nemo_automodel.components.distributed.pipelining.runtime`](/nemo-automodel/nemo_automodel/components/distributed/pipelining/runtime)**

## Package Contents

### Classes

| Name                                                                                          | Description                                                                     |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`AutoPipeline`](#nemo_automodel-components-distributed-pipelining-autopipeline-AutoPipeline) | Orchestrates pipeline-parallel training on top of torch.distributed.pipelining. |

### API

```python
class nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline(
    world_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    pp_axis_name: str = 'pp',
    dp_axis_names: tuple[str, ...] = ('dp',),
    cp_axis_name: str | None = None,
    tp_axis_name: str | None = None,
    ep_axis_name: str | None = None,
    ep_shard_axis_names: tuple[str, ...] | None = None,
    pp_schedule: str | None = '1f1b',
    pp_schedule_csv: str | None = None,
    pp_microbatch_size: int = 1,
    pp_batch_size: int = 1,
    layers_per_stage: int | None = None,
    round_virtual_stages_to_pp_multiple: typing.Literal['up', 'down'] | None = None,
    module_fqns_per_model_part: list[list[str]] | None = None,
    patch_inner_model: bool = True,
    patch_causal_lm_model: bool = True,
    patch_stage_backward_maybe_with_nosync: bool = False,
    defer_fsdp_grad_sync: bool = True,
    device: torch.device | None = None,
    dtype: torch.dtype | None = None,
    scale_grads_in_schedule: bool = False,
    pp_seq_len: int | None = None,
    pp_recv_buffer_pool: bool = False,
    pp_recv_buffer_pool_slack: int = 2
)
```

Orchestrates pipeline-parallel training on top of torch.distributed.pipelining.

**`_device`** `device`

---

**`_info`**

---

**`_pp_current_seq_len`** `int | None = None`

---

**`_pp_runtime_num_tokens`** `int = 0`

---

**`_runtime_initializers`** `list[PipelineRuntimeInitializer] = []`

---

**`device`** `device`

---

**`info`** `PipelineInfo`

---

**`parts`** `list[Module]`

---

**`pp_mesh`** `DeviceMesh = self.world_mesh[pp_axis_name]`

---

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline._count_parameters(
    module: torch.nn.Module,
    trainable_only: bool = False
) -> int
```

staticmethod

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline._get_schedule_kwargs_chunk_spec(
    kwargs: dict[str, typing.Any]
) -> dict[str, typing.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]`.

**Parameters:**

**`kwargs`** `dict[str, Any]`

Mapping passed to the pipeline schedule. Tensor values may
have arbitrary model-defined layouts; the model hook identifies
any nonstandard batch axis.

---

**Returns:** `dict[str, Any] | None`

A chunk-spec mapping with the same nested structure as `kwargs`,

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.build(
    model: torch.nn.Module,
    loss_fn: typing.Callable | None = None,
    parallelize_fn: nemo_automodel.components.distributed.pipelining.functional.ParallelizeFnProtocol | None = None
)
```

Build the pipeline: validate -> init meta -> split -> schedule.

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.debug_summary() -> str
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.get_stage_param_counts(
    trainable_only: bool = False
) -> list[int]
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.get_total_param_count(
    trainable_only: bool = False
) -> int
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.list_stage_modules() -> list[list[str]]
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.log_debug_summary() -> None
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.pretty_print_stages(
    max_modules_per_stage: int = 16,
    trainable_only: bool = False
) -> str
```

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.step(
    model_input: torch.Tensor,
    target: torch.Tensor | None = None,
    losses: list[torch.Tensor] | None = None,
    kwargs: typing.Any = {}
) -> typing.Any
```

Run one pipeline schedule step with model-owned input chunking.

**Parameters:**

**`model_input`** `torch.Tensor`

Tensor of shape \[batch, ...] containing the first
pipeline stage's input. Ignored on ranks without the first stage.

---

**`target`** `torch.Tensor | None` — default: None

Tensor with a model-defined target layout, or `None` on
ranks without the last pipeline stage.

---

**`losses`** `list[torch.Tensor] | None` — default: None

Mutable list populated with scalar loss tensors, or `None`
on ranks without the last pipeline stage.

---

**`**kwargs`** `Any` — default: \{}

Keyword schedule inputs. Tensor values may have arbitrary
model-defined layouts; model-owned metadata identifies any
nonstandard batch axis.

---

**Returns:** `Any`

The value returned by the underlying PyTorch pipeline schedule.

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.update_seq_len(
    seq_len: int
) -> None
```

Reset pipeline stage infrastructure for a new sequence length.

VLM training batches can have wildly different sequence lengths across steps
(image batches vs. text-only batches).  PyTorch's PipelineStage locks in recv
buffer sizes on the first step, causing a shape-mismatch error on later steps
with different seq\_lens.

Call this before every `schedule.step()` to update the stage shapes without
running an expensive forward pass.  Stage reset is a no-op when `seq_len`
has not changed; model-owned runtime initializers are still prepared.

**Parameters:**

**`seq_len`** `int`

Sequence length of the upcoming batch (`input_ids.shape[1]`).

---

```python
nemo_automodel.components.distributed.pipelining.autopipeline.AutoPipeline.visualize_current_schedule(
    filename: str | None = None
) -> None
```

## Docs

- [nemo_automodel.components.distributed.pipelining.autopipeline](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/autopipeline.md)
- [nemo_automodel.components.distributed.pipelining.config](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/config.md)
- [nemo_automodel.components.distributed.pipelining.functional](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/functional.md)
- [nemo_automodel.components.distributed.pipelining.hf_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/hf_utils.md)
- [nemo_automodel.components.distributed.pipelining.recv_buffer_pool](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/recv_buffer_pool.md)
- [nemo_automodel.components.distributed.pipelining.runtime](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/runtime.md)