Skip to content

fix: preserve dtype across pipeline parallel communication #255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JYMiracle305 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: preserve dtype across pipeline parallel communication #255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions example/gpt2/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,17 @@ constexpr char kDtypeBF16[] = "bfloat16";
const std::unordered_set<std::string> 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<std::string, nn::TransformerConfig> kModelToConfigs = {
{"d12",
Expand Down Expand Up @@ -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 /
Expand Down Expand Up @@ -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 =
Expand Down
22 changes: 13 additions & 9 deletions example/llama3/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,17 @@ constexpr char kDtypeBF16[] = "bfloat16";
const std::unordered_set<std::string> 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) {
Expand Down Expand Up @@ -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 /
Expand Down Expand Up @@ -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 =
Expand Down
22 changes: 13 additions & 9 deletions example/qwen3/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,17 @@ constexpr char kDtypeFP32[] = "float32";
constexpr char kDtypeBF16[] = "bfloat16";
const std::unordered_set<std::string> 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) {
Expand Down Expand Up @@ -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 /
Expand Down Expand Up @@ -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 =
Expand Down
6 changes: 6 additions & 0 deletions infini_train/include/nn/parallel/model_parallel_config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

#include <optional>

#include "infini_train/include/datatype.h"

namespace infini_train::nn::parallel {

// Model-parallel topology and communication policy. Runtime-only process
Expand All @@ -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<int> expert_tensor_parallel_size = std::nullopt;
Expand Down
6 changes: 4 additions & 2 deletions infini_train/include/nn/parallel/pp/send_recv.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
#include <memory>
#include <vector>

#include "infini_train/include/datatype.h"

namespace infini_train {
class Tensor;
class Device;
Expand All @@ -15,9 +17,9 @@ namespace infini_train::nn::parallel {
std::vector<std::shared_ptr<Tensor>> ISend(
const std::vector<std::shared_ptr<Tensor>> &input_tensors,
Device target_device, int peer_rank,
const std::vector<std::vector<int64_t>> &shape);
const std::vector<std::vector<int64_t>> &shape, DataType pipeline_dtype);

std::vector<std::shared_ptr<Tensor>> IRecv(
const std::vector<std::shared_ptr<Tensor>> &outputs, Device src_device,
int peer_rank);
int peer_rank, DataType pipeline_dtype, DataType activation_dtype);
} // namespace infini_train::nn::parallel
19 changes: 12 additions & 7 deletions infini_train/src/nn/parallel/pp/pipeline_schedule.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -56,11 +55,11 @@ std::vector<std::shared_ptr<Tensor>> PipelineSchedule::ReceiveFromPrev(
int peer_rank) {
std::vector<std::shared_ptr<Tensor>> 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<Tensor>(shapes[i], DataType::kFLOAT32,
stage_->device());
auto tensor =
std::make_shared<Tensor>(shapes[i], pipeline_dtype, stage_->device());
tensor->set_requires_grad(true);

// Mark as non-leaf to prevent the autograd engine from creating an
Expand All @@ -75,12 +74,18 @@ std::vector<std::shared_ptr<Tensor>> 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<std::shared_ptr<Tensor>> PipelineSchedule::SendToNext(
const std::vector<std::shared_ptr<Tensor>> &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(
Expand Down
Loading
Loading