Bound symmetric all-reduce scratch storage across input sizes

[Sunt-ing]

Description

Symmetric all-reduce retains a staging allocation for every distinct input size. A long-lived dynamic inference engine can therefore accumulate buffers for past prefill lengths, even though it only needs one size at a time.

Round the internal scratch allocation to a power-of-two capacity and reduce only the requested slice. Retaining the capacity buckets keeps old CUDA graph addresses valid while bounding their total storage to less than twice the largest bucket, per dtype, device and process group. The allocation helper's exact-size behavior and the existing zero-copy path remain unchanged.

This is an alternative implementation of the memory reuse proposed in #2037. It retains old allocations so captured graphs can continue to use them after capacity grows.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Reuse geometric scratch capacities inside symmetric all-reduce.
• Add a two-rank regression covering varying input sizes, bounded retained storage, and replay of previously captured graphs after growth.

Reproduction and validation

Base: eb3c860a4d3394b7eb5e21dba3021fa10782a52e.
Candidate: e860467c74fcdd5c6ed2324bee0ee346f4deb0ee.

In a TE-supported PyTorch development environment with two peer-accessible GPUs, save the attached patch as external-change.patch, then run:

git clone https://github.com/NVIDIA/TransformerEngine.git transformer-engine
cd transformer-engine
git checkout eb3c860a4d3394b7eb5e21dba3021fa10782a52e
git submodule update --init --recursive
git apply ../external-change.patch
export NCCL_HOME="$PWD/.validation/nccl"
make -C 3rdparty/nccl-extensions/third_party/nccl -j32 src.build \
  BUILDDIR="$NCCL_HOME" NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90"
export LD_LIBRARY_PATH="$NCCL_HOME/lib:${LD_LIBRARY_PATH:-}"
NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS=90 python -m pip install -e . -v --no-build-isolation --no-deps
python -m pytest -q tests/pytorch/distributed/test_symmetric_memory.py

To run the same test against the original implementation in a disposable checkout, retain the new test and restore only the production file:

git show eb3c860a4d3394b7eb5e21dba3021fa10782a52e:transformer_engine/pytorch/distributed.py > transformer_engine/pytorch/distributed.py
python -m pytest -q tests/pytorch/distributed/test_symmetric_memory.py

The original implementation fails the retained-capacity assertion after 32 different row counts. The test also checks numerical reductions and old graph replay after another capacity is allocated.

The source-built candidate passed the focused test (1 passed in 11.46s). The existing Linear sanity selection passed 486 cases; 954 unsupported parameter combinations were skipped by the upstream test. The Python L0 lint completed with Your code has been rated at 10.00/10; the license check passed.

The candidate's common and PyTorch native extensions were built from source for H200 (SM90), using the repository-pinned NCCL 2.30.7, CUDA 13.1 and PyTorch 2.11.0a0+eb65b36914.nv26.02. Import checks confirmed that Python and both native libraries came from this checkout. Use an isolated environment without an older TE core distribution; the optional DeepEP extension from the original image was excluded because it was linked against another NCCL version and is not used by this dense model. NCCL NVLS was disabled because of the environment's multicast setup. The commands above target SM90 and require the project's build dependencies.

Supplemental model evidence uses Qwen3-32B through Megatron dynamic inference at TP4. The original cache retains 48.75 MiB per rank across 16 natural-text prompt lengths; a second sweep over those lengths adds no further buffers. The candidate retains 15 MiB. All four ranks match the original symmetric-reduction implementation exactly for generated tokens and log probabilities across solo, mixed and reversed requests. No OOM was observed, and this experiment does not establish production shape frequencies.

The full native build also passed Qwen3-32B BF16 TP4 eager inference: all four ranks matched solo/mixed/reversed tokens and log probabilities, with 6912 symmetric reductions per rank. The repeated length scan retained 15 MiB per rank and did not grow on the second sweep. Model-level decode CUDA graphs exercised 42 graph steps per rank; outputs matched each other and the eager run exactly. FlashAttention 4.0.0b35 was used. These are supplemental integration results; the full upstream CI suite has not been run.

Checklist:

• I have read and followed the contributing guidelines.
• The functionality is complete.
• I have commented my code, particularly in hard-to-understand areas.
• I have made corresponding changes to the documentation.
• My changes generate no new warnings.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I, the PR author, have personally reviewed every line of this PR.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The changes appear safe to merge; no new actionable issues were found.

Summary

Bounds retained scratch storage by reusing capacity buckets while keeping old buffers alive for CUDA graphs.

• Symmetric all-reduce reuses scratch buffers across nearby input sizes.

The earlier, unnumbered finding about unsupported machines is fixed: the test checks symmetric-memory support and peer access before starting workers.

_{Reviews (3) · Last reviewed commit: "2x -> 25% memory overhead" · Reviewed by Greptile}

[ksivaman]

/te-ci pytorch L0

[ksivaman]

LGTM

[ksivaman]

Part of the greptile AI comment also applies here: #2037 (comment) where the workload could be using more memory with no changes from user side.