wehub-resource-sync
328 lines
14 KiB
C++
328 lines
14 KiB
C++
/*!
|
|
* \file multi_gpu_loader.cc
|
|
* \brief Implementation of a multi-GPU loader with loading-time sharding.
|
|
*/
|
|
#ifndef MLC_SINGLE_GPU_ONLY
|
|
#include <tvm/ffi/container/array.h>
|
|
#include <tvm/ffi/extra/json.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/optional.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/runtime/device_api.h>
|
|
#include <tvm/runtime/disco/builtin.h>
|
|
#include <tvm/runtime/disco/disco_worker.h>
|
|
#include <tvm/runtime/vm/tensor_cache_support.h>
|
|
|
|
#include <chrono>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <functional>
|
|
#include <numeric>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include "../metadata/model.h"
|
|
#include "../support/progress_bar.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace multi_gpu {
|
|
|
|
using tvm::Device;
|
|
using tvm::runtime::vm::TensorCacheMetadata;
|
|
using namespace tvm::runtime;
|
|
using tvm::ffi::Array;
|
|
using tvm::ffi::Function;
|
|
using tvm::ffi::Object;
|
|
using tvm::ffi::Optional;
|
|
using tvm::ffi::Shape;
|
|
using tvm::ffi::TypedFunction;
|
|
using DurationType = std::chrono::microseconds;
|
|
|
|
class RangeTimer {
|
|
public:
|
|
explicit RangeTimer(DurationType* result)
|
|
: start(std::chrono::high_resolution_clock::now()), result(result) {}
|
|
|
|
~RangeTimer() {
|
|
std::chrono::time_point<std::chrono::high_resolution_clock> end =
|
|
std::chrono::high_resolution_clock::now(); //
|
|
auto duration = end - start;
|
|
(*result) += std::chrono::duration_cast<DurationType>(end - start);
|
|
}
|
|
|
|
private:
|
|
std::chrono::time_point<std::chrono::high_resolution_clock> start;
|
|
DurationType* result;
|
|
};
|
|
|
|
class PreprocessorPool {
|
|
public:
|
|
explicit PreprocessorPool(const ModelMetadata& model_metadata, Module vm_module) {
|
|
for (const ModelMetadata::Param& param : model_metadata.params) {
|
|
for (const ModelMetadata::Param::Preproc& preproc : param.preprocs) {
|
|
const std::string& func_name = preproc.func_name;
|
|
if (Function f = vm_module.defined()
|
|
? vm_module->GetFunction(func_name, true).value_or(Function(nullptr))
|
|
: nullptr;
|
|
f != nullptr) {
|
|
preproc_funcs[func_name] = f;
|
|
} else if (const auto f = Function::GetGlobal(func_name); f.has_value()) {
|
|
preproc_funcs[func_name] = *f;
|
|
} else {
|
|
LOG(FATAL) << "ValueError: Undefined function: " << func_name;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor Apply(Tensor param, const ModelMetadata::Param& param_info) const {
|
|
for (const ModelMetadata::Param::Preproc& preproc : param_info.preprocs) {
|
|
const std::string& func_name = preproc.func_name;
|
|
Tensor param_in = param;
|
|
param = Tensor::Empty(preproc.out_shape, preproc.out_dtype, param->device);
|
|
TVM_FFI_ICHECK(preproc_funcs.count(func_name));
|
|
DLTensor dl_param_in = *param_in.operator->();
|
|
DLTensor dl_param = *param.operator->();
|
|
preproc_funcs.at(func_name)(&dl_param_in, &dl_param);
|
|
}
|
|
return param;
|
|
}
|
|
|
|
private:
|
|
std::unordered_map<std::string, TypedFunction<void(DLTensor*, DLTensor*)>> preproc_funcs;
|
|
};
|
|
|
|
struct ParamInfo {
|
|
const TensorCacheMetadata::FileRecord* file;
|
|
const TensorCacheMetadata::FileRecord::ParamRecord* param;
|
|
};
|
|
|
|
Tensor RecvFromGlobalWorker0(Device device, const ModelMetadata::Param& param_info) {
|
|
Shape shape = param_info.preprocs.empty() ? param_info.shape : param_info.preprocs[0].in_shape;
|
|
Tensor result = Tensor::Empty(shape, param_info.dtype, device);
|
|
RecvFromWorker0(result);
|
|
return result;
|
|
}
|
|
|
|
Tensor BroadcastOrShardAndScatter(Tensor param, const ModelMetadata::Param& param_info,
|
|
int num_shards, const PreprocessorPool& preprocs) {
|
|
bool needs_sharding = !param_info.preprocs.empty();
|
|
if (!needs_sharding) {
|
|
BroadcastFromWorker0(param, /*in_group=*/true, param);
|
|
return param;
|
|
}
|
|
Device device = param->device;
|
|
Shape shape = param_info.preprocs.back().out_shape;
|
|
DLDataType dtype = param_info.preprocs.back().out_dtype;
|
|
TVM_FFI_ICHECK(shape.size() >= 1 && shape[0] == num_shards)
|
|
<< "ValueError: The first dimension of the output shape must be equal to the "
|
|
<< "number of shards, but got: " << shape << " and num_shards = " << num_shards;
|
|
param = preprocs.Apply(param, param_info);
|
|
Tensor result = Tensor::Empty(Shape(shape.begin() + 1, shape.end()), dtype, device);
|
|
ScatterFromWorker0(param, /*in_group=*/true, result);
|
|
return result;
|
|
}
|
|
|
|
Tensor ReceiveBroadcastedOrSharded(Device device, const ModelMetadata::Param& param_info,
|
|
int num_shards) {
|
|
bool needs_sharding = !param_info.preprocs.empty();
|
|
Tensor result;
|
|
if (needs_sharding) {
|
|
Shape shape = param_info.preprocs.back().out_shape;
|
|
DLDataType dtype = param_info.preprocs.back().out_dtype;
|
|
result = Tensor::Empty(Shape(shape.begin() + 1, shape.end()), dtype, device);
|
|
ScatterFromWorker0(std::nullopt, /*in_group=*/true, result);
|
|
} else {
|
|
result = Tensor::Empty(param_info.shape, param_info.dtype, device);
|
|
BroadcastFromWorker0(result, /*in_group=*/true, result);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
std::string FormatDuration(DurationType duration) {
|
|
std::ostringstream os;
|
|
auto float_seconds = std::chrono::duration_cast<std::chrono::duration<float>>(duration).count();
|
|
os << std::fixed << std::setprecision(3) << float_seconds << " s";
|
|
return os.str();
|
|
}
|
|
|
|
Array<Optional<Tensor>> LoadMultiGPU(const std::string& model_path, Module vm_module,
|
|
const std::string& model_config_str) {
|
|
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
|
Device device = worker->default_device;
|
|
int worker_id = worker->worker_id;
|
|
int group_size = worker->num_workers / worker->num_groups;
|
|
int num_shards = group_size;
|
|
int group_id = worker_id / group_size;
|
|
LOG(INFO) << "[Worker #" << worker_id << "] Loading model to device: " << device;
|
|
// Step 0. Initialize metadata and paths
|
|
TensorCacheMetadata tensor_cache_metadata = TensorCacheMetadata::Load(model_path);
|
|
tvm::ffi::json::Value model_config = tvm::ffi::json::Parse(model_config_str);
|
|
ModelMetadata model_metadata =
|
|
ModelMetadata::FromModule(vm_module, model_config.cast<tvm::ffi::json::Object>());
|
|
TVM_FFI_ICHECK_EQ(model_metadata.tensor_parallel_shards, num_shards)
|
|
<< "ValueError: The model is compiled using `--tensor-parallel-shards="
|
|
<< model_metadata.tensor_parallel_shards
|
|
<< "`, but mlc-chat-config.json is configured to use " << num_shards << " GPUs. "
|
|
<< "Please set \"tensor_parallel_shards\" in mlc-chat-config.json to "
|
|
<< model_metadata.tensor_parallel_shards;
|
|
// Step 1. Extract auxiliary information
|
|
PreprocessorPool preprocs(model_metadata, vm_module);
|
|
std::unordered_map<std::string, ModelMetadata::Param> param_name2info;
|
|
for (const ModelMetadata::Param& param : model_metadata.params) {
|
|
param_name2info[param.name] = param;
|
|
}
|
|
// Step 2. Load, preprocess and shard all the parameters
|
|
std::unordered_map<std::string, Tensor> sharded_params;
|
|
if (worker_id == 0) {
|
|
DurationType time_loading(0);
|
|
DurationType time_preproc(0);
|
|
ProgressBar progress_bar(model_metadata.params.size());
|
|
LOG(INFO) << "Loading parameters...";
|
|
for (const TensorCacheMetadata::FileRecord& record : tensor_cache_metadata.records) {
|
|
Array<Tensor> loaded_params;
|
|
{
|
|
RangeTimer _(&time_loading);
|
|
std::string raw_data_buffer;
|
|
loaded_params = record.Load(device, model_path, &raw_data_buffer);
|
|
DeviceAPI::Get(device)->StreamSync(device, nullptr);
|
|
}
|
|
// For each parameter in the shard file, preprocess and shard it
|
|
for (size_t i = 0; i < record.records.size(); ++i, progress_bar.Progress()) {
|
|
RangeTimer _(&time_preproc);
|
|
const std::string& param_name = record.records[i].name;
|
|
const ModelMetadata::Param& param_info = param_name2info.at(param_name);
|
|
for (int group_id : param_info.pipeline_stages) {
|
|
if (group_id == 0) {
|
|
// Broadcast or shard-scatter this parameter to all workers in worker group 0.
|
|
sharded_params[param_name] =
|
|
BroadcastOrShardAndScatter(loaded_params[i], param_info, num_shards, preprocs);
|
|
} else {
|
|
// Send this parameter to the first worker of the worker group of "group_id",
|
|
// and let that first worker to process this parameter.
|
|
SendToWorker(loaded_params[i], /*receiver_id=*/group_id * group_size);
|
|
}
|
|
}
|
|
DeviceAPI::Get(device)->StreamSync(device, nullptr);
|
|
}
|
|
}
|
|
LOG(INFO) << "Loading done. Time used:" << std::fixed << std::setprecision(3) //
|
|
<< " Loading " << FormatDuration(time_loading) << " Preprocessing "
|
|
<< FormatDuration(time_preproc) << ".";
|
|
} else {
|
|
for (const TensorCacheMetadata::FileRecord& record : tensor_cache_metadata.records) {
|
|
for (size_t i = 0; i < record.records.size(); ++i) {
|
|
const std::string& param_name = record.records[i].name;
|
|
const ModelMetadata::Param& param_info = param_name2info.at(param_name);
|
|
if (std::find(param_info.pipeline_stages.begin(), param_info.pipeline_stages.end(),
|
|
group_id) == param_info.pipeline_stages.end()) {
|
|
// This worker group doesn't need to hold a copy of this parameter.
|
|
continue;
|
|
}
|
|
|
|
if (worker_id % group_size == 0) {
|
|
// The worker is the first worker of its worker group (while not the first worker group).
|
|
// Receive the full parameter from the global worker 0.
|
|
Tensor full_param = RecvFromGlobalWorker0(device, param_info);
|
|
// Broadcast or shard-scatter this parameter to all workers in its worker group.
|
|
sharded_params[param_name] =
|
|
BroadcastOrShardAndScatter(full_param, param_info, num_shards, preprocs);
|
|
} else {
|
|
// The worker is not the first worker of its worker group.
|
|
// Receive from the first worker in the its worker group.
|
|
sharded_params[param_name] = ReceiveBroadcastedOrSharded(device, param_info, num_shards);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Step 3. Reorder the sharded parameters according to the order in model_metadata
|
|
Array<Optional<Tensor>> shards;
|
|
shards.reserve(model_metadata.params.size());
|
|
for (const ModelMetadata::Param& param : model_metadata.params) {
|
|
const auto& it = sharded_params.find(param.name);
|
|
shards.push_back(it == sharded_params.end() ? Optional<Tensor>() : it->second);
|
|
}
|
|
return shards;
|
|
}
|
|
|
|
Array<Optional<Tensor>> LoadMultiGPUPresharded(const std::string& model_path, Module vm_module,
|
|
const std::string& model_config_str) {
|
|
DiscoWorker* worker = DiscoWorker::ThreadLocal();
|
|
Device device = worker->default_device;
|
|
int worker_id = worker->worker_id;
|
|
int group_size = worker->num_workers / worker->num_groups;
|
|
int num_shards = group_size;
|
|
int group_id = worker_id / group_size;
|
|
int local_worker_id = worker_id % group_size;
|
|
LOG(INFO) << "[Worker #" << worker_id << "] Loading model to device: " << device;
|
|
// Step 0. Initialize metadata and paths
|
|
TensorCacheMetadata tensor_cache_metadata = TensorCacheMetadata::Load(model_path);
|
|
tvm::ffi::json::Value model_config = tvm::ffi::json::Parse(model_config_str);
|
|
ModelMetadata model_metadata =
|
|
ModelMetadata::FromModule(vm_module, model_config.cast<tvm::ffi::json::Object>());
|
|
|
|
std::unordered_map<std::string, ParamInfo> param_info_map;
|
|
for (const TensorCacheMetadata::FileRecord& file_record : tensor_cache_metadata.records) {
|
|
for (const TensorCacheMetadata::FileRecord::ParamRecord& param_record : file_record.records) {
|
|
const std::string& param_name = param_record.name;
|
|
param_info_map[param_name] = ParamInfo{&file_record, ¶m_record};
|
|
}
|
|
}
|
|
|
|
Array<Optional<Tensor>> params;
|
|
const TensorCacheMetadata::FileRecord* current_file_;
|
|
std::string current_file_stream_;
|
|
params.reserve(model_metadata.params.size());
|
|
DurationType time_loading(0);
|
|
for (const ModelMetadata::Param& param : model_metadata.params) {
|
|
RangeTimer _(&time_loading);
|
|
if (std::find(param.pipeline_stages.begin(), param.pipeline_stages.end(), group_id) ==
|
|
param.pipeline_stages.end()) {
|
|
// This worker group doesn't need to hold a copy of this parameter.
|
|
params.push_back(Optional<Tensor>());
|
|
continue;
|
|
}
|
|
bool needs_sharding = !param.preprocs.empty();
|
|
std::string param_name =
|
|
needs_sharding ? static_cast<const std::stringstream&>(
|
|
std::stringstream() << param.name << "_shard-" << local_worker_id)
|
|
.str()
|
|
: std::string(param.name);
|
|
auto it = param_info_map.find(param_name);
|
|
TVM_FFI_ICHECK(it != param_info_map.end())
|
|
<< "ValueError: Cannot find parameter: " << param_name;
|
|
const ParamInfo& param_info = (*it).second;
|
|
const TensorCacheMetadata::FileRecord::ParamRecord* param_record = param_info.param;
|
|
const TensorCacheMetadata::FileRecord* file_record = param_info.file;
|
|
|
|
if (file_record != current_file_) {
|
|
current_file_ = file_record;
|
|
file_record->Load(device, model_path, ¤t_file_stream_);
|
|
}
|
|
|
|
params.push_back(param_record->Load(device, ¤t_file_stream_));
|
|
}
|
|
SyncWorker();
|
|
if (worker_id == 0) {
|
|
LOG(INFO) << "Loading done. Time used: " << FormatDuration(time_loading) << ".";
|
|
}
|
|
return params;
|
|
}
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef()
|
|
.def("mlc.multi_gpu.LoadMultiGPU", LoadMultiGPU)
|
|
.def("mlc.multi_gpu.LoadMultiGPUPresharded", LoadMultiGPUPresharded);
|
|
}
|
|
|
|
} // namespace multi_gpu
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_SINGLE_GPU_ONLY
|