diff --git a/CMakeLists.txt b/CMakeLists.txt index 567edc1dc698..dd85ebf34faf 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,15 @@ if(USE_ROCM AND USE_RCCL) target_link_libraries(tvm_runtime_extra PRIVATE tvm_rccl_objs rccl) 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) + 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/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) 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/include/tvm/runtime/disco/disco_worker.h b/include/tvm/runtime/disco/disco_worker.h index e5c88e0d215b..9d7a145476c0 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..f851bcba0fd8 100644 --- a/include/tvm/runtime/disco/session.h +++ b/include/tvm/runtime/disco/session.h @@ -76,6 +76,7 @@ #include #include #include +#include #include #include @@ -294,8 +295,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 +308,14 @@ 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 +337,72 @@ 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..2fa0cbfd968a 100644 --- a/python/tvm/exec/disco_worker.py +++ b/python/tvm/exec/disco_worker.py @@ -101,12 +101,21 @@ 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 +124,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/relax/op/ccl/__init__.py b/python/tvm/relax/op/ccl/__init__.py index 47006358f93f..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 +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 e49614f4ded1..d47d640bcf7f 100644 --- a/python/tvm/relax/op/ccl/ccl.py +++ b/python/tvm/relax/op/ccl/ccl.py @@ -106,3 +106,25 @@ 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: + """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) diff --git a/python/tvm/relax/transform/legalize_ops/ccl.py b/python/tvm/relax/transform/legalize_ops/ccl.py index 659b1f7d4397..fcfbd2251472 100644 --- a/python/tvm/relax/transform/legalize_ops/ccl.py +++ b/python/tvm/relax/transform/legalize_ops/ccl.py @@ -123,3 +123,26 @@ 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 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/python/tvm/runtime/disco/process_pool.py b/python/tvm/runtime/disco/process_pool.py index 93444f94c2aa..9fc136c37106 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,44 @@ 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..ad9626f0c47b 100644 --- a/python/tvm/runtime/disco/session.py +++ b/python/tvm/runtime/disco/session.py @@ -322,6 +322,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/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 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..785d541771d1 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..4573bd59e7ae --- /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(); + + 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; + 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(); + + 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); + + 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; + } + + 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..b5e5c0597555 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,130 @@ 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 +189,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,9 +212,40 @@ 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 } int64_t GetNumWorkers() final { return num_nodes_ * num_workers_per_node_; } @@ -181,12 +333,23 @@ 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 +370,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 +386,27 @@ 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,21 +420,48 @@ 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); + local_session_->BroadcastPacked(wrapped); } else { - local_session_->SendPacked(local_worker_id, args); + local_session_->SendPacked(local_worker_id, wrapped); } break; } 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 +471,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,12 +487,26 @@ 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"); 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_; @@ -288,12 +516,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 +541,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) { - auto n = - ffi::make_object(num_nodes, num_workers_per_node, num_groups, host, 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, build_ring); return Session(n); } @@ -315,13 +554,14 @@ 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..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) + 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 +80,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() { @@ -166,30 +181,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) { +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); 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..3a9d2cdd0ef1 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,36 @@ 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..4e0b16d7edc1 100644 --- a/src/support/socket.h +++ b/src/support/socket.h @@ -293,6 +293,19 @@ 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 */ diff --git a/tests/python/disco/test_cpuccl.py b/tests/python/disco/test_cpuccl.py new file mode 100644 index 000000000000..ba94a5bf2827 --- /dev/null +++ b/tests/python/disco/test_cpuccl.py @@ -0,0 +1,483 @@ +# 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 ThreadedSession, ProcessSession and SocketSession. +Reuse the test cases from test_ccl.py""" + +import socket +import subprocess +import sys +import tempfile +import threading + +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.runtime.vm import VirtualMachine +from tvm.script import relax as R + +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.""" + + def __init__(self, num_workers, num_nodes=2, num_groups=1, build_ring=True): + 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: + # Shut down the session first so remote nodes can exit gracefully. + 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): + 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 + + +_all_session_kinds = [di.ThreadedSession, di.ProcessSession, create_socket_session] + + +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] + + def run_test(_sess): + pass + + _run_with_ccl_session(session_kind, devices, run_test) + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_allreduce(session_kind): + devices = [0, 1] + + 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] + + 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(session_kind, use_explicit_output): + devices = [0, 1] + + 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) + + _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(session_kind, use_explicit_output): + devices = [0, 1] + + 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, :, :]) + + _run_with_ccl_session(session_kind, devices, run_test) + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_scatter_with_implicit_reshape(session_kind): + devices = [0, 1] + + 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) + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_gather(session_kind): + devices = [0, 1] + + 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): + 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) + + +@pytest.mark.parametrize("session_kind", _all_session_kinds, ids=_session_id) +def test_mlp(session_kind): # pylint: disable=too-many-locals + devices = [0, 1] + + # 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() + + 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) + + +@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] + + # 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() + + 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) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py index 2e84170940e1..f82fd5081f01 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,12 @@ 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 +126,7 @@ 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 +135,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 +158,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 +169,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 +181,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 +195,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 +207,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 +221,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 +281,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 +351,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 diff --git a/tests/python/disco/test_socket_upload.py b/tests/python/disco/test_socket_upload.py new file mode 100644 index 000000000000..2b4a03396d8d --- /dev/null +++ b/tests/python/disco/test_socket_upload.py @@ -0,0 +1,204 @@ +# 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 + +"""Test upload_vm_module over a multi-node SocketSession.""" + +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): + 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() + # 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}") + + 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() 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()