frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[c10d][nvshmem] add nvshmem build rules and dependency for libtorch_cuda (#159562)

Browse Source 
Summary:
Add guarded build option for nvshmem-related c10d code with `-c fbcode.caffe2_use_nvshmem`

Guarded clause include nvshmem device + host code (static-linked) + these 2 files:
- `torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu`
-    `torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159562
Approved by: https://github.com/Skylion007, https://github.com/kwen2501
This commit is contained in:
Tan Hoang
2025-08-31 12:56:51 +00:00
committed by PyTorch MergeBot
parent 75bc23cfc3
commit 91f0bcf43f
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
build_variables.bzl
View File

@ -758,6 +758,13 @@ libtorch_cuda_distributed_extra_sources = [
"torch/csrc/distributed/rpc/tensorpipe_cuda.cpp",
]
libtorch_nvshmem_sources = [
"torch/csrc/distributed/c10d/cuda/utils.cpp",
"torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp",
"torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu",
"torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu",
]
libtorch_cuda_distributed_sources = libtorch_cuda_distributed_base_sources + libtorch_cuda_distributed_extra_sources
libtorch_cuda_sources = libtorch_cuda_core_sources + libtorch_cuda_distributed_sources + [