wehub-resource-sync
101 lines
3.7 KiB
C++
101 lines
3.7 KiB
C++
/*!
|
|
* \file builtin.cc
|
|
* \brief Multi-GPU builtin functions in MLC LLM.
|
|
*/
|
|
#ifndef MLC_SINGLE_GPU_ONLY
|
|
|
|
#include <tvm/ffi/container/array.h>
|
|
#include <tvm/ffi/container/shape.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/optional.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/runtime/disco/builtin.h>
|
|
#include <tvm/runtime/disco/disco_worker.h>
|
|
#include <tvm/runtime/tensor.h>
|
|
#include <tvm/runtime/vm/vm.h>
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace multi_gpu {
|
|
|
|
using namespace tvm::runtime;
|
|
using tvm::ffi::Array;
|
|
using tvm::ffi::ObjectRef;
|
|
using tvm::ffi::Optional;
|
|
using tvm::ffi::Shape;
|
|
|
|
ObjectRef DispatchFunctionByGroup(tvm::ffi::AnyView vm_arg,
|
|
Array<Array<ObjectRef>> funcs_and_args) {
|
|
using namespace vm;
|
|
VirtualMachine* vm = VirtualMachine::GetContextPtr(vm_arg);
|
|
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
|
int world_size = worker->num_workers;
|
|
int group_size = worker->num_workers / worker->num_groups;
|
|
int num_group = world_size / group_size;
|
|
TVM_FFI_ICHECK_EQ(funcs_and_args.size(), num_group)
|
|
<< "Number of groups mismatches. There are " << num_group
|
|
<< " groups while the function/arg array has " << funcs_and_args.size() << " elements.";
|
|
|
|
int group_id = worker->worker_id / group_size;
|
|
TVM_FFI_ICHECK(!funcs_and_args[group_id].empty())
|
|
<< "No function is provided for group " << group_id;
|
|
VMClosure func = funcs_and_args[group_id][0].as_or_throw<VMClosure>();
|
|
|
|
int num_args = static_cast<int>(funcs_and_args[group_id].size()) - 1;
|
|
std::vector<tvm::ffi::AnyView> packed_args(num_args);
|
|
for (int i = 0; i < num_args; ++i) {
|
|
// NOTE: Need explicily define `arg` so that the argument does not
|
|
// have type code kTVMObjectRValueRefArg.
|
|
packed_args[i] = funcs_and_args[group_id][1 + i];
|
|
}
|
|
|
|
tvm::ffi::Any rv;
|
|
vm->InvokeClosurePacked(funcs_and_args[group_id][0].as_or_throw<VMClosure>(),
|
|
tvm::ffi::PackedArgs(packed_args.data(), packed_args.size()), &rv);
|
|
return rv.cast<ObjectRef>();
|
|
}
|
|
|
|
ObjectRef SendFromLastGroupToWorker0(Tensor send, Optional<Tensor> recv, Shape shape,
|
|
DLDataType dtype) {
|
|
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
|
int worker_id = worker->worker_id;
|
|
int world_size = worker->num_workers;
|
|
int group_size = worker->num_workers / worker->num_groups;
|
|
TVM_FFI_ICHECK_NE(world_size, group_size) << "Cannot perform when there is only one group.";
|
|
int sender_id = world_size - group_size;
|
|
if (worker_id == 0) {
|
|
TVM_FFI_ICHECK(recv.has_value()) << "The receive Tensor is undefined for worker 0.";
|
|
Tensor recv_arr = recv.value().CreateView(shape, dtype);
|
|
RecvFromWorker(recv_arr, sender_id);
|
|
return recv_arr;
|
|
} else if (worker_id == sender_id) {
|
|
TVM_FFI_ICHECK_EQ(send->dtype, dtype)
|
|
<< "The src Tensor has mismatched dtype than the expected dtype.";
|
|
TVM_FFI_ICHECK_EQ(send->ndim, shape.size())
|
|
<< "The src Tensor has mismatched shape than the expected shape.";
|
|
for (int i = 0; i < send->ndim; ++i) {
|
|
TVM_FFI_ICHECK_EQ(send->shape[i], shape[i])
|
|
<< "The src Tensor has mismatched shape than the expected shape.";
|
|
}
|
|
SendToWorker(send, /*receiver_id=*/0);
|
|
return recv.value_or(Tensor(nullptr));
|
|
}
|
|
|
|
// We only process for worker 0 and the first worker of the last group.
|
|
// For other workers, we return the input object.
|
|
return recv.value_or(Tensor(nullptr));
|
|
}
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef()
|
|
.def("mlc.multi_gpu.DispatchFunctionByGroup", DispatchFunctionByGroup)
|
|
.def("mlc.multi_gpu.SendFromLastGroupToWorker0", SendFromLastGroupToWorker0);
|
|
}
|
|
|
|
} // namespace multi_gpu
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_SINGLE_GPU_ONLY
|