Skip to content

init #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

init #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fc3f2c7
init
Chao1Han Nov 27, 2024
0560c8b
update
Chao1Han Dec 2, 2024
31d2bab
update
Chao1Han Dec 2, 2024
ba7a39a
update
Chao1Han Dec 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions caffe2/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1104,6 +1104,10 @@ if(USE_XPU)
message(WARNING "Failed to include ATen XPU implementation target")
else()
target_link_libraries(torch_xpu PRIVATE torch_xpu_ops)
# USE_C10D_XCCL from third_party torch-xpu-ops repository for xccl registration.
if(USE_C10D_XCCL)
Copy link
Collaborator

@zhangxiaoli73 zhangxiaoli73 Dec 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add comments

target_compile_definitions(torch_xpu PUBLIC USE_C10D_XCCL)
endif()
if(MSVC)
# Windows
target_link_options(torch_xpu PRIVATE
Expand Down
8 changes: 8 additions & 0 deletions torch/_C/_distributed_c10d.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,3 +679,11 @@ class _SymmetricMemory:
def stream_write_value32(
tensor: torch.Tensor, offset: int, val: int
) -> torch.Tensor: ...

class ProcessGroupXCCL(Backend):
def __init__(
self,
store: Store,
rank: int,
size: int,
): ...
7 changes: 6 additions & 1 deletion torch/csrc/distributed/c10d/ProcessGroup.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@ class TORCH_API ProcessGroup : public torch::CustomClassHolder {
NCCL = 2,
UCC = 3,
MPI = 4,
CUSTOM = 5,
XCCL = 5,
CUSTOM = 6,
};

static std::string backendTypeToString(const BackendType& type) {
Expand All @@ -86,6 +87,8 @@ class TORCH_API ProcessGroup : public torch::CustomClassHolder {
return "gloo";
case BackendType::NCCL:
return "nccl";
case BackendType::XCCL:
return "xccl";
case BackendType::UCC:
return "ucc";
case BackendType::MPI:
Expand All @@ -106,6 +109,8 @@ class TORCH_API ProcessGroup : public torch::CustomClassHolder {
return BackendType::GLOO;
} else if (backend == "nccl") {
return BackendType::NCCL;
} else if (backend == "xccl") {
return BackendType::XCCL;
} else if (backend == "ucc") {
return BackendType::UCC;
} else if (backend == "mpi") {
Expand Down
22 changes: 22 additions & 0 deletions torch/csrc/distributed/c10d/init.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@
#include <torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp>
#endif

#ifdef USE_C10D_XCCL
#include <torch/csrc/distributed/c10d/ProcessGroupXCCL.hpp>
#endif

#ifdef USE_C10D_NCCL
#include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
#include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
Expand Down Expand Up @@ -2311,6 +2315,7 @@ The hook must have the following signature:
.value("UNDEFINED", ::c10d::ProcessGroup::BackendType::UNDEFINED)
.value("GLOO", ::c10d::ProcessGroup::BackendType::GLOO)
.value("NCCL", ::c10d::ProcessGroup::BackendType::NCCL)
.value("XCCL", ::c10d::ProcessGroup::BackendType::XCCL)
.value("UCC", ::c10d::ProcessGroup::BackendType::UCC)
.value("MPI", ::c10d::ProcessGroup::BackendType::MPI)
.value("CUSTOM", ::c10d::ProcessGroup::BackendType::CUSTOM)
Expand Down Expand Up @@ -2927,6 +2932,23 @@ Example::
py::call_guard<py::gil_scoped_release>());
#endif

#ifdef USE_C10D_XCCL
auto processGroupXCCL =
intrusive_ptr_no_gil_destructor_class_<::c10d::ProcessGroupXCCL>(
module, "ProcessGroupXCCL", backend)
.def(
py::init([](const c10::intrusive_ptr<::c10d::Store>& store,
int rank,
int size) {
return c10::make_intrusive<::c10d::ProcessGroupXCCL>(
store, rank, size);
}),
py::arg("store"),
py::arg("rank"),
py::arg("size"),
py::call_guard<py::gil_scoped_release>());
#endif

#ifdef USE_C10D_UCC
auto processGroupUCC =
intrusive_ptr_no_gil_destructor_class_<::c10d::ProcessGroupUCC>(
Expand Down
30 changes: 28 additions & 2 deletions torch/distributed/distributed_c10d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@
"is_nccl_available",
"is_torchelastic_launched",
"is_ucc_available",
"is_xccl_available",
"isend",
"monitored_barrier",
"new_group",
Expand Down Expand Up @@ -132,6 +133,7 @@
_NCCL_AVAILABLE = True
_GLOO_AVAILABLE = True
_UCC_AVAILABLE = True
_XCCL_AVAILABLE = True

_pickler = pickle.Pickler
_unpickler = pickle.Unpickler
Expand Down Expand Up @@ -195,6 +197,14 @@ def _export_c_types() -> None:
except ImportError:
_UCC_AVAILABLE = False

try:
from torch._C._distributed_c10d import ProcessGroupXCCL

ProcessGroupXCCL.__module__ = "torch.distributed.distributed_c10d"
__all__ += ["ProcessGroupXCCL"]
except ImportError:
_XCCL_AVAILABLE = False

logger = logging.getLogger(__name__)

PG_WRAPPER_STORE_PREFIX = "pg_wrapper"
Expand Down Expand Up @@ -224,7 +234,7 @@ class Backend(str):
"""
An enum-like class for backends.

Available backends: GLOO, NCCL, UCC, MPI, and other registered backends.
Available backends: GLOO, NCCL, UCC, MPI, XCCL, and other registered backends.

The values of this class are lowercase strings, e.g., ``"gloo"``. They can
be accessed as attributes, e.g., ``Backend.NCCL``.
Expand All @@ -244,21 +254,24 @@ class Backend(str):
NCCL = "nccl"
UCC = "ucc"
MPI = "mpi"
XCCL = "xccl"

_BackendPlugin = namedtuple("_BackendPlugin", ["creator_fn", "extended_api"])

_plugins: Dict[str, _BackendPlugin] = {}

backend_list = [UNDEFINED, GLOO, NCCL, UCC, MPI]
backend_list = [UNDEFINED, GLOO, NCCL, XCCL, UCC, MPI]

default_device_backend_map: Dict[str, str] = {
"cpu": GLOO,
"cuda": NCCL,
"xpu": XCCL,
}

backend_capability: Dict[str, List[str]] = {
GLOO: ["cpu", "cuda"],
NCCL: ["cuda"],
XCCL: ["xpu"],
UCC: ["cpu", "cuda"],
MPI: ["cpu", "cuda"],
}
Expand All @@ -267,6 +280,7 @@ class Backend(str):
UNDEFINED: ProcessGroup.BackendType.UNDEFINED,
GLOO: ProcessGroup.BackendType.GLOO,
NCCL: ProcessGroup.BackendType.NCCL,
XCCL: ProcessGroup.BackendType.XCCL,
UCC: ProcessGroup.BackendType.UCC,
MPI: ProcessGroup.BackendType.MPI,
}
Expand Down Expand Up @@ -1185,6 +1199,11 @@ def is_ucc_available() -> bool:
return _UCC_AVAILABLE


def is_xccl_available() -> bool:
"""Check if the XCCL backend is available."""
return _XCCL_AVAILABLE


def is_backend_available(backend: str) -> bool:
"""
Check backend availability.
Expand Down Expand Up @@ -1885,6 +1904,13 @@ def _new_process_group_helper(
backend_prefix_store, group_rank, group_size, timeout=timeout
)
backend_type = ProcessGroup.BackendType.UCC
elif backend_str == Backend.XCCL:
if not is_xccl_available():
raise RuntimeError("Distributed package doesn't have XCCL built in")
backend_class = ProcessGroupXCCL(
backend_prefix_store, group_rank, group_size
)
backend_type = ProcessGroup.BackendType.XCCL
else:
assert (
backend_str.upper() in Backend._plugins
Expand Down