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`	True on rank 0 or when torch.distributed is not initialized (single log line per job).
`_pool_stage_backward_recvs`	Install instance-local allocation and ownership rules before stage setup.
`_pool_stage_forward_recvs`	Copy received inputs into independent leaves for each microbatch.
`_recv_slots`	Color receive lifetimes through the compute action that copies them.
`_ring_size`	Number of real buffer sets for a stage: in-flight depth plus slack.
`install_interleaved_recv_pool`	Pool receives using this interleaved schedule’s actual actions.
`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._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.

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.

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.

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_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

intDefaults to 2

Additional slots above the posted-receive lifetime bound.

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