This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2f889774ec [3rdparty] AUTO mode for custom all-reduce strategy (#16797)
2f889774ec is described below
commit 2f889774ec10b56ebfac89f78698e06eb200db46
Author: Ruihang Lai <[email protected]>
AuthorDate: Wed Mar 27 01:30:09 2024 -0400
[3rdparty] AUTO mode for custom all-reduce strategy (#16797)
This PR adds the automatic mode selection for customized all-reduce
kernels, referring TensorRT-LLM.
Meanwhile, this PR fixes a bug that may cause customized all-reduce
kernel to hang forever. Prior to this PR, each worker resets its
barrier values to 0 *after using all-gather to exchange their
barrier handles*. Afterwards, the customized all-reduce kernels
update the barriers of all workers. So it is possible that, worker 0
updates worker 1's barrier *before* worker 1 resets its barrier to 0.
This lead to the all-reduce kernel hanging forever.
This PR changes the behavior to resetting barriers before all-gather,
and forcing a device synchronization after reset.
---
3rdparty/tensorrt_llm/custom_allreduce_kernels.h | 33 ++++++++++++++++++++++
.../tvm/relax/transform/ipc_allreduce_rewrite.py | 2 --
src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc | 26 +++++++++++------
src/runtime/disco/cuda_ipc/custom_allreduce.cc | 12 ++++++--
tests/python/disco/test_custom_allreduce.py | 4 +++
5 files changed, 63 insertions(+), 14 deletions(-)
diff --git a/3rdparty/tensorrt_llm/custom_allreduce_kernels.h
b/3rdparty/tensorrt_llm/custom_allreduce_kernels.h
index 7fd66e5d10..7c515a03ac 100644
--- a/3rdparty/tensorrt_llm/custom_allreduce_kernels.h
+++ b/3rdparty/tensorrt_llm/custom_allreduce_kernels.h
@@ -25,8 +25,10 @@ constexpr size_t MAX_RANKS_PER_NODE = 8;
constexpr size_t DEFAULT_BLOCK_SIZE = 1024;
enum class AllReduceStrategyType : int8_t {
+ RING = 0,
ONESHOT = 1,
TWOSHOT = 2,
+ AUTO = 3,
};
struct AllReduceParams {
@@ -42,6 +44,37 @@ struct AllReduceParams {
void* local_output_buffer_ptr;
};
+inline size_t GetMaxRequiredWorkspaceSize(int world_size) {
+ if (world_size <= 2) {
+ return 16 * 1000 * 1000;
+ }
+ return 8 * 1000 * 1000;
+}
+
+inline AllReduceStrategyType SelectImplementation(size_t message_size, int
world_size) {
+ const size_t maxWorkspaceSize = GetMaxRequiredWorkspaceSize(world_size);
+
+ if (message_size > maxWorkspaceSize) {
+ return AllReduceStrategyType::RING;
+ }
+
+ if (world_size <= 2) {
+ return AllReduceStrategyType::ONESHOT;
+ }
+
+ if (world_size <= 4) {
+ if (message_size < 1 * 1000 * 1000) {
+ return AllReduceStrategyType::ONESHOT;
+ }
+ return AllReduceStrategyType::TWOSHOT;
+ }
+
+ if (message_size < 500 * 1000) {
+ return AllReduceStrategyType::ONESHOT;
+ }
+ return AllReduceStrategyType::TWOSHOT;
+}
+
void customAllReduce(AllReduceParams& params, void* data, size_t elts,
DLDataType dataType,
AllReduceStrategyType strat, cudaStream_t stream);
diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py
b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
index 3e7b005a60..df40181cb9 100644
--- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py
+++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
@@ -40,8 +40,6 @@ class IPCAllReduceRewrite:
The all-reduce strategy. Only "1" and "2" are supported.
"1" stands for one-shot, and "2" stands for two-shot.
"""
- if allreduce_strategy not in [1, 2]:
- raise ValueError(f"All-reduce strategy {allreduce_strategy} is not
supported.")
self.allreduce_strategy = allreduce_strategy
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext)
-> IRModule:
diff --git a/src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc
b/src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc
index 451c3df0cb..fec5abec86 100644
--- a/src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc
+++ b/src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc
@@ -91,15 +91,13 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
private:
void* DeviceAllocDataSpace(Device dev, size_t size, size_t alignment,
DLDataType type_hint) final {
- auto [data_ptr, data_comm_ptrs] = AllocIPCMemory(dev, size, alignment,
type_hint);
+ auto [data_ptr, data_comm_ptrs] =
+ AllocIPCMemory(dev, size, alignment, type_hint,
/*reset_memory_to_zero=*/false);
int barrier_ptr_size = sizeof(uint32_t) * (MAX_ALL_REDUCE_BLOCKS + 2) *
MAX_RANKS_PER_NODE;
- auto [barrier_in_ptr, barrier_in_comm_ptrs] =
- AllocIPCMemory(dev, barrier_ptr_size, alignment, DataType::UInt(32));
- auto [barrier_out_ptr, barrier_out_comm_ptrs] =
- AllocIPCMemory(dev, barrier_ptr_size, alignment, DataType::UInt(32));
- // Initialize the barrier values to 0 to avoid synchronization issue.
- CUDA_CALL(cudaMemset(barrier_in_ptr, 0, barrier_ptr_size));
- CUDA_CALL(cudaMemset(barrier_out_ptr, 0, barrier_ptr_size));
+ auto [barrier_in_ptr, barrier_in_comm_ptrs] = AllocIPCMemory(
+ dev, barrier_ptr_size, alignment, DataType::UInt(32),
/*reset_memory_to_zero=*/true);
+ auto [barrier_out_ptr, barrier_out_comm_ptrs] = AllocIPCMemory(
+ dev, barrier_ptr_size, alignment, DataType::UInt(32),
/*reset_memory_to_zero=*/true);
// Create the CUDAIPCMemory object.
ObjectPtr<CUDAIPCMemoryObj> ipc_memory = make_object<CUDAIPCMemoryObj>();
@@ -142,12 +140,22 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
* pointer.
*/
std::pair<void*, std::vector<void*>> AllocIPCMemory(Device dev, size_t size,
size_t alignment,
- DLDataType type_hint) {
+ DLDataType type_hint,
+ bool
reset_memory_to_zero) {
// Alloc local buffer
ICHECK(dev.device_type == kDLCUDA);
void* ptr;
CUDA_CALL(cudaSetDevice(dev.device_id));
CUDA_CALL(cudaMalloc(&ptr, size));
+ // Reset allocated memory to zero when required.
+ // We explicitly synchronize after memset, to make sure memset finishes
+ // before using all-gather to exchange IPC handles.
+ // This is important to ensure the memory reset get ordered
+ // before any other peers read the memory.
+ if (reset_memory_to_zero) {
+ CUDA_CALL(cudaMemset(ptr, 0, size));
+ CUDA_CALL(cudaDeviceSynchronize());
+ }
// Create ipc handle
cudaIpcMemHandle_t local_handle;
CUDA_CALL(cudaIpcGetMemHandle(&local_handle, ptr));
diff --git a/src/runtime/disco/cuda_ipc/custom_allreduce.cc
b/src/runtime/disco/cuda_ipc/custom_allreduce.cc
index e9be5973e1..98fd777b83 100644
--- a/src/runtime/disco/cuda_ipc/custom_allreduce.cc
+++ b/src/runtime/disco/cuda_ipc/custom_allreduce.cc
@@ -66,7 +66,15 @@ void CustomAllReduce(DLTensor* send, int strategy, DLTensor*
recv) {
int64_t num_elements = TensorSize(send);
nccl::CCLThreadLocalContext* ctx = nccl::CCLThreadLocalContext::Get();
- if (!CanApplyCustomAllReduce(num_elements, send->dtype)) {
+ tensorrt_llm::AllReduceStrategyType strategy_ =
+ static_cast<tensorrt_llm::AllReduceStrategyType>(strategy);
+ if (strategy_ == tensorrt_llm::AllReduceStrategyType::AUTO) {
+ strategy_ = tensorrt_llm::SelectImplementation(
+ num_elements * ((send->dtype.bits * send->dtype.lanes + 7) / 8),
ctx->worker->num_workers);
+ }
+
+ if (strategy_ == tensorrt_llm::AllReduceStrategyType::RING ||
+ !CanApplyCustomAllReduce(num_elements, send->dtype)) {
// Dispatch to nccl AllReduce if the customized all-reduce cannot apply.
deviceStream_t stream = ctx->GetDefaultStream();
NCCL_CALL(ncclAllReduce(send->data, recv->data, num_elements,
@@ -92,8 +100,6 @@ void CustomAllReduce(DLTensor* send, int strategy, DLTensor*
recv) {
params.peer_barrier_ptrs_out[i] =
reinterpret_cast<uint32_t*>(ipc_memory->barrier_out[i]);
}
- tensorrt_llm::AllReduceStrategyType strategy_ =
- static_cast<tensorrt_llm::AllReduceStrategyType>(strategy);
if (!CanApplyTwoShotAllReduce(num_elements, send->dtype,
ctx->worker->num_workers)) {
// Two-shot all-reduce does not support this case.
// So we fallback to the one-shot strategy.
diff --git a/tests/python/disco/test_custom_allreduce.py
b/tests/python/disco/test_custom_allreduce.py
index 47b5f9590a..4aed32c052 100644
--- a/tests/python/disco/test_custom_allreduce.py
+++ b/tests/python/disco/test_custom_allreduce.py
@@ -29,15 +29,19 @@ from tvm.runtime.disco import Session
class AllReduceStrategyType(enum.IntEnum):
+ RING = 0
ONESHOT = 1
TWOSHOT = 2
+ AUTO = 3
_shapes = [(2, 3), (3, 4), (128, 128)]
_strategies = [
+ AllReduceStrategyType.RING,
AllReduceStrategyType.ONESHOT,
AllReduceStrategyType.TWOSHOT,
+ AllReduceStrategyType.AUTO,
]
_ccl = [ccl for ccl in tvm.get_global_func("runtime.disco.compiled_ccl")() if
ccl == "nccl"]
