> This page is for version 0.5.0 · 26.06.
> 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.components.distributed.pipelining.recv_buffer_pool

Pooled point-to-point recv buffers for torch pipeline stages.

`torch.distributed.pipelining` pre-allocates one full-size recv buffer per
microbatch and per direction (activations forward, gradients backward). At
large microbatch counts this dominates pipeline-stage memory: e.g. Kimi-K3 at
2k-token rows x GBS 4096 / mbs 2 / dp32 needs 64 buffer sets per direction,
\~45 GiB on middle stages — more than the activations themselves — which is
exactly what pushes that shape out of memory on 64 GB200 nodes.

Under the 1F1B schedule only `num_stages - stage_index` microbatches are
ever in flight on a stage: a forward recv buffer is live from irecv-post
until that chunk's backward consumes the stage input (activation-checkpoint
recompute reads it), and a grad recv buffer only until the chunk's backward
returns. This module therefore rebinds the per-chunk recv-info maps onto a
ring of `K = (num_stages - stage_index) + slack` real buffer sets.

Gradient accumulation on reused leaf buffers is already impossible upstream:
`stage_backward()` harvests input grads and sets `val.grad = None` per
chunk.

Opt-in via `install_recv_buffer_pool` (wired to
`PipelineConfig.pp_recv_buffer_pool`); it is a no-op that logs a warning
and leaves stock behavior when the torch internals it adapts are not
recognized. Only schedules with bounded in-flight depth are safe — see
`schedule_supports_recv_pool`; schedules with unbounded in-flight
microbatches (e.g. GPipe) would silently corrupt gradients (verified by the
CPU parity harness in the unit tests).

Native interleaved 1F1B uses a separate, instance-local implementation in
`install_interleaved_recv_pool`. Its pools are colored from the actual lowered schedule. Copies give each
microbatch independent input and gradient storage before receive-slot reuse.

## Module Contents

### Functions

| Name                                                                                                                                | Description                                                                            |
| ----------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- |
| [`_is_rank_zero`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_is_rank_zero)                                 | True on rank 0 or when torch.distributed is not initialized (single log line per job). |
| [`_pool_stage_backward_recvs`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_pool_stage_backward_recvs)       | Install instance-local allocation and ownership rules before stage setup.              |
| [`_pool_stage_forward_recvs`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_pool_stage_forward_recvs)         | Copy received inputs into independent leaves for each microbatch.                      |
| [`_recv_slots`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_recv_slots)                                     | Color receive lifetimes through the compute action that copies them.                   |
| [`_ring_size`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_ring_size)                                       | Number of real buffer sets for a stage: in-flight depth plus slack.                    |
| [`install_interleaved_recv_pool`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-install_interleaved_recv_pool) | Pool receives using this interleaved schedule's actual actions.                        |
| [`install_recv_buffer_pool`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-install_recv_buffer_pool)           | Monkeypatch pipeline-stage recv-buffer setup with ring pooling.                        |
| [`schedule_supports_recv_pool`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-schedule_supports_recv_pool)     | Return True when the schedule's in-flight depth bound makes pooling safe.              |

### Data

[`_INSTALLED`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_INSTALLED)

[`_INSTALLED_LAYOUT`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_INSTALLED_LAYOUT)

[`_SUPPORTED_SCHEDULES`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-_SUPPORTED_SCHEDULES)

[`logger`](#nemo_automodel-components-distributed-pipelining-recv_buffer_pool-logger)

### API

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._is_rank_zero() -> bool
```

True on rank 0 or when torch.distributed is not initialized (single log line per job).

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._pool_stage_backward_recvs(
    stage: torch.distributed.pipelining.PipelineStage,
    slots: tuple[int, ...],
    count: int
) -> None
```

Install instance-local allocation and ownership rules before stage setup.

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._pool_stage_forward_recvs(
    stage: torch.distributed.pipelining.PipelineStage,
    slots: tuple[int, ...],
    count: int
) -> None
```

Copy received inputs into independent leaves for each microbatch.

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._recv_slots(
    actions: list[torch.distributed.pipelining.schedules._Action],
    stage_index: int,
    num_microbatches: int,
    slack: int,
    forward: bool = False
) -> tuple[tuple[int, ...], int]
```

Color receive lifetimes through the compute action that copies them.

The copied forward input remains live through backward independently of
its receive slot. A copied gradient may outlive its backward and send.

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._ring_size(
    stage,
    num_microbatches: int,
    slack: int
) -> int
```

Number of real buffer sets for a stage: in-flight depth plus slack.

**Parameters:**

**`stage`**

torch `PipelineStage` (reads `num_stages` / `stage_index`).

---

**`num_microbatches`** `int`

Microbatches per step (upper bound for the ring).

---

**`slack`** `int`

Extra buffer sets beyond the 1F1B in-flight depth.

---

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool.install_interleaved_recv_pool(
    schedule: torch.distributed.pipelining.schedules.ScheduleInterleaved1F1B,
    slack: int = 2
) -> None
```

Pool receives using this interleaved schedule's actual actions.

Call once after schedule construction and before the first step. Received
activations become independent leaves before forward, while gradients are
copied before backward. These owned copies protect checkpoint recomputation,
escaping parameter gradients and asynchronous sends from slot reuse.

Pooling removes preallocation for every microbatch but retains owned inputs
until their actual consumers finish. The resulting peak memory depends on
the schedule and model, including whether outputs alias stage inputs.

**Parameters:**

**`schedule`** `ScheduleInterleaved1F1B`

Native interleaved 1F1B schedule with lowered P2P actions.

---

**`slack`** `int` — default: 2

Additional slots above the posted-receive lifetime bound.

---

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool.install_recv_buffer_pool(
    slack: int = 2
) -> bool
```

Monkeypatch pipeline-stage recv-buffer setup with ring pooling.

Must run before pipeline schedule/stage infra preparation. Installs
process-wide (class-level) and is idempotent; the first call's `slack`
wins.

**Parameters:**

**`slack`** `int` — default: 2

Extra buffer sets beyond the 1F1B in-flight depth, absorbing
transient recv-ahead (default 2).

---

**Returns:** `bool`

True if installed (or already installed), False if the torch

**Raises:**

* `ValueError`: if `slack` is negative (a ring smaller than the 1F1B
  in-flight depth would silently corrupt gradients).

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool.schedule_supports_recv_pool(
    pp_schedule: str | None
) -> bool
```

Return True when the schedule's in-flight depth bound makes pooling safe.

**Parameters:**

**`pp_schedule`** `str | None`

Schedule name as configured (e.g. `"1f1b"`, `"gpipe"`),
or None (custom CSV schedule).

---

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._INSTALLED = False
```

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._INSTALLED_LAYOUT: str | None = None
```

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool._SUPPORTED_SCHEDULES = frozenset({'1f1b'})
```

```python
nemo_automodel.components.distributed.pipelining.recv_buffer_pool.logger = logging.getLogger(__name__)
```