Skip to content
Markdown

PyTorch Symmetric Memory

Scope: torch.distributed._symmetric_memory, PyTorch's pluggable one-sided GPU-to-GPU communication primitive built on a shared symmetric buffer (same size, same slot on every rank), the rendezvous handshake that maps peer buffers into each rank's address space, and the put/get/multimem/all-to-all ops that run against it instead of a full NCCL collective. Covers the current alpha maturity label, the strict same-order-on-all-ranks rendezvous constraint, and the CUDA/NCCL/NVSHMEM backend requirements as stated in PyTorch's own documentation.

Reference templates on real, currently-shipping APIs (PyTorch 2.13 stable docs, matching the main/nightly docs as of this writing). Code below needs multiple CUDA GPUs with peer-to-peer or RDMA connectivity and is not executed here; the allocation-order invariant that governs rendezvous is validated with a runnable, asserted Python simulation later on this page. Pin the exact PyTorch version before trusting any signature: this module is explicitly under active development and PyTorch's own docs warn API changes are possible.

What it is

PyTorch Symmetric Memory (SymmMem) lets every rank in a process group allocate a buffer of identical size, then perform a collective rendezvous that maps each peer's buffer into the local rank's address space. After rendezvous, a rank can read or write another rank's memory directly, one-sided, without routing the operation through a full NCCL collective launched and synchronized from the host.

The two-call programming model, straight from PyTorch's documentation:

import torch.distributed._symmetric_memory as symm_mem

t = symm_mem.empty(128, device=torch.device("cuda", rank))
hdl = symm_mem.rendezvous(t, group)

symm_mem.empty allocates a tensor backed by a symmetric memory allocation. symm_mem.rendezvous establishes the handshake with peers in the group and returns a handle (hdl) exposing buffer_ptrs (pointers to every peer's symmetric buffer), multicast_ptr (if the hardware supports multicast, e.g. NVLink SHARP), and signal_pad_ptrs (per-allocation synchronization pads used by device-side barriers). PyTorch's documentation states plainly: "The empty and rendezvous functions must be called in the same order on all ranks in the group." That is not a style suggestion, it is the mechanism: rendezvous exchanges one metadata record per call through a collective store, so a rank that calls empty/rendezvous in a different order pairs its N-th exchange with the wrong peer record.

Once rendezvoused, operations run either through pre-built ops in the torch.ops.symm_mem namespace (one_shot_all_reduce, two_shot_all_reduce_, multimem_all_reduce_, multimem_all_gather_out, all_to_all_vdev, tile_reduce, and others), or through hand-written CUDA or Triton kernels that dereference hdl.buffer_ptrs directly and synchronize with the signal pads. A small one-sided get is also exposed directly from Python:

src = symm_mem.empty(1024, device=device)
hdl = symm_mem.rendezvous(src, group)

if dist.get_rank(group) == 0:
    dst = torch.empty((512,), device=device)
    symm_mem.get(dst, hdl, peer=1, offset=512)   # copy peer 1's buffer[512:1024] into dst

SymmMem supports multiple backends selected with symm_mem.set_backend(...): CUDA (P2P via NVLink/PCIe), NCCL (NCCL's own symmetric-memory-aware kernels, including "Copy Engine Collectives" that offload data movement to DMA engines instead of SMs), and NVSHMEM (for cross-node scale-out via NVSHMEM's put/get, exposed to Triton kernels through torch.distributed._symmetric_memory._nvshmem_triton). Which backend is active changes which ops and cross-node behavior are available, so treat set_backend as part of the API surface to pin, not a hidden default.

Why use it

  • Removes the host from the critical path for fine-grained exchanges. A full NCCL collective is a host-launched, synchronized operation; SymmMem's one-sided ops (once rendezvoused) let a kernel read or write a peer's memory as part of its own execution, without a separate collective launch per exchange.
  • SM offload via Copy Engine Collectives. With the NCCL backend, NCCL_CTA_POLICY_ZERO plus symmetric memory tensors moves standard collectives (all_gather_single, all_to_all_single) onto the GPU's copy (DMA) engines instead of streaming multiprocessors, freeing SMs for compute during the same window: the same motivation as comms/compute overlap, but at the collective-implementation level rather than the scheduling level.
  • Free numerical-accuracy upgrade for some collectives. When tensors are allocated with symmetric memory and reduced with the NCCL backend, NCCL's symmetric kernels internally accumulate BF16/FP16 inputs in FP32 before writing back BF16/FP16 outputs, for all_reduce and reduce_scatter within the supported NVLink (and, as of NCCL 2.29, network) domain, with no code change beyond allocating the input as symmetric.
  • Building block for custom fused collectives and MoE dispatch. all_to_all_vdev / all_to_all_vdev_2d (device-side split/offset all-to-all-v over NVSHMEM) and tile_reduce / multi_root_tile_reduce (partial-tile reductions to arbitrary roots) are primitives for hand-rolled MoE token dispatch/combine and custom fused kernels that do not map onto a standard NCCL collective shape.
  • Triton-native. Symmetric buffers, signal pads, and NVSHMEM put/get are all reachable from @triton.jit kernels, so a compute kernel and its communication can live in one kernel body without dropping to raw CUDA.

When to use it (and when not)

Use PyTorch Symmetric Memory when:

  • You are writing a custom fused kernel that needs to read or write a peer GPU's memory as part of its own compute, not as a separate collective step (custom all-reduce variants, tile-wise partial reductions, MoE dispatch/combine shaped as all_to_all_vdev).
  • You want an existing NCCL collective (all_gather_single, all_to_all_single, all_reduce, reduce_scatter) to run on copy engines instead of SMs, and you can meet the NCCL 2.28+ / P2P / zero-CTA-policy requirements below.
  • You are building or extending a research system (an inference engine, an MoE framework, a custom parallelism strategy) where the standard collective shapes in NCCL collectives and algorithm selection do not fit and you are willing to own correctness at the pointer level.

Do not use it when:

  • You need a stable, broadly-supported production primitive today. PyTorch's own documentation states torch.distributed._symmetric_memory "is currently in alpha state and under development. API changes may be possible." Treat every function name and signature in this page as subject to change on upgrade; do not lock a production system to the current API shape without a re-validation step per PyTorch release.
  • A standard bucketed all-reduce or all-gather already saturates your interconnect. NCCL's default collectives (used by DDP and FSDP) are topology-aware and already well overlapped; reach for SymmMem only when you have a specific, measured reason a host-launched collective is the bottleneck, not by default.
  • Your team cannot enforce strict SPMD-style call ordering across every rank (see Failure modes). A framework with rank-dependent control flow around allocation is a bad fit until that is refactored.
  • You need this to run identically across arbitrary GPU vendors/backends. The backends documented are CUDA (NVLink/PCIe P2P), NCCL, and NVSHMEM; each has its own hardware and version prerequisites, none of which are guaranteed present on an arbitrary cluster.

Architecture

flowchart TB
    subgraph Alloc["Symmetric allocation (same call, every rank)"]
        E0["Rank 0: symm_mem.empty(128)"]
        E1["Rank 1: symm_mem.empty(128)"]
        E2["Rank 2: symm_mem.empty(128)"]
    end

    Alloc --> RDV["Collective rendezvous(t, group)<br/>metadata exchange via TCPStore / PG allgather<br/>MUST be called in the same order on all ranks"]

    RDV --> HDL0["hdl on rank 0:<br/>buffer_ptrs, multicast_ptr, signal_pad_ptrs"]
    RDV --> HDL1["hdl on rank 1: same fields"]
    RDV --> HDL2["hdl on rank 2: same fields"]

    HDL0 --> OS["One-sided put/get or multimem op<br/>(symm_mem.get, one_shot_all_reduce, tile_reduce...)<br/>issued from device / current stream, no per-op host collective launch"]
    HDL1 --> OS
    HDL2 --> OS

    NCCL["Traditional NCCL collective<br/>(all_reduce / all_gather / reduce_scatter)<br/>host launches op on every rank, ranks synchronize,<br/>ring/tree algorithm moves data, all ranks return together"]

    OS -.->|"same logical result, different path"| CMP{"Compare"}
    NCCL -.->|"same logical result, different path"| CMP

The left path (symmetric allocation, rendezvous, one-sided op) is SymmMem: allocation and rendezvous are collective (every rank must participate, in the same order), but the data-movement op itself can be one-sided and device-initiated. The right path is a conventional NCCL collective: still collective end to end, host-launched and synchronized on every call. SymmMem's rendezvous replaces per-call collective overhead with a one-time handshake; after that, individual put/get/multimem ops can be issued without a fresh host-side collective launch each time.

How to use it

Per the current PyTorch documentation (docs.pytorch.org/docs/2.13/symmetric_memory.html, matching main). Requires multiple CUDA GPUs with peer-to-peer access (or a supported RDMA/NVSHMEM path for cross-node), and PyTorch's distributed runtime initialized. Not executed on hardware here.

# symm_mem_basic.py  — reference template, needs 2+ CUDA GPUs, run with torchrun
import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem

dist.init_process_group()
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device("cuda", rank)

# 1. Allocate a symmetric tensor: same call, same size, on every rank.
t = symm_mem.empty(4096, device=device)

# 2. Rendezvous: collective handshake exchanging peer buffer/multicast/signal-pad
#    pointers. Must happen in the same order on every rank in `group`.
hdl = symm_mem.rendezvous(t, dist.group.WORLD)

# 3a. A prebuilt one-sided op, issued on the current CUDA stream.
torch.ops.symm_mem.one_shot_all_reduce(t, "sum", dist.group.WORLD.group_name)

# 3b. A basic one-sided get: copy a slice of a peer's symmetric buffer locally.
if rank == 0 and world_size > 1:
    dst = torch.empty((512,), device=device)
    symm_mem.get(dst, hdl, peer=1, offset=0)   # dst <- peer 1's t[0:512]

To use the NCCL-backend Copy Engine Collectives path (offloads standard collectives to DMA engines, requires NCCL 2.28+):

# copy_engine_collectives.py  — reference template, needs NCCL 2.28+, not executed here
opts = dist.ProcessGroupNCCL.Options()
opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
dist.init_process_group(backend="nccl", pg_options=opts, device_id=device)

symm_mem.set_backend("NCCL")
group_name = dist.group.WORLD.group_name

inp = symm_mem.empty(1024 * 1024, device=device)
out = symm_mem.empty(1024 * 1024 * world_size, device=device)
symm_mem.rendezvous(inp, group=group_name)
symm_mem.rendezvous(out, group=group_name)

# Runs on copy engines instead of SMs because inp/out are symmetric and the PG
# is configured with the zero-CTA policy. As of NCCL 2.28, CE collectives
# cannot run on the default stream, so async_op=True (or a side stream) is required.
work = dist.all_gather_single(out, inp, async_op=True)
work.wait()

At world sizes beyond a single NVLink domain (documented range: 8-72 ranks per domain, tested up to 100k total ranks), the default TCPStore-based rendezvous becomes a bottleneck (roughly 3.6 s at 10k ranks with 72-rank groups, roughly 36 s at 100k ranks, per PyTorch's documented measurements). PyTorch documents a faster path using the process group's own NCCL allgather:

opts = dist.ProcessGroupNCCL.Options()
opts.use_pg_for_symm_mem_rendezvous = True
pg = dist.new_group(ranks, pg_options=opts)

t = symm_mem.empty(size, device=device)
hdl = symm_mem.rendezvous(t, group=pg)

Failure modes

  • Mismatched allocation order across ranks. Rendezvous exchanges metadata one call at a time; if rank 0 calls empty/rendezvous for buffer A then buffer B while rank 1 calls them for B then A, each rank's N-th exchange pairs with a different logical buffer on its peer. PyTorch's documentation is explicit that empty and rendezvous "must be called in the same order on all ranks in the group." Depending on how the mismatch lands, the result is a hang (ranks waiting on store keys that never arrive in the expected order) or silent corruption (a rank ends up holding pointers to the wrong peer buffer and reads/writes garbage without any error). This is a strict SPMD constraint: conditional, rank-dependent allocation code is the most common way to violate it. See the executable simulation below for how the same shape of bug looks as a plain order-matching check.
  • Unsupported backend/CUDA/NCCL/hardware combination. Copy Engine Collectives require NCCL 2.28 or later and GPUs with peer-to-peer access, plus the process group configured with NCCL_CTA_POLICY_ZERO (or NCCL_CTA_POLICY=2); as of NCCL 2.28, CE collectives cannot run on the default stream, only on an explicit async/side stream. Higher-precision (FP32-accumulate) reduction over symmetric tensors is scoped to all_reduce/reduce_scatter, within the NVLink domain as of PyTorch 2.9 (NCCL 2.27), extending to NVLink+network for reduce_scatter as of PyTorch 2.11 (NCCL 2.29); other collectives and other domains are silently unaffected, not accelerated. multimem_all_reduce_ and multimem_all_gather_out require hardware multicast support, on NVIDIA GPUs specifically NVLink SHARP; running them without that hardware is unsupported. MemPool support (allocate ordinary-looking tensors from a symmetric pool) is documented as available for the CUDA and NVSHMEM backends as of PyTorch 2.11, with NCCL backend MemPool support still "in progress": check the installed version's docs before assuming a backend/feature pairing works.
  • Treating alpha as production-stable. PyTorch's own documentation states the module "is currently in alpha state and under development. API changes may be possible," and this note is present in both the stable (2.13) and main/nightly documentation builds at time of writing, i.e. it is not a nightly-only caveat that will disappear once you pin a stable release. Function names, return types, and default backends can change between releases without the deprecation cycle a stable API would get; re-validate on every PyTorch upgrade rather than freezing a version and assuming forward compatibility.
  • Rendezvous cost at scale, mistaken for a hang. The default TCPStore-backed rendezvous is a real, non-trivial exchange (documented at roughly 3.6-36 seconds moving from 10k to 100k total ranks); a slow rendezvous at large world size can look like a stall rather than expected overhead. Use use_pg_for_symm_mem_rendezvous at that scale rather than assuming a fixed cost.

Executable check of the ordering invariant (pure Python, no torch/GPU needed; models the rendezvous handshake as a per-step, per-rank tag comparison, which is the same shape of check a strict SPMD rendezvous protocol performs):

def simulate_rendezvous(alloc_orders):
    """alloc_orders: list of per-rank lists of allocation tags, in call order.
    A real rendezvous exchanges one metadata record per empty()/rendezvous() call
    via a collective store op; every rank must issue the same sequence of calls
    (SPMD) or the N-th store get on one rank pairs with the wrong record on
    another. Returns (ok, mismatch_step): the first step, if any, where ranks
    disagree on which buffer they are rendezvousing.
    """
    n_ranks = len(alloc_orders)
    assert n_ranks >= 2, "need at least 2 ranks to rendezvous"
    lengths = {len(o) for o in alloc_orders}
    assert len(lengths) == 1, "ranks called a different number of allocations, definite hang"
    n_steps = lengths.pop()
    for step in range(n_steps):
        tags_this_step = {alloc_orders[r][step] for r in range(n_ranks)}
        if len(tags_this_step) != 1:
            return False, step
    return True, None


# happy path: every rank allocates the same three buffers in the same order
matched = [
    ["expert_buf", "grad_buf", "kv_buf"],
    ["expert_buf", "grad_buf", "kv_buf"],
    ["expert_buf", "grad_buf", "kv_buf"],
]
ok, step = simulate_rendezvous(matched)
assert ok and step is None, (ok, step)

# adversarial: rank 1 swaps the order of its second and third allocation
mismatched = [
    ["expert_buf", "grad_buf", "kv_buf"],
    ["expert_buf", "kv_buf", "grad_buf"],
    ["expert_buf", "grad_buf", "kv_buf"],
]
ok, step = simulate_rendezvous(mismatched)
assert not ok and step == 1, (ok, step)   # detected at the first differing step

# edge case: rank-dependent control flow skips an allocation entirely
conditional_alloc = [
    ["expert_buf", "grad_buf"],
    ["expert_buf", "grad_buf", "kv_buf"],
]
try:
    simulate_rendezvous(conditional_alloc)
    raise SystemExit("expected AssertionError for mismatched call counts")
except AssertionError as e:
    assert "different number" in str(e)

print("rendezvous order-invariant simulation: all asserts passed")

Run with python3; all three assertions (matched, mismatched-at-step-1, and the differing-call-count edge case) pass. This models the shape of the invariant PyTorch's rendezvous depends on; it is not the actual PyTorch/NCCL rendezvous implementation, which additionally handles the store-based metadata exchange, memory mapping, and CUDA IPC handles.

References

  • PyTorch Symmetric Memory, official documentation (stable, PyTorch 2.13): https://docs.pytorch.org/docs/stable/symmetric_memory.html
  • PyTorch Symmetric Memory, official documentation (main/nightly, same content as stable): https://docs.pytorch.org/docs/main/symmetric_memory.html
  • PyTorch Symmetric Memory raw source (exact function signatures, the alpha-state note, and the ordering constraint): https://docs.pytorch.org/docs/main/_sources/symmetric_memory.md.txt
  • GitHub RFC, "[RFC] Support Symmetric Memory programming," pytorch/pytorch#163666 (closed, milestone 2.9.0): https://github.com/pytorch/pytorch/issues/163666
  • GitHub RFC, "[RFC] support symmetric memory in torch.compile," pytorch/pytorch#162859: https://github.com/pytorch/pytorch/issues/162859
  • GitHub issue tracking a real-world rendezvous failure on DGX B200, pytorch/pytorch#167537 (illustrates that rendezvous failures are an active, current-version concern, not a hypothetical): https://github.com/pytorch/pytorch/issues/167537
  • Meta/PyTorch kraken examples of symmetric-memory Triton kernels: https://github.com/meta-pytorch/kraken/blob/main/kraken

Related: NVSHMEM: GPU-Initiated Communication · NCCL Collectives and Algorithm Selection · Communication-Computation Overlap · FSDP · Glossary