From cb61c1a738438574ab2e46300761559fb1315068 Mon Sep 17 00:00:00 2001 From: lzhang2 Date: Mon, 30 Jun 2025 14:25:56 +0800 Subject: [PATCH 01/10] enable symm on xpu --- torch/_C/_distributed_c10d.pyi | 6 ++++ torch/csrc/distributed/c10d/init.cpp | 6 ++++ .../c10d/symm_mem/SymmetricMemory.hpp | 1 + .../distributed/_symmetric_memory/__init__.py | 36 +++++++++++-------- 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/torch/_C/_distributed_c10d.pyi b/torch/_C/_distributed_c10d.pyi index 72fde27d02576..168e4d9dfc841 100644 --- a/torch/_C/_distributed_c10d.pyi +++ b/torch/_C/_distributed_c10d.pyi @@ -804,6 +804,12 @@ class _SymmetricMemory: channel: int = 0, timeout_ms: int = 0, ) -> None: ... + def copy_buffer( + self, + src: torch.Tensor, + dst: torch.Tensor, + size: int, + ) -> None: ... @staticmethod def memset32( tensor: torch.Tensor, offset: int, val: int, count: int = 1 diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp index a0904a814637c..60714977ef4ba 100644 --- a/torch/csrc/distributed/c10d/init.cpp +++ b/torch/csrc/distributed/c10d/init.cpp @@ -1198,6 +1198,12 @@ This class does not support ``__members__`` property.)"); py::arg("src_rank"), py::arg("channel") = 0, py::arg("timeout_ms") = 0) + .def( + "copy_buffer", + &SymmetricMemory::copy_buffer, + py::arg("src"), + py::arg("dst"), + py::arg("size")) // Util functions that are often used together with symmetric memory but // not necessarily directly on symmetric memory. .def_static( diff --git a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp index c2828de04c9b3..cd45572ceaefa 100644 --- a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp +++ b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp @@ -68,6 +68,7 @@ class TORCH_API SymmetricMemory : public c10::intrusive_ptr_target { virtual void barrier(int channel, size_t timeout_ms) = 0; virtual void put_signal(int dst_rank, int channel, size_t timeout_ms) = 0; virtual void wait_signal(int src_rank, int channel, size_t timeout_ms) = 0; + virtual void copy_buffer(at::Tensor src, at::Tensor dst , size_t size) = 0; virtual int get_rank() = 0; virtual int get_world_size() = 0; diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 4b0e9acc19bd7..91a6c83ee0d84 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -10,6 +10,7 @@ from typing import Any, Callable, Literal, Optional import torch +import torch.distributed as dist import torch.distributed._functional_collectives as funcol import torch.distributed.distributed_c10d as c10d from torch._C._autograd import DeviceType @@ -105,7 +106,7 @@ def get_symm_mem_workspace(group_name: str, min_size: int) -> _SymmetricMemory: tensor = _group_name_to_workspace_tensor.get(group_name) size = tensor.numel() * tensor.element_size() if tensor is not None else 0 if tensor is None or size < min_size: - if torch.cuda.is_current_stream_capturing(): + if False: #torch.cuda.is_current_stream_capturing(): curr_size = 0 if tensor is None else tensor.numel() * tensor.element_size() raise RuntimeError( f"get_symm_mem_workspace(): the requested size ({min_size} bytes) " @@ -120,19 +121,19 @@ def get_symm_mem_workspace(group_name: str, min_size: int) -> _SymmetricMemory: (max(size, min_size),), [1], torch.uint8, - torch.device(f"cuda:{torch.cuda.current_device()}"), + torch.device(f"xpu:{torch.xpu.current_device()}"), group_name, ) _group_name_to_workspace_tensor[group_name] = tensor return _SymmetricMemory.rendezvous(tensor) -_backend_streams: dict[int, torch.cuda.Stream] = {} +_backend_streams: dict[int, torch.xpu.Stream] = {} -def _get_backend_stream(priority: int = 0) -> torch.cuda.Stream: +def _get_backend_stream(priority: int = 0) -> torch.xpu.Stream: if priority not in _backend_streams: - _backend_streams[priority] = torch.cuda.Stream(priority=priority) + _backend_streams[priority] = torch.xpu.Stream(priority=priority) return _backend_streams[priority] @@ -167,9 +168,10 @@ def _pipelined_multi_all_gather_and_consume( group_size = symm_mem.world_size rank = symm_mem.rank - symm_mem.barrier(channel=0) + # symm_mem.barrier(channel=0) + dist.barrier() backend_stream = _get_backend_stream() - backend_stream.wait_stream(torch.cuda.current_stream()) + backend_stream.wait_stream(torch.xpu.current_stream()) for x, y in zip(shard, ag_out): assert x.is_contiguous(), ( @@ -185,7 +187,8 @@ def _pipelined_multi_all_gather_and_consume( def copy_shard(dst: list[torch.Tensor], src: list[torch.Tensor]) -> None: for d, s in zip(dst, src): - d.copy_(s) + symm_mem.copy_buffer(s, d, s.numel()) + # d.copy_(s) def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: offset_bytes = 0 @@ -249,8 +252,11 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: # prevent suboptimal scenarios, we are giving up the chance to overlap "mv" # and "b" with the first shard_consumer for now. copy_shard(dst=local_p2p_bufs, src=shard) - symm_mem.barrier(channel=1) - backend_stream.wait_stream(torch.cuda.current_stream()) + #torch.xpu.synchronize() + #print(f"[Python] zl_debug copy shard tensor to local done {local_p2p_bufs} with shape {local_p2p_bufs.shape}", flush=True) + # symm_mem.barrier(channel=1) + dist.barrier() + backend_stream.wait_stream(torch.xpu.current_stream()) # At this point, all ranks have copied their local shard to # their local p2p buffer. Each rank can now copy and consume @@ -259,7 +265,7 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: for step in range(1, group_size): if step % 2 == 0: - stream = torch.cuda.current_stream() + stream = torch.xpu.current_stream() else: stream = backend_stream remote_rank = (step + rank) % group_size @@ -272,14 +278,15 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: # Copy from input to the all-gather output. Opportunistically overlap # it with the last shard_consumer. if group_size % 2 == 0: - stream = torch.cuda.current_stream() + stream = torch.xpu.current_stream() else: stream = backend_stream with stream: copy_shard(dst=shards[rank], src=shard) - torch.cuda.current_stream().wait_stream(backend_stream) - symm_mem.barrier(channel=0) + torch.xpu.current_stream().wait_stream(backend_stream) + # symm_mem.barrier(channel=0) + dist.barrier() def _pipelined_all_gather_and_consume( @@ -642,6 +649,7 @@ def _fused_all_gather_matmul_fallback( @torch.library.impl(lib, "fused_all_gather_matmul", "CUDA") +@torch.library.impl(lib, "fused_all_gather_matmul", "XPU") def _fused_all_gather_matmul( A_shard: torch.Tensor, Bs: list[torch.Tensor], From 27a2420009471797b505be0e09fa794a05d18884 Mon Sep 17 00:00:00 2001 From: lzhang2 Date: Thu, 3 Jul 2025 11:24:07 +0800 Subject: [PATCH 02/10] enable fused matmul+reducescatter --- .../distributed/_symmetric_memory/__init__.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 91a6c83ee0d84..d310f36a0da3d 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -339,9 +339,10 @@ def _pipelined_produce_and_all2all( group_size = symm_mem.world_size rank = symm_mem.rank - symm_mem.barrier(channel=0) + # symm_mem.barrier(channel=0) + dist.barrier() backend_stream = _get_backend_stream() - backend_stream.wait_stream(torch.cuda.current_stream()) + backend_stream.wait_stream(torch.xpu.current_stream()) def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: assert idx in (0, 1) @@ -359,7 +360,7 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: for step in range(1, group_size): remote_rank = (rank - step) % group_size if step % 2 == 0: - stream = torch.cuda.current_stream() + stream = torch.xpu.current_stream() p2p_buf = local_p2p_buf_1 remote_p2p_buf = get_p2p_buf(remote_rank, 1) else: @@ -415,22 +416,26 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: # the correct order, there's very little room for the scheduling # order of subsequent kernels to be inconsistent across ranks. if step == 2: - torch.cuda._sleep(100) + torch.xpu._sleep(100) # todo: no this API right now chunk_producer((rank + step) % group_size, p2p_buf) - symm_mem.barrier(channel=step % 2) - out_chunks[remote_rank].copy_(remote_p2p_buf) + # symm_mem.barrier(channel=step % 2) + dist.barrier() + # replaced with copy_buffer + # out_chunks[remote_rank].copy_(remote_p2p_buf) + symm_mem.copy_buffer(remote_p2p_buf, out_chunks[remote_rank], remote_p2p_buf.numel()) # The local P2P buffer can only be overwritten by the next # chunk_producer after all peers have finished reading from it. - symm_mem.barrier(channel=step % 2) + # symm_mem.barrier(channel=step % 2) + dist.barrier() # If the sleep wasn't issued in the above loop, do it now. if group_size == 2: - torch.cuda._sleep(100) + torch.xpu._sleep(100) chunk_producer(rank, out_chunks[rank]) - torch.cuda.current_stream().wait_stream(backend_stream) - symm_mem.barrier(channel=0) - + torch.xpu.current_stream().wait_stream(backend_stream) + # symm_mem.barrier(channel=0) + dist.barrier() lib = torch.library.Library("symm_mem", "DEF") # noqa: TOR901 lib.define( @@ -1001,6 +1006,7 @@ def restride_A_shard_for_fused_all_gather_matmul( @torch.library.impl(lib, "fused_matmul_reduce_scatter", "CUDA") +@torch.library.impl(lib, "fused_matmul_reduce_scatter", "XPU") def _fused_matmul_reduce_scatter( A: torch.Tensor, B: torch.Tensor, From 7bb8e78394758a276dad1b4572d89f897c33c426 Mon Sep 17 00:00:00 2001 From: lzhang2 Date: Thu, 10 Jul 2025 10:23:49 +0800 Subject: [PATCH 03/10] refine fused matmul and reducescatter --- .../distributed/_symmetric_memory/__init__.py | 78 +++---------------- 1 file changed, 11 insertions(+), 67 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index d310f36a0da3d..010335380c18b 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -334,7 +334,7 @@ def _pipelined_produce_and_all2all( dist.all_to_all_single(output=output, input=torch.cat(chunks)) """ out_chunks = output.chunk(c10d._get_group_size_by_name(group_name)) - p2p_workspace_size_req = out_chunks[0].numel() * out_chunks[0].element_size() * 2 + p2p_workspace_size_req = out_chunks[0].numel() * out_chunks[0].element_size() * dist.get_world_size() symm_mem = get_symm_mem_workspace(group_name, min_size=p2p_workspace_size_req) group_size = symm_mem.world_size rank = symm_mem.rank @@ -345,79 +345,25 @@ def _pipelined_produce_and_all2all( backend_stream.wait_stream(torch.xpu.current_stream()) def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: - assert idx in (0, 1) - offset = 0 if idx == 0 else out_chunks[0].numel() + offset = out_chunks[0].numel() * idx return symm_mem.get_buffer( rank, out_chunks[0].shape, out_chunks[0].dtype, offset ) - # Prepare two local p2p buffers, so that a remote rank can pull the result - # of step [i] in one p2p buffer while the local rank can compute the - # result of step [i+1] and write it directly the other p2p buffer. - local_p2p_buf_0 = get_p2p_buf(rank, 0) - local_p2p_buf_1 = get_p2p_buf(rank, 1) - for step in range(1, group_size): remote_rank = (rank - step) % group_size + producer_rank = (rank + step) % group_size + p2p_buf = get_p2p_buf(rank, producer_rank) + remote_p2p_buf = get_p2p_buf(remote_rank, rank + ) if step % 2 == 0: stream = torch.xpu.current_stream() - p2p_buf = local_p2p_buf_1 - remote_p2p_buf = get_p2p_buf(remote_rank, 1) else: stream = backend_stream - p2p_buf = local_p2p_buf_0 - remote_p2p_buf = get_p2p_buf(remote_rank, 0) with stream: - # Parallelization strategy: every rank issues independent compute - # -> barrier -> p2p copy sequences on two streams. In addition to - # computation/communication overlapping, the strategy allows for - # computation/computation overlapping, greatly reducing - # quantization inefficiency. - # - # Ideally, stream activities would look like this ("b" for - # barriers, "cp" for p2p copies): - # - # [rank 0] - # stream 0: [ chunk_producer ][b][ cp ][ chunk_producer ][b][ cp ] - # stream 1: [ chunk_producer ][b][ cp ][ chunk_producer ][b][ cp ] - # - # [rank 1] - # stream 0: [ chunk_producer ][b][ cp ][ chunk_producer ][b][ cp ] - # stream 1: [ chunk_producer ][b][ cp ][ chunk_producer ][b][ cp ] - # - # Note that the barriers synchronize streams with the same ID - # across ranks. They don't synchronize streams on the same rank. - # - # Since the work on both streams is independent, there's no - # guarantee that the chunk_producer from stream 0 or stream 1 will - # be scheduled first. If there is a scheduling mismatch across - # ranks, the barrier forces all ranks to wait for the slowest. - # - # When scheduling mismatches occur among ranks, the stream - # activities might look like this (note that p2p copies from - # different streams cannot overlap with each other): - # - # [rank 0] - # stream 0: [ chunk_producer ][b ][ cp ][ chunk_producer ][b ][ cp ] - # stream 1: [ chunk_producer ][b] [ cp ][ chunk_producer ][b] [ cp ] - # - # [rank 1] - # stream 0: [ chunk_producer ][b] [ cp ][ chunk_producer ][b] [ cp ] - # stream 1: [ chunk_producer ][b ][ cp ][ chunk_producer ][b ][ cp ] - # - # To prevent this, we need to ensure that the chunk_producer on - # stream 1 gets scheduled first on every rank. Without access to - # the underlying kernels, CUDA offers no API to control the - # scheduling order of two independent, overlapping kernels. Our - # solution is to issue a small sleep kernel in stream 0. The sleep - # duration is insignificant, but having an extra task in stream 0 - # will almost guarantee that the chunk_producer on stream 1 gets - # scheduled first. Once the first chunk_producer is scheduled in - # the correct order, there's very little room for the scheduling - # order of subsequent kernels to be inconsistent across ranks. - if step == 2: - torch.xpu._sleep(100) # todo: no this API right now - chunk_producer((rank + step) % group_size, p2p_buf) + # if step == 2: + # torch.xpu._sleep(100) + chunk_producer(producer_rank, p2p_buf) # symm_mem.barrier(channel=step % 2) dist.barrier() # replaced with copy_buffer @@ -426,12 +372,10 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: # The local P2P buffer can only be overwritten by the next # chunk_producer after all peers have finished reading from it. # symm_mem.barrier(channel=step % 2) - dist.barrier() # If the sleep wasn't issued in the above loop, do it now. - if group_size == 2: - torch.xpu._sleep(100) - + # if group_size == 2: + # torch.xpu._sleep(100) chunk_producer(rank, out_chunks[rank]) torch.xpu.current_stream().wait_stream(backend_stream) # symm_mem.barrier(channel=0) From 0339c94d24f007f3b5493c44ffdeade2a990b330 Mon Sep 17 00:00:00 2001 From: lzhang2 Date: Thu, 17 Jul 2025 16:56:37 +0800 Subject: [PATCH 04/10] fp8 scaled support --- torch/distributed/_symmetric_memory/__init__.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 010335380c18b..32437271dd8c8 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -501,6 +501,7 @@ def unflatten(t: torch.Tensor) -> torch.Tensor: scale_mode = _check_and_verify_fp8_all_gather_scale_mode( shard=A_shard, scale=A_scale, gather_dim=gather_dim, group_size=group.size() ) + print(f"zl_debug get scaled mode = {scale_mode} of allgather+matmul", flush=True) # Computing block-wise matmul along the first dim of A if scale_mode == _ScaleMode.ROW_WISE_SHARDED: @@ -843,6 +844,7 @@ def scaled_matmul( @torch.library.impl(lib, "fused_all_gather_scaled_matmul", "CUDA") +@torch.library.impl(lib, "fused_all_gather_scaled_matmul", "XPU") def _fused_all_gather_scaled_matmul( A_shard: torch.Tensor, Bs: list[torch.Tensor], @@ -898,7 +900,8 @@ def _fused_all_gather_scaled_matmul( with torch.profiler.record_function("fused_all_gather_scaled_matmul"): A, res = _fused_all_gather_matmul_impl( - torch.ops.aten._scaled_mm.out, + # torch.ops.aten._scaled_mm.out, + torch.ops.vllm.fp8_gemm.out, A_shard, Bs, A_scale, @@ -1056,6 +1059,7 @@ def chunk_producer(rank: int, out: torch.Tensor) -> None: @torch.library.impl(lib, "fused_scaled_matmul_reduce_scatter", "CUDA") +@torch.library.impl(lib, "fused_scaled_matmul_reduce_scatter", "XPU") def _fused_scaled_matmul_reduce_scatter( A: torch.Tensor, B: torch.Tensor, @@ -1089,7 +1093,8 @@ def _fused_scaled_matmul_reduce_scatter( ) with torch.profiler.record_function("fused_scaled_matmul_reduce_scatter"): return _fused_scaled_matmul_reduce_scatter_impl( - mm_out_op=torch.ops.aten._scaled_mm.out, + # mm_out_op=torch.ops.aten._scaled_mm.out, + mm_out_op=torch.ops.vllm.fp8_gemm.out, A=A, B=B, A_scale=A_scale, From 1b514dcfae74b480f9137dc646f9ac848e4ced6b Mon Sep 17 00:00:00 2001 From: "Han, Chao1" Date: Thu, 31 Jul 2025 16:00:56 +0800 Subject: [PATCH 05/10] format --- .../distributed/_symmetric_memory/__init__.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 32437271dd8c8..70906b366010c 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -106,7 +106,7 @@ def get_symm_mem_workspace(group_name: str, min_size: int) -> _SymmetricMemory: tensor = _group_name_to_workspace_tensor.get(group_name) size = tensor.numel() * tensor.element_size() if tensor is not None else 0 if tensor is None or size < min_size: - if False: #torch.cuda.is_current_stream_capturing(): + if False: # torch.cuda.is_current_stream_capturing(): curr_size = 0 if tensor is None else tensor.numel() * tensor.element_size() raise RuntimeError( f"get_symm_mem_workspace(): the requested size ({min_size} bytes) " @@ -252,8 +252,8 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: # prevent suboptimal scenarios, we are giving up the chance to overlap "mv" # and "b" with the first shard_consumer for now. copy_shard(dst=local_p2p_bufs, src=shard) - #torch.xpu.synchronize() - #print(f"[Python] zl_debug copy shard tensor to local done {local_p2p_bufs} with shape {local_p2p_bufs.shape}", flush=True) + # torch.xpu.synchronize() + # print(f"[Python] zl_debug copy shard tensor to local done {local_p2p_bufs} with shape {local_p2p_bufs.shape}", flush=True) # symm_mem.barrier(channel=1) dist.barrier() backend_stream.wait_stream(torch.xpu.current_stream()) @@ -334,7 +334,9 @@ def _pipelined_produce_and_all2all( dist.all_to_all_single(output=output, input=torch.cat(chunks)) """ out_chunks = output.chunk(c10d._get_group_size_by_name(group_name)) - p2p_workspace_size_req = out_chunks[0].numel() * out_chunks[0].element_size() * dist.get_world_size() + p2p_workspace_size_req = ( + out_chunks[0].numel() * out_chunks[0].element_size() * dist.get_world_size() + ) symm_mem = get_symm_mem_workspace(group_name, min_size=p2p_workspace_size_req) group_size = symm_mem.world_size rank = symm_mem.rank @@ -354,8 +356,7 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: remote_rank = (rank - step) % group_size producer_rank = (rank + step) % group_size p2p_buf = get_p2p_buf(rank, producer_rank) - remote_p2p_buf = get_p2p_buf(remote_rank, rank - ) + remote_p2p_buf = get_p2p_buf(remote_rank, rank) if step % 2 == 0: stream = torch.xpu.current_stream() else: @@ -368,7 +369,9 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: dist.barrier() # replaced with copy_buffer # out_chunks[remote_rank].copy_(remote_p2p_buf) - symm_mem.copy_buffer(remote_p2p_buf, out_chunks[remote_rank], remote_p2p_buf.numel()) + symm_mem.copy_buffer( + remote_p2p_buf, out_chunks[remote_rank], remote_p2p_buf.numel() + ) # The local P2P buffer can only be overwritten by the next # chunk_producer after all peers have finished reading from it. # symm_mem.barrier(channel=step % 2) @@ -381,6 +384,7 @@ def get_p2p_buf(rank: int, idx: int) -> torch.Tensor: # symm_mem.barrier(channel=0) dist.barrier() + lib = torch.library.Library("symm_mem", "DEF") # noqa: TOR901 lib.define( "fused_all_gather_matmul(" From 97b55d14461a555ab9c723aede51f9cd30bb2b7a Mon Sep 17 00:00:00 2001 From: "Yu, Guangye" <106960996+guangyey@users.noreply.github.com> Date: Thu, 10 Jul 2025 00:28:11 -0700 Subject: [PATCH 06/10] Enable XPU built with Level Zero (#321) --- CMakeLists.txt | 1 + cmake/Dependencies.cmake | 1 + cmake/Modules/FindLevelZero.cmake | 59 +++++++++++++++++++++++++++++++ cmake/Summary.cmake | 4 +++ cmake/public/xpu.cmake | 28 ++++++++++++++- 5 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 cmake/Modules/FindLevelZero.cmake diff --git a/CMakeLists.txt b/CMakeLists.txt index ad7368e192983..77f24490d56f2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -236,6 +236,7 @@ option(USE_ASAN "Use Address+Undefined Sanitizers" OFF) option(USE_TSAN "Use Thread Sanitizer" OFF) option(USE_CUDA "Use CUDA" ON) option(USE_XPU "Use XPU" ON) +cmake_dependent_option(USE_LEVEL_ZERO "Use Level Zero for XPU backend." ON "USE_XPU" OFF) cmake_dependent_option( BUILD_LAZY_CUDA_LINALG "Build cuda linalg ops as separate library" ON "USE_CUDA AND LINUX AND BUILD_PYTHON" OFF) diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 944c7821f6676..409fad5d378f9 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -95,6 +95,7 @@ if(USE_XPU) message(WARNING "Not compiling with XPU. Could NOT find SYCL. " "Suppress this warning with -DUSE_XPU=OFF.") caffe2_update_option(USE_XPU OFF) + caffe2_update_option(USE_LEVEL_ZERO OFF) endif() foreach(flag ${XPU_HOST_CXX_FLAGS}) add_definitions(${flag}) diff --git a/cmake/Modules/FindLevelZero.cmake b/cmake/Modules/FindLevelZero.cmake new file mode 100644 index 0000000000000..c8b9720f84432 --- /dev/null +++ b/cmake/Modules/FindLevelZero.cmake @@ -0,0 +1,59 @@ +# This will define the following variables: +# LevelZero_FOUND : True if the system has the LevelZero library. +# LevelZero_INCLUDE_DIR : Level Zero include directory. +# LevelZero_LIBRARY : Level Zero library directory. + +include(FindPackageHandleStandardArgs) + +# Find Level Zero include directory. +if(CMAKE_SYSTEM_NAME MATCHES "Linux") + find_path( + LevelZero_INCLUDE_DIR + NAMES level_zero/ze_api.h + PATH_SUFFIXES include + ) +elseif(CMAKE_SYSTEM_NAME MATCHES "Windows") + find_path( + LevelZero_INCLUDE_DIR + NAMES level_zero/ze_api.h + HINTS $ENV{LEVEL_ZERO_V1_SDK_PATH} + PATH_SUFFIXES include + ) +endif() + +if(NOT LevelZero_INCLUDE_DIR) + set(LevelZero_FOUND False) + set(LevelZero_REASON_FAILURE "Level Zero include directory not found!!") + set(LevelZero_NOT_FOUND_MESSAGE "${LevelZero_REASON_FAILURE}") + return() +endif() + +# Find Level Zero library directory. +if(CMAKE_SYSTEM_NAME MATCHES "Linux") + find_library( + LevelZero_LIBRARY + NAMES ze_loader + PATH_SUFFIXES x86_64_linux_gnu lib lib/x64 lib64 + ) +elseif(CMAKE_SYSTEM_NAME MATCHES "Windows") + find_library( + LevelZero_LIBRARY + NAMES ze_loader + HINTS $ENV{LEVEL_ZERO_V1_SDK_PATH} + PATH_SUFFIXES lib lib64 + ) +endif() + +if(NOT LevelZero_LIBRARY) + set(LevelZero_FOUND False) + set(LevelZero_REASON_FAILURE "Level Zero library not found!!") + set(LevelZero_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}") + return() +endif() + +find_package_handle_standard_args( + LevelZero + FOUND_VAR LevelZero_FOUND + REQUIRED_VARS LevelZero_INCLUDE_DIR LevelZero_LIBRARY + REASON_FAILURE_MESSAGE "${LevelZero_REASON_FAILURE}" + ) diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake index 745d9ea058687..1f10d8ed9badd 100644 --- a/cmake/Summary.cmake +++ b/cmake/Summary.cmake @@ -124,6 +124,10 @@ function(caffe2_print_configuration_summary) message(STATUS " SYCL compiler version : ${SYCL_COMPILER_VERSION}") message(STATUS " SYCL include path : ${SYCL_INCLUDE_DIR}") message(STATUS " SYCL library : ${SYCL_LIBRARY}") + if(${USE_LEVEL_ZERO}) + message(STATUS " Level Zero include path : ${LevelZero_INCLUDE_DIR}") + message(STATUS " Level Zero library : ${LevelZero_LIBRARY}") + endif() endif() message(STATUS " USE_ROCM : ${USE_ROCM}") if(${USE_ROCM}) diff --git a/cmake/public/xpu.cmake b/cmake/public/xpu.cmake index b39e31d0ade8a..b85ca74466dbc 100644 --- a/cmake/public/xpu.cmake +++ b/cmake/public/xpu.cmake @@ -14,6 +14,14 @@ if(NOT SYCL_FOUND) # Exit early to avoid populating XPU_HOST_CXX_FLAGS. return() endif() +# Find LevelZero library. +if(USE_LEVEL_ZERO) + find_package(LevelZero REQUIRED) + if(NOT LevelZero_FOUND) + set(PYTORCH_FOUND_XPU FALSE) + return() + endif() +endif() set(PYTORCH_FOUND_XPU TRUE) # SYCL library interface @@ -26,11 +34,29 @@ set_property( TARGET torch::sycl PROPERTY INTERFACE_LINK_LIBRARIES ${SYCL_LIBRARY}) +# LevelZero library interface +if(USE_LEVEL_ZERO) + add_library(torch::level_zero INTERFACE IMPORTED) + + set_property( + TARGET torch::level_zero PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${LevelZero_INCLUDE_DIR}) + set_property( + TARGET torch::level_zero PROPERTY INTERFACE_LINK_LIBRARIES + ${LevelZero_LIBRARY}) +endif() + # xpurt add_library(torch::xpurt INTERFACE IMPORTED) + +set(xpurt_deps torch::sycl) +if(USE_LEVEL_ZERO) + list(APPEND xpurt_deps torch::level_zero) +endif() + set_property( TARGET torch::xpurt PROPERTY INTERFACE_LINK_LIBRARIES - torch::sycl) + "${xpurt_deps}") # setting xpu arch flags torch_xpu_get_arch_list(XPU_ARCH_FLAGS) From f4b119e5e835dcfedbdc19a9235a61206ec1f67b Mon Sep 17 00:00:00 2001 From: "Han, Chao1" Date: Tue, 5 Aug 2025 15:44:42 +0800 Subject: [PATCH 07/10] stage < TP --- .../distributed/_symmetric_memory/__init__.py | 33 +++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 70906b366010c..fca426f5e62ef 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -263,16 +263,29 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: # remote shards. shard_consumer(shard, rank) - for step in range(1, group_size): - if step % 2 == 0: - stream = torch.xpu.current_stream() - else: - stream = backend_stream - remote_rank = (step + rank) % group_size - remote_p2p_bufs = get_p2p_bufs(remote_rank) + # for step in range(1, group_size): + # if step % 2 == 0: + # stream = torch.xpu.current_stream() + # else: + # stream = backend_stream + # remote_rank = (step + rank) % group_size + # remote_p2p_bufs = get_p2p_bufs(remote_rank) + # with stream: + # copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) + # shard_consumer(shards[remote_rank], remote_rank) + + stage_size = 2 + for stage in range(1, group_size, stage_size): + stream = torch.xpu.current_stream() if (stage // stage_size) % 2 == 0 else backend_stream with stream: - copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) - shard_consumer(shards[remote_rank], remote_rank) + for i in range(stage_size): + step = stage + i + if step >= group_size: + break + remote_rank = (step + rank) % group_size + remote_p2p_bufs = get_p2p_bufs(remote_rank) + copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) + shard_consumer(shards[remote_rank], remote_rank) if ag_out_needed: # Copy from input to the all-gather output. Opportunistically overlap @@ -1722,7 +1735,7 @@ def is_nvshmem_available() -> bool: return _is_nvshmem_available() -def set_backend(name: Literal["NVSHMEM", "CUDA", "NCCL"]) -> None: +def set_backend(name: Literal["NVSHMEM", "CUDA", "NCCL", "XCCL"]) -> None: r""" Set the backend for symmetric memory allocation. This is a global setting and affects all subsequent calls to From 9bd58a610fdd61ec8912f0f93dc3a2fdc463f762 Mon Sep 17 00:00:00 2001 From: Han Chao Date: Wed, 13 Aug 2025 13:26:11 +0800 Subject: [PATCH 08/10] register op --- torch/distributed/_symmetric_memory/__init__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index fca426f5e62ef..5881b01e47032 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -1376,7 +1376,7 @@ def _maybe_convert_scalar_types_to_dtypes( class Work(_Work): def __init__(self) -> None: super().__init__() - self.event = torch.cuda.Event() + self.event = torch.xpu.Event() self.event.record() def wait(self, timeout: timedelta = timedelta(seconds=0)) -> bool: @@ -1421,7 +1421,7 @@ def _low_contention_all_gather_meta( group_size = c10d._get_group_size_by_name(group_name) return tensor.new_empty(tensor.shape[0] * group_size, *tensor.shape[1:]) - +@torch.library.impl(lib, "_low_contention_all_gather", "XPU") @torch.library.impl(lib, "_low_contention_all_gather", "CUDA") def _low_contention_all_gather( tensor: torch.Tensor, @@ -1454,7 +1454,7 @@ def _low_contention_all_gather( output = tensor.new_empty(tensor.shape[0] * world_size, *tensor.shape[1:]) chunks = output.chunk(world_size) - _get_backend_stream().wait_stream(torch.cuda.current_stream()) + _get_backend_stream().wait_stream(torch.xpu.current_stream()) with _get_backend_stream(): if not input_is_symm_mem: local_buf = symm_mem.get_buffer(rank, tensor.shape, tensor.dtype) @@ -1492,7 +1492,7 @@ def _low_contention_reduce_scatter_with_symm_mem_input( a2a_res = torch.empty_like(tensor) chunks = a2a_res.chunk(world_size) - _get_backend_stream().wait_stream(torch.cuda.current_stream()) + _get_backend_stream().wait_stream(torch.xpu.current_stream()) with _get_backend_stream(): # pull + offline reduction symm_mem.barrier() @@ -1529,7 +1529,7 @@ def _low_contention_reduce_scatter_with_workspace( assert tensor.shape[0] % world_size == 0 chunks = tensor.chunk(world_size) - _get_backend_stream().wait_stream(torch.cuda.current_stream()) + _get_backend_stream().wait_stream(torch.xpu.current_stream()) with _get_backend_stream(): # push + offline reduction workspace.barrier() @@ -1552,7 +1552,7 @@ def _low_contention_reduce_scatter_with_workspace( torch._C._distributed_c10d._register_work(ret, Work()) return ret - +@torch.library.impl(lib, "_low_contention_reduce_scatter", "XPU") @torch.library.impl(lib, "_low_contention_reduce_scatter", "CUDA") def _low_contention_reduce_scatter( tensor: torch.Tensor, From c1d75fb783bae4bf1a650b45667315a3d985e1ad Mon Sep 17 00:00:00 2001 From: Han Chao Date: Mon, 25 Aug 2025 14:23:27 +0800 Subject: [PATCH 09/10] update --- test/distributed/test_symmetric_memory.py | 188 +++++++++--------- .../distributed/_symmetric_memory/__init__.py | 53 +++-- 2 files changed, 118 insertions(+), 123 deletions(-) diff --git a/test/distributed/test_symmetric_memory.py b/test/distributed/test_symmetric_memory.py index 8327a5611ef4c..c5e0c8315b66c 100644 --- a/test/distributed/test_symmetric_memory.py +++ b/test/distributed/test_symmetric_memory.py @@ -24,7 +24,7 @@ from torch.testing._internal.common_cuda import _get_torch_cuda_version, SM90OrLater from torch.testing._internal.common_device_type import e4m3_type from torch.testing._internal.common_distributed import ( - MultiProcContinuousTest, + MultiProcContinousTest, MultiProcessTestCase, requires_multicast_support, skip_if_lt_x_gpu, @@ -46,20 +46,20 @@ test_contexts = [nullcontext, _test_mode] # So that tests are written in device-agnostic way -device_type = "cuda" +device_type = "xpu" device_module = torch.get_device_module(device_type) @instantiate_parametrized_tests -@requires_cuda_p2p_access() -class SymmetricMemoryTest(MultiProcContinuousTest): +# @requires_cuda_p2p_access() +class SymmetricMemoryTest(MultiProcContinousTest): @property def device(self) -> torch.device: return torch.device(device_type, self.rank) def _init_process(self, set_device: bool = True): if set_device: - torch.cuda.set_device(self.device) + torch.xpu.set_device(self.device) torch.manual_seed(42 + self.rank) def test_has_multicast_support(self) -> None: @@ -67,12 +67,12 @@ def test_has_multicast_support(self) -> None: self.assertFalse(_SymmetricMemory.has_multicast_support(DeviceType.CPU, 0)) # NOTE: DeviceType.CUDA is implicitly tested through @requires_multicast_support - @skipIfRocm - @skip_if_lt_x_gpu(2) + # @skipIfRocm + # @skip_if_lt_x_gpu(2) def test_get_backend(self) -> None: - backend = symm_mem.get_backend(torch.device("cuda")) + backend = symm_mem.get_backend(torch.device("xpu")) self.assertIsNotNone(backend) - backend = symm_mem.get_backend("cuda") + backend = symm_mem.get_backend("xpu") self.assertIsNotNone(backend) @skipIfRocm @@ -83,13 +83,13 @@ def test_cuda_nvlink_connectivity_detection(self) -> None: connectivity = _detect_dma_connectivity(DeviceType.CUDA, "nvlink") self.assertEqual(connectivity.device_type, DeviceType.CUDA) self.assertEqual(connectivity.connection_type, "nvlink") - self.assertEqual(len(connectivity.matrix), torch.cuda.device_count()) + self.assertEqual(len(connectivity.matrix), torch.xpu.device_count()) for row in connectivity.matrix: - self.assertEqual(len(row), torch.cuda.device_count()) + self.assertEqual(len(row), torch.xpu.device_count()) - @runOnRocmArch(MI300_ARCH) + # @runOnRocmArch(MI300_ARCH) def test_large_alloc(self) -> None: - t = symm_mem.empty(2 * 1024**3, dtype=torch.uint8, device="cuda") + t = symm_mem.empty(2 * 1024**3, dtype=torch.uint8, device="xpu") self.assertEqual(t.numel() * t.element_size(), 2 * 1024**3) def _get_test_alloc_args(self): @@ -127,8 +127,8 @@ def _verify_symmetric_memory(self, symm_mem_hdl): symm_mem_hdl.barrier() - @runOnRocmArch(MI300_ARCH) - @skip_if_lt_x_gpu(2) + # @runOnRocmArch(MI300_ARCH) + # @skip_if_lt_x_gpu(2) @parametrize("set_device", [True, False]) def test_empty_strided_p2p(self, set_device: bool) -> None: self._init_process(set_device) @@ -145,8 +145,8 @@ def test_empty_strided_p2p(self, set_device: bool) -> None: del t self._verify_symmetric_memory(symm_mem_hdl) - @skipIfRocm # started failing during ROCm 6.4 CI upgrade - @skip_if_lt_x_gpu(2) + # @skipIfRocm # started failing during ROCm 6.4 CI upgrade + # @skip_if_lt_x_gpu(2) @parametrize("set_device", [True, False]) def test_empty_strided_p2p_persistent(self, set_device: bool) -> None: self._init_process(set_device) @@ -172,12 +172,12 @@ def test_empty_strided_p2p_persistent(self, set_device: bool) -> None: symm_mem_hdl = _SymmetricMemory.rendezvous(t) self._verify_symmetric_memory(symm_mem_hdl) - @runOnRocmArch(MI300_ARCH) - @skip_if_lt_x_gpu(2) + # @runOnRocmArch(MI300_ARCH) + # @skip_if_lt_x_gpu(2) def test_get_signal_pad(self) -> None: self._init_process() - t = symm_mem.empty(1, device="cuda") + t = symm_mem.empty(1, device="xpu") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) peer_rank = (self.rank + 1) % self.world_size @@ -206,18 +206,18 @@ def test_get_signal_pad(self) -> None: self.assertEqual(signal_pad.numel(), 64) # Sanity check that writes to buffer doesn't corrupt signal_pad - t = symm_mem.empty(0, device="cuda") + t = symm_mem.empty(0, device="xpu") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) signal_pad = symm_mem_hdl.get_signal_pad(self.rank) signal_pad.fill_(42) t.fill_(0) self.assertTrue(signal_pad.eq(42).all()) - @runOnRocmArch(MI300_ARCH) - @requires_cuda + # @runOnRocmArch(MI300_ARCH) + # @requires_cuda def test_allow_overlapping_devices(self) -> None: os.environ["TORCH_SYMM_MEM_ALLOW_OVERLAPPING_DEVICES"] = "1" - t = symm_mem.empty(64, device="cuda:0") + t = symm_mem.empty(64, device="xpu:0") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) self.assertEqual(symm_mem_hdl.rank, self.rank) @@ -232,8 +232,8 @@ def test_allow_overlapping_devices(self) -> None: os.environ["TORCH_SYMM_MEM_ALLOW_OVERLAPPING_DEVICES"] = "0" - @runOnRocmArch(MI300_ARCH) - @skip_if_lt_x_gpu(2) + # @runOnRocmArch(MI300_ARCH) + # @skip_if_lt_x_gpu(2) @parametrize("gather_dim", [0, 1]) def test_fused_all_gather_matmul(self, gather_dim: int) -> None: self._init_process() @@ -246,8 +246,8 @@ def test_fused_all_gather_matmul(self, gather_dim: int) -> None: rank = self.rank torch.manual_seed(42 + rank) - A_shard = torch.rand(BATCH, M // self.world_size, K, device="cuda") - Bs = [torch.rand(K, N, device="cuda") for _ in range(3)] + A_shard = torch.rand(BATCH, M // self.world_size, K, device="xpu") + Bs = [torch.rand(K, N, device="xpu") for _ in range(3)] ag_output_0, mm_outputs_0 = _fused_all_gather_matmul_fallback( A_shard, Bs, gather_dim=gather_dim, group_name=group.group_name @@ -262,12 +262,12 @@ def test_fused_all_gather_matmul(self, gather_dim: int) -> None: assert torch.allclose(mm_output_0, mm_output_1) assert mm_output_0.stride(), mm_output_1.stride() - @skipIfRocm # this requires async_input_mm support - @skipIf( - not SM90OrLater, - "_fused_all_gather_matmul_native currently only supports sm>=90", - ) - @skip_if_lt_x_gpu(2) + # @skipIfRocm # this requires async_input_mm support + # @skipIf( + # not SM90OrLater, + # "_fused_all_gather_matmul_native currently only supports sm>=90", + # ) + # @skip_if_lt_x_gpu(2) @parametrize("symm_mem_input", [True, False]) @parametrize("is_b_row_major", [True, False]) def test_fused_all_gather_matmul_native( @@ -293,20 +293,20 @@ def test_fused_all_gather_matmul_native( ).normal_() else: A_shard = torch.rand( - M // self.world_size, K, dtype=torch.bfloat16, device="cuda" + M // self.world_size, K, dtype=torch.bfloat16, device="xpu" ) if is_b_row_major: - B = torch.rand(K, N, dtype=torch.bfloat16, device="cuda") + B = torch.rand(K, N, dtype=torch.bfloat16, device="xpu") else: - B = torch.rand(N, K, dtype=torch.bfloat16, device="cuda").t() + B = torch.rand(N, K, dtype=torch.bfloat16, device="xpu").t() ag_baseline, mm_baseline = _fused_all_gather_matmul_fallback( A_shard, [B], gather_dim=0, group_name=group_name ) with torch.profiler.profile( activities=[ - torch.profiler.ProfilerActivity.CUDA, + torch.profiler.ProfilerActivity.XPU, ], ) as prof: ag_target, mm_target = torch.ops.symm_mem.fused_all_gather_matmul( @@ -335,10 +335,10 @@ def test_multimem_all_gather_matmul(self) -> None: torch.manual_seed(42 + self.rank) A_shard = torch.rand( - M // self.world_size, K, dtype=torch.bfloat16, device="cuda" + M // self.world_size, K, dtype=torch.bfloat16, device="xpu" ) - B = torch.rand(K, N, dtype=torch.bfloat16, device="cuda") + B = torch.rand(K, N, dtype=torch.bfloat16, device="xpu") ag_baseline, mm_baseline = _fused_all_gather_matmul_fallback( A_shard, [B], gather_dim=0, group_name=group_name, return_A=False @@ -359,8 +359,8 @@ def test_multimem_all_gather_matmul(self) -> None: torch.testing.assert_close(ag_target, ag_baseline) torch.testing.assert_close(mm_target[0], mm_baseline[0]) - @runOnRocmArch(MI300_ARCH) - @skip_if_lt_x_gpu(2) + # @runOnRocmArch(MI300_ARCH) + # @skip_if_lt_x_gpu(2) @parametrize("gather_dim", [0, 1]) @parametrize( "scale_mode", ["tensor-wise", "row-wise-replicated", "row-wise-sharded"] @@ -386,20 +386,20 @@ def test_fused_all_gather_scaled_matmul( torch.manual_seed(42 + rank) - A_shard = torch.rand(*leading_dims, K, device="cuda").to(e4m3_type) - Bs = [torch.rand(N, K, device="cuda").to(e4m3_type).T for _ in range(3)] + A_shard = torch.rand(*leading_dims, K, device="xpu").to(e4m3_type) + Bs = [torch.rand(N, K, device="xpu").to(e4m3_type).T for _ in range(3)] if scale_mode == "tensor-wise": - A_scale = torch.tensor(0.1, device="cuda") - B_scales = [torch.tensor(0.1, device="cuda") for _ in range(3)] + A_scale = torch.tensor(0.1, device="xpu") + B_scales = [torch.tensor(0.1, device="xpu") for _ in range(3)] out_dtypes = [None, torch.bfloat16, torch.float32] elif scale_mode == "row-wise-sharded": - A_scale = torch.full((*leading_dims, 1), 0.1, device="cuda") - B_scales = [torch.full((1, N), 0.1, device="cuda") for _ in range(3)] + A_scale = torch.full((*leading_dims, 1), 0.1, device="xpu") + B_scales = [torch.full((1, N), 0.1, device="xpu") for _ in range(3)] out_dtypes = [torch.bfloat16] * 3 elif scale_mode == "row-wise-replicated": - A_scale = torch.full((BATCH, M, 1), 0.1, device="cuda") - B_scales = [torch.full((1, N), 0.1, device="cuda") for _ in range(3)] + A_scale = torch.full((BATCH, M, 1), 0.1, device="xpu") + B_scales = [torch.full((1, N), 0.1, device="xpu") for _ in range(3)] out_dtypes = [torch.bfloat16] * 3 else: raise AssertionError(f"Invalid scale_mode: {scale_mode}") @@ -445,7 +445,7 @@ def test_fused_all_gather_scaled_matmul( self.assertEqual(mm_output_0.stride(), mm_output_1.stride()) self.assertEqual(mm_output_0.dtype, mm_output_1.dtype) - @runOnRocmArch(MI300_ARCH) + # @runOnRocmArch(MI300_ARCH) @skip_if_lt_x_gpu(2) @parametrize("scatter_dim", [0, 1]) def test_fused_matmul_reduce_scatter(self, scatter_dim: int) -> None: @@ -459,8 +459,8 @@ def test_fused_matmul_reduce_scatter(self, scatter_dim: int) -> None: rank = self.rank torch.manual_seed(42 + rank) - A = torch.rand(BATCH, M, K, device="cuda") - B = torch.rand(K, N, device="cuda") + A = torch.rand(BATCH, M, K, device="xpu") + B = torch.rand(K, N, device="xpu") output_0 = _fused_matmul_reduce_scatter_fallback( A, B, "avg", scatter_dim=scatter_dim, group_name=group.group_name @@ -472,7 +472,7 @@ def test_fused_matmul_reduce_scatter(self, scatter_dim: int) -> None: assert torch.allclose(output_0, output_1) assert output_0.stride() == output_1.stride() - @skipIfRocm # AsyncTP support changed _fused_scaled_matmul_reduce_scatter_fallback API, need more changes + # @skipIfRocm # AsyncTP support changed _fused_scaled_matmul_reduce_scatter_fallback API, need more changes @skip_if_lt_x_gpu(2) @parametrize("scatter_dim", [0, 1]) @parametrize("rowwise", [True, False]) @@ -489,15 +489,15 @@ def test_fused_scaled_matmul_reduce_scatter( rank = self.rank torch.manual_seed(42 + rank) - A = torch.rand(BATCH, M, K, device="cuda").to(e4m3_type) - B = torch.rand(N, K, device="cuda").to(e4m3_type).T + A = torch.rand(BATCH, M, K, device="xpu").to(e4m3_type) + B = torch.rand(N, K, device="xpu").to(e4m3_type).T if rowwise: - A_scale = torch.full((BATCH, M, 1), 0.1, device="cuda") - B_scale = torch.full((1, N), 0.1, device="cuda") + A_scale = torch.full((BATCH, M, 1), 0.1, device="xpu") + B_scale = torch.full((1, N), 0.1, device="xpu") else: - A_scale = torch.tensor(0.1, device="cuda") - B_scale = torch.tensor(0.1, device="cuda") + A_scale = torch.tensor(0.1, device="xpu") + B_scale = torch.tensor(0.1, device="xpu") output_shape = [*A.shape[:-1], B.shape[1]] @@ -522,7 +522,7 @@ def test_fused_scaled_matmul_reduce_scatter( assert outputs[0].stride() == outputs[1].stride() assert torch.allclose(outputs[0], outputs[1]), (outputs[0], outputs[1]) - @runOnRocmArch(MI300_ARCH) + # @runOnRocmArch(MI300_ARCH) @parametrize("dim", [0, 1, 2]) def test_optimal_layout(self, dim: int) -> None: t = torch.rand(8, 64, 32, 16) @@ -535,8 +535,8 @@ def test_optimal_layout(self, dim: int) -> None: self.assertTrue(x.movedim(dim, 0).is_contiguous()) self.assertTrue(torch.allclose(x, t)) - @runOnRocmArch(MI300_ARCH) - @skip_if_lt_x_gpu(2) + # @runOnRocmArch(MI300_ARCH) + # @skip_if_lt_x_gpu(2) @parametrize("symm_mem_input", [True, False]) def test_low_contention_all_gather(self, symm_mem_input: bool) -> None: self._init_process() @@ -560,7 +560,7 @@ def test_low_contention_all_gather(self, symm_mem_input: bool) -> None: for r in range(self.world_size): self.assertTrue(chunks[r].eq(r).all()) - @runOnRocmArch(MI300_ARCH) + # @runOnRocmArch(MI300_ARCH) @skip_if_lt_x_gpu(2) @parametrize("reduce_op", ["sum", "avg"]) @parametrize("symm_mem_input", [True, False]) @@ -608,7 +608,7 @@ def test_subgroup(self) -> None: world = dist.group.WORLD subgroup = subgroup_0 if world.rank() < world.size() // 2 else subgroup_1 - t = symm_mem.empty(64, device="cuda") + t = symm_mem.empty(64, device="xpu") symm_mem_world = symm_mem.rendezvous(t, group=world) symm_mem_subgroup = symm_mem.rendezvous(t, group=subgroup) @@ -636,8 +636,8 @@ def test_subgroup(self) -> None: # This Test class is used to test the error handling of SymmetricMemory APIs. # Since a process restart is often needed after each test, we use the -# MultiProcessTestCase instead of MultiProcContinuousTest. -@requires_cuda_p2p_access() +# MultiProcessTestCase instead of MultiProcContinousTest. +# @requires_cuda_p2p_access() class SymmMemNegativeTest(MultiProcessTestCase): def setUp(self) -> None: super().setUp() @@ -652,10 +652,10 @@ def device(self) -> torch.device: return torch.device(device_type, self.rank) def _init_process(self): - torch.cuda.set_device(self.device) + torch.xpu.set_device(self.device) store = dist.FileStore(self.file_name, self.world_size) dist.init_process_group( - backend="nccl", + backend="xccl", world_size=self.world_size, rank=self.rank, store=store, @@ -672,15 +672,15 @@ def _init_process(self): def test_barrier_timeout(self) -> None: self._init_process() - t = symm_mem.empty(1, device="cuda") + t = symm_mem.empty(1, device="xpu") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) if self.rank == 0: with self.assertRaises(RuntimeError): symm_mem_hdl.barrier(timeout_ms=1000) - torch.cuda.synchronize() + torch.xpu.synchronize() else: - torch.cuda.synchronize() + torch.xpu.synchronize() # The device-side timeout triggers a __trap() that causes all # subsequent host/device interactions to result in an "unspecified @@ -698,7 +698,7 @@ def test_barrier_timeout(self) -> None: def test_put_signal_timeout(self) -> None: self._init_process() - t = symm_mem.empty(1, device="cuda") + t = symm_mem.empty(1, device="xpu") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) if self.rank == 0: @@ -707,9 +707,9 @@ def test_put_signal_timeout(self) -> None: # doesn't wait on this signal, the subsequent put will timeout. symm_mem_hdl.put_signal(dst_rank=1) symm_mem_hdl.put_signal(dst_rank=1, timeout_ms=1000) - torch.cuda.synchronize() + torch.xpu.synchronize() else: - torch.cuda.synchronize() + torch.xpu.synchronize() # The device-side timeout triggers a __trap() that causes all # subsequent host/device interactions to result in an "unspecified @@ -727,15 +727,15 @@ def test_put_signal_timeout(self) -> None: def test_wait_signal_timeout(self) -> None: self._init_process() - t = symm_mem.empty(1, device="cuda") + t = symm_mem.empty(1, device="xpu") symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD) if self.rank == 0: with self.assertRaises(RuntimeError): symm_mem_hdl.wait_signal(src_rank=1, timeout_ms=1000) - torch.cuda.synchronize() + torch.xpu.synchronize() else: - torch.cuda.synchronize() + torch.xpu.synchronize() # The device-side timeout triggers a __trap() that causes all # subsequent host/device interactions to result in an "unspecified @@ -745,14 +745,14 @@ def test_wait_signal_timeout(self) -> None: @instantiate_parametrized_tests -@requires_cuda_p2p_access() -class SymmMemCollectiveTest(MultiProcContinuousTest): +# @requires_cuda_p2p_access() +class SymmMemCollectiveTest(MultiProcContinousTest): @property def device(self) -> torch.device: return torch.device(device_type, self.rank) def _init_process(self): - torch.cuda.set_device(self.device) + torch.xpu.set_device(self.device) torch.manual_seed(42 + self.rank) @skip_if_lt_x_gpu(4) @@ -992,10 +992,10 @@ def test_multimem_all_gather(self, align_bytes: int) -> None: @instantiate_parametrized_tests -@requires_cuda_p2p_access() -class LoweringTest(MultiProcContinuousTest): +# @requires_cuda_p2p_access() +class LoweringTest(MultiProcContinousTest): def _init_process(self) -> None: - torch.cuda.set_device(self.device) + torch.xpu.set_device(self.device) enable_symm_mem_for_group(dist.group.WORLD.group_name) torch.manual_seed(42 + self.rank) torch._inductor.config._collective.auto_select = True @@ -1060,15 +1060,15 @@ def func_3(x): class SymmMemSingleProcTest(TestCase): - @requires_cuda - @skipIf( - not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 0), - "stream_write_value32 currently only supports cuda version>=12.0", - ) - @runOnRocmArch(MI300_ARCH) + # @requires_cuda + # @skipIf( + # not TEST_WITH_ROCM and _get_torch_cuda_version() < (12, 0), + # "stream_write_value32 currently only supports xpu version>=12.0", + # ) + # @runOnRocmArch(MI300_ARCH) def test_stream_write_value32(self): - tensor = torch.zeros(4, dtype=torch.uint32, device="cuda") - expect = torch.tril(torch.ones(4, 4, device="cuda")).to(torch.uint32) + tensor = torch.zeros(4, dtype=torch.uint32, device="xpu") + expect = torch.tril(torch.ones(4, 4, device="xpu")).to(torch.uint32) for i in range(4): _SymmetricMemory.stream_write_value32(tensor, i, 1) @@ -1080,14 +1080,14 @@ def test_stream_write_value32(self): with self.assertRaises(RuntimeError): _SymmetricMemory.stream_write_value32(tensor, offset=0, val=4294967296) - @requires_cuda + # @requires_cuda @runOnRocmArch(MI300_ARCH) def test_memset32(self): t = _SymmetricMemory.empty_strided_p2p( (64,), (1,), dtype=torch.uint32, - device=torch.device("cuda:0"), + device=torch.device("xpu:0"), group_name="0", ).fill_(0) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 5881b01e47032..c3f54db0b8e9b 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -168,8 +168,8 @@ def _pipelined_multi_all_gather_and_consume( group_size = symm_mem.world_size rank = symm_mem.rank - # symm_mem.barrier(channel=0) - dist.barrier() + symm_mem.barrier(channel=0) + # dist.barrier() backend_stream = _get_backend_stream() backend_stream.wait_stream(torch.xpu.current_stream()) @@ -263,29 +263,29 @@ def get_p2p_bufs(remote_rank: int) -> list[torch.Tensor]: # remote shards. shard_consumer(shard, rank) - # for step in range(1, group_size): - # if step % 2 == 0: - # stream = torch.xpu.current_stream() - # else: - # stream = backend_stream - # remote_rank = (step + rank) % group_size - # remote_p2p_bufs = get_p2p_bufs(remote_rank) - # with stream: - # copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) - # shard_consumer(shards[remote_rank], remote_rank) - - stage_size = 2 - for stage in range(1, group_size, stage_size): - stream = torch.xpu.current_stream() if (stage // stage_size) % 2 == 0 else backend_stream + for step in range(1, group_size): + if step % 2 == 0: + stream = torch.xpu.current_stream() + else: + stream = backend_stream + remote_rank = (step + rank) % group_size + remote_p2p_bufs = get_p2p_bufs(remote_rank) with stream: - for i in range(stage_size): - step = stage + i - if step >= group_size: - break - remote_rank = (step + rank) % group_size - remote_p2p_bufs = get_p2p_bufs(remote_rank) - copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) - shard_consumer(shards[remote_rank], remote_rank) + copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) + shard_consumer(shards[remote_rank], remote_rank) + + # stage_size = 2 + # for stage in range(1, group_size, stage_size): + # stream = torch.xpu.current_stream() if (stage // stage_size) % 2 == 0 else backend_stream + # with stream: + # for i in range(stage_size): + # step = stage + i + # if step >= group_size: + # break + # remote_rank = (step + rank) % group_size + # remote_p2p_bufs = get_p2p_bufs(remote_rank) + # copy_shard(dst=shards[remote_rank], src=remote_p2p_bufs) + # shard_consumer(shards[remote_rank], remote_rank) if ag_out_needed: # Copy from input to the all-gather output. Opportunistically overlap @@ -1250,11 +1250,6 @@ def _fused_scaled_matmul_reduce_scatter_impl( .flatten(0, -2) ) A_scale_shards = list(A_scale.chunk(group.size())) - # cuBLAS's row-wise kernel requires scales to be aligned to 16 bytes. - # When we slice them we might break this and need to reallocate them. - A_scale_shards = [ - t if t.data_ptr() % 16 == 0 else t.clone() for t in A_scale_shards - ] else: raise ValueError("A_scale cannot be none for scaled_mm") From 3cc755e4e641493304485324249cd55699bbe0d3 Mon Sep 17 00:00:00 2001 From: Han Chao Date: Wed, 10 Sep 2025 13:41:41 +0800 Subject: [PATCH 10/10] Revert "Enable XPU built with Level Zero (#321)" This reverts commit 97b55d14461a555ab9c723aede51f9cd30bb2b7a. --- CMakeLists.txt | 1 - cmake/Dependencies.cmake | 1 - cmake/Modules/FindLevelZero.cmake | 59 ------------------------------- cmake/Summary.cmake | 4 --- cmake/public/xpu.cmake | 28 +-------------- 5 files changed, 1 insertion(+), 92 deletions(-) delete mode 100644 cmake/Modules/FindLevelZero.cmake diff --git a/CMakeLists.txt b/CMakeLists.txt index 77f24490d56f2..ad7368e192983 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -236,7 +236,6 @@ option(USE_ASAN "Use Address+Undefined Sanitizers" OFF) option(USE_TSAN "Use Thread Sanitizer" OFF) option(USE_CUDA "Use CUDA" ON) option(USE_XPU "Use XPU" ON) -cmake_dependent_option(USE_LEVEL_ZERO "Use Level Zero for XPU backend." ON "USE_XPU" OFF) cmake_dependent_option( BUILD_LAZY_CUDA_LINALG "Build cuda linalg ops as separate library" ON "USE_CUDA AND LINUX AND BUILD_PYTHON" OFF) diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 409fad5d378f9..944c7821f6676 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -95,7 +95,6 @@ if(USE_XPU) message(WARNING "Not compiling with XPU. Could NOT find SYCL. " "Suppress this warning with -DUSE_XPU=OFF.") caffe2_update_option(USE_XPU OFF) - caffe2_update_option(USE_LEVEL_ZERO OFF) endif() foreach(flag ${XPU_HOST_CXX_FLAGS}) add_definitions(${flag}) diff --git a/cmake/Modules/FindLevelZero.cmake b/cmake/Modules/FindLevelZero.cmake deleted file mode 100644 index c8b9720f84432..0000000000000 --- a/cmake/Modules/FindLevelZero.cmake +++ /dev/null @@ -1,59 +0,0 @@ -# This will define the following variables: -# LevelZero_FOUND : True if the system has the LevelZero library. -# LevelZero_INCLUDE_DIR : Level Zero include directory. -# LevelZero_LIBRARY : Level Zero library directory. - -include(FindPackageHandleStandardArgs) - -# Find Level Zero include directory. -if(CMAKE_SYSTEM_NAME MATCHES "Linux") - find_path( - LevelZero_INCLUDE_DIR - NAMES level_zero/ze_api.h - PATH_SUFFIXES include - ) -elseif(CMAKE_SYSTEM_NAME MATCHES "Windows") - find_path( - LevelZero_INCLUDE_DIR - NAMES level_zero/ze_api.h - HINTS $ENV{LEVEL_ZERO_V1_SDK_PATH} - PATH_SUFFIXES include - ) -endif() - -if(NOT LevelZero_INCLUDE_DIR) - set(LevelZero_FOUND False) - set(LevelZero_REASON_FAILURE "Level Zero include directory not found!!") - set(LevelZero_NOT_FOUND_MESSAGE "${LevelZero_REASON_FAILURE}") - return() -endif() - -# Find Level Zero library directory. -if(CMAKE_SYSTEM_NAME MATCHES "Linux") - find_library( - LevelZero_LIBRARY - NAMES ze_loader - PATH_SUFFIXES x86_64_linux_gnu lib lib/x64 lib64 - ) -elseif(CMAKE_SYSTEM_NAME MATCHES "Windows") - find_library( - LevelZero_LIBRARY - NAMES ze_loader - HINTS $ENV{LEVEL_ZERO_V1_SDK_PATH} - PATH_SUFFIXES lib lib64 - ) -endif() - -if(NOT LevelZero_LIBRARY) - set(LevelZero_FOUND False) - set(LevelZero_REASON_FAILURE "Level Zero library not found!!") - set(LevelZero_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}") - return() -endif() - -find_package_handle_standard_args( - LevelZero - FOUND_VAR LevelZero_FOUND - REQUIRED_VARS LevelZero_INCLUDE_DIR LevelZero_LIBRARY - REASON_FAILURE_MESSAGE "${LevelZero_REASON_FAILURE}" - ) diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake index 1f10d8ed9badd..745d9ea058687 100644 --- a/cmake/Summary.cmake +++ b/cmake/Summary.cmake @@ -124,10 +124,6 @@ function(caffe2_print_configuration_summary) message(STATUS " SYCL compiler version : ${SYCL_COMPILER_VERSION}") message(STATUS " SYCL include path : ${SYCL_INCLUDE_DIR}") message(STATUS " SYCL library : ${SYCL_LIBRARY}") - if(${USE_LEVEL_ZERO}) - message(STATUS " Level Zero include path : ${LevelZero_INCLUDE_DIR}") - message(STATUS " Level Zero library : ${LevelZero_LIBRARY}") - endif() endif() message(STATUS " USE_ROCM : ${USE_ROCM}") if(${USE_ROCM}) diff --git a/cmake/public/xpu.cmake b/cmake/public/xpu.cmake index b85ca74466dbc..b39e31d0ade8a 100644 --- a/cmake/public/xpu.cmake +++ b/cmake/public/xpu.cmake @@ -14,14 +14,6 @@ if(NOT SYCL_FOUND) # Exit early to avoid populating XPU_HOST_CXX_FLAGS. return() endif() -# Find LevelZero library. -if(USE_LEVEL_ZERO) - find_package(LevelZero REQUIRED) - if(NOT LevelZero_FOUND) - set(PYTORCH_FOUND_XPU FALSE) - return() - endif() -endif() set(PYTORCH_FOUND_XPU TRUE) # SYCL library interface @@ -34,29 +26,11 @@ set_property( TARGET torch::sycl PROPERTY INTERFACE_LINK_LIBRARIES ${SYCL_LIBRARY}) -# LevelZero library interface -if(USE_LEVEL_ZERO) - add_library(torch::level_zero INTERFACE IMPORTED) - - set_property( - TARGET torch::level_zero PROPERTY INTERFACE_INCLUDE_DIRECTORIES - ${LevelZero_INCLUDE_DIR}) - set_property( - TARGET torch::level_zero PROPERTY INTERFACE_LINK_LIBRARIES - ${LevelZero_LIBRARY}) -endif() - # xpurt add_library(torch::xpurt INTERFACE IMPORTED) - -set(xpurt_deps torch::sycl) -if(USE_LEVEL_ZERO) - list(APPEND xpurt_deps torch::level_zero) -endif() - set_property( TARGET torch::xpurt PROPERTY INTERFACE_LINK_LIBRARIES - "${xpurt_deps}") + torch::sycl) # setting xpu arch flags torch_xpu_get_arch_list(XPU_ARCH_FLAGS)