From 9fc6016c1e6d57df597b5d65061db22c76375bee Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 13:54:32 +0800 Subject: [PATCH 01/10] [RUNTIME][Disco] Add CPU collective communication support [RUNTIME][Disco] Add CPU collective communication support [FixBugs] Add build_ring and UploadModule() --- CMakeLists.txt | 12 +- include/tvm/runtime/disco/disco_worker.h | 7 +- include/tvm/runtime/disco/session.h | 64 +++- python/tvm/exec/disco_worker.py | 17 +- python/tvm/runtime/disco/process_pool.py | 51 ++- python/tvm/runtime/disco/session.py | 37 +- src/runtime/extra/disco/bcast_session.cc | 5 + src/runtime/extra/disco/bcast_session.h | 24 ++ src/runtime/extra/disco/builtin.cc | 21 ++ src/runtime/extra/disco/cpuccl/cpuccl.cc | 324 ++++++++++++++++++ .../extra/disco/distributed/socket_session.cc | 252 ++++++++++++-- src/runtime/extra/disco/process_session.cc | 47 ++- src/runtime/extra/disco/threaded_session.cc | 35 +- src/support/socket.h | 14 + 14 files changed, 866 insertions(+), 44 deletions(-) create mode 100644 src/runtime/extra/disco/cpuccl/cpuccl.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index 567edc1dc698..1d1fd96ce18f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -46,7 +46,7 @@ tvm_option(USE_OPENCL "Build with OpenCL" OFF) tvm_option(USE_OPENCL_ENABLE_HOST_PTR "Enable OpenCL memory object access to host" OFF) tvm_option(USE_OPENCL_GTEST "Path to OpenCL specific gtest version for runtime cpp tests." /path/to/opencl/gtest) tvm_option(USE_VULKAN "Build with Vulkan" OFF) - +tvm_option(USE_CPU_CCL "Build with CPU CCL" OFF) # Whether to use spirv-tools and SPIRV-Headers from Khronos GitHub or GitLab. # @@ -360,7 +360,7 @@ tvm_file_glob(GLOB RUNTIME_SRCS src/runtime/vm/*.cc src/runtime/memory/*.cc src/runtime/rpc/minrpc/*.cc -) + ) # Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra. # Note: src/backend/*/runtime sources move to per-backend DSOs. set(TVM_RUNTIME_EXT_OBJS "") @@ -524,6 +524,14 @@ if(USE_ROCM AND USE_RCCL) target_link_libraries(tvm_runtime_extra PRIVATE tvm_rccl_objs rccl) endif() +# CPUCCL. +if(USE_CPU_CCL) + tvm_file_glob(GLOB _cpu_ccl_srcs src/runtime/extra/disco/cpuccl/*.cc) + add_library(tvm_cpu_ccl_objs OBJECT ${_cpu_ccl_srcs}) + target_link_libraries(tvm_cpu_ccl_objs PRIVATE tvm_runtime_extra_defs) + target_link_libraries(tvm_runtime_extra PRIVATE tvm_cpu_ccl_objs) +endif() + target_link_libraries(tvm_runtime_extra PUBLIC tvm_runtime) # If disco/cuda_ipc is included, link the CUDA DSO. diff --git a/include/tvm/runtime/disco/disco_worker.h b/include/tvm/runtime/disco/disco_worker.h index e5c88e0d215b..68b0e6e8f5fa 100644 --- a/include/tvm/runtime/disco/disco_worker.h +++ b/include/tvm/runtime/disco/disco_worker.h @@ -58,7 +58,9 @@ class DiscoWorker { default_device(Device{DLDeviceType::kDLCPU, 0}), worker_zero_data(worker_zero_data), channel(channel), - register_file{} {} + register_file{}, + ring_in(nullptr), + ring_out(nullptr){} /*! \brief Main loop of the worker */ void MainLoop(); @@ -96,6 +98,9 @@ class DiscoWorker { struct Impl; friend struct DiscoWorker::Impl; + /*! \brief CPU communication ring endpoints assigned by the session during BuildRing(). */ + DiscoRingChannel* ring_in; + DiscoRingChannel* ring_out; }; /*! * \brief A threadlocal wrapper of DiscoWorker. diff --git a/include/tvm/runtime/disco/session.h b/include/tvm/runtime/disco/session.h index 99098606b3f7..6ad7e2ddd204 100644 --- a/include/tvm/runtime/disco/session.h +++ b/include/tvm/runtime/disco/session.h @@ -82,6 +82,8 @@ #include #include +#include + namespace tvm { namespace runtime { @@ -294,8 +296,10 @@ class Session : public ffi::ObjectRef { * \brief Create a session backed by a thread pool of workers * \param num_workers The number of workers. * \param num_groups The number of worker groups. + * \param build_ring Whether to wire the workers into a unidirectional pipe ring + * (worker `i` -> worker `i + 1`), used as the data plane for CPU collectives. */ - TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups); + TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups, bool build_ring); /*! * \brief Create a session backed by pipe-based multiprocessing * \param num_workers The number of workers. @@ -305,12 +309,15 @@ class Session : public ffi::ObjectRef { * When `worker-id` is 0, it shuts down the process pool; Otherwise, it retursn a tuple * (read_fd, writefd) used to communicate with the corresponding worker. * \param entrypoint The entrypoint of DiscoWorker main worker function. + * \param build_ring Whether to wire the workers into a unidirectional pipe ring + * (worker `i` -> worker `i + 1`), used as the data plane for CPU collectives. * \note Worker-0 is always co-located with the controler as a separate thread, and therefore * worker-0 does not exist in the process pool. */ TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups, ffi::String process_pool_creator, - ffi::String entrypoint); + ffi::String entrypoint, + bool build_ring); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj); }; @@ -332,6 +339,59 @@ class DiscoChannel { virtual ffi::PackedArgs RecvReply() = 0; }; +class DiscoRingChannel { + public: + explicit DiscoRingChannel(int fd) : fd_(fd) { + TVM_FFI_ICHECK_GE(fd_, 0) << "DiscoRingChannel: invalid fd " << fd_; + } + ~DiscoRingChannel() { Close(); } + + DiscoRingChannel(const DiscoRingChannel&) = delete; + DiscoRingChannel& operator=(const DiscoRingChannel&) = delete; + DiscoRingChannel(DiscoRingChannel&& other) noexcept : fd_(other.fd_) { other.fd_ = -1; } + + void Send(const void* data, size_t size) { + const char* write_data = static_cast(data); + while (size > 0) { + ssize_t n; + do { n = ::write(fd_, write_data, size); } while (n < 0 && errno == EINTR); + TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Send failed: " << std::strerror(errno); + write_data += n; size -= n; + } + } + + void Recv(void* data, size_t size) { + char* read_data = static_cast(data); + while (size > 0) { + ssize_t n; + do { n = ::read(fd_, read_data, size); } while (n < 0 && errno == EINTR); + TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Recv failed: " << std::strerror(errno); + read_data += n; size -= n; + } + } + + + ssize_t ReadSome(void* data, size_t max_size) { + if (fd_ < 0) return 0; + ssize_t n; + do { n = ::read(fd_, data, max_size); } while (n < 0 && errno == EINTR); + return n; + } + + ssize_t WriteSome(const void* data, size_t size) { + if (fd_ < 0) return -1; + ssize_t n; + do { n = ::write(fd_, data, size); } while (n < 0 && errno == EINTR); + return n; + } + + void Close() { if (fd_ >= 0) { ::close(fd_); fd_ = -1; } } + + private: + int fd_; +}; + + /*! * \brief A special communication channel between controler and worker-0, * assuming they are always collocated in the same process. diff --git a/python/tvm/exec/disco_worker.py b/python/tvm/exec/disco_worker.py index 2052ae7b04ca..6330e12fd435 100644 --- a/python/tvm/exec/disco_worker.py +++ b/python/tvm/exec/disco_worker.py @@ -101,12 +101,20 @@ def fget_item(param_name: str, param_index: int) -> Tensor: def main(): """Main worker function""" - if len(sys.argv) != 6: - print("Usage: ") + argc = len(sys.argv) + if argc not in (6, 8): + print( + "Usage: " + " [ ]") return worker_id = int(sys.argv[1]) num_workers = int(sys.argv[2]) num_groups = int(sys.argv[3]) + + ring_out_fd = -1 + ring_in_fd = -1 + has_ring = (argc == 8) + if sys.platform == "win32": import msvcrt # pylint: disable=import-outside-toplevel,import-error @@ -115,9 +123,12 @@ def main(): else: reader = int(sys.argv[4]) writer = int(sys.argv[5]) + if has_ring: + ring_in_fd = int(sys.argv[6]) + ring_out_fd = int(sys.argv[7]) worker_func = get_global_func("runtime.disco.WorkerProcess") - worker_func(worker_id, num_workers, num_groups, reader, writer) + worker_func(worker_id, num_workers, num_groups, reader, writer, ring_in_fd, ring_out_fd) if __name__ == "__main__": diff --git a/python/tvm/runtime/disco/process_pool.py b/python/tvm/runtime/disco/process_pool.py index 93444f94c2aa..4020424fa5d4 100644 --- a/python/tvm/runtime/disco/process_pool.py +++ b/python/tvm/runtime/disco/process_pool.py @@ -56,6 +56,8 @@ def __init__( # pylint: disable=too-many-arguments entrypoint: str = "tvm.exec.disco_worker", stdout=None, stderr=None, + ring_in_fd: int = -1, + ring_out_fd: int = -1, ): self.worker_id = worker_id self.num_workers = num_workers @@ -64,6 +66,8 @@ def __init__( # pylint: disable=too-many-arguments self._proc = None self._stdout = stdout self._stderr = stderr + self.ring_out_fd = ring_out_fd + self.ring_in_fd = ring_in_fd def __del__(self): try: @@ -117,6 +121,8 @@ def start(self): main_read, worker_write = os.pipe() worker_read, main_write = os.pipe() + has_ring = self.ring_out_fd >= 0 and self.ring_in_fd >= 0 + cmd = [ sys.executable, "-m", @@ -133,6 +139,8 @@ def start(self): os.set_handle_inheritable(worker_read_handle, True) os.set_handle_inheritable(worker_write_handle, True) cmd += [str(worker_read_handle), str(worker_write_handle)] + #CPU ring is not implemented on Windows. + cmd += ["-1", "-1"] self._proc = subprocess.Popen( cmd, close_fds=False, @@ -141,9 +149,17 @@ def start(self): ) else: cmd += [str(worker_read), str(worker_write)] + pass_fds = [worker_read, worker_write] + + if has_ring: + cmd += [str(self.ring_in_fd), str(self.ring_out_fd)] + pass_fds += [self.ring_in_fd, self.ring_out_fd] + else: + cmd += ["-1", "-1"] + self._proc = subprocess.Popen( # pylint: disable=consider-using-with cmd, - pass_fds=(worker_read, worker_write), + pass_fds=tuple(pass_fds), stdout=self._stdout, stderr=self._stderr, ) @@ -151,6 +167,9 @@ def start(self): # close worker side of the pipe os.close(worker_read) os.close(worker_write) + if has_ring: + os.close(self.ring_out_fd) + os.close(self.ring_in_fd) return main_read, main_write @@ -178,12 +197,36 @@ def _kill_child_processes(pid): @register_global_func("runtime.disco.create_process_pool") -def _create_process_pool(num_workers: int, num_groups: int, entrypoint: str): +def _create_process_pool(num_workers: int, num_groups: int, entrypoint: str, build_ring: bool): """Create a process pool where the workers' are [1, num_workers).""" - pool = [DiscoPopenWorker(i, num_workers, num_groups, entrypoint) for i in range(1, num_workers)] + ring_pipes = [] + if build_ring: + for _ in range(num_workers): + ring_pipes.append(os.pipe()) + + pool = [] + for i in range(1, num_workers): + ring_out_fd = -1 + ring_in_fd = -1 + if build_ring: + ring_out_fd = ring_pipes[i][1] # write end of pipe[i] + ring_in_fd = ring_pipes[(i - 1) % num_workers][0] # read end of pipe[i-1] + + pool.append(DiscoPopenWorker( + i, num_workers, num_groups, entrypoint, + ring_in_fd=ring_in_fd, ring_out_fd=ring_out_fd, + )) def result_func(worker_id: int): - nonlocal pool + nonlocal pool, ring_pipes + + # Special ID (-1): Return the ring FDs of worker 0 for the controller. + if worker_id == -1: + assert build_ring, "Ring not enabled; cannot query worker 0 ring fds" + return Shape([ + ring_pipes[num_workers - 1][0], # W0 ring_in (read of pipe[N-1]) + ring_pipes[0][1], # W0 ring_out (write of pipe[0]) + ]) if worker_id != 0: read_fd, write_fd = pool[worker_id - 1].start() return Shape([read_fd, write_fd]) diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py index 0a8125090067..445f088c1687 100644 --- a/python/tvm/runtime/disco/session.py +++ b/python/tvm/runtime/disco/session.py @@ -321,6 +321,30 @@ def load_vm_module( """ func = self._get_cached_method("runtime.disco.load_vm_module") return DModule(func(path, device), self) + + def upload_vm_module(self, path: str) -> str: + """Broadcast a module file to every worker and write it locally on each. + + This reads the file on the controller and broadcasts its bytes over the session, + so each worker writes its own copy at `path` (parent directories are created as needed). + No shared filesystem or scp is required. + + Parameters + ---------- + path : str + The path of the module file. It is read locally on the controller and written to the + same path on every worker; pass this same path to `load_vm_module`. + + Returns + ------- + path : str + The path written on every worker (identical to the argument), for chaining into + `load_vm_module`. + """ + with open(path, "rb") as f: + blob = bytearray(f.read()) + self.get_global_func("runtime.disco.upload_module")(blob, path) + return path def init_ccl(self, ccl: str, *device_ids): """Initialize the underlying communication collective library. @@ -336,7 +360,7 @@ def init_ccl(self, ccl: str, *device_ids): *device_ids : int The device IDs to be used by the underlying communication library. """ - assert ccl in ("nccl", "rccl"), f"Unsupported CCL backend: {ccl}" + assert ccl in ("nccl", "rccl", "cpuccl"), f"Unsupported CCL backend: {ccl}" _ffi_api.SessionInitCCL(self, ccl, Shape(device_ids)) # type: ignore # pylint: disable=no-member self._clear_ipc_memory_pool() @@ -545,12 +569,13 @@ def _clear_ipc_memory_pool(self): class ThreadedSession(Session): """A Disco session backed by multi-threading.""" - def __init__(self, num_workers: int, num_groups: int = 1) -> None: + def __init__(self, num_workers: int, num_groups: int = 1, build_ring: bool = False) -> None: """Create a disco session backed by multiple threads in the same process.""" self.__init_handle_by_constructor__( _ffi_api.SessionThreaded, # type: ignore # pylint: disable=no-member num_workers, num_groups, + build_ring, ) @@ -563,6 +588,7 @@ def __init__( num_workers: int, num_groups: int = 1, entrypoint: str = "tvm.exec.disco_worker", + build_ring: bool = False, ) -> None: self.__init_handle_by_constructor__( _ffi_api.SessionProcess, # type: ignore # pylint: disable=no-member @@ -570,6 +596,7 @@ def __init__( num_groups, "runtime.disco.create_process_pool", entrypoint, + build_ring, ) self._configure_structlog() @@ -597,9 +624,9 @@ def _configure_structlog(self) -> None: @register_global_func("runtime.disco.create_socket_session_local_workers") -def _create_socket_session_local_workers(num_workers) -> Session: +def _create_socket_session_local_workers(num_workers, build_ring) -> Session: """Create the local session for each distributed node over socket session.""" - return ProcessSession(num_workers) + return ProcessSession(num_workers, build_ring=build_ring) @register_object("runtime.disco.SocketSession") @@ -613,6 +640,7 @@ def __init__( num_groups: int, host: str, port: int, + build_ring: bool = False, ) -> None: self.__init_handle_by_constructor__( _ffi_api.SocketSession, # type: ignore # pylint: disable=no-member @@ -621,6 +649,7 @@ def __init__( num_groups, host, port, + build_ring, ) diff --git a/src/runtime/extra/disco/bcast_session.cc b/src/runtime/extra/disco/bcast_session.cc index e6a326eeb544..a70582c1176f 100644 --- a/src/runtime/extra/disco/bcast_session.cc +++ b/src/runtime/extra/disco/bcast_session.cc @@ -65,6 +65,8 @@ void BcastSessionObj::CopyToWorker0(const Tensor& host_array, const DRef& remote } void BcastSessionObj::Shutdown() { + if (shutdown_) return; + shutdown_ = true; BcastSessionObj::Internal::BroadcastUnpacked(this, DiscoAction::kShutDown, 0); } @@ -103,6 +105,9 @@ DRef BcastSessionObj::CallWithPacked(const ffi::PackedArgs& args) { } void BcastSessionObj::DeallocReg(int reg_id) { + // After shutdown the workers (and their pipes) are gone; broadcasting would raise SIGPIPE. + // A DRef may be garbage-collected after the session has been shut down, so guard here. + if (shutdown_) return; BcastSessionObj::Internal::BroadcastUnpacked(this, DiscoAction::kKillReg, reg_id); this->free_regs_.push_back(reg_id); } diff --git a/src/runtime/extra/disco/bcast_session.h b/src/runtime/extra/disco/bcast_session.h index dc9760ee813a..d2b13441b4bf 100644 --- a/src/runtime/extra/disco/bcast_session.h +++ b/src/runtime/extra/disco/bcast_session.h @@ -84,17 +84,41 @@ class BcastSessionObj : public SessionObj { */ virtual ffi::PackedArgs RecvReplyPacked(int worker_id) = 0; + /*! + * \brief Atomically swap W_0's ring_in channel. SocketSession uses this + * to redirect cross-node ring inbound through a proxy pipe. + * Default no-op; ProcessSession overrides. + * \param new_ch New channel to install; ownership taken. + * \return Previously installed channel; ownership returned to caller. + */ + virtual std::unique_ptr RerouteRingIn(std::unique_ptr new_ch) { + return nullptr; + } + + /*! + * \brief Close worker-0's ring channels now, instead of waiting for the object destructor. + * SocketSession calls this on its local session during Shutdown, before joining its proxy + * threads: the send proxy reads the ring pipe whose write end is worker-0's ring_out. With + * a single worker per node that write end lives in-process and is only released at + * destruction, so without closing it here the proxy-thread join deadlocks. Default no-op; + * ProcessSession overrides. + */ + virtual void CloseRing() {} + + /*! \brief A side channel to communicate with worker-0 */ WorkerZeroData worker_zero_data_; /*! \brief Number of registers used, including those in `free_regs_` */ int reg_count_ = 1; /*! \brief The regsiter ids that have been deallocated */ std::vector free_regs_; + bool shutdown_ = false; struct Internal; friend struct Internal; friend class SocketSessionObj; friend class RemoteSocketSession; + friend class RingProxyEndpoint; }; /*! diff --git a/src/runtime/extra/disco/builtin.cc b/src/runtime/extra/disco/builtin.cc index d9d5fc132768..1aaee0a375fb 100644 --- a/src/runtime/extra/disco/builtin.cc +++ b/src/runtime/extra/disco/builtin.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include #include #include @@ -25,7 +26,11 @@ #include #include +#include +#include #include +#include +#include #include "./utils.h" @@ -71,6 +76,21 @@ ffi::Module LoadVMModule(std::string path, ffi::Optional device) { return mod; } +static void MakeParentDirs(const std::string& path) { + for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { + ::mkdir(path.substr(0, pos).c_str(), 0755); + } +} + +void UploadModule(ffi::Bytes data, std::string path) { + if (DiscoWorker::ThreadLocal()->local_worker_id != 0) return; + MakeParentDirs(path); + std::ofstream fs(path, std::ios::binary | std::ios::trunc); + TVM_FFI_ICHECK(fs.is_open()) << "disco.upload_module: Cannot open " << path; + fs.write(data.data(), static_cast(data.size())); + TVM_FFI_ICHECK(fs.good()) << "disco.upload_module: Write failed for " << path; +} + Tensor DiscoEmptyTensor(ffi::Shape shape, DLDataType dtype, ffi::Optional device) { return Tensor::Empty(shape, dtype, UseDefaultDeviceIfNone(device)); } @@ -130,6 +150,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("runtime.disco.load_vm_module", LoadVMModule) + .def("runtime.disco.upload_module", UploadModule) .def("runtime.disco.empty", [](ffi::Shape shape, DLDataType dtype, ffi::Optional device, bool worker0_only, bool in_group) -> ffi::Optional { diff --git a/src/runtime/extra/disco/cpuccl/cpuccl.cc b/src/runtime/extra/disco/cpuccl/cpuccl.cc new file mode 100644 index 000000000000..ab2f6db0f43e --- /dev/null +++ b/src/runtime/extra/disco/cpuccl/cpuccl.cc @@ -0,0 +1,324 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + + +namespace tvm { +namespace runtime { +namespace cpuccl { + +#ifndef TVM_DISCO_DEVICE_NAME +#define TVM_DISCO_DEVICE_NAME "cpu" +#endif +#ifndef TVM_DISCO_CCL_NAME +#define TVM_DISCO_CCL_NAME "cpuccl" +#endif + +/*! \brief Number of bytes taken by a single element of `dtype`. */ +inline int64_t DTypeBytes(DLDataType dtype) { return (dtype.bits * dtype.lanes + 7) / 8; } + +template +static void DivideInPlace(T* data, int64_t n, int div) { + for (int64_t i = 0; i < n; ++i) data[i] /= static_cast(div); +} + +static void DivideInPlaceBytes(void* data, int64_t numel, DLDataType dtype, int div) { + if (dtype == DLDataType{kDLFloat, 32, 1}) + DivideInPlace(static_cast(data), numel, div); + else if (dtype == DLDataType{kDLFloat, 64, 1}) + DivideInPlace(static_cast(data), numel, div); + else if (dtype == DLDataType{kDLInt, 32, 1}) + DivideInPlace(static_cast(data), numel, div); + else if (dtype == DLDataType{kDLInt, 64, 1}) + DivideInPlace(static_cast(data), numel, div); + else + TVM_FFI_THROW(ValueError) << "CPU collective: kAvg unsupported dtype " << dtype; +} + +template +static void ReduceInPlace(T* dst, const T* src, int64_t n, ReduceKind op) { + switch (op) { + case ReduceKind::kSum: + for (int64_t i = 0; i < n; ++i) dst[i] += src[i]; + break; + case ReduceKind::kMax: + for (int64_t i = 0; i < n; ++i) dst[i] = std::max(dst[i], src[i]); + break; + case ReduceKind::kMin: + for (int64_t i = 0; i < n; ++i) dst[i] = std::min(dst[i], src[i]); + break; + case ReduceKind::kProd: + for (int64_t i = 0; i < n; ++i) dst[i] *= src[i]; + break; + default: + TVM_FFI_THROW(ValueError) << "CPU collective: unsupported ReduceKind " + << static_cast(op); + } +} + +static void ReduceBytes(void* dst, const void* src, int64_t numel, DLDataType dtype, + ReduceKind op) { + if (dtype == DLDataType{kDLFloat, 32, 1}) + ReduceInPlace(static_cast(dst), static_cast(src), numel, op); + else if (dtype == DLDataType{kDLFloat, 64, 1}) + ReduceInPlace(static_cast(dst), static_cast(src), numel, op); + else if (dtype == DLDataType{kDLInt, 32, 1}) + ReduceInPlace(static_cast(dst), static_cast(src), numel, op); + else if (dtype == DLDataType{kDLInt, 64, 1}) + ReduceInPlace(static_cast(dst), static_cast(src), numel, op); + else + TVM_FFI_THROW(ValueError) << "CPU collective: unsupported dtype " << dtype; +} + +void InitCCL(Session sess, ffi::Shape device_ids) { + DRef func = sess->GetGlobalFunc("runtime.disco." TVM_DISCO_CCL_NAME ".init_ccl_per_worker"); + DLOG(INFO) << "Initializing " TVM_DISCO_CCL_NAME " with devices: " << device_ids; + sess->CallPacked(func, device_ids); +} + +void InitCCLPerWorker(ffi::Shape device_ids) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + TVM_FFI_ICHECK(worker != nullptr) << "Must be called on a DiscoWorker thread"; + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "CPU ring not built; create session with build_ring=true"; + } + worker->ccl = TVM_DISCO_CCL_NAME; +} + +void SyncWorker() {} + +void AllReduce(Tensor send, ReduceKind reduce_kind, bool /*in_group*/, Tensor recv) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + } + + DLDataType dtype = send->dtype; + int64_t dtype_bytes = DTypeBytes(dtype); + int num_workers = worker->num_workers; + int rank = worker->worker_id; + int64_t numel = send.Shape().Product(); + int64_t bytes = numel * dtype_bytes; + + std::memcpy(recv->data, send->data, static_cast(bytes)); + char* buf = static_cast(recv->data); + + ReduceKind acc_kind = (reduce_kind == ReduceKind::kAvg) ? ReduceKind::kSum : reduce_kind; + + if (num_workers == 1) return; + + TVM_FFI_CHECK_EQ(numel % num_workers, 0, ValueError) + << "cpuccl AllReduce: element count (" << numel << ") must be divisible by num_workers (" + << num_workers << "). A non-divisible size leaves empty chunks in the ring reduce-scatter " + << "and deadlocks. Pad/reshape so numel % num_workers == 0."; + + int64_t chunk_elem = (numel + num_workers - 1) / num_workers; + int64_t chunk_bytes = chunk_elem * dtype_bytes; + std::vector tmp(static_cast(chunk_bytes)); + + // Phase A: reduce-scatter + for (int r = 0; r < num_workers - 1; ++r) { + int si = ((rank - r) % num_workers + num_workers) % num_workers; + int ri = ((rank - r - 1) % num_workers + num_workers) % num_workers; + int64_t s_off = si * chunk_bytes; + int64_t r_off = ri * chunk_bytes; + int64_t s_size = std::min(chunk_bytes, bytes - s_off); + int64_t r_size = std::min(chunk_bytes, bytes - r_off); + if (s_size <= 0 || r_size <= 0) continue; + + worker->ring_out->Send(buf + s_off, static_cast(s_size)); + worker->ring_in->Recv(tmp.data(), static_cast(r_size)); + ReduceBytes(buf + r_off, tmp.data(), r_size / dtype_bytes, dtype, acc_kind); + } + + // Phase B: allgather + for (int r = 0; r < num_workers - 1; ++r) { + int si = ((rank + 1 - r) % num_workers + num_workers) % num_workers; + int ri = ((rank - r) % num_workers + num_workers) % num_workers; + int64_t s_off = si * chunk_bytes; + int64_t r_off = ri * chunk_bytes; + int64_t s_size = std::min(chunk_bytes, bytes - s_off); + int64_t r_size = std::min(chunk_bytes, bytes - r_off); + if (s_size <= 0 || r_size <= 0) continue; + + worker->ring_out->Send(buf + s_off, static_cast(s_size)); + worker->ring_in->Recv(buf + r_off, static_cast(r_size)); + } + + if (reduce_kind == ReduceKind::kAvg) { + DivideInPlaceBytes(buf, numel, dtype, num_workers); + } +} + +void AllGather(Tensor send, bool /*in_group*/, Tensor recv) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + + int num_workers = worker->num_workers; + int rank = worker->worker_id; + int64_t chunk_bytes = send.Shape().Product() * DTypeBytes(send->dtype); + char* buf = static_cast(recv->data); + + std::memcpy(buf + rank * chunk_bytes, send->data, static_cast(chunk_bytes)); + + if (num_workers == 1) return; + + for (int r = 0; r < num_workers - 1; ++r) { + int si = ((rank - r) % num_workers + num_workers) % num_workers; + int ri = ((rank - r - 1) % num_workers + num_workers) % num_workers; + + worker->ring_out->Send(buf + si * chunk_bytes, static_cast(chunk_bytes)); + worker->ring_in->Recv(buf + ri * chunk_bytes, static_cast(chunk_bytes)); + } +} + +void BroadcastFromWorker0(ffi::Optional send, bool /*in_group*/, Tensor recv) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + } + + int rank = worker->worker_id; + int num_workers = worker->num_workers; + int64_t bytes = recv.Shape().Product() * DTypeBytes(recv->dtype); + + if (rank == 0) { + TVM_FFI_CHECK(send.has_value(), ValueError) + << "Worker 0 must provide send buffer for broadcast"; + std::memcpy(recv->data, send.value()->data, static_cast(bytes)); + + if (num_workers > 1) { + worker->ring_out->Send(recv->data, static_cast(bytes)); + } + } else { + worker->ring_in->Recv(recv->data, static_cast(bytes)); + + if (rank < num_workers - 1) { + worker->ring_out->Send(recv->data, static_cast(bytes)); + } + } +} + +void ScatterFromWorker0(ffi::Optional send, bool /*in_group*/, Tensor recv) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + } + + int rank = worker->worker_id; + int num_workers = worker->num_workers; + int64_t recv_bytes = recv.Shape().Product() * DTypeBytes(recv->dtype); + + if (num_workers == 1) { + TVM_FFI_CHECK(send.has_value(), ValueError) + << "Worker 0 must provide send buffer for scatter"; + std::memcpy(recv->data, send.value()->data, static_cast(recv_bytes)); + return; + } + + if (rank == 0) { + TVM_FFI_CHECK(send.has_value(), ValueError) + << "Worker 0 must provide send buffer for scatter"; + char* src = static_cast(send.value()->data); + std::memcpy(recv->data, src, static_cast(recv_bytes)); + int64_t tail = recv_bytes * (num_workers - 1); + worker->ring_out->Send(src + recv_bytes, static_cast(tail)); + } else { + int64_t incoming = recv_bytes * (num_workers - rank); + std::vector buf(static_cast(incoming)); + worker->ring_in->Recv(buf.data(), static_cast(incoming)); + std::memcpy(recv->data, buf.data(), static_cast(recv_bytes)); + + if (rank < num_workers - 1) { + int64_t forward_bytes = incoming - recv_bytes; + worker->ring_out->Send(buf.data() + recv_bytes, static_cast(forward_bytes)); + } + } +} + +void GatherToWorker0(Tensor send, bool /*in_group*/, ffi::Optional recv) { + DiscoWorker* worker = DiscoWorker::ThreadLocal(); + + int rank = worker->worker_id; + int num_workers = worker->num_workers; + int64_t chunk_bytes = send.Shape().Product() * DTypeBytes(send->dtype); + + if (num_workers == 1) { + TVM_FFI_CHECK(recv.has_value(), ValueError) + << "Worker 0 must provide recv buffer for gather"; + std::memcpy(recv.value()->data, send->data, static_cast(chunk_bytes)); + return; + } + + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + + if (rank == 0) { + TVM_FFI_CHECK(recv.has_value(), ValueError) + << "Worker 0 must provide recv buffer for gather"; + char* dst = static_cast(recv.value()->data); + std::memcpy(dst, send->data, static_cast(chunk_bytes)); + int64_t incoming = chunk_bytes * (num_workers - 1); + worker->ring_in->Recv(dst + chunk_bytes, static_cast(incoming)); + } else { + int64_t prior = chunk_bytes * (rank - 1); + std::vector buf(static_cast(prior + chunk_bytes)); + if (rank > 1) { + worker->ring_in->Recv(buf.data(), static_cast(prior)); + } + std::memcpy(buf.data() + prior, send->data, static_cast(chunk_bytes)); + worker->ring_out->Send(buf.data(), buf.size()); + } +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef() + .def("runtime.disco." TVM_DISCO_CCL_NAME ".init_ccl", InitCCL) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".init_ccl_per_worker", InitCCLPerWorker) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".allreduce", AllReduce) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".allgather", AllGather) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".broadcast_from_worker0", BroadcastFromWorker0) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".scatter_from_worker0", ScatterFromWorker0) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".gather_to_worker0", GatherToWorker0) + .def("runtime.disco." TVM_DISCO_CCL_NAME ".sync_worker", SyncWorker); +} + +} // namespace cpuccl +} // namespace runtime +} // namespace tvm diff --git a/src/runtime/extra/disco/distributed/socket_session.cc b/src/runtime/extra/disco/distributed/socket_session.cc index 3fb14914cab1..368358a5043e 100644 --- a/src/runtime/extra/disco/distributed/socket_session.cc +++ b/src/runtime/extra/disco/distributed/socket_session.cc @@ -21,6 +21,7 @@ #include #include +#include #include "../../../../support/socket.h" #include "../bcast_session.h" @@ -54,16 +55,120 @@ class DiscoSocketChannel : public DiscoChannel { DiscoStreamMessageQueue message_queue_; }; + +// Proxy thread helpers — both controller and remote share these. +// send: pipe → TCP +// recv: TCP → pipe +constexpr size_t kProxyBufSize = 64 * 1024; + +void ProxyLoop(DiscoRingChannel* src, DiscoRingChannel* dst, std::string tag) { + + LOG(INFO) << "[Proxy " << tag << "] entering, src=" << src << " dst=" << dst; + std::vector buf(kProxyBufSize); + + while (true) { + ssize_t n = src->ReadSome(buf.data(), buf.size()); + if (n <= 0) { LOG(INFO) << "[Proxy " << tag << "] src closed (n=" << n << "), exit"; return; } + LOG(INFO) << "[Proxy " << tag << "] transfer " << n << " bytes"; + + ssize_t written = 0; + while (written < n) { + ssize_t w = dst->WriteSome(buf.data() + written, n - written); + if (w <= 0) { LOG(INFO) << "[Proxy " << tag << "] dst failed (w=" << w << "), exit"; return; } + written += w; + } + } +} + +class RingProxyEndpoint { + public: + void Connect(int node_id, int num_nodes, int base_ring_port, + const std::vector& node_hosts, const BcastSession& local_session, + const std::string& tag) { + const int next_node_id = (node_id + 1) % num_nodes; + const int recv_port = base_ring_port + node_id; + const int send_port = base_ring_port + next_node_id; + + LOG(INFO) << "[Ring " << tag << "] node_id=" << node_id << " listen=:" << recv_port + << " connect_to=" << node_hosts[next_node_id].c_str() << ":" << send_port; + + // Inbound: accept the previous node's connection on our recv_port. + TCPSocket recv_sock; + std::thread accept_thread([recv_port, &recv_sock, &tag, &node_hosts, node_id]() { + TCPSocket listen_sock; + listen_sock.Create(); + listen_sock.SetKeepAlive(true); + listen_sock.Bind(SockAddr(node_hosts[node_id].c_str(), recv_port)); + listen_sock.Listen(); + SockAddr tmp_addr; + recv_sock = listen_sock.Accept(&tmp_addr); + LOG(INFO) << "[Ring " << tag << "] inbound from " << tmp_addr.AsString(); + listen_sock.Close(); + }); + + // Outbound: connect to the next node, retrying until it is listening. + TCPSocket send_sock; + send_sock.Create(); + send_sock.SetKeepAlive(true); + bool ok = false; + for (int i = 0; i < 100; ++i) { + if (send_sock.Connect(SockAddr(node_hosts[next_node_id].c_str(), send_port))) { + ok = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + TVM_FFI_ICHECK(ok) << "[Ring " << tag << "] failed to connect to next node " + << node_hosts[next_node_id].c_str() << ":" << send_port; + accept_thread.join(); + LOG(INFO) << "[Ring " << tag << "] ring TCP ready"; + + tcp_in_ch_ = std::make_unique(recv_sock.Release()); + tcp_out_ch_ = std::make_unique(send_sock.Release()); + + // Splice the local session's W0 ring pipe onto the TCP ring via a pipe + two proxy threads. + int sess_fds[2]; // 0:read 1:write + TVM_FFI_ICHECK_EQ(::pipe(sess_fds), 0); + proxy_out_to_tcp_ = local_session->RerouteRingIn(std::make_unique(sess_fds[0])); + TVM_FFI_ICHECK(proxy_out_to_tcp_ != nullptr); + proxy_in_from_tcp_ = std::make_unique(sess_fds[1]); + + send_thread_ = std::thread(ProxyLoop, proxy_out_to_tcp_.get(), tcp_out_ch_.get(), tag + "-send"); + recv_thread_ = std::thread(ProxyLoop, tcp_in_ch_.get(), proxy_in_from_tcp_.get(), tag + "-recv"); + LOG(INFO) << "[Ring " << tag << "] send/recv proxy threads started"; + } + + void Close() { + if (proxy_out_to_tcp_) proxy_out_to_tcp_->Close(); + if (proxy_in_from_tcp_) proxy_in_from_tcp_->Close(); + if (tcp_in_ch_) tcp_in_ch_->Close(); + if (tcp_out_ch_) tcp_out_ch_->Close(); + } + + void Join() { + if (send_thread_.joinable()) send_thread_.join(); + if (recv_thread_.joinable()) recv_thread_.join(); + } + + private: + std::unique_ptr proxy_out_to_tcp_; + std::unique_ptr proxy_in_from_tcp_; + std::unique_ptr tcp_in_ch_; + std::unique_ptr tcp_out_ch_; + std::thread send_thread_; + std::thread recv_thread_; +}; + class SocketSessionObj : public BcastSessionObj { public: explicit SocketSessionObj(int num_nodes, int num_workers_per_node, int num_groups, - const ffi::String& host, int port) - : num_nodes_(num_nodes), num_workers_per_node_(num_workers_per_node) { + const ffi::String& host, int port, bool build_ring) + : num_nodes_(num_nodes), num_workers_per_node_(num_workers_per_node), build_ring_(build_ring) { const auto f_create_local_session = tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers"); TVM_FFI_ICHECK(f_create_local_session.has_value()) << "Cannot find function runtime.disco.create_socket_session_local_workers"; - local_session_ = ((*f_create_local_session)(num_workers_per_node)).cast(); + local_session_ = ((*f_create_local_session)(num_workers_per_node, build_ring_)).cast(); DRef f_init_workers = local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers"); local_session_->CallPacked(f_init_workers, num_nodes_, /*node_id=*/0, num_groups, @@ -74,12 +179,18 @@ class SocketSessionObj : public BcastSessionObj { socket_.SetKeepAlive(true); socket_.Bind(SockAddr(host.c_str(), port)); socket_.Listen(); - LOG(INFO) << "SocketSession controller listening on " << host << ":" << port; + LOG(INFO) << "[Host] SocketSession controller listening on " << host << ":" << port; - AnyView packed_args[4]; + AnyView packed_args[5]; packed_args[0] = num_nodes; packed_args[1] = num_workers_per_node; packed_args[2] = num_groups; + packed_args[4] = build_ring; + + if (build_ring_) { + node_hosts_.reserve(num_nodes); + node_hosts_.push_back(host); + } for (int i = 0; i + 1 < num_nodes; ++i) { SockAddr addr; @@ -91,8 +202,38 @@ class SocketSessionObj : public BcastSessionObj { // - num_workers_per_node // - num_groups // - node_id - remote_channels_.back()->Send(ffi::PackedArgs(packed_args, 4)); - LOG(INFO) << "Remote node " << addr.AsString() << " connected"; + // - build_ring + remote_channels_.back()->Send(ffi::PackedArgs(packed_args, 5)); + LOG(INFO) << "[Host] Remote node " << addr.AsString() << " connected"; + + if (build_ring_) { + std::string full = addr.AsString(); + ffi::String ip(full.substr(0, full.find(':'))); + node_hosts_.push_back(ip); + LOG(INFO) << "[Host] IP table : node " << (i + 1) << " ip=" << ip.c_str(); + } + + if( build_ring_ && num_nodes_ > 1) { + + const int base_ring_port = port + 1; + + // Broadcast base_ring_port + the IP table so every remote node can wire its ring. + { + std::vector baseport_ip(num_nodes_ + 1); + baseport_ip[0] = base_ring_port; + for (int i = 0; i < num_nodes_; ++i) baseport_ip[1 + i] = node_hosts_[i]; + for (auto& ch : remote_channels_) { + ch->Send(ffi::PackedArgs(baseport_ip.data(), baseport_ip.size())); + } + LOG(INFO) << "[Host] IP and port broadcast completed."; + } + + ring_.Connect(/*node_id=*/0, num_nodes_, base_ring_port, node_hosts_, local_session_,"controller"); + + const int num_workers = num_nodes_ * num_workers_per_node_; + for (int worker_id = 1; worker_id < num_workers; ++worker_id) { this->SyncWorker(worker_id); } + LOG(INFO) << "[Host] Controller: All workers are ready. Initialization complete."; + }// build ring } } @@ -181,12 +322,24 @@ class SocketSessionObj : public BcastSessionObj { } void Shutdown() final { + + if (shutdown_) return; + shutdown_ = true; + // local session will be implicitly shutdown by its destructor std::vector packed_args(2); ffi::PackedArgs::Fill(packed_args.data(), static_cast(DiscoSocketAction::kShutdown), -1); for (auto& channel : remote_channels_) { channel->Send(ffi::PackedArgs(packed_args.data(), packed_args.size())); } + + local_session_->Shutdown(); + if (build_ring_) { + local_session_->CloseRing(); + ring_.Close(); + ring_.Join(); + } + for (auto& socket : remote_sockets_) { socket.Close(); } @@ -207,6 +360,12 @@ class SocketSessionObj : public BcastSessionObj { std::vector remote_sockets_; std::vector> remote_channels_; BcastSession local_session_{nullptr}; + + std::vector node_hosts_; + bool build_ring_ = false; + + private: + RingProxyEndpoint ring_; }; class RemoteSocketSession { @@ -217,17 +376,24 @@ class RemoteSocketSession { socket_.SetKeepAlive(true); SockAddr server_addr{server_host.c_str(), server_port}; Socket::Startup(); - if (!socket_.Connect(server_addr)) { + + bool connected = false; + for (int i = 0; i < 100; ++i) { + if (socket_.Connect(server_addr)) { connected = true; break; } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + if (!connected) { TVM_FFI_THROW(InternalError) << "Failed to connect to server " << server_addr.AsString() << ", errno = " << Socket::GetLastErrorCode(); } channel_ = std::make_unique(socket_); ffi::PackedArgs metadata = channel_->Recv(); - TVM_FFI_ICHECK_EQ(metadata.size(), 4); + TVM_FFI_ICHECK_EQ(metadata.size(), 5); num_nodes_ = metadata[0].cast(); num_workers_per_node_ = metadata[1].cast(); num_groups_ = metadata[2].cast(); node_id_ = metadata[3].cast(); + build_ring = metadata[4].cast(); TVM_FFI_ICHECK_GE(num_local_workers, num_workers_per_node_); InitLocalSession(); } @@ -241,6 +407,19 @@ class RemoteSocketSession { switch (action) { case DiscoSocketAction::kSend: { args = args.Slice(2); + + std::vector keep_alive; + std::vector fwd(args.size()); + for (int i = 0; i < args.size(); ++i) { + if (args[i].type_index() == ffi::TypeIndex::kTVMFFITensor) { + keep_alive.push_back(DiscoDebugObject::Wrap(args[i])); + fwd[i] = keep_alive.back(); + } else { + fwd[i] = args[i]; + } + } + ffi::PackedArgs wrapped(fwd.data(), static_cast(fwd.size())); + if (worker_id == -1) { local_session_->BroadcastPacked(args); } else { @@ -250,12 +429,26 @@ class RemoteSocketSession { } case DiscoSocketAction::kReceive: { args = local_session_->RecvReplyPacked(local_worker_id); - channel_->Reply(args); + + std::vector keep_alive; + std::vector fwd(args.size()); + for (int i = 0; i < args.size(); ++i) { + if (args[i].type_index() == ffi::TypeIndex::kTVMFFITensor) { + keep_alive.push_back(DiscoDebugObject::Wrap(args[i])); + fwd[i] = keep_alive.back(); + } else { + fwd[i] = args[i]; + } + } + + channel_->Reply(ffi::PackedArgs(fwd.data(), static_cast(fwd.size()))); break; } case DiscoSocketAction::kShutdown: { + LOG(INFO) << "[Remote] node_id=" << node_id_ << " recv controller cmd kShutdown"; local_session_->Shutdown(); - LOG(INFO) << "Connection closed by remote controller."; + LOG(INFO) << "[Remote] node_id=" << node_id_ + << " connection closed by controller, shutting down."; return; } default: @@ -265,7 +458,15 @@ class RemoteSocketSession { } ~RemoteSocketSession() { - socket_.Close(); + if (local_session_.defined()) { + local_session_->Shutdown(); + if (build_ring) { + local_session_->CloseRing(); + ring_.Close(); + ring_.Join(); + } + } + if (!socket_.IsClosed()) socket_.Close(); Socket::Finalize(); } @@ -273,7 +474,7 @@ class RemoteSocketSession { void InitLocalSession() { const auto f_create_local_session = tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers"); - local_session_ = ((*f_create_local_session)(num_workers_per_node_)).cast(); + local_session_ = ((*f_create_local_session)(num_workers_per_node_, build_ring)).cast(); DRef f_init_workers = local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers"); @@ -288,12 +489,23 @@ class RemoteSocketSession { int node_id_{-1}; int num_groups_{-1}; int num_workers_per_node_{-1}; + + bool build_ring = false; + RingProxyEndpoint ring_; }; void RemoteSocketSessionEntryPoint(const ffi::String& server_host, int server_port, int num_local_workers) { - RemoteSocketSession proxy(server_host, server_port, num_local_workers); - proxy.MainLoop(); + try { + RemoteSocketSession proxy(server_host, server_port, num_local_workers); + proxy.MainLoop(); + LOG(INFO) << "[Remote] " << server_host << ":" << server_port + << " main loop exited cleanly, process ending."; + + } catch (const std::exception& e) { + LOG(INFO) << "[Remote] node exiting due to exception (controller lost?): " << e.what(); + throw; + } } TVM_FFI_STATIC_INIT_BLOCK() { @@ -302,9 +514,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Session SocketSession(int num_nodes, int num_workers_per_node, int num_groups, - const ffi::String& host, int port) { + const ffi::String& host, int port, bool build_ring) { auto n = - ffi::make_object(num_nodes, num_workers_per_node, num_groups, host, port); + ffi::make_object(num_nodes, num_workers_per_node, num_groups, host, port, build_ring); return Session(n); } @@ -315,13 +527,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("runtime.disco.SocketSession", SocketSession) .def("runtime.disco.socket_session_init_workers", [](int num_nodes, int node_id, int num_groups, int num_workers_per_node) { - LOG(INFO) << "Initializing worker group with " << num_nodes << " nodes, " - << num_workers_per_node << " workers per node, and " << num_groups - << " groups."; DiscoWorker* worker = DiscoWorker::ThreadLocal(); worker->num_groups = num_groups; worker->worker_id = worker->worker_id + node_id * num_workers_per_node; worker->num_workers = num_nodes * num_workers_per_node; + LOG(INFO) << "[Disco_worker_" << worker->worker_id << "]" << " Initializing worker group with " << num_nodes << " nodes, " + << num_workers_per_node << " workers per node, and " << num_groups + << " groups."; }); } diff --git a/src/runtime/extra/disco/process_session.cc b/src/runtime/extra/disco/process_session.cc index 41888caeced0..21d7154a48d3 100644 --- a/src/runtime/extra/disco/process_session.cc +++ b/src/runtime/extra/disco/process_session.cc @@ -60,7 +60,7 @@ class DiscoProcessChannel final : public DiscoChannel { class ProcessSessionObj final : public BcastSessionObj { public: - explicit ProcessSessionObj(int num_workers, int num_groups, ffi::Function process_pool) + explicit ProcessSessionObj(int num_workers, int num_groups, ffi::Function process_pool, bool build_ring) : process_pool_(process_pool), worker_0_( std::make_unique(0, num_workers, num_groups, &worker_zero_data_)) { @@ -79,6 +79,20 @@ class ProcessSessionObj final : public BcastSessionObj { for (int i = 0; i < num_workers - 1; ++i) { workers_.emplace_back(std::make_unique(write_fds[i], read_fds[i])); } + + if (build_ring) { + ffi::Shape w0_fds = process_pool_(-1).cast(); + TVM_FFI_CHECK_EQ(w0_fds.size(), 2, ValueError) + << "process_pool(-1) should return a tuple of size 2 (worker_0's ring fds), " + << "but got a tuple of size " << w0_fds.size() << "."; + int64_t w0_ring_in_fd = w0_fds[0]; + int64_t w0_ring_out_fd = w0_fds[1]; + + ring_in_w0_ = std::make_unique(w0_ring_in_fd); + ring_out_w0_ = std::make_unique(w0_ring_out_fd); + worker_0_->worker->ring_in = ring_in_w0_.get(); + worker_0_->worker->ring_out = ring_out_w0_.get(); + } } void Kill() { @@ -165,31 +179,54 @@ class ProcessSessionObj final : public BcastSessionObj { } return workers_.at(worker_id - 1).get(); } + + std::unique_ptr RerouteRingIn(std::unique_ptr new_ch) override { + auto old = std::move(ring_in_w0_); + ring_in_w0_ = std::move(new_ch); + worker_0_->worker->ring_in = ring_in_w0_.get(); + return old; + } + + void CloseRing() override { + ring_out_w0_.reset(); + ring_in_w0_.reset(); + } ffi::Function process_pool_; std::unique_ptr worker_0_; std::vector> workers_; + std::unique_ptr ring_in_w0_; + std::unique_ptr ring_out_w0_; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("runtime.disco.ProcessSession", ProcessSessionObj, SessionObj); }; Session Session::ProcessSession(int num_workers, int num_group, ffi::String process_pool_creator, - ffi::String entrypoint) { + ffi::String entrypoint, bool build_ring) { TVM_FFI_ICHECK_EQ(num_workers % num_group, 0) << "The number of workers should be divisible by the number of worker group."; const auto pf = tvm::ffi::Function::GetGlobal(process_pool_creator); TVM_FFI_CHECK(pf, ValueError) << "Cannot find function " << process_pool_creator << " in the registry. Please check if it is registered."; - auto process_pool = (*pf)(num_workers, num_group, entrypoint).cast(); - auto n = ffi::make_object(num_workers, num_group, process_pool); + auto process_pool = (*pf)(num_workers, num_group, entrypoint, build_ring).cast(); + auto n = ffi::make_object(num_workers, num_group, process_pool, build_ring); return Session(n); } void WorkerProcess(int worker_id, int num_workers, int num_group, int64_t read_fd, - int64_t write_fd) { + int64_t write_fd, int64_t ring_in_fd, int64_t ring_out_fd) { TVM_FFI_ICHECK_EQ(num_workers % num_group, 0) << "The number of workers should be divisible by the number of worker group."; DiscoProcessChannel channel(read_fd, write_fd); DiscoWorker worker(worker_id, num_workers, num_group, nullptr, &channel); + + std::unique_ptr ring_in_ch, ring_out_ch; + if (ring_in_fd >= 0 && ring_out_fd >= 0) { + ring_in_ch = std::make_unique(ring_in_fd); + ring_out_ch = std::make_unique(ring_out_fd); + worker.ring_in = ring_in_ch.get(); + worker.ring_out = ring_out_ch.get(); + } + worker.MainLoop(); } diff --git a/src/runtime/extra/disco/threaded_session.cc b/src/runtime/extra/disco/threaded_session.cc index 83de74766bb7..7b4996cb08cf 100644 --- a/src/runtime/extra/disco/threaded_session.cc +++ b/src/runtime/extra/disco/threaded_session.cc @@ -143,16 +143,21 @@ DiscoWorkerThread::DiscoWorkerThread(int worker_id, int num_workers, int num_gro class ThreadedSessionObj final : public BcastSessionObj { public: - explicit ThreadedSessionObj(int num_workers, int num_groups) { + explicit ThreadedSessionObj(int num_workers, int num_groups, bool build_ring) { for (int i = 0; i < num_workers; ++i) { WorkerZeroData* data = (i == 0) ? &worker_zero_data_ : nullptr; workers_.emplace_back(i, num_workers, num_groups, data); } + + if (build_ring && num_workers > 1) { + BuildRing(num_workers); + } } ~ThreadedSessionObj() { this->Shutdown(); workers_.clear(); + ring_channels_.clear(); } int64_t GetNumWorkers() { return workers_.size(); } @@ -184,13 +189,37 @@ class ThreadedSessionObj final : public BcastSessionObj { SessionObj); std::vector workers_; + std::vector> ring_channels_; + + private: + void BuildRing(int num_workers) { + + std::vector read_fds(num_workers), write_fds(num_workers); + for (int i = 0; i < num_workers; ++i) { + int fds[2]; + TVM_FFI_ICHECK_EQ(::pipe(fds), 0) << "Failed to create ring pipe for worker " << i; + read_fds[i] = fds[0]; + write_fds[i] = fds[1]; + } + + ring_channels_.reserve(num_workers * 2); + for (int i = 0; i < num_workers; ++i) { + + auto channel_out = std::make_unique(write_fds[i]); + auto channel_in = std::make_unique(read_fds[(i - 1 + num_workers) % num_workers]); + workers_[i].worker->ring_out = channel_out.get(); + workers_[i].worker->ring_in = channel_in.get(); + ring_channels_.push_back(std::move(channel_in)); + ring_channels_.push_back(std::move(channel_out)); + } + } }; -Session Session::ThreadedSession(int num_workers, int num_group) { +Session Session::ThreadedSession(int num_workers, int num_group, bool build_ring) { TVM_FFI_ICHECK_EQ(num_workers % num_group, 0) << "The number of workers should be divisible by the number of worker group."; ffi::ObjectPtr n = - ffi::make_object(num_workers, num_group); + ffi::make_object(num_workers, num_group, build_ring); return Session(std::move(n)); } diff --git a/src/support/socket.h b/src/support/socket.h index 364705449765..7830852bedcc 100644 --- a/src/support/socket.h +++ b/src/support/socket.h @@ -293,6 +293,20 @@ class Socket { if (err == EBADF || err == EINTR) return true; return false; } + + /*! + * \brief Relinquish ownership of the underlying descriptor without closing it, and return it. + * The caller becomes responsible for closing the returned fd (e.g. when handing it to + * another owner). Mirrors std::unique_ptr::release. Prevents a later double-close. + * \return The descriptor previously held; INVALID_SOCKET afterwards. + */ + SockType Release() { + SockType fd = sockfd; + sockfd = INVALID_SOCKET; + return fd; + } + + /*! \brief check if socket is already closed */ bool IsClosed() const { return sockfd == INVALID_SOCKET; } /*! \brief close the socket */ From 8af26b714a8681a153959e951b4845ca40d94af2 Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 15:33:28 +0800 Subject: [PATCH 02/10] [CMAKE] Add set(USE_CPU_CCL) option and mesg --- CMakeLists.txt | 1 + cmake/config.cmake | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 1d1fd96ce18f..dd85ebf34faf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -526,6 +526,7 @@ endif() # CPUCCL. if(USE_CPU_CCL) + message(STATUS "Build with CPU CCL support...") tvm_file_glob(GLOB _cpu_ccl_srcs src/runtime/extra/disco/cpuccl/*.cc) add_library(tvm_cpu_ccl_objs OBJECT ${_cpu_ccl_srcs}) target_link_libraries(tvm_cpu_ccl_objs PRIVATE tvm_runtime_extra_defs) diff --git a/cmake/config.cmake b/cmake/config.cmake index 068b17fbf09b..40b1ce8896b2 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -322,3 +322,8 @@ SET(CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE "x64") # Enable Qualcomm OpenCL extension support set(USE_OPENCL_EXTN_QCOM OFF) + +# Whether to enable CPU CCL support: +# - ON: enable CPU collectives over a pipe ring between disco workers +# - OFF: disable CPU CCL +set(USE_CPU_CCL OFF) From 8e12af94bc461a8cc4f91f66439a3e7c570a1cbf Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 16:52:42 +0800 Subject: [PATCH 03/10] [FixBugs] RemoteConnect and fix multiple broadcast issues in iptables [FixBugs] RemoteConnect [FixBugs] Fix multiple broadcast issues in iptables --- .../extra/disco/distributed/socket_session.cc | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/runtime/extra/disco/distributed/socket_session.cc b/src/runtime/extra/disco/distributed/socket_session.cc index 368358a5043e..97ace9aa1324 100644 --- a/src/runtime/extra/disco/distributed/socket_session.cc +++ b/src/runtime/extra/disco/distributed/socket_session.cc @@ -212,8 +212,9 @@ class SocketSessionObj : public BcastSessionObj { node_hosts_.push_back(ip); LOG(INFO) << "[Host] IP table : node " << (i + 1) << " ip=" << ip.c_str(); } + } - if( build_ring_ && num_nodes_ > 1) { + if( build_ring_ && num_nodes_ > 1) { const int base_ring_port = port + 1; @@ -227,14 +228,12 @@ class SocketSessionObj : public BcastSessionObj { } LOG(INFO) << "[Host] IP and port broadcast completed."; } - ring_.Connect(/*node_id=*/0, num_nodes_, base_ring_port, node_hosts_, local_session_,"controller"); const int num_workers = num_nodes_ * num_workers_per_node_; for (int worker_id = 1; worker_id < num_workers; ++worker_id) { this->SyncWorker(worker_id); } LOG(INFO) << "[Host] Controller: All workers are ready. Initialization complete."; }// build ring - } } int64_t GetNumWorkers() final { return num_nodes_ * num_workers_per_node_; } @@ -480,6 +479,19 @@ class RemoteSocketSession { local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers"); local_session_->CallPacked(f_init_workers, num_nodes_, node_id_, num_groups_, num_workers_per_node_); + if( build_ring && num_nodes_ > 1) { + // The controller broadcasts base_ring_port + the IP table on our control channel. + ffi::PackedArgs host_info = channel_->Recv(); + TVM_FFI_ICHECK_EQ(host_info.size(), num_nodes_ + 1); + const int nNodes = host_info.size() - 1; + const int base_ring_port = host_info[0].cast(); + std::vector node_hosts; + node_hosts.reserve(nNodes); + for (int i = 0; i < nNodes; ++i) node_hosts.push_back(host_info[1 + i].cast()); + + ring_.Connect(node_id_, nNodes, base_ring_port, node_hosts, local_session_, + "node_" + std::to_string(node_id_)); + }//build ring } TCPSocket socket_; From 9fe84e42a0b527a67bebd33ab96b7540fce39161 Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 19:17:21 +0800 Subject: [PATCH 04/10] [Relax][Op] Add gather_to_worker0 --- include/tvm/relax/attrs/ccl.h | 17 +++++++++ python/tvm/relax/op/ccl/__init__.py | 2 +- python/tvm/relax/op/ccl/ccl.py | 21 ++++++++++ .../tvm/relax/transform/legalize_ops/ccl.py | 25 ++++++++++++ src/relax/op/ccl/ccl.cc | 38 +++++++++++++++++++ src/relax/op/ccl/ccl.h | 3 ++ 6 files changed, 105 insertions(+), 1 deletion(-) diff --git a/include/tvm/relax/attrs/ccl.h b/include/tvm/relax/attrs/ccl.h index 031a1de49311..bbc0b7118b0d 100644 --- a/include/tvm/relax/attrs/ccl.h +++ b/include/tvm/relax/attrs/ccl.h @@ -85,6 +85,23 @@ struct ScatterCollectiveAttrs : public tvm::AttrsNode { AttrsNode); }; // struct ScatterCollectiveAttrs +/*! \brief Attributes used in gather operators */ +struct GatherCollectiveAttrs : public tvm::AttrsNode { + int num_workers; + bool in_group; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("num_workers", &GatherCollectiveAttrs::num_workers, + "The number of workers to gather data from.") + .def_ro("in_group", &GatherCollectiveAttrs::in_group, + "Whether the gather operation performs in group or globally as default."); + } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherCollectiveAttrs", GatherCollectiveAttrs, + AttrsNode); +}; // struct GatherCollectiveAttrs + } // namespace relax } // namespace tvm diff --git a/python/tvm/relax/op/ccl/__init__.py b/python/tvm/relax/op/ccl/__init__.py index 47006358f93f..7a66e7405e9b 100644 --- a/python/tvm/relax/op/ccl/__init__.py +++ b/python/tvm/relax/op/ccl/__init__.py @@ -17,4 +17,4 @@ # under the License. """CCL related operators.""" -from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0 +from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0, gather_to_worker0 diff --git a/python/tvm/relax/op/ccl/ccl.py b/python/tvm/relax/op/ccl/ccl.py index e49614f4ded1..47613213ee92 100644 --- a/python/tvm/relax/op/ccl/ccl.py +++ b/python/tvm/relax/op/ccl/ccl.py @@ -106,3 +106,24 @@ def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr: Chunked Tensor received by different workers. """ return _ffi_api.scatter_from_worker0(x, num_workers, axis) + +def gather_to_worker0(x: Expr, num_workers: int, in_group: bool = True) -> Expr: # pylint: disable=invalid-name + """Gather data from all workers to worker-0. + + Parameters + ---------- + x : relax.Expr + The tensor to be gathered from each worker. + + num_workers : int + The number of workers to gather data from. + + in_group : bool + Whether the gather operation performs globally or in group as default. + + Returns + ------- + result : relax.Expr + The concatenated tensor received by worker-0 only. + """ + return _ffi_api.gather_to_worker0(x, num_workers, in_group) # type: ignore # pylint: disable=no-member diff --git a/python/tvm/relax/transform/legalize_ops/ccl.py b/python/tvm/relax/transform/legalize_ops/ccl.py index 659b1f7d4397..e3bf81a95c85 100644 --- a/python/tvm/relax/transform/legalize_ops/ccl.py +++ b/python/tvm/relax/transform/legalize_ops/ccl.py @@ -123,3 +123,28 @@ def _scatter_from_worker0(_bb: BlockBuilder, call: Call) -> Expr: vdevice=call.args[0].ty.vdevice, ), ) + + +@register_legalize("relax.ccl.gather_to_worker0") +def _gather_to_worker0(_bb: BlockBuilder, call: Call) -> Expr: + output_shape = [] + arg_ty = call.args[0].ty + assert isinstance( + arg_ty, TensorType + ), "The input type of gather_to_worker0 should be TensorType." + assert isinstance(arg_ty.shape.ty, ShapeType) + arg_shape = arg_ty.shape.ty + for i, shape_value in enumerate(arg_shape.values): + if i == 0: + output_shape.append(shape_value * call.attrs.num_workers) + else: + output_shape.append(shape_value) + return call_dps_packed( + "runtime.disco.gather_to_worker0", + [call.args[0], call.attrs.in_group], + out_ty=TensorType( + shape=output_shape, + dtype=arg_ty.dtype, + vdevice=arg_ty.vdevice, + ), + ) diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc index c140ba275ba4..fa22436c6bfb 100644 --- a/src/relax/op/ccl/ccl.cc +++ b/src/relax/op/ccl/ccl.cc @@ -34,6 +34,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { AllReduceAttrs::RegisterReflection(); AllGatherAttrs::RegisterReflection(); ScatterCollectiveAttrs::RegisterReflection(); + GatherCollectiveAttrs::RegisterReflection(); } Expr allreduce(Expr x, ffi::String op_type, bool in_group) { @@ -173,5 +174,42 @@ TVM_REGISTER_OP("relax.ccl.scatter_from_worker0") .set_attr("FInferType", InferTypeScatter) .set_attr("FPurity", true); +/* relax.ccl.gather_to_worker0 */ +Expr gather_to_worker0(Expr x, int num_workers, bool in_group) { + ffi::ObjectPtr attrs = ffi::make_object(); + attrs->num_workers = std::move(num_workers); + attrs->in_group = std::move(in_group); + static const Op& op = Op::Get("relax.ccl.gather_to_worker0"); + return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("relax.op.ccl.gather_to_worker0", gather_to_worker0); +} + +Type InferTypeGatherToWorker0(const Call& call, const BlockBuilder& ctx) { + TensorType input_ty = GetUnaryInputTensorType(call, ctx); + + const auto* attrs = call->attrs.as(); + int num_workers = attrs->num_workers; + + ffi::Optional output_dtype = input_ty->dtype; + auto input_shape = input_ty->GetShape(); + if (!input_shape.has_value()) { + return input_ty; + } + ffi::Array output_shape = input_shape.value(); + output_shape.Set(0, output_shape[0] * num_workers); + return TensorType(ShapeExpr(output_shape), output_dtype, input_ty->vdevice); +} + +TVM_REGISTER_OP("relax.ccl.gather_to_worker0") + .set_num_inputs(1) + .add_argument("x", "Tensor", "Input to be gathered to worker-0.") + .set_attrs_type() + .set_attr("FInferType", InferTypeGatherToWorker0) + .set_attr("FPurity", true); + } // namespace relax } // namespace tvm diff --git a/src/relax/op/ccl/ccl.h b/src/relax/op/ccl/ccl.h index 1d049382d0ae..210e1f405c2f 100644 --- a/src/relax/op/ccl/ccl.h +++ b/src/relax/op/ccl/ccl.h @@ -44,6 +44,9 @@ Expr broadcast_from_worker0(Expr data); /*! \brief Perform a scatter operation from worker-0, chunking the given buffer into equal parts. */ Expr scatter_from_worker0(Expr data, int num_workers, int axis); +/*! \brief Gather data from all workers to worker-0. */ +Expr gather_to_worker0(Expr data, int num_workers, bool in_group); + } // namespace relax } // namespace tvm From 87c2d2c85ad6e82f608b73ebb397b786efa6d3c2 Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 19:43:31 +0800 Subject: [PATCH 05/10] [FixBugs] Re-wrap Tensor args when relaying kSend to local workers [FixBugs] Re-wrap Tensor args when relaying kSend to local workers Fix ruff line length and import order for gather_to_worker0 --- python/tvm/relax/op/ccl/__init__.py | 8 +++++++- python/tvm/relax/op/ccl/ccl.py | 5 +++-- python/tvm/relax/transform/legalize_ops/ccl.py | 4 +--- src/runtime/extra/disco/distributed/socket_session.cc | 4 ++-- 4 files changed, 13 insertions(+), 8 deletions(-) diff --git a/python/tvm/relax/op/ccl/__init__.py b/python/tvm/relax/op/ccl/__init__.py index 7a66e7405e9b..e737346fc263 100644 --- a/python/tvm/relax/op/ccl/__init__.py +++ b/python/tvm/relax/op/ccl/__init__.py @@ -17,4 +17,10 @@ # under the License. """CCL related operators.""" -from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0, gather_to_worker0 +from .ccl import ( + allgather, + allreduce, + broadcast_from_worker0, + gather_to_worker0, + scatter_from_worker0, +) diff --git a/python/tvm/relax/op/ccl/ccl.py b/python/tvm/relax/op/ccl/ccl.py index 47613213ee92..d47d640bcf7f 100644 --- a/python/tvm/relax/op/ccl/ccl.py +++ b/python/tvm/relax/op/ccl/ccl.py @@ -107,7 +107,8 @@ def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr: """ return _ffi_api.scatter_from_worker0(x, num_workers, axis) -def gather_to_worker0(x: Expr, num_workers: int, in_group: bool = True) -> Expr: # pylint: disable=invalid-name + +def gather_to_worker0(x: Expr, num_workers: int, in_group: bool = True) -> Expr: """Gather data from all workers to worker-0. Parameters @@ -126,4 +127,4 @@ def gather_to_worker0(x: Expr, num_workers: int, in_group: bool = True) -> Expr: result : relax.Expr The concatenated tensor received by worker-0 only. """ - return _ffi_api.gather_to_worker0(x, num_workers, in_group) # type: ignore # pylint: disable=no-member + return _ffi_api.gather_to_worker0(x, num_workers, in_group) diff --git a/python/tvm/relax/transform/legalize_ops/ccl.py b/python/tvm/relax/transform/legalize_ops/ccl.py index e3bf81a95c85..fcfbd2251472 100644 --- a/python/tvm/relax/transform/legalize_ops/ccl.py +++ b/python/tvm/relax/transform/legalize_ops/ccl.py @@ -129,9 +129,7 @@ def _scatter_from_worker0(_bb: BlockBuilder, call: Call) -> Expr: def _gather_to_worker0(_bb: BlockBuilder, call: Call) -> Expr: output_shape = [] arg_ty = call.args[0].ty - assert isinstance( - arg_ty, TensorType - ), "The input type of gather_to_worker0 should be TensorType." + assert isinstance(arg_ty, TensorType), "The input of gather_to_worker0 should be TensorType." assert isinstance(arg_ty.shape.ty, ShapeType) arg_shape = arg_ty.shape.ty for i, shape_value in enumerate(arg_shape.values): diff --git a/src/runtime/extra/disco/distributed/socket_session.cc b/src/runtime/extra/disco/distributed/socket_session.cc index 97ace9aa1324..cd2b1d4da906 100644 --- a/src/runtime/extra/disco/distributed/socket_session.cc +++ b/src/runtime/extra/disco/distributed/socket_session.cc @@ -420,9 +420,9 @@ class RemoteSocketSession { ffi::PackedArgs wrapped(fwd.data(), static_cast(fwd.size())); if (worker_id == -1) { - local_session_->BroadcastPacked(args); + local_session_->BroadcastPacked(wrapped); } else { - local_session_->SendPacked(local_worker_id, args); + local_session_->SendPacked(local_worker_id, wrapped); } break; } From 74b2834cfbd09fae9d8f9fa13d94acb38c9f2871 Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 19:55:16 +0800 Subject: [PATCH 06/10] [FixBug] check ring only worker->num_workers > 1 --- src/runtime/extra/disco/cpuccl/cpuccl.cc | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/runtime/extra/disco/cpuccl/cpuccl.cc b/src/runtime/extra/disco/cpuccl/cpuccl.cc index ab2f6db0f43e..679e87c39197 100644 --- a/src/runtime/extra/disco/cpuccl/cpuccl.cc +++ b/src/runtime/extra/disco/cpuccl/cpuccl.cc @@ -182,8 +182,11 @@ void AllReduce(Tensor send, ReduceKind reduce_kind, bool /*in_group*/, Tensor re void AllGather(Tensor send, bool /*in_group*/, Tensor recv) { DiscoWorker* worker = DiscoWorker::ThreadLocal(); - TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) - << "Ring not initialized"; + + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + } int num_workers = worker->num_workers; int rank = worker->worker_id; @@ -274,6 +277,11 @@ void ScatterFromWorker0(ffi::Optional send, bool /*in_group*/, Tensor re void GatherToWorker0(Tensor send, bool /*in_group*/, ffi::Optional recv) { DiscoWorker* worker = DiscoWorker::ThreadLocal(); + if (worker->num_workers > 1) { + TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) + << "Ring not initialized"; + } + int rank = worker->worker_id; int num_workers = worker->num_workers; int64_t chunk_bytes = send.Shape().Product() * DTypeBytes(send->dtype); @@ -285,9 +293,6 @@ void GatherToWorker0(Tensor send, bool /*in_group*/, ffi::Optional recv) return; } - TVM_FFI_ICHECK(worker->ring_in != nullptr && worker->ring_out != nullptr) - << "Ring not initialized"; - if (rank == 0) { TVM_FFI_CHECK(recv.has_value(), ValueError) << "Worker 0 must provide recv buffer for gather"; From 3fa0d38bcabfbc9057607ba84faf39e9cf4701f6 Mon Sep 17 00:00:00 2001 From: yuyi Date: Tue, 28 Jul 2026 21:40:25 +0800 Subject: [PATCH 07/10] [Test][Disco] Cover build_ring in session tests and fix test_shape_tuple hang --- tests/python/disco/test_session.py | 56 ++++++++++++++++++------------ 1 file changed, 34 insertions(+), 22 deletions(-) diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py index 2e84170940e1..946b31f50c2b 100644 --- a/tests/python/disco/test_session.py +++ b/tests/python/disco/test_session.py @@ -60,7 +60,7 @@ class SocketSessionTester: launched with the current Python interpreter. """ - def __init__(self, num_workers, num_nodes=2, num_groups=1): + def __init__(self, num_workers, num_nodes=2, num_groups=1, build_ring=False): # Initialize the attributes used by __del__ first, so that teardown is # safe even when __init__ raises below. self.sess = None @@ -74,7 +74,7 @@ def __init__(self, num_workers, num_nodes=2, num_groups=1): def start_server(): try: self.sess = di.SocketSession( - num_nodes, num_workers_per_node, num_groups, server_host, server_port + num_nodes, num_workers_per_node, num_groups, server_host, server_port, build_ring ) except Exception as exc: # pylint: disable=broad-except server_exc.append(exc) @@ -121,7 +121,8 @@ def __del__(self): node.wait() -def create_socket_session(num_workers): + +def create_socket_session(num_workers, build_ring=False): """Create a socket session backed by one local and one remote node. The tester is kept alive in a module-level global so that the session @@ -130,7 +131,7 @@ def create_socket_session(num_workers): global _SOCKET_SESSION_TESTER # Rebind (not `del`) so the global stays defined if the constructor raises. _SOCKET_SESSION_TESTER = None - _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers) + _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers, build_ring=build_ring) assert _SOCKET_SESSION_TESTER.sess is not None return _SOCKET_SESSION_TESTER.sess @@ -153,9 +154,10 @@ def _numpy_from_worker_0(sess: di.Session, remote_array, shape, dtype): @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_int(session_kind): # pylint: disable=invalid-name +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_int(session_kind, build_ring_): # pylint: disable=invalid-name num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one") result: di.DRef = func(1) for i in range(num_workers): @@ -163,9 +165,10 @@ def test_int(session_kind): # pylint: disable=invalid-name @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_float(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_float(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one_float") result: di.DRef = func(1.5) @@ -174,9 +177,10 @@ def test_float(session_kind): @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_tensor(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_tensor(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) device = tvm.cpu(0) x_np = np.arange(6).astype("float32").reshape([2, 3]) y_np = np.arange(6).astype("float32").reshape([2, 3]) + 1 @@ -187,9 +191,10 @@ def test_tensor(session_kind): @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_string(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_string(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) func: di.DPackedFunc = sess.get_global_func("tests.disco.str") result: di.DRef = func("hello") @@ -198,9 +203,10 @@ def test_string(session_kind): @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_string_obj(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_string_obj(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) func: di.DPackedFunc = sess.get_global_func("tests.disco.str_obj") result: di.DRef = func(String("hello")) @@ -211,21 +217,25 @@ def test_string_obj(session_kind): @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_shape_tuple(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_shape_tuple(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) func: di.DPackedFunc = sess.get_global_func("tests.disco.shape_tuple") result: di.DRef = func(Shape([1, 2, 3])) for i in range(num_workers): value = result.debug_get_from_remote(i) assert isinstance(value, Shape) assert list(value) == [1, 2, 3, 4, 5] + # Without this the returned Shape is freed on a worker thread and deadlocks on the GIL. + sess.shutdown() @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_vm_module(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_vm_module(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) # pylint: disable=invalid-name @I.ir_module(s_tir=True) @@ -267,9 +277,10 @@ def main(A: R.Tensor((8, 16), dtype="float32")) -> R.Tensor((16, 8), dtype="floa @pytest.mark.parametrize("session_kind", _all_session_kinds) -def test_vm_multi_func(session_kind): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_vm_multi_func(session_kind, build_ring_): num_workers = 4 - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) # pylint: disable=invalid-name @I.ir_module(s_tir=True) @@ -336,10 +347,11 @@ def transpose_2(A: R.Tensor((16, 8), dtype="float32")) -> R.Tensor( @pytest.mark.parametrize("session_kind", _all_session_kinds) @pytest.mark.parametrize("num_workers", [1, 2, 4]) -def test_num_workers(session_kind, num_workers): +@pytest.mark.parametrize("build_ring_", [True, False]) +def test_num_workers(session_kind, num_workers, build_ring_): if session_kind == create_socket_session and num_workers < 2: return - sess = session_kind(num_workers=num_workers) + sess = session_kind(num_workers=num_workers, build_ring=build_ring_) assert sess.num_workers == num_workers From da0ad3b0c60df2b7a473436aa0975199cd4e1012 Mon Sep 17 00:00:00 2001 From: yuyi Date: Wed, 29 Jul 2026 13:34:11 +0800 Subject: [PATCH 08/10] [Test][Disco] Add tests for cpuccl and up_load_vm_module() --- tests/python/disco/test_session.py | 7 +- tests/python/disco/test_socket_cpuccl.py | 471 +++++++++++++++++++++++ tests/python/disco/test_socket_upload.py | 198 ++++++++++ 3 files changed, 675 insertions(+), 1 deletion(-) create mode 100644 tests/python/disco/test_socket_cpuccl.py create mode 100644 tests/python/disco/test_socket_upload.py diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py index 946b31f50c2b..f7916ef0db0a 100644 --- a/tests/python/disco/test_session.py +++ b/tests/python/disco/test_session.py @@ -74,7 +74,12 @@ def __init__(self, num_workers, num_nodes=2, num_groups=1, build_ring=False): def start_server(): try: self.sess = di.SocketSession( - num_nodes, num_workers_per_node, num_groups, server_host, server_port, build_ring + num_nodes, + num_workers_per_node, + num_groups, + server_host, + server_port, + build_ring, ) except Exception as exc: # pylint: disable=broad-except server_exc.append(exc) diff --git a/tests/python/disco/test_socket_cpuccl.py b/tests/python/disco/test_socket_cpuccl.py new file mode 100644 index 000000000000..74fe6c26a72f --- /dev/null +++ b/tests/python/disco/test_socket_cpuccl.py @@ -0,0 +1,471 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=missing-docstring +"""Collective (cpuccl) tests over a multi-node SocketSession. + +This file is intentionally separate from ``test_ccl.py`` (which targets nccl/rccl on +ThreadedSession/ProcessSession). + +The socket path needs its own multi-node harness -- a remote subprocess per extra node, plus +careful startup ordering -- which does not belong in the upstream nccl test matrix. cpuccl +always builds the ring (``build_ring=True``) and does not support group (``in_group``) +semantics, so only the global collectives are covered. +""" +import socket +import subprocess +import sys +import tempfile +import threading + +import numpy as np +import pytest + +import tvm +import tvm.testing +import tvm.script +from tvm import relax as rx +from tvm.runtime import disco as di +from tvm.runtime.vm import VirtualMachine +from tvm.script import relax as R +from tvm.exec import disco_worker as _ # pylint: disable=unused-import + +if di is None: + pytest.skip("disco runtime is not available", allow_module_level=True) + +_CCL = "cpuccl" +_SOCKET_SESSION_TESTER = None + + +def _get_free_port(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +class SocketSessionTester: + """Run a disco SocketSession with one local node and remote nodes. + + Each remote node is a `tvm.exec.disco_remote_socket_session` subprocess launched with the + current Python interpreter. Unlike test_session.py, `build_ring` defaults to True because + cpuccl uses the ring as its data plane. + """ + + def __init__(self, num_workers, num_nodes=2, num_groups=1, build_ring=True): + # Initialize the attributes used by __del__ first, so that teardown is + # safe even when __init__ raises below. + self.sess = None + self.remote_nodes = [] + assert num_workers % num_nodes == 0 + num_workers_per_node = num_workers // num_nodes + server_host = "localhost" + server_port = _get_free_port() + server_exc = [] + + def start_server(): + try: + self.sess = di.SocketSession( + num_nodes, + num_workers_per_node, + num_groups, + server_host, + server_port, + build_ring, + ) + except Exception as exc: # pylint: disable=broad-except + server_exc.append(exc) + + thread = threading.Thread(target=start_server) + thread.start() + + cmd = "tvm.exec.disco_remote_socket_session" + for _i in range(num_nodes - 1): + self.remote_nodes.append( + subprocess.Popen( + [ + sys.executable, + "-m", + cmd, + server_host, + str(server_port), + str(num_workers_per_node), + ], + stdout=sys.stdout, + stderr=sys.stderr, + ) + ) + + thread.join() + if server_exc: + raise server_exc[0] + + # Bound at class creation: module globals may already be cleared when + # __del__ runs during interpreter shutdown. + _TIMEOUT_EXPIRED = subprocess.TimeoutExpired + + def __del__(self): + try: + if self.sess is not None: + self.sess.shutdown() + finally: + for node in self.remote_nodes: + try: + node.wait(timeout=10) + except self._TIMEOUT_EXPIRED: + node.kill() + node.wait() + + +def create_socket_session(num_workers, build_ring=True): + """Create a socket session backed by one local and one remote node. + + The tester is kept alive in a module-level global so that the session + survives until the next call (or interpreter exit) replaces it. + """ + global _SOCKET_SESSION_TESTER + # Rebind (not `del`) so the global stays defined if the constructor raises. + _SOCKET_SESSION_TESTER = None + _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers, build_ring=build_ring) + assert _SOCKET_SESSION_TESTER.sess is not None + return _SOCKET_SESSION_TESTER.sess + + + +def test_init(): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + +def test_allreduce(): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array_1 = np.arange(12, dtype="float32").reshape(3, 4) + array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4) + d_array = sess.empty((3, 4), "float32") + d_array.debug_copy_from(0, array_1) + d_array.debug_copy_from(1, array_2) + for op, np_op in [ # pylint: disable=invalid-name + ("sum", np.add), + ("prod", np.multiply), + ("min", np.minimum), + ("max", np.maximum), + ("avg", lambda a, b: (a + b) * 0.5), + ]: + dst_array = sess.empty((3, 4), "float32") + sess.allreduce(d_array, dst_array, op=op) + result = dst_array.debug_get_from_remote(0).numpy() + expected = np_op(array_1, array_2) + np.testing.assert_equal(result, expected) + + +def test_allgather(): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array = np.arange(36, dtype="float32") + d_src = sess.empty((3, 3, 2), "float32") + d_dst = sess.empty((3, 4, 3), "float32") + d_src.debug_copy_from(0, array[:18]) + d_src.debug_copy_from(1, array[18:]) + sess.allgather(d_src, d_dst) + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) + np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array.reshape(3, 4, 3)) + + +@pytest.mark.parametrize("use_explicit_output", [True, False]) +def test_broadcast(use_explicit_output): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array = np.arange(12, dtype="float32").reshape(3, 4) + if use_explicit_output: + src_array = sess.empty((3, 4), "float32", worker0_only=True) + src_array.debug_copy_from(0, array) + dst_array = sess.empty((3, 4), "float32") + sess.broadcast_from_worker0(src_array, dst_array) + else: + dst_array = sess.broadcast(array) + + result = dst_array.debug_get_from_remote(1).numpy() + np.testing.assert_equal(result, array) + + +@pytest.mark.parametrize("use_explicit_output", [True, False]) +def test_scatter(use_explicit_output): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array = np.arange(36, dtype="float32").reshape(2, 6, 3) + if use_explicit_output: + d_src = sess.empty((2, 6, 3), "float32", worker0_only=True) + d_dst = sess.empty((6, 3), "float32") + d_src.debug_copy_from(0, array) + sess.scatter_from_worker0(d_src, d_dst) + else: + d_dst = sess.scatter(array) + + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array[0, :, :]) + np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array[1, :, :]) + + +def test_scatter_with_implicit_reshape(): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array = np.arange(36, dtype="float32").reshape(3, 4, 3) + d_src = sess.empty((3, 4, 3), "float32", worker0_only=True) + d_dst = sess.empty((3, 3, 2), "float32") + d_src.debug_copy_from(0, array) + sess.scatter_from_worker0(d_src, d_dst) + + np.testing.assert_equal( + d_dst.debug_get_from_remote(0).numpy(), array.flat[:18].reshape(3, 3, 2) + ) + np.testing.assert_equal( + d_dst.debug_get_from_remote(1).numpy(), array.flat[18:].reshape(3, 3, 2) + ) + + + + +def test_gather(): + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + array = np.arange(36, dtype="float32") + d_src = sess.empty((3, 3, 2), "float32") + d_dst = sess.empty((3, 4, 3), "float32", worker0_only=True) + d_src.debug_copy_from(0, array[:18]) + d_src.debug_copy_from(1, array[18:]) + sess.gather_to_worker0(d_src, d_dst) + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) + + + + +def relax_build(mod): + target = tvm.target.Target("llvm") + with target: + mod = rx.get_pipeline("zero")(mod) # pylint: disable=no-value-for-parameter + return tvm.compile(mod, target=target) + + +def test_mlp(): # pylint: disable=too-many-locals + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + # pylint: disable=invalid-name + @tvm.script.ir_module + class MLP: # pylint: disable=too-few-public-methods + @R.function + def main( + x: R.Tensor((128, 128), "float32"), + W1: R.Tensor((128, 128), "float32"), + W2: R.Tensor((128, 128), "float32"), + ) -> R.Tensor((128, 128), "float32"): + R.func_attr({"global_symbol": "main"}) + with R.dataflow(): + lv0: R.Tensor((128, 128), "float32") = R.matmul(x, W1) + lv1: R.Tensor((128, 128), "float32") = R.nn.gelu(lv0) + lv2: R.Tensor((128, 128), "float32") = R.matmul(lv1, W2) + R.output(lv2) + return lv2 + + @tvm.script.ir_module + class ShardedMLP: # pylint: disable=too-few-public-methods + @R.function + def main( + x: R.Tensor((128, 128), "float32"), + W1: R.Tensor((128, 64), "float32"), # shard along axis 1 + W2: R.Tensor((64, 128), "float32"), # shard along axis 0 + ) -> R.Tensor((128, 128), "float32"): + R.func_attr({"global_symbol": "main"}) + with R.dataflow(): + broadcast_x: R.Tensor((128, 128), "float32") = R.ccl.broadcast_from_worker0(x) + lv0: R.Tensor((128, 64), "float32") = R.matmul(broadcast_x, W1) + lv1: R.Tensor((128, 64), "float32") = R.nn.gelu(lv0) + lv2: R.Tensor((128, 128), "float32") = R.matmul(lv1, W2) + lv3: R.Tensor((128, 128), "float32") = R.ccl.allreduce(lv2, "sum") + R.output(lv3) + return lv3 + + dev = tvm.cpu(0) + X = np.random.randn(128, 128).astype("float32") + W1 = np.random.randn(128, 128).astype("float32") + W2 = np.random.randn(128, 128).astype("float32") + Y_expected = VirtualMachine(relax_build(MLP), device=dev)["main"]( + tvm.runtime.tensor(X, device=dev), + tvm.runtime.tensor(W1, device=dev), + tvm.runtime.tensor(W2, device=dev), + ).numpy() + + with tempfile.TemporaryDirectory() as tmpdir: + path = tmpdir + "/test.so" + relax_build(ShardedMLP).export_library(path) + + mod = sess.load_vm_module(path) + + d_X = sess.empty((128, 128), "float32") + d_W1 = sess.empty((128, 64), "float32") + d_W2 = sess.empty((64, 128), "float32") + + d_X.debug_copy_from(0, X) + d_W1.debug_copy_from(0, W1[:, :64]) + d_W1.debug_copy_from(1, W1[:, 64:]) + d_W2.debug_copy_from(0, W2[:64, :]) + d_W2.debug_copy_from(1, W2[64:, :]) + d_Y = mod["main"](d_X, d_W1, d_W2) + Y_result = tvm.runtime.empty((128, 128), "float32", device=dev) + sess.copy_from_worker_0(Y_result, d_Y) + sess.sync_worker_0() + Y_result = Y_result.numpy() + # pylint: enable=invalid-name + np.testing.assert_allclose(Y_result, Y_expected, rtol=1e-4, atol=1e-4) + + +def test_attention(): # pylint: disable=too-many-locals,too-many-statements + devices = [0, 1] + sess = create_socket_session(num_workers=len(devices)) + sess.init_ccl(_CCL, *devices) + + # pylint: disable=invalid-name + @tvm.script.ir_module + class Attention: # pylint: disable=too-few-public-methods + @R.function + def main( # pylint: disable=too-many-locals + x: R.Tensor((1, 10, 128), "float32"), + Wq: R.Tensor((128, 512), "float32"), + Wk: R.Tensor((128, 512), "float32"), + Wv: R.Tensor((128, 512), "float32"), + Wo: R.Tensor((512, 128), "float32"), + ) -> R.Tensor((128, 128), "float32"): + R.func_attr({"global_symbol": "main"}) + with R.dataflow(): + lv0: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wq) + lv1: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv0, [1, 10, 8, 64]) + lv2: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv1, [0, 2, 1, 3]) + lv3: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wk) + lv4: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv3, [1, 10, 8, 64]) + lv5: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv4, [0, 2, 1, 3]) + lv6: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wv) + lv7: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv6, [1, 10, 8, 64]) + lv8: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv7, [0, 2, 1, 3]) + lv9: R.Tensor((1, 8, 64, 10), "float32") = R.permute_dims(lv5, [0, 1, 3, 2]) + lv10: R.Tensor((1, 8, 10, 10), "float32") = R.matmul(lv2, lv9) + lv11: R.Tensor((1, 8, 10, 10), "float32") = R.multiply( + lv10, R.const(1 / 8, "float32") + ) + lv12: R.Tensor((1, 8, 10, 10), "float32") = R.nn.softmax(lv11, axis=-1) + lv13: R.Tensor((1, 8, 10, 64), "float32") = R.matmul(lv12, lv8) + lv14: R.Tensor((1, 10, 8, 64), "float32") = R.permute_dims(lv13, [0, 2, 1, 3]) + lv15: R.Tensor((1, 10, 512), "float32") = R.reshape(lv14, [1, 10, 512]) + lv16: R.Tensor((1, 10, 128), "float32") = R.matmul(lv15, Wo) + R.output(lv16) + return lv16 + + @tvm.script.ir_module + class ShardedAttention: # pylint: disable=too-few-public-methods + @R.function + def main( # pylint: disable=too-many-locals + x: R.Tensor((1, 10, 128), "float32"), + Wq: R.Tensor((128, 256), "float32"), # shard along axis 1 + Wk: R.Tensor((128, 256), "float32"), # shard along axis 1 + Wv: R.Tensor((128, 256), "float32"), # shard along axis 1 + Wo: R.Tensor((256, 128), "float32"), # shard along axis 0 + ) -> R.Tensor((128, 128), "float32"): + R.func_attr({"global_symbol": "main"}) + with R.dataflow(): + broadcast_x: R.Tensor((1, 10, 128), "float32") = R.ccl.broadcast_from_worker0(x) + lv0: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wq) + lv1: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv0, [1, 10, 4, 64]) + lv2: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv1, [0, 2, 1, 3]) + lv3: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wk) + lv4: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv3, [1, 10, 4, 64]) + lv5: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv4, [0, 2, 1, 3]) + lv6: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wv) + lv7: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv6, [1, 10, 4, 64]) + lv8: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv7, [0, 2, 1, 3]) + lv9: R.Tensor((1, 4, 64, 10), "float32") = R.permute_dims(lv5, [0, 1, 3, 2]) + lv10: R.Tensor((1, 4, 10, 10), "float32") = R.matmul(lv2, lv9) + lv11: R.Tensor((1, 4, 10, 10), "float32") = R.multiply( + lv10, R.const(1 / 8, "float32") + ) + lv12: R.Tensor((1, 4, 10, 10), "float32") = R.nn.softmax(lv11, axis=-1) + lv13: R.Tensor((1, 4, 10, 64), "float32") = R.matmul(lv12, lv8) + lv14: R.Tensor((1, 10, 4, 64), "float32") = R.permute_dims(lv13, [0, 2, 1, 3]) + lv15: R.Tensor((1, 10, 256), "float32") = R.reshape(lv14, [1, 10, 256]) + lv16: R.Tensor((1, 10, 128), "float32") = R.matmul(lv15, Wo) + lv17: R.Tensor((1, 10, 128), "float32") = R.ccl.allreduce(lv16, "sum") + R.output(lv17) + return lv17 + + dev = tvm.cpu(0) + X = np.random.randn(1, 10, 128).astype("float32") + Wq = np.random.randn(128, 512).astype("float32") + Wk = np.random.randn(128, 512).astype("float32") + Wv = np.random.randn(128, 512).astype("float32") + Wo = np.random.randn(512, 128).astype("float32") + Y_expected = VirtualMachine(relax_build(Attention), device=dev)["main"]( + tvm.runtime.tensor(X, device=dev), + tvm.runtime.tensor(Wq, device=dev), + tvm.runtime.tensor(Wk, device=dev), + tvm.runtime.tensor(Wv, device=dev), + tvm.runtime.tensor(Wo, device=dev), + ).numpy() + + with tempfile.TemporaryDirectory() as tmpdir: + path = tmpdir + "/test.so" + relax_build(ShardedAttention).export_library(path) + + mod = sess.load_vm_module(path) + + d_X = sess.empty((1, 10, 128), "float32") + d_Wq = sess.empty((128, 256), "float32") + d_Wk = sess.empty((128, 256), "float32") + d_Wv = sess.empty((128, 256), "float32") + d_Wo = sess.empty((256, 128), "float32") + + d_X.debug_copy_from(0, X) + d_Wq.debug_copy_from(0, Wq[:, :256]) + d_Wq.debug_copy_from(1, Wq[:, 256:]) + d_Wk.debug_copy_from(0, Wk[:, :256]) + d_Wk.debug_copy_from(1, Wk[:, 256:]) + d_Wv.debug_copy_from(0, Wv[:, :256]) + d_Wv.debug_copy_from(1, Wv[:, 256:]) + d_Wo.debug_copy_from(0, Wo[:256, :]) + d_Wo.debug_copy_from(1, Wo[256:, :]) + d_Y = mod["main"](d_X, d_Wq, d_Wk, d_Wv, d_Wo) + Y_result = tvm.runtime.empty((1, 10, 128), "float32", device=dev) + sess.copy_from_worker_0(Y_result, d_Y) + sess.sync_worker_0() + Y_result = Y_result.numpy() + # pylint: enable=invalid-name + np.testing.assert_allclose(Y_result, Y_expected, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/disco/test_socket_upload.py b/tests/python/disco/test_socket_upload.py new file mode 100644 index 000000000000..245c035b9889 --- /dev/null +++ b/tests/python/disco/test_socket_upload.py @@ -0,0 +1,198 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=missing-docstring + +import hashlib +import pathlib +import socket +import subprocess +import sys +import tempfile + +import numpy as np +import pytest + +import tvm +import tvm.script +import tvm.testing +from tvm import relax as rx +from tvm.runtime import disco as di +from tvm.script import relax as R + +if di is None: + pytest.skip("disco runtime is not available", allow_module_level=True) + +_NUM_NODES = 4 +_NUM_WORKERS = 4 +_REL_PATH = "./mod.so" + + +def _node_dir_name(node_id): + return f"upload_node{node_id}" + + +def _get_free_port(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def _sha256(path): + with open(path, "rb") as file: + return hashlib.sha256(file.read()).hexdigest() + + +@tvm.script.ir_module +class Mod: # pylint: disable=too-few-public-methods + @R.function + def double(x: R.Tensor((4,), "float32")) -> R.Tensor((4,), "float32"): + R.func_attr({"global_symbol": "double"}) + with R.dataflow(): + y: R.Tensor((4,), "float32") = R.add(x, x) + R.output(y) + return y + + +class SocketSessionTester: + + def __init__(self, num_workers, node_cwds, num_groups=1, build_ring=True): + # Initialize the attributes used by __del__ first, so that teardown is + # safe even when __init__ raises below. + self.sess = None + self.remote_nodes = [] + num_nodes = len(node_cwds) + assert num_workers % num_nodes == 0 + num_workers_per_node = num_workers // num_nodes + server_host = "localhost" + server_port = _get_free_port() + + cmd = "tvm.exec.disco_remote_socket_session" + for cwd in node_cwds[1:]: # node 0 is this process + self.remote_nodes.append( + subprocess.Popen( + [ + sys.executable, + "-m", + cmd, + server_host, + str(server_port), + str(num_workers_per_node), + ], + cwd=str(cwd), + stdout=sys.stdout, + stderr=sys.stderr, + ) + ) + + self.sess = di.SocketSession( + num_nodes, + num_workers_per_node, + num_groups, + server_host, + server_port, + build_ring, + ) + + # Bound at class creation: module globals may already be cleared when + # __del__ runs during interpreter shutdown. + _TIMEOUT_EXPIRED = subprocess.TimeoutExpired + + def close(self): + try: + # Shut down the session first so remote nodes can exit gracefully. + if self.sess is not None: + self.sess.shutdown() + self.sess = None + finally: + for node in self.remote_nodes: + try: + node.wait(timeout=10) + except self._TIMEOUT_EXPIRED: + node.kill() + node.wait() + self.remote_nodes = [] + + def __del__(self): + self.close() + + +def test_upload_vm_module(monkeypatch): + + with tempfile.TemporaryDirectory(prefix="disco_upload_") as tmp_root: + root = pathlib.Path(tmp_root) + node_dirs = [root / _node_dir_name(i) for i in range(_NUM_NODES)] + for node_dir in node_dirs: + node_dir.mkdir() + print(f"\n[setup] created {_NUM_NODES} node directories under {root}") + for node_dir in node_dirs: + print(f" {node_dir}") + + # The controller is this process, so give it upload_node0 too -- otherwise the relative + # path would resolve against the pytest rootdir and leak the artifact out of the sandbox. + monkeypatch.chdir(node_dirs[0]) + + target = tvm.target.Target("llvm") + tvm.compile(rx.get_pipeline("zero")(Mod), target=target).export_library(_REL_PATH) + expected_sha = _sha256(_REL_PATH) + built = node_dirs[0] / _REL_PATH + print(f"[setup] compiled {built} ({built.stat().st_size} bytes, sha256={expected_sha[:16]})") + + # Precondition: only the controller has the artifact. + for node_dir in node_dirs[1:]: + assert not (node_dir / _REL_PATH).exists() + print(f"[setup] upload_node1..{_NUM_NODES - 1} are empty, as expected") + + tester = SocketSessionTester(_NUM_WORKERS, node_dirs) + try: + sess = tester.sess + sess._sync_all() + sess.upload_vm_module(_REL_PATH) + sess._sync_all() # pylint: disable=protected-access + + print("[upload] per-node state after upload_vm_module:") + for i, node_dir in enumerate(node_dirs): + written = node_dir / _REL_PATH + assert written.is_file(), f"{_node_dir_name(i)} never received the module" + digest = _sha256(written) + print( + f" {_node_dir_name(i)}: {written.stat().st_size} bytes," + f" sha256={digest[:16]}" + ) + assert digest == expected_sha, f"{_node_dir_name(i)} got a corrupt copy" + + # The bytes arrived intact; now show they are loadable and runnable on every node. + mod = sess.load_vm_module(_REL_PATH) + print(f"[load] load_vm_module({_REL_PATH!r}) succeeded on all {_NUM_NODES} nodes") + + d_x = sess.empty((4,), "float32") + for i in range(_NUM_WORKERS): + d_x.debug_copy_from(i, np.full((4,), i + 1, dtype="float32")) + d_y = mod["double"](d_x) + for i in range(_NUM_WORKERS): + got = d_y.debug_get_from_remote(i).numpy() + np.testing.assert_equal(got, np.full((4,), 2 * (i + 1), dtype="float32")) + print(f" worker {i}: double({i + 1}) -> {got}") + finally: + tester.close() + + print(f"[cleanup] removed {tmp_root}") + + +if __name__ == "__main__": + tvm.testing.main() From fa40b6090809d5dd2b5343a44824a75d5a8eeeba Mon Sep 17 00:00:00 2001 From: yuyi Date: Wed, 29 Jul 2026 13:36:24 +0800 Subject: [PATCH 09/10] [Test][Relax][Op] Add test for gather_to_worker0 --- tests/python/relax/test_op_ccl.py | 61 +++++++++++++++++++ .../relax/test_transform_legalize_ops_ccl.py | 23 +++++++ 2 files changed, 84 insertions(+) diff --git a/tests/python/relax/test_op_ccl.py b/tests/python/relax/test_op_ccl.py index e0cb8396bf3b..415651f49533 100644 --- a/tests/python/relax/test_op_ccl.py +++ b/tests/python/relax/test_op_ccl.py @@ -29,6 +29,7 @@ def test_op_correctness(): assert relax.op.ccl.allreduce(x).op == Op.get("relax.ccl.allreduce") assert relax.op.ccl.broadcast_from_worker0(x).op == Op.get("relax.ccl.broadcast_from_worker0") assert relax.op.ccl.allgather(x, 2).op == Op.get("relax.ccl.allgather") + assert relax.op.ccl.gather_to_worker0(x, 2).op == Op.get("relax.ccl.gather_to_worker0") def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type): @@ -255,5 +256,65 @@ def test_scatter_from_worker0_infer_ty_more_input_dtype(): ) +# gather_to_worker0 concatenates each worker's shard along axis 0, so its inferred type matches +# allgather's; only worker-0 actually holds the result at runtime. +def test_gather_to_worker0_infer_ty(): + bb = relax.BlockBuilder() + x0 = relax.Var("x", R.Tensor((2, 3), "float32")) + x1 = relax.Var("x", R.Tensor("float32", ndim=3)) + x2 = relax.Var("x", R.Tensor("float32", ndim=-1)) + x3 = relax.Var("x", R.Tensor((2, 3))) + x4 = relax.Var("x", R.Tensor()) + x5 = relax.Var("x", R.Tensor((3, 4))) + + _check_inference(bb, relax.op.ccl.gather_to_worker0(x0, 2), relax.TensorType((4, 3), "float32")) + _check_inference( + bb, relax.op.ccl.gather_to_worker0(x1, 2), relax.TensorType(dtype="float32", ndim=3) + ) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x2, 2), relax.TensorType(dtype="float32")) + _check_inference( + bb, relax.op.ccl.gather_to_worker0(x3, 2), relax.TensorType((4, 3), dtype=None) + ) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x4, 2), relax.TensorType(dtype=None)) + _check_inference( + bb, relax.op.ccl.gather_to_worker0(x5, 2), relax.TensorType((6, 4), dtype=None) + ) + + +def test_gather_to_worker0_infer_ty_shape_symbolic(): + bb = relax.BlockBuilder() + m = tirx.Var("m", "int64") + n = tirx.Var("n", "int64") + x0 = relax.Var("x", R.Tensor((m, n), "float32")) + x1 = relax.Var("x", R.Tensor((4, n), "float32")) + + _check_inference( + bb, relax.op.ccl.gather_to_worker0(x0, 2), relax.TensorType((m * 2, n), "float32") + ) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x1, 2), relax.TensorType((8, n), "float32")) + + +def test_gather_to_worker0_infer_ty_shape_var(): + bb = relax.BlockBuilder() + s0 = relax.Var("s", relax.ShapeType(ndim=2)) + s1 = relax.Var("s", relax.ShapeType()) + x0 = relax.Var("x", relax.TensorType(s0, "float32")) + x1 = relax.Var("x", relax.TensorType(s1, "float32")) + + _check_inference(bb, relax.op.ccl.gather_to_worker0(x0, 2), relax.TensorType(s0, "float32")) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x1, 2), relax.TensorType(s1, "float32")) + + +def test_gather_to_worker0_infer_ty_more_input_dtype(): + bb = relax.BlockBuilder() + x0 = relax.Var("x", R.Tensor((2, 3), "float64")) + x1 = relax.Var("x", R.Tensor((2, 3), "int8")) + x2 = relax.Var("x", R.Tensor((2, 3), "int64")) + + _check_inference(bb, relax.op.ccl.gather_to_worker0(x0, 2), relax.TensorType((4, 3), "float64")) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x1, 2), relax.TensorType((4, 3), "int8")) + _check_inference(bb, relax.op.ccl.gather_to_worker0(x2, 2), relax.TensorType((4, 3), "int64")) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_transform_legalize_ops_ccl.py b/tests/python/relax/test_transform_legalize_ops_ccl.py index b6459e3ae694..5830fb657bd0 100644 --- a/tests/python/relax/test_transform_legalize_ops_ccl.py +++ b/tests/python/relax/test_transform_legalize_ops_ccl.py @@ -143,5 +143,28 @@ def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 5), dtype="flo tvm.ir.assert_structural_equal(mod, Expected) +def test_gather_to_worker0(): + # fmt: off + @tvm.script.ir_module + class GatherToWorker0: + @R.function + def main(x: R.Tensor((10, 10), "float32")) -> R.Tensor((10, 10), "float32"): + gv0: R.Tensor((20, 10), "float32") = R.ccl.gather_to_worker0(x, 2) + gv1 = R.ccl.gather_to_worker0(x, 2) + return x + + @I.ir_module(s_tir=True) + class Expected: + @R.function + def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"): + gv0: R.Tensor((20, 10), dtype="float32") = R.call_dps_packed("runtime.disco.gather_to_worker0", [x, True], out_ty=R.Tensor((20, 10), dtype="float32")) + gv1: R.Tensor((20, 10), dtype="float32") = R.call_dps_packed("runtime.disco.gather_to_worker0", [x, True], out_ty=R.Tensor((20, 10), dtype="float32")) + return x + # fmt: on + + mod = LegalizeOps()(GatherToWorker0) + tvm.ir.assert_structural_equal(mod, Expected) + + if __name__ == "__main__": tvm.testing.main() From 0cebf05b7618a555077df3cda95fa529327486fd Mon Sep 17 00:00:00 2001 From: yuyi Date: Wed, 29 Jul 2026 14:22:52 +0800 Subject: [PATCH 10/10] [Disco][Test] Rename test_socket_cpuccl.py and add ProcessSess and ThreadSess test [Disco][Test] Rename test_socket_cpuccl.py and add ProcessSess and ThreadSess test [Disco][Test] Adjust the comment format Fix pre-commit formatting issues --- include/tvm/runtime/disco/disco_worker.h | 2 +- include/tvm/runtime/disco/session.h | 41 ++- python/tvm/exec/disco_worker.py | 11 +- python/tvm/runtime/disco/process_pool.py | 34 +- python/tvm/runtime/disco/session.py | 2 +- src/runtime/extra/disco/bcast_session.h | 6 +- src/runtime/extra/disco/cpuccl/cpuccl.cc | 13 +- .../extra/disco/distributed/socket_session.cc | 106 +++--- src/runtime/extra/disco/process_session.cc | 22 +- src/runtime/extra/disco/threaded_session.cc | 15 +- src/support/socket.h | 3 +- .../{test_socket_cpuccl.py => test_cpuccl.py} | 342 +++++++++--------- tests/python/disco/test_session.py | 1 - tests/python/disco/test_socket_upload.py | 22 +- 14 files changed, 334 insertions(+), 286 deletions(-) rename tests/python/disco/{test_socket_cpuccl.py => test_cpuccl.py} (60%) diff --git a/include/tvm/runtime/disco/disco_worker.h b/include/tvm/runtime/disco/disco_worker.h index 68b0e6e8f5fa..9d7a145476c0 100644 --- a/include/tvm/runtime/disco/disco_worker.h +++ b/include/tvm/runtime/disco/disco_worker.h @@ -60,7 +60,7 @@ class DiscoWorker { channel(channel), register_file{}, ring_in(nullptr), - ring_out(nullptr){} + ring_out(nullptr) {} /*! \brief Main loop of the worker */ void MainLoop(); diff --git a/include/tvm/runtime/disco/session.h b/include/tvm/runtime/disco/session.h index 6ad7e2ddd204..f851bcba0fd8 100644 --- a/include/tvm/runtime/disco/session.h +++ b/include/tvm/runtime/disco/session.h @@ -76,14 +76,13 @@ #include #include #include +#include #include #include #include #include -#include - namespace tvm { namespace runtime { @@ -316,8 +315,7 @@ class Session : public ffi::ObjectRef { */ TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups, ffi::String process_pool_creator, - ffi::String entrypoint, - bool build_ring); + ffi::String entrypoint, bool build_ring); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj); }; @@ -354,9 +352,12 @@ class DiscoRingChannel { const char* write_data = static_cast(data); while (size > 0) { ssize_t n; - do { n = ::write(fd_, write_data, size); } while (n < 0 && errno == EINTR); + do { + n = ::write(fd_, write_data, size); + } while (n < 0 && errno == EINTR); TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Send failed: " << std::strerror(errno); - write_data += n; size -= n; + write_data += n; + size -= n; } } @@ -364,34 +365,44 @@ class DiscoRingChannel { char* read_data = static_cast(data); while (size > 0) { ssize_t n; - do { n = ::read(fd_, read_data, size); } while (n < 0 && errno == EINTR); + do { + n = ::read(fd_, read_data, size); + } while (n < 0 && errno == EINTR); TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Recv failed: " << std::strerror(errno); - read_data += n; size -= n; + read_data += n; + size -= n; } } - ssize_t ReadSome(void* data, size_t max_size) { if (fd_ < 0) return 0; ssize_t n; - do { n = ::read(fd_, data, max_size); } while (n < 0 && errno == EINTR); + do { + n = ::read(fd_, data, max_size); + } while (n < 0 && errno == EINTR); return n; } ssize_t WriteSome(const void* data, size_t size) { if (fd_ < 0) return -1; ssize_t n; - do { n = ::write(fd_, data, size); } while (n < 0 && errno == EINTR); + do { + n = ::write(fd_, data, size); + } while (n < 0 && errno == EINTR); return n; } - void Close() { if (fd_ >= 0) { ::close(fd_); fd_ = -1; } } + void Close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + } - private: - int fd_; + private: + int fd_; }; - /*! * \brief A special communication channel between controler and worker-0, * assuming they are always collocated in the same process. diff --git a/python/tvm/exec/disco_worker.py b/python/tvm/exec/disco_worker.py index 6330e12fd435..2fa0cbfd968a 100644 --- a/python/tvm/exec/disco_worker.py +++ b/python/tvm/exec/disco_worker.py @@ -104,16 +104,17 @@ def main(): argc = len(sys.argv) if argc not in (6, 8): print( - "Usage: " - " [ ]") + "Usage: " + " [ ]" + ) return worker_id = int(sys.argv[1]) num_workers = int(sys.argv[2]) num_groups = int(sys.argv[3]) ring_out_fd = -1 - ring_in_fd = -1 - has_ring = (argc == 8) + ring_in_fd = -1 + has_ring = argc == 8 if sys.platform == "win32": import msvcrt # pylint: disable=import-outside-toplevel,import-error @@ -124,7 +125,7 @@ def main(): reader = int(sys.argv[4]) writer = int(sys.argv[5]) if has_ring: - ring_in_fd = int(sys.argv[6]) + ring_in_fd = int(sys.argv[6]) ring_out_fd = int(sys.argv[7]) worker_func = get_global_func("runtime.disco.WorkerProcess") diff --git a/python/tvm/runtime/disco/process_pool.py b/python/tvm/runtime/disco/process_pool.py index 4020424fa5d4..9fc136c37106 100644 --- a/python/tvm/runtime/disco/process_pool.py +++ b/python/tvm/runtime/disco/process_pool.py @@ -139,7 +139,7 @@ def start(self): os.set_handle_inheritable(worker_read_handle, True) os.set_handle_inheritable(worker_write_handle, True) cmd += [str(worker_read_handle), str(worker_write_handle)] - #CPU ring is not implemented on Windows. + # CPU ring is not implemented on Windows. cmd += ["-1", "-1"] self._proc = subprocess.Popen( cmd, @@ -207,15 +207,21 @@ def _create_process_pool(num_workers: int, num_groups: int, entrypoint: str, bui pool = [] for i in range(1, num_workers): ring_out_fd = -1 - ring_in_fd = -1 + ring_in_fd = -1 if build_ring: - ring_out_fd = ring_pipes[i][1] # write end of pipe[i] - ring_in_fd = ring_pipes[(i - 1) % num_workers][0] # read end of pipe[i-1] - - pool.append(DiscoPopenWorker( - i, num_workers, num_groups, entrypoint, - ring_in_fd=ring_in_fd, ring_out_fd=ring_out_fd, - )) + ring_out_fd = ring_pipes[i][1] # write end of pipe[i] + ring_in_fd = ring_pipes[(i - 1) % num_workers][0] # read end of pipe[i-1] + + pool.append( + DiscoPopenWorker( + i, + num_workers, + num_groups, + entrypoint, + ring_in_fd=ring_in_fd, + ring_out_fd=ring_out_fd, + ) + ) def result_func(worker_id: int): nonlocal pool, ring_pipes @@ -223,10 +229,12 @@ def result_func(worker_id: int): # Special ID (-1): Return the ring FDs of worker 0 for the controller. if worker_id == -1: assert build_ring, "Ring not enabled; cannot query worker 0 ring fds" - return Shape([ - ring_pipes[num_workers - 1][0], # W0 ring_in (read of pipe[N-1]) - ring_pipes[0][1], # W0 ring_out (write of pipe[0]) - ]) + return Shape( + [ + ring_pipes[num_workers - 1][0], # W0 ring_in (read of pipe[N-1]) + ring_pipes[0][1], # W0 ring_out (write of pipe[0]) + ] + ) if worker_id != 0: read_fd, write_fd = pool[worker_id - 1].start() return Shape([read_fd, write_fd]) diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py index 445f088c1687..ad9626f0c47b 100644 --- a/python/tvm/runtime/disco/session.py +++ b/python/tvm/runtime/disco/session.py @@ -321,7 +321,7 @@ def load_vm_module( """ func = self._get_cached_method("runtime.disco.load_vm_module") return DModule(func(path, device), self) - + def upload_vm_module(self, path: str) -> str: """Broadcast a module file to every worker and write it locally on each. diff --git a/src/runtime/extra/disco/bcast_session.h b/src/runtime/extra/disco/bcast_session.h index d2b13441b4bf..785d541771d1 100644 --- a/src/runtime/extra/disco/bcast_session.h +++ b/src/runtime/extra/disco/bcast_session.h @@ -84,14 +84,15 @@ class BcastSessionObj : public SessionObj { */ virtual ffi::PackedArgs RecvReplyPacked(int worker_id) = 0; - /*! + /*! * \brief Atomically swap W_0's ring_in channel. SocketSession uses this * to redirect cross-node ring inbound through a proxy pipe. * Default no-op; ProcessSession overrides. * \param new_ch New channel to install; ownership taken. * \return Previously installed channel; ownership returned to caller. */ - virtual std::unique_ptr RerouteRingIn(std::unique_ptr new_ch) { + virtual std::unique_ptr RerouteRingIn( + std::unique_ptr new_ch) { return nullptr; } @@ -105,7 +106,6 @@ class BcastSessionObj : public SessionObj { */ virtual void CloseRing() {} - /*! \brief A side channel to communicate with worker-0 */ WorkerZeroData worker_zero_data_; /*! \brief Number of registers used, including those in `free_regs_` */ diff --git a/src/runtime/extra/disco/cpuccl/cpuccl.cc b/src/runtime/extra/disco/cpuccl/cpuccl.cc index 679e87c39197..4573bd59e7ae 100644 --- a/src/runtime/extra/disco/cpuccl/cpuccl.cc +++ b/src/runtime/extra/disco/cpuccl/cpuccl.cc @@ -29,7 +29,6 @@ #include #include - namespace tvm { namespace runtime { namespace cpuccl { @@ -248,15 +247,13 @@ void ScatterFromWorker0(ffi::Optional send, bool /*in_group*/, Tensor re int64_t recv_bytes = recv.Shape().Product() * DTypeBytes(recv->dtype); if (num_workers == 1) { - TVM_FFI_CHECK(send.has_value(), ValueError) - << "Worker 0 must provide send buffer for scatter"; + TVM_FFI_CHECK(send.has_value(), ValueError) << "Worker 0 must provide send buffer for scatter"; std::memcpy(recv->data, send.value()->data, static_cast(recv_bytes)); return; } if (rank == 0) { - TVM_FFI_CHECK(send.has_value(), ValueError) - << "Worker 0 must provide send buffer for scatter"; + TVM_FFI_CHECK(send.has_value(), ValueError) << "Worker 0 must provide send buffer for scatter"; char* src = static_cast(send.value()->data); std::memcpy(recv->data, src, static_cast(recv_bytes)); int64_t tail = recv_bytes * (num_workers - 1); @@ -287,15 +284,13 @@ void GatherToWorker0(Tensor send, bool /*in_group*/, ffi::Optional recv) int64_t chunk_bytes = send.Shape().Product() * DTypeBytes(send->dtype); if (num_workers == 1) { - TVM_FFI_CHECK(recv.has_value(), ValueError) - << "Worker 0 must provide recv buffer for gather"; + TVM_FFI_CHECK(recv.has_value(), ValueError) << "Worker 0 must provide recv buffer for gather"; std::memcpy(recv.value()->data, send->data, static_cast(chunk_bytes)); return; } if (rank == 0) { - TVM_FFI_CHECK(recv.has_value(), ValueError) - << "Worker 0 must provide recv buffer for gather"; + TVM_FFI_CHECK(recv.has_value(), ValueError) << "Worker 0 must provide recv buffer for gather"; char* dst = static_cast(recv.value()->data); std::memcpy(dst, send->data, static_cast(chunk_bytes)); int64_t incoming = chunk_bytes * (num_workers - 1); diff --git a/src/runtime/extra/disco/distributed/socket_session.cc b/src/runtime/extra/disco/distributed/socket_session.cc index cd2b1d4da906..b5e5c0597555 100644 --- a/src/runtime/extra/disco/distributed/socket_session.cc +++ b/src/runtime/extra/disco/distributed/socket_session.cc @@ -55,26 +55,30 @@ class DiscoSocketChannel : public DiscoChannel { DiscoStreamMessageQueue message_queue_; }; - // Proxy thread helpers — both controller and remote share these. // send: pipe → TCP // recv: TCP → pipe constexpr size_t kProxyBufSize = 64 * 1024; void ProxyLoop(DiscoRingChannel* src, DiscoRingChannel* dst, std::string tag) { - LOG(INFO) << "[Proxy " << tag << "] entering, src=" << src << " dst=" << dst; std::vector buf(kProxyBufSize); while (true) { ssize_t n = src->ReadSome(buf.data(), buf.size()); - if (n <= 0) { LOG(INFO) << "[Proxy " << tag << "] src closed (n=" << n << "), exit"; return; } + if (n <= 0) { + LOG(INFO) << "[Proxy " << tag << "] src closed (n=" << n << "), exit"; + return; + } LOG(INFO) << "[Proxy " << tag << "] transfer " << n << " bytes"; ssize_t written = 0; while (written < n) { ssize_t w = dst->WriteSome(buf.data() + written, n - written); - if (w <= 0) { LOG(INFO) << "[Proxy " << tag << "] dst failed (w=" << w << "), exit"; return; } + if (w <= 0) { + LOG(INFO) << "[Proxy " << tag << "] dst failed (w=" << w << "), exit"; + return; + } written += w; } } @@ -86,8 +90,8 @@ class RingProxyEndpoint { const std::vector& node_hosts, const BcastSession& local_session, const std::string& tag) { const int next_node_id = (node_id + 1) % num_nodes; - const int recv_port = base_ring_port + node_id; - const int send_port = base_ring_port + next_node_id; + const int recv_port = base_ring_port + node_id; + const int send_port = base_ring_port + next_node_id; LOG(INFO) << "[Ring " << tag << "] node_id=" << node_id << " listen=:" << recv_port << " connect_to=" << node_hosts[next_node_id].c_str() << ":" << send_port; @@ -123,26 +127,29 @@ class RingProxyEndpoint { accept_thread.join(); LOG(INFO) << "[Ring " << tag << "] ring TCP ready"; - tcp_in_ch_ = std::make_unique(recv_sock.Release()); + tcp_in_ch_ = std::make_unique(recv_sock.Release()); tcp_out_ch_ = std::make_unique(send_sock.Release()); // Splice the local session's W0 ring pipe onto the TCP ring via a pipe + two proxy threads. int sess_fds[2]; // 0:read 1:write TVM_FFI_ICHECK_EQ(::pipe(sess_fds), 0); - proxy_out_to_tcp_ = local_session->RerouteRingIn(std::make_unique(sess_fds[0])); + proxy_out_to_tcp_ = + local_session->RerouteRingIn(std::make_unique(sess_fds[0])); TVM_FFI_ICHECK(proxy_out_to_tcp_ != nullptr); proxy_in_from_tcp_ = std::make_unique(sess_fds[1]); - send_thread_ = std::thread(ProxyLoop, proxy_out_to_tcp_.get(), tcp_out_ch_.get(), tag + "-send"); - recv_thread_ = std::thread(ProxyLoop, tcp_in_ch_.get(), proxy_in_from_tcp_.get(), tag + "-recv"); + send_thread_ = + std::thread(ProxyLoop, proxy_out_to_tcp_.get(), tcp_out_ch_.get(), tag + "-send"); + recv_thread_ = + std::thread(ProxyLoop, tcp_in_ch_.get(), proxy_in_from_tcp_.get(), tag + "-recv"); LOG(INFO) << "[Ring " << tag << "] send/recv proxy threads started"; } void Close() { - if (proxy_out_to_tcp_) proxy_out_to_tcp_->Close(); + if (proxy_out_to_tcp_) proxy_out_to_tcp_->Close(); if (proxy_in_from_tcp_) proxy_in_from_tcp_->Close(); - if (tcp_in_ch_) tcp_in_ch_->Close(); - if (tcp_out_ch_) tcp_out_ch_->Close(); + if (tcp_in_ch_) tcp_in_ch_->Close(); + if (tcp_out_ch_) tcp_out_ch_->Close(); } void Join() { @@ -163,12 +170,15 @@ class SocketSessionObj : public BcastSessionObj { public: explicit SocketSessionObj(int num_nodes, int num_workers_per_node, int num_groups, const ffi::String& host, int port, bool build_ring) - : num_nodes_(num_nodes), num_workers_per_node_(num_workers_per_node), build_ring_(build_ring) { + : num_nodes_(num_nodes), + num_workers_per_node_(num_workers_per_node), + build_ring_(build_ring) { const auto f_create_local_session = tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers"); TVM_FFI_ICHECK(f_create_local_session.has_value()) << "Cannot find function runtime.disco.create_socket_session_local_workers"; - local_session_ = ((*f_create_local_session)(num_workers_per_node, build_ring_)).cast(); + local_session_ = + ((*f_create_local_session)(num_workers_per_node, build_ring_)).cast(); DRef f_init_workers = local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers"); local_session_->CallPacked(f_init_workers, num_nodes_, /*node_id=*/0, num_groups, @@ -210,30 +220,32 @@ class SocketSessionObj : public BcastSessionObj { std::string full = addr.AsString(); ffi::String ip(full.substr(0, full.find(':'))); node_hosts_.push_back(ip); - LOG(INFO) << "[Host] IP table : node " << (i + 1) << " ip=" << ip.c_str(); + LOG(INFO) << "[Host] IP table : node " << (i + 1) << " ip=" << ip.c_str(); } } - if( build_ring_ && num_nodes_ > 1) { - - const int base_ring_port = port + 1; + if (build_ring_ && num_nodes_ > 1) { + const int base_ring_port = port + 1; - // Broadcast base_ring_port + the IP table so every remote node can wire its ring. - { - std::vector baseport_ip(num_nodes_ + 1); - baseport_ip[0] = base_ring_port; - for (int i = 0; i < num_nodes_; ++i) baseport_ip[1 + i] = node_hosts_[i]; - for (auto& ch : remote_channels_) { - ch->Send(ffi::PackedArgs(baseport_ip.data(), baseport_ip.size())); - } - LOG(INFO) << "[Host] IP and port broadcast completed."; + // Broadcast base_ring_port + the IP table so every remote node can wire its ring. + { + std::vector baseport_ip(num_nodes_ + 1); + baseport_ip[0] = base_ring_port; + for (int i = 0; i < num_nodes_; ++i) baseport_ip[1 + i] = node_hosts_[i]; + for (auto& ch : remote_channels_) { + ch->Send(ffi::PackedArgs(baseport_ip.data(), baseport_ip.size())); } - ring_.Connect(/*node_id=*/0, num_nodes_, base_ring_port, node_hosts_, local_session_,"controller"); + LOG(INFO) << "[Host] IP and port broadcast completed."; + } + ring_.Connect(/*node_id=*/0, num_nodes_, base_ring_port, node_hosts_, local_session_, + "controller"); - const int num_workers = num_nodes_ * num_workers_per_node_; - for (int worker_id = 1; worker_id < num_workers; ++worker_id) { this->SyncWorker(worker_id); } - LOG(INFO) << "[Host] Controller: All workers are ready. Initialization complete."; - }// build ring + const int num_workers = num_nodes_ * num_workers_per_node_; + for (int worker_id = 1; worker_id < num_workers; ++worker_id) { + this->SyncWorker(worker_id); + } + LOG(INFO) << "[Host] Controller: All workers are ready. Initialization complete."; + } // build ring } int64_t GetNumWorkers() final { return num_nodes_ * num_workers_per_node_; } @@ -321,7 +333,6 @@ class SocketSessionObj : public BcastSessionObj { } void Shutdown() final { - if (shutdown_) return; shutdown_ = true; @@ -363,8 +374,8 @@ class SocketSessionObj : public BcastSessionObj { std::vector node_hosts_; bool build_ring_ = false; - private: - RingProxyEndpoint ring_; + private: + RingProxyEndpoint ring_; }; class RemoteSocketSession { @@ -378,7 +389,10 @@ class RemoteSocketSession { bool connected = false; for (int i = 0; i < 100; ++i) { - if (socket_.Connect(server_addr)) { connected = true; break; } + if (socket_.Connect(server_addr)) { + connected = true; + break; + } std::this_thread::sleep_for(std::chrono::milliseconds(100)); } if (!connected) { @@ -473,13 +487,14 @@ class RemoteSocketSession { void InitLocalSession() { const auto f_create_local_session = tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers"); - local_session_ = ((*f_create_local_session)(num_workers_per_node_, build_ring)).cast(); + local_session_ = + ((*f_create_local_session)(num_workers_per_node_, build_ring)).cast(); DRef f_init_workers = local_session_->GetGlobalFunc("runtime.disco.socket_session_init_workers"); local_session_->CallPacked(f_init_workers, num_nodes_, node_id_, num_groups_, num_workers_per_node_); - if( build_ring && num_nodes_ > 1) { + if (build_ring && num_nodes_ > 1) { // The controller broadcasts base_ring_port + the IP table on our control channel. ffi::PackedArgs host_info = channel_->Recv(); TVM_FFI_ICHECK_EQ(host_info.size(), num_nodes_ + 1); @@ -491,7 +506,7 @@ class RemoteSocketSession { ring_.Connect(node_id_, nNodes, base_ring_port, node_hosts, local_session_, "node_" + std::to_string(node_id_)); - }//build ring + } // build ring } TCPSocket socket_; @@ -513,10 +528,10 @@ void RemoteSocketSessionEntryPoint(const ffi::String& server_host, int server_po proxy.MainLoop(); LOG(INFO) << "[Remote] " << server_host << ":" << server_port << " main loop exited cleanly, process ending."; - + } catch (const std::exception& e) { LOG(INFO) << "[Remote] node exiting due to exception (controller lost?): " << e.what(); - throw; + throw; } } @@ -527,8 +542,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { Session SocketSession(int num_nodes, int num_workers_per_node, int num_groups, const ffi::String& host, int port, bool build_ring) { - auto n = - ffi::make_object(num_nodes, num_workers_per_node, num_groups, host, port, build_ring); + auto n = ffi::make_object(num_nodes, num_workers_per_node, num_groups, host, + port, build_ring); return Session(n); } @@ -543,7 +558,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { worker->num_groups = num_groups; worker->worker_id = worker->worker_id + node_id * num_workers_per_node; worker->num_workers = num_nodes * num_workers_per_node; - LOG(INFO) << "[Disco_worker_" << worker->worker_id << "]" << " Initializing worker group with " << num_nodes << " nodes, " + LOG(INFO) << "[Disco_worker_" << worker->worker_id << "]" + << " Initializing worker group with " << num_nodes << " nodes, " << num_workers_per_node << " workers per node, and " << num_groups << " groups."; }); diff --git a/src/runtime/extra/disco/process_session.cc b/src/runtime/extra/disco/process_session.cc index 21d7154a48d3..5476877b0af7 100644 --- a/src/runtime/extra/disco/process_session.cc +++ b/src/runtime/extra/disco/process_session.cc @@ -60,7 +60,8 @@ class DiscoProcessChannel final : public DiscoChannel { class ProcessSessionObj final : public BcastSessionObj { public: - explicit ProcessSessionObj(int num_workers, int num_groups, ffi::Function process_pool, bool build_ring) + explicit ProcessSessionObj(int num_workers, int num_groups, ffi::Function process_pool, + bool build_ring) : process_pool_(process_pool), worker_0_( std::make_unique(0, num_workers, num_groups, &worker_zero_data_)) { @@ -85,12 +86,12 @@ class ProcessSessionObj final : public BcastSessionObj { TVM_FFI_CHECK_EQ(w0_fds.size(), 2, ValueError) << "process_pool(-1) should return a tuple of size 2 (worker_0's ring fds), " << "but got a tuple of size " << w0_fds.size() << "."; - int64_t w0_ring_in_fd = w0_fds[0]; + int64_t w0_ring_in_fd = w0_fds[0]; int64_t w0_ring_out_fd = w0_fds[1]; - ring_in_w0_ = std::make_unique(w0_ring_in_fd); + ring_in_w0_ = std::make_unique(w0_ring_in_fd); ring_out_w0_ = std::make_unique(w0_ring_out_fd); - worker_0_->worker->ring_in = ring_in_w0_.get(); + worker_0_->worker->ring_in = ring_in_w0_.get(); worker_0_->worker->ring_out = ring_out_w0_.get(); } } @@ -179,8 +180,9 @@ class ProcessSessionObj final : public BcastSessionObj { } return workers_.at(worker_id - 1).get(); } - - std::unique_ptr RerouteRingIn(std::unique_ptr new_ch) override { + + std::unique_ptr RerouteRingIn( + std::unique_ptr new_ch) override { auto old = std::move(ring_in_w0_); ring_in_w0_ = std::move(new_ch); worker_0_->worker->ring_in = ring_in_w0_.get(); @@ -212,8 +214,8 @@ Session Session::ProcessSession(int num_workers, int num_group, ffi::String proc return Session(n); } -void WorkerProcess(int worker_id, int num_workers, int num_group, int64_t read_fd, - int64_t write_fd, int64_t ring_in_fd, int64_t ring_out_fd) { +void WorkerProcess(int worker_id, int num_workers, int num_group, int64_t read_fd, int64_t write_fd, + int64_t ring_in_fd, int64_t ring_out_fd) { TVM_FFI_ICHECK_EQ(num_workers % num_group, 0) << "The number of workers should be divisible by the number of worker group."; DiscoProcessChannel channel(read_fd, write_fd); @@ -221,9 +223,9 @@ void WorkerProcess(int worker_id, int num_workers, int num_group, int64_t read_f std::unique_ptr ring_in_ch, ring_out_ch; if (ring_in_fd >= 0 && ring_out_fd >= 0) { - ring_in_ch = std::make_unique(ring_in_fd); + ring_in_ch = std::make_unique(ring_in_fd); ring_out_ch = std::make_unique(ring_out_fd); - worker.ring_in = ring_in_ch.get(); + worker.ring_in = ring_in_ch.get(); worker.ring_out = ring_out_ch.get(); } diff --git a/src/runtime/extra/disco/threaded_session.cc b/src/runtime/extra/disco/threaded_session.cc index 7b4996cb08cf..3a9d2cdd0ef1 100644 --- a/src/runtime/extra/disco/threaded_session.cc +++ b/src/runtime/extra/disco/threaded_session.cc @@ -157,7 +157,7 @@ class ThreadedSessionObj final : public BcastSessionObj { ~ThreadedSessionObj() { this->Shutdown(); workers_.clear(); - ring_channels_.clear(); + ring_channels_.clear(); } int64_t GetNumWorkers() { return workers_.size(); } @@ -191,24 +191,23 @@ class ThreadedSessionObj final : public BcastSessionObj { std::vector workers_; std::vector> ring_channels_; - private: + private: void BuildRing(int num_workers) { - std::vector read_fds(num_workers), write_fds(num_workers); for (int i = 0; i < num_workers; ++i) { int fds[2]; TVM_FFI_ICHECK_EQ(::pipe(fds), 0) << "Failed to create ring pipe for worker " << i; - read_fds[i] = fds[0]; - write_fds[i] = fds[1]; + read_fds[i] = fds[0]; + write_fds[i] = fds[1]; } ring_channels_.reserve(num_workers * 2); for (int i = 0; i < num_workers; ++i) { - auto channel_out = std::make_unique(write_fds[i]); - auto channel_in = std::make_unique(read_fds[(i - 1 + num_workers) % num_workers]); + auto channel_in = + std::make_unique(read_fds[(i - 1 + num_workers) % num_workers]); workers_[i].worker->ring_out = channel_out.get(); - workers_[i].worker->ring_in = channel_in.get(); + workers_[i].worker->ring_in = channel_in.get(); ring_channels_.push_back(std::move(channel_in)); ring_channels_.push_back(std::move(channel_out)); } diff --git a/src/support/socket.h b/src/support/socket.h index 7830852bedcc..4e0b16d7edc1 100644 --- a/src/support/socket.h +++ b/src/support/socket.h @@ -294,7 +294,7 @@ class Socket { return false; } - /*! + /*! * \brief Relinquish ownership of the underlying descriptor without closing it, and return it. * The caller becomes responsible for closing the returned fd (e.g. when handing it to * another owner). Mirrors std::unique_ptr::release. Prevents a later double-close. @@ -306,7 +306,6 @@ class Socket { return fd; } - /*! \brief check if socket is already closed */ bool IsClosed() const { return sockfd == INVALID_SOCKET; } /*! \brief close the socket */ diff --git a/tests/python/disco/test_socket_cpuccl.py b/tests/python/disco/test_cpuccl.py similarity index 60% rename from tests/python/disco/test_socket_cpuccl.py rename to tests/python/disco/test_cpuccl.py index 74fe6c26a72f..ba94a5bf2827 100644 --- a/tests/python/disco/test_socket_cpuccl.py +++ b/tests/python/disco/test_cpuccl.py @@ -15,16 +15,9 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-docstring -"""Collective (cpuccl) tests over a multi-node SocketSession. +"""Collective (cpuccl) tests over ThreadedSession, ProcessSession and SocketSession. +Reuse the test cases from test_ccl.py""" -This file is intentionally separate from ``test_ccl.py`` (which targets nccl/rccl on -ThreadedSession/ProcessSession). - -The socket path needs its own multi-node harness -- a remote subprocess per extra node, plus -careful startup ordering -- which does not belong in the upstream nccl test matrix. cpuccl -always builds the ring (``build_ring=True``) and does not support group (``in_group``) -semantics, so only the global collectives are covered. -""" import socket import subprocess import sys @@ -35,13 +28,12 @@ import pytest import tvm -import tvm.testing import tvm.script +import tvm.testing from tvm import relax as rx from tvm.runtime import disco as di from tvm.runtime.vm import VirtualMachine from tvm.script import relax as R -from tvm.exec import disco_worker as _ # pylint: disable=unused-import if di is None: pytest.skip("disco runtime is not available", allow_module_level=True) @@ -59,16 +51,9 @@ def _get_free_port(): class SocketSessionTester: - """Run a disco SocketSession with one local node and remote nodes. - - Each remote node is a `tvm.exec.disco_remote_socket_session` subprocess launched with the - current Python interpreter. Unlike test_session.py, `build_ring` defaults to True because - cpuccl uses the ring as its data plane. - """ + """Run a disco SocketSession with one local node and remote nodes.""" def __init__(self, num_workers, num_nodes=2, num_groups=1, build_ring=True): - # Initialize the attributes used by __del__ first, so that teardown is - # safe even when __init__ raises below. self.sess = None self.remote_nodes = [] assert num_workers % num_nodes == 0 @@ -120,6 +105,7 @@ def start_server(): def __del__(self): try: + # Shut down the session first so remote nodes can exit gracefully. if self.sess is not None: self.sess.shutdown() finally: @@ -132,11 +118,6 @@ def __del__(self): def create_socket_session(num_workers, build_ring=True): - """Create a socket session backed by one local and one remote node. - - The tester is kept alive in a module-level global so that the session - survives until the next call (or interpreter exit) replaces it. - """ global _SOCKET_SESSION_TESTER # Rebind (not `del`) so the global stays defined if the constructor raises. _SOCKET_SESSION_TESTER = None @@ -145,125 +126,152 @@ def create_socket_session(num_workers, build_ring=True): return _SOCKET_SESSION_TESTER.sess +_all_session_kinds = [di.ThreadedSession, di.ProcessSession, create_socket_session] + -def test_init(): +def _session_id(session_kind): + return session_kind.__name__ + + +def _run_with_ccl_session(session_kind, devices, func): + sess = session_kind(num_workers=len(devices), build_ring=True) + try: + sess.init_ccl(_CCL, *devices) + return func(sess) + finally: + sess.shutdown() + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_init(session_kind): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) + def run_test(_sess): + pass + + _run_with_ccl_session(session_kind, devices, run_test) -def test_allreduce(): + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_allreduce(session_kind): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - - array_1 = np.arange(12, dtype="float32").reshape(3, 4) - array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4) - d_array = sess.empty((3, 4), "float32") - d_array.debug_copy_from(0, array_1) - d_array.debug_copy_from(1, array_2) - for op, np_op in [ # pylint: disable=invalid-name - ("sum", np.add), - ("prod", np.multiply), - ("min", np.minimum), - ("max", np.maximum), - ("avg", lambda a, b: (a + b) * 0.5), - ]: - dst_array = sess.empty((3, 4), "float32") - sess.allreduce(d_array, dst_array, op=op) - result = dst_array.debug_get_from_remote(0).numpy() - expected = np_op(array_1, array_2) - np.testing.assert_equal(result, expected) - - -def test_allgather(): + + def run_test(sess): + array_1 = np.arange(12, dtype="float32").reshape(3, 4) + array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4) + d_array = sess.empty((3, 4), "float32") + d_array.debug_copy_from(0, array_1) + d_array.debug_copy_from(1, array_2) + for op, np_op in [ # pylint: disable=invalid-name + ("sum", np.add), + ("prod", np.multiply), + ("min", np.minimum), + ("max", np.maximum), + ("avg", lambda a, b: (a + b) * 0.5), + ]: + dst_array = sess.empty((3, 4), "float32") + sess.allreduce(d_array, dst_array, op=op) + result = dst_array.debug_get_from_remote(0).numpy() + expected = np_op(array_1, array_2) + np.testing.assert_equal(result, expected) + + _run_with_ccl_session(session_kind, devices, run_test) + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_allgather(session_kind): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - array = np.arange(36, dtype="float32") - d_src = sess.empty((3, 3, 2), "float32") - d_dst = sess.empty((3, 4, 3), "float32") - d_src.debug_copy_from(0, array[:18]) - d_src.debug_copy_from(1, array[18:]) - sess.allgather(d_src, d_dst) - np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) - np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array.reshape(3, 4, 3)) + def run_test(sess): + array = np.arange(36, dtype="float32") + d_src = sess.empty((3, 3, 2), "float32") + d_dst = sess.empty((3, 4, 3), "float32") + d_src.debug_copy_from(0, array[:18]) + d_src.debug_copy_from(1, array[18:]) + sess.allgather(d_src, d_dst) + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) + np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array.reshape(3, 4, 3)) + + _run_with_ccl_session(session_kind, devices, run_test) +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) @pytest.mark.parametrize("use_explicit_output", [True, False]) -def test_broadcast(use_explicit_output): +def test_broadcast(session_kind, use_explicit_output): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - array = np.arange(12, dtype="float32").reshape(3, 4) - if use_explicit_output: - src_array = sess.empty((3, 4), "float32", worker0_only=True) - src_array.debug_copy_from(0, array) - dst_array = sess.empty((3, 4), "float32") - sess.broadcast_from_worker0(src_array, dst_array) - else: - dst_array = sess.broadcast(array) + def run_test(sess): + array = np.arange(12, dtype="float32").reshape(3, 4) + if use_explicit_output: + src_array = sess.empty((3, 4), "float32", worker0_only=True) + src_array.debug_copy_from(0, array) + dst_array = sess.empty((3, 4), "float32") + sess.broadcast_from_worker0(src_array, dst_array) + else: + dst_array = sess.broadcast(array) - result = dst_array.debug_get_from_remote(1).numpy() - np.testing.assert_equal(result, array) + result = dst_array.debug_get_from_remote(1).numpy() + np.testing.assert_equal(result, array) + _run_with_ccl_session(session_kind, devices, run_test) + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) @pytest.mark.parametrize("use_explicit_output", [True, False]) -def test_scatter(use_explicit_output): +def test_scatter(session_kind, use_explicit_output): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - array = np.arange(36, dtype="float32").reshape(2, 6, 3) - if use_explicit_output: - d_src = sess.empty((2, 6, 3), "float32", worker0_only=True) - d_dst = sess.empty((6, 3), "float32") - d_src.debug_copy_from(0, array) - sess.scatter_from_worker0(d_src, d_dst) - else: - d_dst = sess.scatter(array) + def run_test(sess): + array = np.arange(36, dtype="float32").reshape(2, 6, 3) + if use_explicit_output: + d_src = sess.empty((2, 6, 3), "float32", worker0_only=True) + d_dst = sess.empty((6, 3), "float32") + d_src.debug_copy_from(0, array) + sess.scatter_from_worker0(d_src, d_dst) + else: + d_dst = sess.scatter(array) - np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array[0, :, :]) - np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array[1, :, :]) + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array[0, :, :]) + np.testing.assert_equal(d_dst.debug_get_from_remote(1).numpy(), array[1, :, :]) + _run_with_ccl_session(session_kind, devices, run_test) -def test_scatter_with_implicit_reshape(): - devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - array = np.arange(36, dtype="float32").reshape(3, 4, 3) - d_src = sess.empty((3, 4, 3), "float32", worker0_only=True) - d_dst = sess.empty((3, 3, 2), "float32") - d_src.debug_copy_from(0, array) - sess.scatter_from_worker0(d_src, d_dst) +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_scatter_with_implicit_reshape(session_kind): + devices = [0, 1] - np.testing.assert_equal( - d_dst.debug_get_from_remote(0).numpy(), array.flat[:18].reshape(3, 3, 2) - ) - np.testing.assert_equal( - d_dst.debug_get_from_remote(1).numpy(), array.flat[18:].reshape(3, 3, 2) - ) + def run_test(sess): + array = np.arange(36, dtype="float32").reshape(3, 4, 3) + d_src = sess.empty((3, 4, 3), "float32", worker0_only=True) + d_dst = sess.empty((3, 3, 2), "float32") + d_src.debug_copy_from(0, array) + sess.scatter_from_worker0(d_src, d_dst) + np.testing.assert_equal( + d_dst.debug_get_from_remote(0).numpy(), array.flat[:18].reshape(3, 3, 2) + ) + np.testing.assert_equal( + d_dst.debug_get_from_remote(1).numpy(), array.flat[18:].reshape(3, 3, 2) + ) + _run_with_ccl_session(session_kind, devices, run_test) -def test_gather(): +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_gather(session_kind): devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) - - array = np.arange(36, dtype="float32") - d_src = sess.empty((3, 3, 2), "float32") - d_dst = sess.empty((3, 4, 3), "float32", worker0_only=True) - d_src.debug_copy_from(0, array[:18]) - d_src.debug_copy_from(1, array[18:]) - sess.gather_to_worker0(d_src, d_dst) - np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) + def run_test(sess): + array = np.arange(36, dtype="float32") + d_src = sess.empty((3, 3, 2), "float32") + d_dst = sess.empty((3, 4, 3), "float32", worker0_only=True) + d_src.debug_copy_from(0, array[:18]) + d_src.debug_copy_from(1, array[18:]) + sess.gather_to_worker0(d_src, d_dst) + np.testing.assert_equal(d_dst.debug_get_from_remote(0).numpy(), array.reshape(3, 4, 3)) + _run_with_ccl_session(session_kind, devices, run_test) def relax_build(mod): @@ -273,10 +281,9 @@ def relax_build(mod): return tvm.compile(mod, target=target) -def test_mlp(): # pylint: disable=too-many-locals +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_mlp(session_kind): # pylint: disable=too-many-locals devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) # pylint: disable=invalid-name @tvm.script.ir_module @@ -323,34 +330,36 @@ def main( tvm.runtime.tensor(W2, device=dev), ).numpy() - with tempfile.TemporaryDirectory() as tmpdir: - path = tmpdir + "/test.so" - relax_build(ShardedMLP).export_library(path) - - mod = sess.load_vm_module(path) - - d_X = sess.empty((128, 128), "float32") - d_W1 = sess.empty((128, 64), "float32") - d_W2 = sess.empty((64, 128), "float32") - - d_X.debug_copy_from(0, X) - d_W1.debug_copy_from(0, W1[:, :64]) - d_W1.debug_copy_from(1, W1[:, 64:]) - d_W2.debug_copy_from(0, W2[:64, :]) - d_W2.debug_copy_from(1, W2[64:, :]) - d_Y = mod["main"](d_X, d_W1, d_W2) - Y_result = tvm.runtime.empty((128, 128), "float32", device=dev) - sess.copy_from_worker_0(Y_result, d_Y) - sess.sync_worker_0() - Y_result = Y_result.numpy() + def run_test(sess): + with tempfile.TemporaryDirectory() as tmpdir: + path = tmpdir + "/test.so" + relax_build(ShardedMLP).export_library(path) + + mod = sess.load_vm_module(path) + + d_X = sess.empty((128, 128), "float32") + d_W1 = sess.empty((128, 64), "float32") + d_W2 = sess.empty((64, 128), "float32") + + d_X.debug_copy_from(0, X) + d_W1.debug_copy_from(0, W1[:, :64]) + d_W1.debug_copy_from(1, W1[:, 64:]) + d_W2.debug_copy_from(0, W2[:64, :]) + d_W2.debug_copy_from(1, W2[64:, :]) + d_Y = mod["main"](d_X, d_W1, d_W2) + Y_result = tvm.runtime.empty((128, 128), "float32", device=dev) + sess.copy_from_worker_0(Y_result, d_Y) + sess.sync_worker_0() + return Y_result.numpy() + + Y_result = _run_with_ccl_session(session_kind, devices, run_test) # pylint: enable=invalid-name np.testing.assert_allclose(Y_result, Y_expected, rtol=1e-4, atol=1e-4) -def test_attention(): # pylint: disable=too-many-locals,too-many-statements +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_attention(session_kind): # pylint: disable=too-many-locals,too-many-statements devices = [0, 1] - sess = create_socket_session(num_workers=len(devices)) - sess.init_ccl(_CCL, *devices) # pylint: disable=invalid-name @tvm.script.ir_module @@ -437,32 +446,35 @@ def main( # pylint: disable=too-many-locals tvm.runtime.tensor(Wo, device=dev), ).numpy() - with tempfile.TemporaryDirectory() as tmpdir: - path = tmpdir + "/test.so" - relax_build(ShardedAttention).export_library(path) - - mod = sess.load_vm_module(path) - - d_X = sess.empty((1, 10, 128), "float32") - d_Wq = sess.empty((128, 256), "float32") - d_Wk = sess.empty((128, 256), "float32") - d_Wv = sess.empty((128, 256), "float32") - d_Wo = sess.empty((256, 128), "float32") - - d_X.debug_copy_from(0, X) - d_Wq.debug_copy_from(0, Wq[:, :256]) - d_Wq.debug_copy_from(1, Wq[:, 256:]) - d_Wk.debug_copy_from(0, Wk[:, :256]) - d_Wk.debug_copy_from(1, Wk[:, 256:]) - d_Wv.debug_copy_from(0, Wv[:, :256]) - d_Wv.debug_copy_from(1, Wv[:, 256:]) - d_Wo.debug_copy_from(0, Wo[:256, :]) - d_Wo.debug_copy_from(1, Wo[256:, :]) - d_Y = mod["main"](d_X, d_Wq, d_Wk, d_Wv, d_Wo) - Y_result = tvm.runtime.empty((1, 10, 128), "float32", device=dev) - sess.copy_from_worker_0(Y_result, d_Y) - sess.sync_worker_0() - Y_result = Y_result.numpy() + def run_test(sess): + with tempfile.TemporaryDirectory() as tmpdir: + path = tmpdir + "/test.so" + relax_build(ShardedAttention).export_library(path) + + mod = sess.load_vm_module(path) + + d_X = sess.empty((1, 10, 128), "float32") + d_Wq = sess.empty((128, 256), "float32") + d_Wk = sess.empty((128, 256), "float32") + d_Wv = sess.empty((128, 256), "float32") + d_Wo = sess.empty((256, 128), "float32") + + d_X.debug_copy_from(0, X) + d_Wq.debug_copy_from(0, Wq[:, :256]) + d_Wq.debug_copy_from(1, Wq[:, 256:]) + d_Wk.debug_copy_from(0, Wk[:, :256]) + d_Wk.debug_copy_from(1, Wk[:, 256:]) + d_Wv.debug_copy_from(0, Wv[:, :256]) + d_Wv.debug_copy_from(1, Wv[:, 256:]) + d_Wo.debug_copy_from(0, Wo[:256, :]) + d_Wo.debug_copy_from(1, Wo[256:, :]) + d_Y = mod["main"](d_X, d_Wq, d_Wk, d_Wv, d_Wo) + Y_result = tvm.runtime.empty((1, 10, 128), "float32", device=dev) + sess.copy_from_worker_0(Y_result, d_Y) + sess.sync_worker_0() + return Y_result.numpy() + + Y_result = _run_with_ccl_session(session_kind, devices, run_test) # pylint: enable=invalid-name np.testing.assert_allclose(Y_result, Y_expected, rtol=1e-3, atol=1e-3) diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py index f7916ef0db0a..f82fd5081f01 100644 --- a/tests/python/disco/test_session.py +++ b/tests/python/disco/test_session.py @@ -126,7 +126,6 @@ def __del__(self): node.wait() - def create_socket_session(num_workers, build_ring=False): """Create a socket session backed by one local and one remote node. diff --git a/tests/python/disco/test_socket_upload.py b/tests/python/disco/test_socket_upload.py index 245c035b9889..2b4a03396d8d 100644 --- a/tests/python/disco/test_socket_upload.py +++ b/tests/python/disco/test_socket_upload.py @@ -16,6 +16,8 @@ # under the License. # pylint: disable=missing-docstring +"""Test upload_vm_module over a multi-node SocketSession.""" + import hashlib import pathlib import socket @@ -37,7 +39,7 @@ pytest.skip("disco runtime is not available", allow_module_level=True) _NUM_NODES = 4 -_NUM_WORKERS = 4 +_NUM_WORKERS = 4 _REL_PATH = "./mod.so" @@ -70,10 +72,7 @@ def double(x: R.Tensor((4,), "float32")) -> R.Tensor((4,), "float32"): class SocketSessionTester: - def __init__(self, num_workers, node_cwds, num_groups=1, build_ring=True): - # Initialize the attributes used by __del__ first, so that teardown is - # safe even when __init__ raises below. self.sess = None self.remote_nodes = [] num_nodes = len(node_cwds) @@ -133,29 +132,33 @@ def __del__(self): def test_upload_vm_module(monkeypatch): - with tempfile.TemporaryDirectory(prefix="disco_upload_") as tmp_root: root = pathlib.Path(tmp_root) node_dirs = [root / _node_dir_name(i) for i in range(_NUM_NODES)] + for node_dir in node_dirs: node_dir.mkdir() + # Please add `-s` to print the directory for verification. print(f"\n[setup] created {_NUM_NODES} node directories under {root}") + for node_dir in node_dirs: print(f" {node_dir}") - # The controller is this process, so give it upload_node0 too -- otherwise the relative - # path would resolve against the pytest rootdir and leak the artifact out of the sandbox. monkeypatch.chdir(node_dirs[0]) target = tvm.target.Target("llvm") tvm.compile(rx.get_pipeline("zero")(Mod), target=target).export_library(_REL_PATH) expected_sha = _sha256(_REL_PATH) built = node_dirs[0] / _REL_PATH - print(f"[setup] compiled {built} ({built.stat().st_size} bytes, sha256={expected_sha[:16]})") + + print( + f"[setup] compiled {built} ({built.stat().st_size} bytes, sha256={expected_sha[:16]})" + ) # Precondition: only the controller has the artifact. for node_dir in node_dirs[1:]: assert not (node_dir / _REL_PATH).exists() + print(f"[setup] upload_node1..{_NUM_NODES - 1} are empty, as expected") tester = SocketSessionTester(_NUM_WORKERS, node_dirs) @@ -166,6 +169,7 @@ def test_upload_vm_module(monkeypatch): sess._sync_all() # pylint: disable=protected-access print("[upload] per-node state after upload_vm_module:") + for i, node_dir in enumerate(node_dirs): written = node_dir / _REL_PATH assert written.is_file(), f"{_node_dir_name(i)} never received the module" @@ -181,8 +185,10 @@ def test_upload_vm_module(monkeypatch): print(f"[load] load_vm_module({_REL_PATH!r}) succeeded on all {_NUM_NODES} nodes") d_x = sess.empty((4,), "float32") + for i in range(_NUM_WORKERS): d_x.debug_copy_from(i, np.full((4,), i + 1, dtype="float32")) + d_y = mod["double"](d_x) for i in range(_NUM_WORKERS): got = d_y.debug_get_from_remote(i).numpy()