wehub-resource-sync
136 lines
4.6 KiB
C++
136 lines
4.6 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include "z1.h"
|
|
#include "deepcompile.h"
|
|
|
|
namespace dc {
|
|
|
|
class Z1CustomOpExecutor : public CustomOpExecutor {
|
|
public:
|
|
Z1CustomOpExecutor(c10::intrusive_ptr<c10d::ProcessGroup> process_group,
|
|
std::shared_ptr<DSParamRegistry> param_registry,
|
|
std::shared_ptr<DoubleBufferedReduceBucket> reduce_buckets,
|
|
std::vector<long> ds_ids,
|
|
ncclComm_t nccl_comm,
|
|
at::cuda::CUDAStream rs_stream,
|
|
at::cuda::CUDAStream copy_stream,
|
|
bool pre_div_reduce)
|
|
: CustomOpExecutor(process_group,
|
|
param_registry,
|
|
reduce_buckets,
|
|
ds_ids,
|
|
nccl_comm,
|
|
rs_stream,
|
|
copy_stream,
|
|
pre_div_reduce)
|
|
{
|
|
}
|
|
~Z1CustomOpExecutor() {}
|
|
|
|
at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id) override
|
|
{
|
|
if (!hasKey(grad_tensors_, ds_id)) {
|
|
grad_tensors_[ds_id] = grad_tensor;
|
|
} else {
|
|
grad_tensors_[ds_id].add_(grad_tensor);
|
|
}
|
|
|
|
if (param_updated_) {
|
|
CustomOpExecutor::reduceGrad(grad_tensors_[ds_id], ds_id);
|
|
grad_tensors_.erase(ds_id);
|
|
}
|
|
|
|
return at::Tensor();
|
|
}
|
|
|
|
void flushReduceBucket(at::ScalarType scalar_type) override
|
|
{
|
|
if (!hasKey(reduce_tasks_, scalar_type)) { return; }
|
|
|
|
blockCopyEvents(scalar_type);
|
|
applyPreDivision(scalar_type);
|
|
|
|
// NCCL AllReduce operation
|
|
ncclGroupStart();
|
|
for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) {
|
|
ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(),
|
|
t.getSendBuf().data_ptr(),
|
|
t.getSendBuf().numel(),
|
|
get_nccl_data_type(scalar_type),
|
|
getReductionOp(),
|
|
nccl_comm_,
|
|
rs_stream_);
|
|
if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); }
|
|
}
|
|
ncclGroupEnd();
|
|
|
|
// Copy results to gradient buffers
|
|
{
|
|
at::cuda::CUDAStreamGuard guard(rs_stream_);
|
|
for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) {
|
|
auto param = param_registry_->getParam(t.getDSId());
|
|
auto grad_buf = param.getGradBuffer().flatten();
|
|
|
|
if (grad_buf.numel() == 0) { continue; }
|
|
|
|
int64_t offset = param.getOffset();
|
|
auto recv_buf = t.getSendBuf().flatten().index(
|
|
{torch::indexing::Slice(offset, offset + grad_buf.numel())});
|
|
grad_buf.copy_(recv_buf);
|
|
}
|
|
}
|
|
|
|
performCleanup(scalar_type);
|
|
}
|
|
|
|
protected:
|
|
std::unordered_map<long, at::Tensor> grad_tensors_;
|
|
};
|
|
|
|
namespace {
|
|
|
|
at::cuda::CUDAStream get_rs_stream()
|
|
{
|
|
static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true);
|
|
return rs_stream;
|
|
}
|
|
|
|
at::cuda::CUDAStream get_copy_stream()
|
|
{
|
|
static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true);
|
|
return copy_stream;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void register_graph_z1(long graph_id, const std::vector<long>& ds_ids)
|
|
{
|
|
executors[graph_id] = std::make_shared<Z1CustomOpExecutor>(process_group,
|
|
param_registry,
|
|
reduce_buckets,
|
|
ds_ids,
|
|
nccl_comm,
|
|
get_rs_stream(),
|
|
get_copy_stream(),
|
|
pre_div_reduce);
|
|
}
|
|
|
|
void register_param(long ds_id,
|
|
const std::vector<int64_t>& ds_shape,
|
|
at::Tensor ds_tensor,
|
|
at::Tensor grad_buffer,
|
|
int64_t offset)
|
|
{
|
|
param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, false, offset, false);
|
|
}
|
|
|
|
void update_param_grad_buffer(long ds_id, at::Tensor grad_buffer, int64_t offset)
|
|
{
|
|
param_registry->updateGradBuffer(ds_id, grad_buffer, offset);
|
|
}
|
|
|
|
} // namespace dc
|