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`	True on rank 0 or when torch.distributed is not initialized (single log line per job).
`_ring_size`	Number of real buffer sets for a stage: in-flight depth plus slack.
`install_recv_buffer_pool`	Monkeypatch pipeline-stage recv-buffer setup with ring pooling.
`schedule_supports_recv_pool`	Return True when the schedule’s in-flight depth bound makes pooling safe.

Data

_INSTALLED

_INSTALLED_LAYOUT

_SUPPORTED_SCHEDULES

logger

API

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).

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.

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

intDefaults to 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).

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).

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

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

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

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