Skip to content
Markdown

PyTorch Symmetric Memory

Scope: torch.distributed._symmetric_memory, PyTorch's alpha API for symmetric GPU allocations, one-sided peer access, and device-oriented collective kernels. Covers rendezvous ordering, backend requirements, and the distinction between one-sided get and collective torch.ops.symm_mem operations.

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.

Independently verified here, one level below GPU execution: torch==2.13.0+cpu was installed in a clean venv (no GPU needed for this check) and torch.distributed._symmetric_memory was introspected directly via inspect.signature and inspect.getsource, not re-read from documentation. This caught a real error the documentation-sourced draft of this page carried: symm_mem.get(dst, hdl, peer=1, offset=...) does not exist, at the module level or as a method, in this PyTorch build. The real one-sided read is a method on the rendezvous handle, hdl.get_remote_tensor(peer: int, sizes: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor (confirmed from its C++-bound docstring), alongside hdl.get_buffer(rank, sizes, dtype, storage_offset=0) for viewing a local rank's own symmetric buffer and hdl.put_signal/hdl.wait_signal/hdl.barrier for the application-level readiness synchronization a one-sided read needs and this page's Failure modes section already warns is the caller's responsibility. empty(*size, dtype=None, device=None), rendezvous(tensor, group), and set_backend(name: Literal['NVSHMEM','CUDA','NCCL']) all matched their documented signatures exactly. torch.ops.symm_mem (one_shot_all_reduce and the rest) is populated only by a CUDA-enabled build registering its kernels at import time; on the CPU-only wheel used for this check the namespace is empty, which is expected, not a defect, and does not confirm or refute those op signatures. No rendezvous, peer memory access, or collective was run; this is a source-level API check, not a multi-GPU execution.

Re-verified independently on 2026-07-17 in a separate clean venv (torch==2.13.0+cpu, installed fresh via pip install --index-url https://download.pytorch.org/whl/cpu torch, no GPU present in this environment): inspect.getdoc/inspect.signature on the installed _SymmetricMemory class reproduced the exact same handle methods and signatures described above -- get_remote_tensor(peer, sizes, dtype), get_buffer(rank, sizes, dtype, storage_offset=0), put_signal(dst_rank, channel=0, timeout_ms=0), wait_signal(src_rank, channel=0, timeout_ms=0), barrier(channel=0, timeout_ms=0) -- plus a separate module-level put_signal(src, hdl, peer)/wait_signal(hdl, peer) pair (data-plus-signal in one call) that this page does not otherwise use. torch.ops.symm_mem was confirmed empty on this CPU build again, consistent with the prior finding. This pass also located and fixed a second instance of the one-sided-read defect the prior pass left open: the "How to use it" example's step 3b called hdl.get_remote_tensor with a comment admitting it "omits [readiness sync] for brevity," which is exactly the failure mode this page's own Failure modes section warns about; it now pairs put_signal/wait_signal the same way the example earlier on this page already does, so as written it would not read before the peer has produced data if actually run under torchrun with 2 processes. As before, no rendezvous, peer memory access, or collective was executed; no multi-GPU hardware is available in this environment, so this remains a source-level API check and a logical-correctness fix, not new execution evidence.

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 peer buffers into each rank's address space. After rendezvous, a custom kernel can access peer memory and the handle's get_remote_tensor method can copy from one peer. Operations such as one_shot_all_reduce remain collectives even though their implementation uses symmetric mappings.

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 peer handshake and returns a handle (hdl) exposing buffer_ptrs, multicast_ptr when supported, and signal_pad_ptrs. PyTorch requires empty and rendezvous to be called in the same order on every rank in the group; rank-dependent allocation order violates that documented collective contract.

Once rendezvoused, collective operations run through the torch.ops.symm_mem namespace (one_shot_all_reduce, two_shot_all_reduce_, multimem_all_reduce_, all_to_all_vdev, tile_reduce, and others), or through custom CUDA or Triton kernels that dereference hdl.buffer_ptrs and synchronize with signal pads. One-sided peer reads are methods on the rendezvous handle itself, not a standalone module-level function; hdl.get_remote_tensor reads a peer's buffer, and hdl.wait_signal/hdl.put_signal are the handle's own readiness synchronization, verified directly from the installed package's docstrings (inspect.getsource, not the docs site):

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

if dist.get_rank(group) == 1:
    # Producer side, shown explicitly: an earlier revision of this example
    # omitted this branch entirely, so wait_signal below had no peer that
    # could ever unblock it. Fill this rank's own symmetric buffer, THEN
    # signal rank 0 that it is safe to read -- skip this call and rank 0's
    # wait_signal blocks (or times out) with no producer to unblock it.
    src.fill_(compute_value())
    hdl.put_signal(dst_rank=0)

if dist.get_rank(group) == 0:
    # Consumer side: hdl.wait_signal(src_rank, channel=0, timeout_ms=...) blocks
    # until peer 1 calls hdl.put_signal(dst_rank=0, ...) above; skipping this
    # is exactly the "read before the peer has produced it" failure mode below.
    hdl.wait_signal(src_rank=1)
    dst = hdl.get_remote_tensor(peer=1, sizes=(512,), dtype=src.dtype)   # peer 1's buffer, as a tensor

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

  • Supports device-oriented fine-grained exchanges. Custom CUDA or Triton kernels can dereference mapped peer buffers and use signal pads for synchronization. The handle's get_remote_tensor/get_buffer reads are one-sided; all-reduce, all-to-all, and tile-reduce operations still require their documented collective participation.
  • 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.
  • Higher-precision accumulation for supported reductions. NCCL symmetric kernels can accumulate BF16/FP16 inputs in FP32 before casting outputs back for supported all_reduce and reduce_scatter paths. This does not apply to every symmetric-memory operation; verify the installed PyTorch and NCCL release.
  • 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["Peer access or SymmMem operation<br/>get is one-sided; all-reduce and tile-reduce are collective"]
    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

Symmetric allocation and rendezvous are collective and ordered. After rendezvous, get is one-sided, custom kernels may access mapped peer buffers, and SymmMem all-reduce/all-to-all/tile-reduce operations retain collective semantics. A conventional NCCL collective is shown for comparison, not as proof that every SymmMem operation is one-sided.

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 os

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()
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device("cuda", local_rank)
torch.cuda.set_device(device)

# 1. Allocate a symmetric tensor: same call, same size, on every rank.
t = symm_mem.empty(4096, device=device)
t.fill_(float(rank + 1))   # deterministic, rank-distinguishable content, not garbage

# 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 collective SymmMem op. Every rank in the group must call it.
reduced = torch.ops.symm_mem.one_shot_all_reduce(
    t, "sum", dist.group.WORLD.group_name)
expected = sum(range(1, world_size + 1))   # 1 + 2 + ... + world_size
assert torch.allclose(reduced, torch.full_like(reduced, float(expected))), \
    (reduced, expected)  # only meaningful once run on real hardware; not executed here

# 3b. A basic one-sided read: hdl.get_remote_tensor is a method on the handle,
#     not a standalone symm_mem function. One-sided reads need explicit
#     readiness sync -- skip it and this is exactly the "read before the
#     peer has produced it" failure mode in Failure modes below. Rank 1
#     signals once its buffer is safe to read; rank 0 waits for that signal
#     before reading, same pairing as the put_signal/wait_signal example above.
if rank == 1 and world_size > 1:
    hdl.put_signal(dst_rank=0)   # this rank's t is ready; unblocks rank 0's wait below

if rank == 0 and world_size > 1:
    hdl.wait_signal(src_rank=1)   # blocks until rank 1's put_signal above
    dst = hdl.get_remote_tensor(peer=1, sizes=(512,), dtype=t.dtype)  # 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)

How to develop with it

Failure cases to encode in tests

  • Mismatched allocation order across ranks. If rank 0 allocates buffer A then B while rank 1 allocates B then A, the ranks violate PyTorch's same-order requirement. Test rank-dependent branches and loops explicitly; do not infer a particular failure symptom from an alpha implementation.
  • Unsupported backend/CUDA/NCCL/hardware combination. Copy Engine Collectives require NCCL 2.28 or later, peer-to-peer access, zero-CTA policy, symmetric allocations, rendezvous, and a non-default stream path. Higher-precision reduction, multimem, and MemPool each have narrower release and hardware scopes; verify the installed release rather than extrapolating from another operation.
  • 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.

This pure-Python test checks the documented ordering invariant. It does not simulate PyTorch's metadata exchange or failure behavior:

def simulate_rendezvous(alloc_orders):
    """Return the first step where per-rank allocation tags differ.

    This validates SPMD call ordering, not the PyTorch rendezvous implementation.
    """
    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")

Executed output (.venv/bin/python, this exact script, 2026-07-17):

rendezvous order-invariant simulation: all asserts passed

The matched sequence, reordered sequence, and differing-call-count case all pass their asserts; the single print line above is real captured stdout, not a paraphrase.

How to maintain it

Pin PyTorch and NCCL together. On each upgrade, re-check the alpha API signatures, supported backends, CE-collective scope, stream restriction, and return schemas the way this page's own verification did: pip install torch --index-url https://download.pytorch.org/whl/cpu into a scratch venv (no GPU needed for this check) and inspect.signature/inspect.getsource the real torch.distributed._symmetric_memory module and its _SymmetricMemory class, rather than trusting the docs site or a prior release's remembered shape; this is exactly the check that caught symm_mem.get never having existed as documented here. Run a two-rank smoke test for every operation used by the application and a negative test that deliberately changes allocation order on one rank under a timeout.

How to run it in production

Set the CUDA device from LOCAL_RANK, then initialize one process per GPU. Record rendezvous latency separately from transfer or collective latency: the default TCPStore-backed rendezvous measures in the seconds range at scale (roughly 3.6 s at 10k ranks, roughly 36 s at 100k ranks, per PyTorch's own documented figures cited above), so a rendezvous that takes that long is expected overhead rather than a stall; switch to use_pg_for_symm_mem_rendezvous before that startup cost matters to the deployment. Gate rollout on numerical comparison with the standard collective, timeout every collective sequence, and retain a standard NCCL implementation as the rollback path. Treat a rendezvous that hangs or raises, rather than one that completes slowly, as the separate failure class documented in pytorch/pytorch#167537 (Failure modes above): keep TORCH_SYMM_MEM_DISABLE_MULTICAST=1 and a node-reboot runbook entry available for on-call before assuming an application-level bug.

PyTorch's symmetric-memory documentation (the raw source cited in References) does not publish a per-topology support matrix, a metrics or telemetry API, or a documented fallback ladder beyond the use_pg_for_symm_mem_rendezvous rendezvous path and the multicast-disable switch above; this was checked directly against that source, not inferred from its absence on this page. Do not present a more detailed operational matrix than that as vendor-documented, and build any additional dashboarding (rendezvous duration, collective timeout counts, fallback-path activation) as operator-side instrumentation around these calls, not as something SymmMem exposes itself.

Failure modes

  • Rank-dependent allocation or rendezvous order violates the collective contract.
  • A one-sided get_remote_tensor/get_buffer read without pairing wait_signal/put_signal (or an equivalent barrier) can read data before the peer has produced it; these are real methods on the handle, verified above, not a documentation convenience this page invented.
  • Citing a symmetric-memory API surface from memory or from a prior PyTorch release instead of introspecting the installed package: this page shipped symm_mem.get(dst, hdl, peer, offset) as a real function for some time, and it is not one; re-verify against inspect.getsource/inspect.signature on the exact pinned version before trusting any signature on this page, including the corrected ones.
  • Treating one_shot_all_reduce, tile-reduce, or all-to-all as one-sided omits required participants.
  • Unsupported backend, topology, stream, or NCCL-version combinations can fail or miss the intended optimization.
  • A rendezvous that hangs or raises rather than completing slowly is a distinct, real failure class, not a hypothetical: pytorch/pytorch#167537 (cited in References) reports rendezvous() failing with RuntimeError: CUDA driver error: system not yet initialized, or hanging outright, on DGX B200 (Blackwell) systems. PyTorch maintainers traced it to multicast rendezvous interacting with a bad fabric-manager/driver/kernel state rather than an application bug; the only reported workarounds at issue-close time were setting TORCH_SYMM_MEM_DISABLE_MULTICAST=1 (which disables the multicast-dependent ops, including some torch.ops.symm_mem reductions) or rebooting the affected node.
  • The API is alpha; signatures and backend coverage can change between PyTorch releases.

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