> 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

Ring-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 :func:`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
:func:`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).

## 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). |
| [`_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_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._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_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__)
```