PyTorch enables symmetric memory lowering for ROCm 10.1 and newer
Key Takeaway
PyTorch now supports Inductor lowering for one_shot_all_reduce using symmetric memory on ROCm versions 10.1 and above.
Summary
A new update to PyTorch enables symmetric memory support for Inductor lowering when running on ROCm 10.1 or higher. This change specifically affects the one_shot_all_reduce operation and its associated empty_strided_p2p buffer. The update also enables corresponding unit tests within the Continuous Integration (CI) pipeline.
Symmetric memory in PyTorch Inductor
- Enables one_shot_all_reduce Inductor lowering for ROCm >= 10.1.
- Utilizes the registered symm-mem buffer (empty_strided_p2p).
- Updates CI to include LoweringTest unit tests following the ROCm 10.1 upgrade.
- Verified on gfx942 hardware using
test_symmetric_memory.py.
Why it matters
This improves distributed training performance and stability for users running PyTorch on modern AMD ROCm platforms.
Original reference: github.com →