perf(distributed): ring-pooled PP recv buffers and Kimi-K3 2k benchmark shape

[yisongbetter]

What does this PR do ?

Opt-in ring pooling of torch.distributed.pipelining P2P recv buffers under
1F1B (distributed.pipeline.pp_recv_buffer_pool), which removes the
per-microbatch buffer pre-allocation that OOMs Kimi-K3 at 2k rows × GBS 4096
on 64 GB200 nodes, plus the new kimi_k3_gb200_2k.yaml benchmark shape it
unlocks.

torch.distributed.pipelining pre-allocates one full-size P2P recv buffer
per microbatch and per direction (activations forward, gradients backward).
At large microbatch counts this dominates stage memory: Kimi-K3 at 2k rows x
GBS 4096 / mbs 2 / dp32 means 64 buffer sets per direction, ~52 GiB on middle
stages — more than the activations — and the shape OOMs during step 1 on
64 GB200 nodes (by only 56 MiB, measured).

Under 1F1B only num_stages - stage_index microbatches are in flight per
stage: a forward recv buffer is live from irecv-post until that chunk's
backward consumes the stage input (AC recompute reads it); a grad recv buffer
only until the chunk's backward returns. New module
components/distributed/pipelining/recv_buffer_pool.py rebinds the per-chunk
recv-info maps onto a ring of K = inflight + slack real buffer sets
(~45 GiB saved on middle stages at the shape above).

• Opt-in: distributed.pipeline.pp_recv_buffer_pool: true
  (+ pp_recv_buffer_pool_slack, default 2), threaded
  PipelineConfig -> AutoPipeline.build().
• Schedule-gated: only 1F1B (bounded in-flight depth proof). Other
  schedules log a warning and keep stock behavior — GPipe's unbounded
  in-flight depth silently corrupts gradients through reused buffers
  (verified negative on the CPU parity harness during development).
• torch-version adaptive, fail-open: adapts both the torch >= 2.13
  layout (_setup_forward/backward_recv_info) and the 2.12 line in current
  NGC containers (_prepare_forward/backward_infra); anything else logs a
  warning and leaves stock behavior.
• Gradient safety: grad accumulation on reused leaf buffers is impossible
  upstream — stage_backward() harvests input grads and sets
  val.grad = None per chunk.

Changelog

• nemo_automodel/components/distributed/pipelining/recv_buffer_pool.py (new):
  ring-pooled recv buffers, schedule_supports_recv_pool() allowlist
  ({"1f1b"}), torch-version probe with fail-open install.
• nemo_automodel/components/distributed/pipelining/config.py:
  PipelineConfig.pp_recv_buffer_pool (default False),
  pp_recv_buffer_pool_slack (default 2), documented in the class docstring.
• nemo_automodel/components/distributed/pipelining/autopipeline.py: thread the
  two fields into AutoPipeline.build() and install the pool before stages build.
• examples/llm_benchmark/kimi/kimi_k3_gb200_2k.yaml (new): the measured 2k-row
  champion config with the pool enabled; header states the measured shape.
• tests/unit_tests/distributed/pipelining/test_recv_buffer_pool.py (new): CPU
  gloo parity harness, schedule gating, fail-open probe.
• The perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 commits appear in this diff until perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 merges (stacked PR).

Validation

New benchmark config kimi_k3_gb200_2k.yaml — the pool unlocks 2k rows x
GBS 4096 at 64 nodes, the new best-efficiency K3 shape:

shape	tok/s/GPU	MFU	peak mem
1k x GBS4096 (kimi_k3_gb200.yaml + #3779 flags)	495.7	13.7%	~135 GiB
2k x GBS4096 + recv pool (this file)	615.9	17.1%	139.6 GiB

Mechanism: the workload is launch-bound at this scale, so doubling tokens
per row halves the per-token launch tax; K3's KDA linear attention does not
penalize longer rows. PP bubble 9.9% (64 microbatches). MFU basis
tok/s/GPU x 6 x 104e9 / 2.25e15. The file ships the exact measured
champion config (compile_norm off there; the 2k x compile_norm stack is
pending measurement and only changes numbers, not code).

Measurement provenance: the 17.1% figure was measured before #3792, when
from_config built a 56-head MLA / 18432-wide dense FFN model (~101.5B active
matmul parameters instead of 104.0B). On the checkpoint-aligned shape this
file builds today expect ~2% fewer tok/s at about the same MFU and ~1 GiB more
per GPU (44 GiB headroom measured). A 64-node re-run on the aligned shape is
queued (w53p rung = this exact file); numbers here and in the header will be
replaced when it lands. With #3792 the recipe prints MFU on the config-derived
basis via nemo_automodel/recipes/llm/benchmark.py --config ….

Correctness evidence:

• CPU gloo parity (in-PR unit test): pp4 x 1F1B x 16 microbatches x 3
  steps, stock vs pooled — loss trajectory and per-stage parameter sums
  bitwise identical. The test catches ring-too-small corruption because
  linear weight grads read the saved stage input, which IS the recv buffer.
• pp2 ring-activation gate at cluster (dev evidence): ring active
  (4/8 sets), finite-loss trajectory bitwise-identical over 8 steps,
  throughput neutral at small scale. Re-verified on the current PR tree
  with BenchmarkingRecipe (pp2×ep4, seq 2048, 1F1B): 12/12 iterations,
  recv_buffer_pool: installed (slack=2), recipe MFU within 1% of wall-clock.
• 256-GPU production evidence: w53 rescue run rc=0, mem 139.6 GiB
  (44 GiB headroom), loss pattern identical to non-pooled baseline regime.

Review pre-answers:

1. Interleaved PP / other schedules? Gated off by
   schedule_supports_recv_pool() — allowlist, currently {"1f1b"}.
2. torch private API drift? Version-adaptive probe with fail-open: install
   returns False and logs; stock allocation is kept, nothing breaks.
3. Why monkeypatch instead of subclass? Stages are constructed inside
   torch's pipeline_model flow; class-level rebinding before build is the
   minimal seam that needs no fork of the stage class. Idempotent + opt-in.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (test_recv_buffer_pool.py, listed above)
• Did you add or update any necessary documentation? (PipelineConfig docstring for both fields; yaml header documents the shape and provenance. No docs page enumerates PipelineConfig fields.)

Additional Information

• Part 2 of 2, stacked on perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779. Its three commits (34f28f9, 22be971,
  b5900fc) appear in this diff until perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 merges; review the commits
  after b5900fc. Merge order: perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 first.
• Rebased on main after fix(kimi-k3): correct MFU accounting (FLOPs formula, config defaults) #3792 (K3 FLOPs formula + KimiK3TextConfig
  defaults aligned with the checkpoint).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yisongbetter]

/ok to test 8a58634