From 9bb0b9c92ff7389580382fbdc7c126ba04d61932 Mon Sep 17 00:00:00 2001 From: JYMiracle305 <604951424@qq.com> Date: Thu, 5 Nov 2026 10:16:33 +0000 Subject: [PATCH] feat: support pipeline communication dtype --- example/gpt2/main.cc | 22 +-- example/llama3/main.cc | 22 +-- example/qwen3/main.cc | 22 +-- .../nn/parallel/model_parallel_config.h | 6 + .../include/nn/parallel/pp/send_recv.h | 6 +- .../src/nn/parallel/pp/pipeline_schedule.cc | 19 ++- infini_train/src/nn/parallel/pp/send_recv.cc | 130 ++++++++++++++---- 7 files changed, 166 insertions(+), 61 deletions(-) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 37ce5024..2ea5bdbf 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -183,6 +183,17 @@ constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles = { "none", "constant", "linear", "cosine", "inverse-square-root"}; +DataType ParseDtype(const std::string &dtype) { + if (dtype == kDtypeFP32) { + return DataType::kFLOAT32; + } + if (dtype == kDtypeBF16) { + return DataType::kBFLOAT16; + } + LOG(FATAL) << "Datatype " << dtype << " not supported."; + return DataType::kFLOAT32; +} + // const std::unordered_map kModelToConfigs = { {"d12", @@ -365,15 +376,7 @@ void Train(const nn::parallel::Rank &rank) { // select the data type // TODO(lzm): change to solely rely on the weight file info for determining // the dtype when autocast is supported - DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { - dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { - dtype = DataType::kBFLOAT16; - } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype - << " not supported."; - } + const DataType dtype = ParseDtype(FLAGS_dtype); auto num_micro_batches = FLAGS_total_batch_size / @@ -755,6 +758,7 @@ int main(int argc, char *argv[]) { auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::ModelParallelConfig model_parallel_config; + model_parallel_config.pipeline_dtype = ParseDtype(FLAGS_dtype); model_parallel_config.tensor_model_parallel_size = FLAGS_tensor_parallel; model_parallel_config.pipeline_model_parallel_size = FLAGS_pipeline_parallel; model_parallel_config.virtual_pipeline_model_parallel_size = diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 2143ad98..7a7abb30 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -182,6 +182,17 @@ constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles = { "none", "constant", "linear", "cosine", "inverse-square-root"}; +DataType ParseDtype(const std::string &dtype) { + if (dtype == kDtypeFP32) { + return DataType::kFLOAT32; + } + if (dtype == kDtypeBF16) { + return DataType::kBFLOAT16; + } + LOG(FATAL) << "Datatype " << dtype << " not supported."; + return DataType::kFLOAT32; +} + } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { @@ -341,15 +352,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; - DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { - dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { - dtype = DataType::kBFLOAT16; - } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype - << " not supported."; - } + const DataType dtype = ParseDtype(FLAGS_dtype); auto num_micro_batches = FLAGS_total_batch_size / @@ -730,6 +733,7 @@ int main(int argc, char *argv[]) { auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::ModelParallelConfig model_parallel_config; + model_parallel_config.pipeline_dtype = ParseDtype(FLAGS_dtype); model_parallel_config.tensor_model_parallel_size = FLAGS_tensor_parallel; model_parallel_config.pipeline_model_parallel_size = FLAGS_pipeline_parallel; model_parallel_config.virtual_pipeline_model_parallel_size = diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index 8759713f..7f393312 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -182,6 +182,17 @@ constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles = { "none", "constant", "linear", "cosine", "inverse-square-root"}; + +DataType ParseDtype(const std::string &dtype) { + if (dtype == kDtypeFP32) { + return DataType::kFLOAT32; + } + if (dtype == kDtypeBF16) { + return DataType::kBFLOAT16; + } + LOG(FATAL) << "Datatype " << dtype << " not supported."; + return DataType::kFLOAT32; +} } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { @@ -339,15 +350,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; - DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { - dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { - dtype = DataType::kBFLOAT16; - } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype - << " not supported."; - } + const DataType dtype = ParseDtype(FLAGS_dtype); auto num_micro_batches = FLAGS_total_batch_size / @@ -728,6 +731,7 @@ int main(int argc, char *argv[]) { auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::ModelParallelConfig model_parallel_config; + model_parallel_config.pipeline_dtype = ParseDtype(FLAGS_dtype); model_parallel_config.tensor_model_parallel_size = FLAGS_tensor_parallel; model_parallel_config.pipeline_model_parallel_size = FLAGS_pipeline_parallel; model_parallel_config.virtual_pipeline_model_parallel_size = diff --git a/infini_train/include/nn/parallel/model_parallel_config.h b/infini_train/include/nn/parallel/model_parallel_config.h index d1e8e92b..966747a4 100644 --- a/infini_train/include/nn/parallel/model_parallel_config.h +++ b/infini_train/include/nn/parallel/model_parallel_config.h @@ -2,6 +2,8 @@ #include +#include "infini_train/include/datatype.h" + namespace infini_train::nn::parallel { // Model-parallel topology and communication policy. Runtime-only process @@ -10,6 +12,10 @@ struct ModelParallelConfig { int tensor_model_parallel_size = 1; int pipeline_model_parallel_size = 1; int virtual_pipeline_model_parallel_size = 1; + // Dtype used by pipeline payloads on the communication wire. + DataType pipeline_dtype = DataType::kFLOAT32; + // Logical model-boundary dtype restored after receiving a wire payload. + DataType pipeline_activation_dtype = DataType::kFLOAT32; bool sequence_parallel = false; int expert_model_parallel_size = 1; std::optional expert_tensor_parallel_size = std::nullopt; diff --git a/infini_train/include/nn/parallel/pp/send_recv.h b/infini_train/include/nn/parallel/pp/send_recv.h index 38a21a5e..b7244666 100644 --- a/infini_train/include/nn/parallel/pp/send_recv.h +++ b/infini_train/include/nn/parallel/pp/send_recv.h @@ -5,6 +5,8 @@ #include #include +#include "infini_train/include/datatype.h" + namespace infini_train { class Tensor; class Device; @@ -15,9 +17,9 @@ namespace infini_train::nn::parallel { std::vector> ISend( const std::vector> &input_tensors, Device target_device, int peer_rank, - const std::vector> &shape); + const std::vector> &shape, DataType pipeline_dtype); std::vector> IRecv( const std::vector> &outputs, Device src_device, - int peer_rank); + int peer_rank, DataType pipeline_dtype, DataType activation_dtype); } // namespace infini_train::nn::parallel diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index b40a18de..30663e54 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -12,7 +12,6 @@ #include "glog/logging.h" #include "infini_train/include/autocast.h" -#include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/parallel/global.h" @@ -56,11 +55,11 @@ std::vector> PipelineSchedule::ReceiveFromPrev( int peer_rank) { std::vector> recv_tensors; auto &shapes = stage_->recv_shape(); + const auto pipeline_dtype = + global::GetModelParallelConfig().pipeline_dtype; for (size_t i = 0; i < shapes.size(); ++i) { - // FIXME(jym): The data type between stages is not float32, which will cause - // a crash - auto tensor = std::make_shared(shapes[i], DataType::kFLOAT32, - stage_->device()); + auto tensor = + std::make_shared(shapes[i], pipeline_dtype, stage_->device()); tensor->set_requires_grad(true); // Mark as non-leaf to prevent the autograd engine from creating an @@ -75,12 +74,18 @@ std::vector> PipelineSchedule::ReceiveFromPrev( recv_tensors.push_back(tensor); } - return IRecv(recv_tensors, stage_->device(), peer_rank); + const auto activation_dtype = + global::GetModelParallelConfig().pipeline_activation_dtype; + return IRecv(recv_tensors, stage_->device(), peer_rank, pipeline_dtype, + activation_dtype); } std::vector> PipelineSchedule::SendToNext( const std::vector> &tensors, int peer_rank) { - return ISend(tensors, stage_->device(), peer_rank, stage_->recv_shape()); + const auto pipeline_dtype = + global::GetModelParallelConfig().pipeline_dtype; + return ISend(tensors, stage_->device(), peer_rank, stage_->recv_shape(), + pipeline_dtype); } PipelineParallelScheduler::Task PipelineParallelScheduler::CreateTask( diff --git a/infini_train/src/nn/parallel/pp/send_recv.cc b/infini_train/src/nn/parallel/pp/send_recv.cc index 78b563e1..c61286c9 100644 --- a/infini_train/src/nn/parallel/pp/send_recv.cc +++ b/infini_train/src/nn/parallel/pp/send_recv.cc @@ -1,11 +1,11 @@ #include "infini_train/include/nn/parallel/pp/send_recv.h" -#include #include #include +#include "glog/logging.h" + #include "infini_train/include/autograd/function.h" -#include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/parallel/process_group.h" #include "infini_train/include/nn/parallel/rank.h" @@ -20,15 +20,21 @@ class ISend : public autograd::Function { static constexpr char kType[] = "ISendFunction"; explicit ISend(Device target_device, int peer_rank, - const std::vector> &shape) + const std::vector> &shape, + DataType pipeline_dtype) : autograd::Function(kType), target_device_(target_device), peer_rank_(peer_rank), - shapes_(shape) {} + shapes_(shape), + pipeline_dtype_(pipeline_dtype) {} std::vector> Forward( const std::vector> &input_tensors) override; + void SetupContext( + const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward( const std::vector> &grad_outputs) override; @@ -37,16 +43,23 @@ class ISend : public autograd::Function { Device input_device_; int peer_rank_ = -1; const std::vector> &shapes_; + DataType pipeline_dtype_ = DataType::kFLOAT32; + std::vector> wire_tensors_; + std::vector> wire_grad_tensors_; + std::vector input_dtypes_; }; class IRecv : public autograd::Function { public: static constexpr char kType[] = "IRecvFunction"; - explicit IRecv(Device src_device, int peer_rank) + explicit IRecv(Device src_device, int peer_rank, DataType pipeline_dtype, + DataType activation_dtype) : autograd::Function(kType), src_device_(src_device), - peer_rank_(peer_rank) {} + peer_rank_(peer_rank), + pipeline_dtype_(pipeline_dtype), + activation_dtype_(activation_dtype) {} std::vector> Forward( const std::vector> &input_tensors) override; @@ -62,40 +75,83 @@ class IRecv : public autograd::Function { Device src_device_; Device cur_device_; int peer_rank_ = -1; + DataType pipeline_dtype_ = DataType::kFLOAT32; + DataType activation_dtype_ = DataType::kFLOAT32; + std::vector output_dtypes_; + std::vector> wire_grad_tensors_; }; std::vector> ISend::Forward( const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), shapes_.size()) + << "Pipeline send tensor count does not match the configured receive " + "shapes"; + + CHECK(!input_tensors.empty()); const auto &input = input_tensors[0]; + CHECK(input != nullptr); input_device_ = input->GetDevice(); + wire_tensors_.clear(); + wire_tensors_.reserve(input_tensors.size()); + for (const auto &tensor : input_tensors) { + CHECK(tensor != nullptr); + if (tensor->Dtype() == pipeline_dtype_) { + wire_tensors_.push_back(tensor); + } else { + wire_tensors_.push_back( + std::make_shared(tensor->To(pipeline_dtype_))); + } + CHECK(wire_tensors_.back()->Dtype() == pipeline_dtype_); + } + auto pp_group = ProcessGroupFactory::Instance(input_device_.type()) ->Get(GetPipelineParallelProcessGroupName( input_device_.Rank().GlobalRank())); - pp_group->Send(input_tensors, peer_rank_, false); + pp_group->Send(wire_tensors_, peer_rank_, false); return input_tensors; } +void ISend::SetupContext( + const std::vector> &input_tensors, + const std::vector> &output_tensors) { + input_dtypes_.clear(); + input_dtypes_.reserve(input_tensors.size()); + for (const auto &tensor : input_tensors) { + CHECK(tensor != nullptr); + input_dtypes_.push_back(tensor->Dtype()); + } +} + std::vector> ISend::Backward( const std::vector> &grad_outputs) { - std::vector> recv_tensors; - for (int shape_i = 0; shape_i < shapes_.size(); ++shape_i) { - // FIXME(jym): The data type between stages is not float32, which will cause - // a crash - auto r_tensor = std::make_shared(shapes_[shape_i], - DataType::kFLOAT32, input_device_); - recv_tensors.push_back(r_tensor); + wire_grad_tensors_.clear(); + CHECK_EQ(input_dtypes_.size(), shapes_.size()); + for (size_t shape_i = 0; shape_i < shapes_.size(); ++shape_i) { + auto r_tensor = std::make_shared(shapes_[shape_i], pipeline_dtype_, + input_device_); + wire_grad_tensors_.push_back(r_tensor); } auto pp_group = ProcessGroupFactory::Instance(input_device_.type()) ->Get(GetPipelineParallelProcessGroupName( input_device_.Rank().GlobalRank())); - pp_group->Recv(recv_tensors, peer_rank_, false); - - return recv_tensors; + pp_group->Recv(wire_grad_tensors_, peer_rank_, false); + + std::vector> input_grads; + input_grads.reserve(wire_grad_tensors_.size()); + for (size_t i = 0; i < wire_grad_tensors_.size(); ++i) { + if (wire_grad_tensors_[i]->Dtype() == input_dtypes_[i]) { + input_grads.push_back(wire_grad_tensors_[i]); + } else { + input_grads.push_back( + std::make_shared(wire_grad_tensors_[i]->To(input_dtypes_[i]))); + } + } + return input_grads; } std::vector> IRecv::Forward( @@ -105,7 +161,17 @@ std::vector> IRecv::Forward( src_device_.Rank().GlobalRank())); pp_group->Recv(recv_tensors, peer_rank_, false); - return recv_tensors; + if (pipeline_dtype_ == activation_dtype_) { + return recv_tensors; + } + + std::vector> activation_tensors; + activation_tensors.reserve(recv_tensors.size()); + for (const auto &tensor : recv_tensors) { + activation_tensors.push_back( + std::make_shared(tensor->To(activation_dtype_))); + } + return activation_tensors; } void IRecv::SetupContext( @@ -115,6 +181,11 @@ void IRecv::SetupContext( return; } cur_device_ = output_tensors[0]->GetDevice(); + output_dtypes_.clear(); + output_dtypes_.reserve(output_tensors.size()); + for (const auto &tensor : output_tensors) { + output_dtypes_.push_back(tensor->Dtype()); + } } std::vector> IRecv::Backward( @@ -123,25 +194,34 @@ std::vector> IRecv::Backward( ->Get(GetPipelineParallelProcessGroupName( cur_device_.Rank().GlobalRank())); - pp_group->Send(grad_outputs, peer_rank_, false); + CHECK_EQ(grad_outputs.size(), output_dtypes_.size()); + wire_grad_tensors_ = grad_outputs; + for (size_t i = 0; i < wire_grad_tensors_.size(); ++i) { + if (wire_grad_tensors_[i]->Dtype() != pipeline_dtype_) { + wire_grad_tensors_[i] = + std::make_shared(wire_grad_tensors_[i]->To(pipeline_dtype_)); + } + } + pp_group->Send(wire_grad_tensors_, peer_rank_, false); - return grad_outputs; + return wire_grad_tensors_; } } // namespace functions std::vector> ISend( const std::vector> &input_tensors, Device target_device, int peer_rank, - const std::vector> &shape) { - auto func = - std::make_shared(target_device, peer_rank, shape); + const std::vector> &shape, DataType pipeline_dtype) { + auto func = std::make_shared( + target_device, peer_rank, shape, pipeline_dtype); return func->Apply(input_tensors); } std::vector> IRecv( const std::vector> &outputs, Device src_device, - int peer_rank) { - auto func = std::make_shared(src_device, peer_rank); + int peer_rank, DataType pipeline_dtype, DataType activation_dtype) { + auto func = std::make_shared( + src_device, peer_rank, pipeline_dtype, activation_dtype); return func->Apply(outputs); } } // namespace infini_train::nn::parallel