Repository navigation
[experimental] PyTorch symmetric memory as an allocation provider - #551
nirvedhmeshram wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
Adds an experimental allocation provider backed by torch.distributed._symmetric_memory, enabling Iris device kernels (e.g., iris.store) to operate on torch-allocated symmetric memory without changing kernel code.
Changes:
- Introduces
TorchSymmMemProviderto allocate torch symmetric tensors and expose Iris-compatible per-allocation peer-base tables / descriptors. - Extracts
SymmetricAddressMapinto a sharediris/experimental/symmetric_memory.pymodule and re-exports it from the rocSHMEM provider. - Adds distributed unit tests and updates experimental docs describing the new provider, constraints, and run instructions.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
iris/experimental/torch_symmmem_provider.py |
New torch symmetric-memory provider implementation producing Iris-compatible peer-base tables and descriptors. |
iris/experimental/symmetric_memory.py |
New shared SymmetricAddressMap descriptor used by multiple providers. |
iris/experimental/rocshmem_provider.py |
Switches to importing SymmetricAddressMap from the shared module and re-exports it for backwards-compatible imports. |
tests/unittests/test_torch_symmmem_provider.py |
New distributed tests validating invariants, per-allocation mapping, reachability, and iris.store integration. |
iris/experimental/README.md |
Documents the torch symmetric memory provider behavior, limitations, and verification steps. |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| is reused. Rendezvousing again here would be a collective call made from | ||
| whichever rank happened to ask. | ||
| """ | ||
| handle = self._handles.get(tensor.data_ptr()) |
There was a problem hiding this comment.
Do we want the providers to support tensors not allocated through the provider? I am thinking of a pattern like:
t = symm_mem.empty(*shape, dtype=dtype, device=self._device)
h = symm_mem.rendezvous(tensor, group=self._group)
provider.symmetric_address_map(t, h)
We can do a different PR later if needed. @drprajap let me know your thoughts too.
There was a problem hiding this comment.
yes, we want to support already rendezvoused pytorch tensors, this would enhance support for usecase where we already own PyTorch symmetric tensor, which may have already called symm_mem.empty() and collectively rendezvoused it before Iris is involved. We should use exisitng peer pointers in that case instead of allocating second tensor and copying data.
okay to follow-up with separate PR
| is reused. Rendezvousing again here would be a collective call made from | ||
| whichever rank happened to ask. | ||
| """ | ||
| handle = self._handles.get(tensor.data_ptr()) |
There was a problem hiding this comment.
yes, we want to support already rendezvoused pytorch tensors, this would enhance support for usecase where we already own PyTorch symmetric tensor, which may have already called symm_mem.empty() and collectively rendezvoused it before Iris is involved. We should use exisitng peer pointers in that case instead of allocating second tensor and copying data.
okay to follow-up with separate PR
| | `2.10.0+rocm7.2.1` | works | | ||
| | `2.13.0+rocm7.2` | works | | ||
| | `2.9.1+rocm7.1` | no — lacks `set_signal_pad_size`, and `rendezvous` fails with `HIP error: invalid argument` | | ||
| | `2.13.0+rocm7.15.0a` (nightly) | no — `hipErrorOutOfMemory` on a 4 KB allocation | |
There was a problem hiding this comment.
My recent check on the torch 2.15.0.dev20260930+rocm10.0 nightly had symmetric memory working, but only with set_backend("NVSHMEM"), which is rocSHMEM on ROCm
There was a problem hiding this comment.
Thanks — that's a build we couldn't run, and it shows the README went too far. "Do not call set_backend" and "forcing any other name fails at allocation" generalised from forcing NCCL on the 2.13/rocm7.2 wheels; we never tried NVSHMEM.
The provider never chose a backend, so b0797f8 just says what is true: the backend is the caller's choice, and the provider reads buffer_ptrs off whatever handle comes back. The one real trap stays: get_backend() reports 'CUDA', but set_backend('CUDA') is rejected. Your build is now a row in the table, marked "reported, not run here", and the example shows set_backend("NVSHMEM") as the opt-in. I also scoped two related claims to the default backend: "offsets are not shared between allocations" and the intra-node scope. A rocSHMEM-backed heap would plausibly behave like the rocSHMEM provider on both.
One question: did you run these provider tests on that nightly, or a standalone allocate? If the tests pass there, I'd drop the qualifier.
|
LGTM overall, correctness issue I mentioned before, I checked locally and it is not applicable, I thought handle may have stale address and didn't know holding handle doesn't just keep memory alive but keeps address range mapped as well. In PyTorch allocator,this handle would hold current ranks' own allocation, and that only would get unmapped when last reference goes away. so we will never run into stale handle, or pointing to different tensor. Nit: one thing to clarify in the comments from lifetime standpoint, free() does more than forget the handle. The stored rendezvous handle keeps the allocation and its peer mappings alive, so memory isn't released until free() is called (or the provider is dropped). Could the docstring and README say callers should call free()? |
A second provider needs the same descriptor, and importing it from rocshmem_provider would make that provider depend on rocshmem4py. The new module imports nothing but torch, so no provider drags in another's runtime. Re-exported from rocshmem_provider so existing imports keep working.
Lets Iris device kernels operate on tensors from torch.distributed._symmetric_memory. No Iris device code changes: iris.store, load and copy translate against a peer-base table, and torch's rendezvous handle exposes buffer_ptrs, which already is that table with the local rank's entry equal to the tensor's own data_ptr. The provider reads it off the handle and checks the invariant rather than computing anything. Unlike rocSHMEM, each allocation is a separate IPC mapping rather than a window onto one linear heap, so peer offsets are not shared between allocations and a table must only translate pointers into the allocation it describes. symmetric_address_map refuses tensors this provider did not allocate: rendezvous is collective, so re-deriving a handle on demand would hang the ranks that did not ask. Validated on gfx950 at 2, 4 and 8 ranks: 5 passed.
Covers that it needs no install beyond a torch build whose symmetric memory can allocate, which torch versions those are, why the tests probe by allocating rather than importing, the set_backend trap, and how its per-allocation mapping differs from rocSHMEM's uniform heap offsets.
…c all ranks Rename torch_symmmem_provider to torch_symm_mem_provider. The triple m reads as a typo; torch itself abbreviates torch.distributed._symmetric_memory as symm_mem, so follow that rather than the symmem spelling suggested in review. Validate that handle.buffer_ptrs has one entry per rank. Iris indexes this table by rank inside the kernel, so a short table reads past the end of the allocation instead of raising. Observed builds return exactly world_size ints, so this is insurance rather than a fix. Synchronise on every rank after the barrier in the store test, not only on the writer before it. The barrier orders the ranks; it does not drain the reading rank's stream, so inbound writes were not strictly ordered against work queued next on that rank.
…was measured Register the torch provider in test_provider_unified.py (ROCm#552). Its probe and skips move into a session fixture in conftest so both modules share them, and the two torch tests the unified test now covers -- the table invariant and per-allocation anchoring -- are dropped. The backend guidance overreached. "Do not call set_backend" and "forcing any other name fails at allocation" generalised from forcing NCCL on 2.13 builds; on the 2.15 rocm10.0 nightly symmetric memory reportedly works only with set_backend("NVSHMEM"), i.e. rocSHMEM. The provider never chose a backend anyway, so the docs now say the choice is the caller's, keep the one real trap (set_backend('CUDA') is rejected even though get_backend reports it), and record the nightly report as reported, not measured. Likewise "peer offsets are not shared between allocations" is a property of the default IPC backend, not of torch symmetric memory; reworded as "need not be". Fixes the README saying allocation is collective while the docstring said it is local: only rendezvous is, on the default backend. The module docstring is cut to the points that matter when reading the code; the rest already lives in the README. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
ad31c4c to
b0797f8
Compare
Summary
Adds
torch.distributed._symmetric_memoryas a second allocation provider, alongside the rocSHMEM one from #550. No changes to Iris kernels are needed.Motivation
#546 lists Torch Symmetric Memory and rocSHMEM as target providers of the same allocator-agnostic boundary. #550 did rocSHMEM; this is the other one, and it is the case that tests the boundary harder — rocSHMEM's mapping is uniform across its symmetric heap, so almost any table works. Torch's is not, so the per-allocation descriptor #546 specifies is load-bearing here rather than merely tidy.
As with #550, the whole integration is host-side. The test kernel calls unmodified
iris.storeon memory allocated entirely by torch.Technical details
iris/experimental/torch_symm_mem_provider.pyaddsTorchSymmMemProviderwith the same surface as the rocSHMEM provider —allocate_symmetric,allocate_symmetric_map,symmetric_address_map— matchingIris.allocate_symmetricfrom #549.The provider computes nothing. Torch's rendezvous handle exposes
buffer_ptrs: a per-peer list of pointers into this process's address space, whose local-rank entry already equals the tensor's owndata_ptr(). That is precisely the table Iris translates against, so the provider reads it off the handle and asserts the invariant rather than deriving it. Where the rocSHMEM provider has to queryrocshmem_ptrand cache per-peer offsets, this one has no pointer arithmetic at all.The first commit moves
SymmetricAddressMapout ofrocshmem_provider.pyintoiris/experimental/symmetric_memory.py, which imports nothing but torch. Without that, using the descriptor would dragrocshmem4pyinto a provider that has no use for it. It is re-exported fromrocshmem_providerso existing imports keep working.How it differs from the rocSHMEM provider
Peer offsets need not be shared between allocations. On the default backend each torch allocation is a separate IPC mapping, not a window onto one linear heap, so a table only translates pointers into the allocation it describes. This is the opposite of rocSHMEM, where #550 documents and relies on tables being interchangeable.
This is not theoretical. An earlier revision of the test passed one allocation's table while storing to a different allocation — copied from the rocSHMEM test, where it is legitimate — and the kernel faulted.
The practical consequence is narrow:
iris.storeandloadtranslate one pointer, so you pass the table of the allocation being addressed and a kernel can take several.iris.copytakes a singleheap_basesand translates bothsrcanddstagainst it, so under this provider it is only valid when both live in the same allocation.symmetric_address_maprejects foreign tensors. Rendezvous is collective; deriving a handle on demand would hang whichever ranks did not ask. Handles are recorded at allocation time instead.Tests run in CI
Unlike #550, no install and no image change are needed — torch is already in the CI images, so the tests execute rather than skip with nothing added to
apptainer/iris.deforinstall_rocshmem.sh.That said, symmetric memory working on ROCm is not a given, and the module import is not a useful test of it:
torch.distributed._symmetric_memoryimports successfully on builds where allocation then fails. Measured on gfx950:2.10.0+rocm7.2.1— the Apptainer CI base2.13.0+rocm7.22.9.1+rocm7.1— the Docker baseset_signal_pad_size, andrendezvousfails withHIP error: invalid argument2.13.0+rocm7.15.0anightlyhipErrorOutOfMemoryon a 4 KB allocationThe 2.15.0.dev20260930+rocm10.0 nightly was reported in review to work only with
set_backend("NVSHMEM")(rocSHMEM on ROCm); that has not been run here. The rows above are all the default backend.So the fixture probes by attempting a throwaway allocation and skips on failure. CI builds with Apptainer, so the tests run; on the Docker fallback they would skip cleanly rather than fail. As in #550,
importorskipis inside the fixture rather than at module scope — a module-level one collects zero items, pytest returns exit code 5 (NO_TESTS_COLLECTED), and torchrun reports that as a child failure that fails the whole job.The backend is the caller's choice; the provider never calls
symm_mem.set_backend(). One trap worth recording, since it cost a day: on ROCmget_backend()reports the default as'CUDA'— the HIP IPC path — butset_backend('CUDA')is rejected with "SymmetricMemory does not find allocation backend CUDA". On the 2.13/rocm7.2 wheels, forcingNCCLinstead then fails at allocation, which reads as "symmetric memory is unsupported on ROCm" when the default works.Test plan
The provider is registered in the unified test from #552, which covers the table invariant and per-allocation anchoring (two tensors, two tables, one kernel; Triton and Gluon, bf16 and fp32). The torch module covers the rest: the reachability mask, rejection of foreign tensors, and unmodified
iris.storeover torch-allocated memory.Test results
On gfx950, on the same image CI uses (
rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0, torch2.10.0+rocm7.2.1):Both files together, per rank. The 4 skips are the unified test's rocSHMEM entries, because this image has no
rocshmem4py.All three rank counts were checked rather than just 2, because
symm_mem.rendezvousis collective across the whole group, so rank count is a real variable and CI exercises 1, 2, 4 and 8.Not covered
'CUDA', is the HIP IPC path on ROCm, and is intra-node.NVSHMEM(rocSHMEM) is untested here.iris.copyis not exercised. It is valid within a single allocation and invalid across two, per the mapping difference above — the restriction is documented but untested.iris.storeis exercised.load,put,getand the atomics are single-translation and should behave identically.