Laiyyi opened a new pull request, #20077:
URL: https://github.com/apache/tvm/pull/20077
## Summary
Currently, TVM Disco supports multi-node distributed execution through NCCL
and RCCL. However, these backends require all participating devices to have
GPUs. In some environments where only CPU resources are available, existing
collective communication backends cannot be directly used.
Therefore, this PR introduces CPUCCL, enabling TVM Disco to perform
multi-node distributed execution in CPU-only environments.
- Added `cpuccl.cc` to implement the CPU-based collective communication
backend.
- Added `upload_vm_module` to support VM module transfer across nodes.
- In real distributed environments, only the host node usually contains
the compiled `.so` module.
- The host node can upload the module to remote nodes through
`upload_vm_module`.
- Remote nodes can then load the module through `load_vm_module` and
participate in distributed execution.
- Implemented `allreduce` using a ring-based algorithm. These functions
interface follows the implementation style of `nccl.cc`.
- Added `build_ring` option to `ThreadedSession`,
`ProcessSession`,`SocketSession` and `process_pool.py` to enable optional ring
construction for CPUCCL communication.
- Updated `disco_worker.py` with corresponding changes to support CPUCCL
communication.
- Added `gather_to_worker0` support in Relax.
## SocketSession Cross-node Communication
In `SocketSession`, CPUCCL cross-node communication is implemented using
local rings combined with TCP forwarding.
Assume two nodes, A and B, each containing two workers. Each node first
creates its own local ring:
Node A:
W0 ─A0→ W1 ─A1→ W0
Node B:
W0 ─B0→ W1 ─B1→ W0
Each session additionally creates two threads responsible for cross-node TCP
communication.
When Node A's W1 needs to send data to Node B's W0, it only writes data into
the local ring channel `A1`. The send thread monitors `A1` and forwards the
data to Node B through TCP.
On Node B, `SocketSession` creates a proxy channel. The recv thread receives
data from Node A and writes it into `proxy_out_to_tcp_`. Since Node B W0's
`ring_in` is rerouted through `RerouteRingIn`, W0 on node B can correctly
receive data from Node A W1.
## Bug Fix
- Fixed the garbage collection issue in `test_session.py` caused by
duplicated shutdown messages during session destruction. Added shutdown state
checking in `bcast_session.cc` to prevent repeated shutdown handling.
## Build Configuration
- Added `USE_CPU_CCL` option to allow users to choose whether to compile
`cpuccl.cc`.
- Updated related `CMakeLists.txt` files to include CPUCCL in the build
process.
## Tests
- Added `test_cpuccl.py`, reusing NCCL test cases to validate CPUCCL
collective communication.
- Added build-ring and non-build-ring test cases in `test_session.py`.
- Added `test_socket_upload` to validate the `upload_vm_module`
functionality.
- Added tests in `test_op_ccl.py` and `test_transform_legalize_ops_ccl.py`
to verify Relax `gather_to_worker0` functionality.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
