diff --git a/example/mnist/net.cc b/example/mnist/net.cc index 501fee7ef..75ad72e2b 100644 --- a/example/mnist/net.cc +++ b/example/mnist/net.cc @@ -18,8 +18,8 @@ MNIST::MNIST() { std::vector> layers; layers.push_back(std::make_shared(784, 30)); layers.push_back(std::make_shared()); - modules_["sequential"] = std::make_shared(std::move(layers)); - modules_["linear2"] = std::make_shared(30, 10); + RegisterModule("sequential", std::make_shared(std::move(layers))); + RegisterModule("linear2", std::make_shared(30, 10)); } std::vector> diff --git a/infini_train/include/nn/modules/container.h b/infini_train/include/nn/modules/container.h index 28bceccf1..814af70e2 100644 --- a/infini_train/include/nn/modules/container.h +++ b/infini_train/include/nn/modules/container.h @@ -1,7 +1,8 @@ #pragma once #include -#include +#include +#include #include #include "infini_train/include/nn/modules/module.h" @@ -23,8 +24,8 @@ class Sequential : public CloneableModule { class ModuleDict : public CloneableModule { public: static constexpr char kType[] = "ModuleDict"; - // TODO(dcj): in torch, there is a dict with the order of insertion - explicit ModuleDict(std::unordered_map> modules); + using Item = std::pair>; + explicit ModuleDict(std::vector modules); std::vector> Forward(const std::vector> &input_tensors) override; }; diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h index 1d42b2acc..f83c769f5 100644 --- a/infini_train/include/nn/modules/module.h +++ b/infini_train/include/nn/modules/module.h @@ -47,10 +47,16 @@ class Module : public std::enable_shared_from_this { const std::string &type() const; - virtual std::vector> Parameters() const; + // PyTorch-style registration APIs. Names are unique across parameters, buffers, and child modules. + // Re-registering an existing name in the same registry replaces its value without changing insertion order. + std::shared_ptr RegisterParameter(const std::string &name, std::shared_ptr parameter); + std::shared_ptr RegisterBuffer(const std::string &name, std::shared_ptr buffer, + bool persistent = true); + std::shared_ptr RegisterModule(const std::string &name, std::shared_ptr module); - // InfiniTrain's NamedParameters returns results ordered by full parameter name. - // TODO: Align with PyTorch's ordering in the future. + virtual std::vector> Parameters(bool recurse = true) const; + + // Results follow parameter and module registration order, matching torch.nn.Module.named_parameters(). std::vector>> NamedParameters(const std::string &prefix = "", bool recurse = true, bool remove_duplicate = true) const; bool has_parameter(const std::string &name) const; @@ -106,6 +112,10 @@ class Module : public std::enable_shared_from_this { std::unordered_map> modules_; std::unordered_map> parameters_; std::unordered_map> buffers_; + std::vector module_order_; + std::vector parameter_order_; + std::vector buffer_order_; + std::unordered_set non_persistent_buffers_; std::vector forward_pre_hooks_; std::vector forward_post_hooks_; diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h index 823ae82b5..a3ee502dc 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h +++ b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h @@ -40,6 +40,7 @@ class DistributedDataParallel : public nn::Module { private: void BuildParamAndGradBuffers(); void RegisterBackwardHooks(); + void RegisterForwardPreHooks(); void OnGradReady(const std::shared_ptr ¶m); private: diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h index 729456ce3..a9717ab56 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h +++ b/infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h @@ -57,7 +57,7 @@ class DistributedDataParallelConfig { // TODO(zbl): Unused by now, to be implemented in ParamAndGradBucketGroup int num_distributed_optimizer_instances = 1; - // Maximum number of parameters in each ParamAndGradBucket. + // Target maximum number of elements in each ParamAndGradBucket. // NOTE(zbl): This is distinct from DDP Reducer's MB-based bucket caps. size_t bucket_size_in_elements = 1000000; diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h index d7cea198e..29494d486 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h +++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h @@ -56,6 +56,7 @@ class DistributedOptimizer final : public infini_train::Optimizer { // Inherit from DDP model std::vector> param_grad_buffers_; std::vector> bucket_groups_; + std::vector> first_param_sync_bucket_groups_; // DP info size_t ddp_world_size_; diff --git a/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h b/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h index 4af99d818..f189f236f 100644 --- a/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h +++ b/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h @@ -113,6 +113,9 @@ class ParamAndGradBucketGroup { // Wait for parameter all-gather to complete void FinishParamSync(bool skip_next_bucket_dispatch = false); + // Drain any previous parameter all-gather before the optimizer writes the next parameter version. + void PrepareParamSyncForNextStep(); + // TODO(zbl): For PP, set the next bucket group used for parameter all-gather. void SetNextParamGatherBucketGroup(std::shared_ptr next_group); diff --git a/infini_train/src/nn/lora/lora_linear.cc b/infini_train/src/nn/lora/lora_linear.cc index 0a2d1098d..208b69254 100644 --- a/infini_train/src/nn/lora/lora_linear.cc +++ b/infini_train/src/nn/lora/lora_linear.cc @@ -35,11 +35,11 @@ LoRALinear::LoRALinear(std::shared_ptr base_linear, const LoRAConfig } // Transfer weight from base linear (overwrite base-created one) - parameters_[kParamWeightName] = base_linear->parameter(kParamWeightName); + RegisterParameter(kParamWeightName, base_linear->parameter(kParamWeightName)); // Transfer bias if exists if (has_bias()) { - parameters_[kParamBiasName] = base_linear->parameter(kParamBiasName); + RegisterParameter(kParamBiasName, base_linear->parameter(kParamBiasName)); } // Initialize LoRA weights @@ -52,9 +52,9 @@ LoRALinear::LoRALinear(std::shared_ptr base_linear, const LoRAConfig void LoRALinear::InitLoRAWeights() { // A matrix: [rank, in_features] // Initialize with Kaiming uniform (or normal based on config) - parameters_[kParamLoraAName] - = std::make_shared(std::vector{config_.rank, in_features_}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraAName, std::make_shared(std::vector{config_.rank, in_features_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (config_.use_kaiming_a) { init::KaimingUniform(parameters_[kParamLoraAName], config_.kaiming_a_param); @@ -64,9 +64,9 @@ void LoRALinear::InitLoRAWeights() { // B matrix: [out_features, rank] // Initialize with zeros (ensures LoRA starts as identity transformation) - parameters_[kParamLoraBName] - = std::make_shared(std::vector{out_features_, config_.rank}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraBName, std::make_shared(std::vector{out_features_, config_.rank}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); init::Zeros(parameters_[kParamLoraBName]); } diff --git a/infini_train/src/nn/lora/lora_parallel_linear.cc b/infini_train/src/nn/lora/lora_parallel_linear.cc index 9b038e2d4..7d01d94f4 100644 --- a/infini_train/src/nn/lora/lora_parallel_linear.cc +++ b/infini_train/src/nn/lora/lora_parallel_linear.cc @@ -34,7 +34,7 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptrparameter(kParamWeightName)->GetDevice(); // Transfer weight from base module (overwrite base-created one) - parameters_[kParamWeightName] = base_module->parameter(kParamWeightName); + RegisterParameter(kParamWeightName, base_module->parameter(kParamWeightName)); // Get dimensions from weight shape [out_features_per_partition, in_features] const auto &weight_dims = parameters_[kParamWeightName]->Dims(); @@ -42,7 +42,7 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptrhas_parameter(kParamBiasName)) { - parameters_[kParamBiasName] = base_module->parameter(kParamBiasName); + RegisterParameter(kParamBiasName, base_module->parameter(kParamBiasName)); } // Initialize LoRA weights @@ -66,7 +66,7 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptrparameter(kParamWeightName)->GetDevice(); // Transfer weight from base module (overwrite base-created one) - parameters_[kParamWeightName] = base_module->parameter(kParamWeightName); + RegisterParameter(kParamWeightName, base_module->parameter(kParamWeightName)); // Get dimensions from weight shape [out_features_per_partition, in_features] const auto &weight_dims = parameters_[kParamWeightName]->Dims(); @@ -79,7 +79,7 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptrhas_parameter(kParamBiasName)) { - parameters_[kParamBiasName] = base_module->parameter(kParamBiasName); + RegisterParameter(kParamBiasName, base_module->parameter(kParamBiasName)); } // Initialize LoRA weights @@ -92,9 +92,9 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptr(std::vector{config_.rank, in_features_}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraAName, std::make_shared(std::vector{config_.rank, in_features_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (parallel::global::GetTensorParallelSize() > 1) { const auto global_rank = device_.Rank().GlobalRank(); @@ -122,10 +122,10 @@ void LoRAColumnParallelLinear::InitLoRAWeights() { } } - parameters_[kParamLoraBName] - = std::make_shared(std::vector{out_features_per_partition_, config_.rank}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraBName, + std::make_shared(std::vector{out_features_per_partition_, config_.rank}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); init::Zeros(parameters_[kParamLoraBName]); } @@ -243,7 +243,7 @@ LoRARowParallelLinear::LoRARowParallelLinear(std::shared_ptrparameter(kParamWeightName)->GetDevice(); // Transfer weight from base module (overwrite base-created one) - parameters_[kParamWeightName] = base_module->parameter(kParamWeightName); + RegisterParameter(kParamWeightName, base_module->parameter(kParamWeightName)); // Get dimensions from weight shape [out_features, in_features_per_partition] const auto &weight_dims = parameters_[kParamWeightName]->Dims(); @@ -251,7 +251,7 @@ LoRARowParallelLinear::LoRARowParallelLinear(std::shared_ptrhas_parameter(kParamBiasName)) { - parameters_[kParamBiasName] = base_module->parameter(kParamBiasName); + RegisterParameter(kParamBiasName, base_module->parameter(kParamBiasName)); } // Initialize LoRA weights @@ -275,7 +275,7 @@ LoRARowParallelLinear::LoRARowParallelLinear(std::shared_ptrparameter(kParamWeightName)->GetDevice(); // Transfer weight from base module (overwrite base-created one) - parameters_[kParamWeightName] = base_module->parameter(kParamWeightName); + RegisterParameter(kParamWeightName, base_module->parameter(kParamWeightName)); // Get dimensions from weight shape [out_features, in_features_per_partition] const auto &weight_dims = parameters_[kParamWeightName]->Dims(); @@ -288,7 +288,7 @@ LoRARowParallelLinear::LoRARowParallelLinear(std::shared_ptrhas_parameter(kParamBiasName)) { - parameters_[kParamBiasName] = base_module->parameter(kParamBiasName); + RegisterParameter(kParamBiasName, base_module->parameter(kParamBiasName)); } // Initialize LoRA weights @@ -303,10 +303,10 @@ void LoRARowParallelLinear::InitLoRAWeights() { // lora_B: [out_features, rank] - replicated // lora_A: [rank, in_features_per_partition] - parameters_[kParamLoraAName] - = std::make_shared(std::vector{config_.rank, in_features_per_partition_}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraAName, + std::make_shared(std::vector{config_.rank, in_features_per_partition_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (parallel::global::GetTensorParallelSize() > 1) { const auto global_rank = device_.Rank().GlobalRank(); auto *tp_group = parallel::ProcessGroupFactory::Instance(device_.type()) @@ -334,9 +334,9 @@ void LoRARowParallelLinear::InitLoRAWeights() { } // lora_B: [out_features, rank] - parameters_[kParamLoraBName] - = std::make_shared(std::vector{out_features_, config_.rank}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamLoraBName, std::make_shared(std::vector{out_features_, config_.rank}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); init::Zeros(parameters_[kParamLoraBName]); } diff --git a/infini_train/src/nn/modules/container.cc b/infini_train/src/nn/modules/container.cc index 33707b636..0e6c1c0c8 100644 --- a/infini_train/src/nn/modules/container.cc +++ b/infini_train/src/nn/modules/container.cc @@ -10,7 +10,7 @@ namespace infini_train::nn { Sequential::Sequential(std::vector> &&layers) : CloneableModule(kType) { int idx = 0; for (auto &layer : layers) { - modules_[std::to_string(idx)] = std::move(layer); + RegisterModule(std::to_string(idx), std::move(layer)); ++idx; } } @@ -21,8 +21,8 @@ std::vector> Sequential::Forward(const std::vector> modules) : CloneableModule(kType) { - for (auto &[name, layer] : modules) { modules_[name] = std::move(layer); } +ModuleDict::ModuleDict(std::vector modules) : CloneableModule(kType) { + for (auto &[name, layer] : modules) { RegisterModule(name, std::move(layer)); } } std::vector> ModuleDict::Forward(const std::vector> &input_tensors) { @@ -33,7 +33,7 @@ ModuleList::ModuleList(std::vector> &&layers) : CloneableModule(kType), module_list_(std::move(layers)) { int idx = 0; for (auto &layer : module_list_) { - modules_[std::to_string(idx)] = layer; + RegisterModule(std::to_string(idx), layer); ++idx; } } diff --git a/infini_train/src/nn/modules/linear.cc b/infini_train/src/nn/modules/linear.cc index 7b93fa944..b1f40b615 100644 --- a/infini_train/src/nn/modules/linear.cc +++ b/infini_train/src/nn/modules/linear.cc @@ -14,12 +14,13 @@ Linear::Linear(int64_t in_features, int64_t out_features, bool bias, Device devi : CloneableModule(kType), bias_(bias) { device_ = device; - parameters_[kParamWeightName] - = std::make_shared(std::vector{out_features, in_features}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, std::make_shared(std::vector{out_features, in_features}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (bias) { - parameters_[kParamBiasName] - = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device_)->RequiresGrad(); + RegisterParameter( + kParamBiasName, + std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device_)->RequiresGrad()); } ResetParameters(); } diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc index 9475d49fe..ceefb6ddb 100644 --- a/infini_train/src/nn/modules/module.cc +++ b/infini_train/src/nn/modules/module.cc @@ -27,9 +27,59 @@ Module::Module(const std::string &type) : type_(type), device_(Device()) {} const std::string &Module::type() const { return type_; } -std::vector> Module::Parameters() const { - const auto &named_parameters = NamedParameters(); +namespace { +void CheckRegistrationName(const std::string &name) { + CHECK(!name.empty()) << "Module registration name cannot be empty."; + CHECK(name.find('.') == std::string::npos) << "Module registration name cannot contain '.', got: " << name; +} +} // namespace + +std::shared_ptr Module::RegisterParameter(const std::string &name, std::shared_ptr parameter) { + CheckRegistrationName(name); + CHECK(!modules_.contains(name) && !buffers_.contains(name)) + << "Cannot register parameter '" << name << "': the name is already used by another registry."; + CHECK(!parameter || parameter->is_leaf()) + << "Cannot register parameter '" << name << "': parameters must be leaf tensors."; + + if (!parameters_.contains(name)) { + parameter_order_.push_back(name); + } + parameters_[name] = std::move(parameter); + return parameters_.at(name); +} + +std::shared_ptr Module::RegisterBuffer(const std::string &name, std::shared_ptr buffer, + bool persistent) { + CheckRegistrationName(name); + CHECK(!modules_.contains(name) && !parameters_.contains(name)) + << "Cannot register buffer '" << name << "': the name is already used by another registry."; + + if (!buffers_.contains(name)) { + buffer_order_.push_back(name); + } + buffers_[name] = std::move(buffer); + if (persistent) { + non_persistent_buffers_.erase(name); + } else { + non_persistent_buffers_.insert(name); + } + return buffers_.at(name); +} + +std::shared_ptr Module::RegisterModule(const std::string &name, std::shared_ptr module) { + CheckRegistrationName(name); + CHECK(!parameters_.contains(name) && !buffers_.contains(name)) + << "Cannot register module '" << name << "': the name is already used by another registry."; + + if (!modules_.contains(name)) { + module_order_.push_back(name); + } + modules_[name] = std::move(module); + return modules_.at(name); +} +std::vector> Module::Parameters(bool recurse) const { + const auto &named_parameters = NamedParameters(/*prefix=*/"", recurse); std::vector> params; params.reserve(named_parameters.size()); @@ -53,19 +103,11 @@ Module::NamedParameters(const std::string &prefix, bool recurse, bool remove_dup } for (const auto &[module_prefix, module] : named_modules) { - std::vector>> local_parameters; - local_parameters.reserve(module->parameters_.size()); - - for (const auto &[name, parameter] : module->parameters_) { - if (parameter != nullptr) { - local_parameters.emplace_back(name, parameter); + for (const auto &name : module->parameter_order_) { + const auto ¶meter = module->parameters_.at(name); + if (!parameter) { + continue; } - } - - std::sort(local_parameters.begin(), local_parameters.end(), - [](const auto &lhs, const auto &rhs) { return lhs.first < rhs.first; }); - - for (const auto &[name, parameter] : local_parameters) { if (remove_duplicate && !visited_parameters.insert(parameter.get()).second) { continue; } @@ -93,8 +135,17 @@ const std::shared_ptr &Module::parameter(const std::string &name) const std::vector> Module::Buffers() const { std::vector> buffers; - for (auto &[_, buffer] : buffers_) { buffers.push_back(buffer); } - for (auto &[_, module] : modules_) { + for (const auto &name : buffer_order_) { + const auto &buffer = buffers_.at(name); + if (buffer) { + buffers.push_back(buffer); + } + } + for (const auto &name : module_order_) { + const auto &module = modules_.at(name); + if (!module) { + continue; + } for (auto &buffer : module->Buffers()) { buffers.push_back(buffer); } } return buffers; @@ -137,19 +188,12 @@ Module::NamedModules(std::unordered_set *memory, const std::string &pr // Emit self first (pre-order) named_modules.emplace_back(prefix, shared_from_this()); - // Collect children then sort by key for stable order - std::vector>> children; - children.reserve(modules_.size()); - for (const auto &[name, module] : modules_) { + // Recurse in registration order, matching torch.nn.Module.named_modules(). + for (const auto &name : module_order_) { + const auto &module = modules_.at(name); if (!module) { continue; } - children.emplace_back(name, module); - } - std::sort(children.begin(), children.end(), [](const auto &a, const auto &b) { return a.first < b.first; }); - - // Recurse in sorted order - for (const auto &[name, module] : children) { const auto submodule_prefix = (prefix.empty() ? "" : prefix + ".") + name; auto sub = module->NamedModules(memory, submodule_prefix, remove_duplicate); named_modules.insert(named_modules.end(), sub.begin(), sub.end()); @@ -167,9 +211,23 @@ const Module &Module::module(const std::string &name) const { std::unordered_map> Module::StateDict() const { std::unordered_map> state; - for (auto &[name, param] : parameters_) { state.emplace(name, param); } - for (auto &[name, buffer] : buffers_) { state.emplace(name, buffer); } - for (auto &[name, module] : modules_) { + for (const auto &name : parameter_order_) { + const auto ¶m = parameters_.at(name); + if (param) { + state.emplace(name, param); + } + } + for (const auto &name : buffer_order_) { + const auto &buffer = buffers_.at(name); + if (buffer && !non_persistent_buffers_.contains(name)) { + state.emplace(name, buffer); + } + } + for (const auto &name : module_order_) { + const auto &module = modules_.at(name); + if (!module) { + continue; + } if (name.starts_with("__pp")) { continue; } @@ -304,32 +362,52 @@ void Module::To(Device device) { std::unordered_map> new_parameters; std::unordered_map> new_buffers; - for (auto &[name, param] : parameters_) { - new_parameters.emplace(name, std::make_shared(param->To(device))); + for (const auto &name : parameter_order_) { + const auto ¶m = parameters_.at(name); + new_parameters.emplace(name, param ? std::make_shared(param->To(device)) : nullptr); + } + for (const auto &name : buffer_order_) { + const auto &buffer = buffers_.at(name); + new_buffers.emplace(name, buffer ? std::make_shared(buffer->To(device)) : nullptr); } - for (auto &[name, buffer] : buffers_) { new_buffers.emplace(name, std::make_shared(buffer->To(device))); } parameters_ = std::move(new_parameters); buffers_ = std::move(new_buffers); device_ = device; - for (auto &[_, module] : modules_) { module->To(device); } + for (const auto &name : module_order_) { + if (const auto &module = modules_.at(name); module) { + module->To(device); + } + } } void Module::To(DataType dtype) { std::unordered_map> new_parameters; std::unordered_map> new_buffers; - for (auto &[name, param] : parameters_) { - new_parameters.emplace(name, std::make_shared(param->To(dtype))); + for (const auto &name : parameter_order_) { + const auto ¶m = parameters_.at(name); + new_parameters.emplace(name, param ? std::make_shared(param->To(dtype)) : nullptr); + } + for (const auto &name : buffer_order_) { + const auto &buffer = buffers_.at(name); + new_buffers.emplace(name, buffer ? std::make_shared(buffer->To(dtype)) : nullptr); } - for (auto &[name, buffer] : buffers_) { new_buffers.emplace(name, std::make_shared(buffer->To(dtype))); } parameters_ = std::move(new_parameters); buffers_ = std::move(new_buffers); - for (auto &[_, layer] : modules_) { layer->To(dtype); } + for (const auto &name : module_order_) { + if (const auto &module = modules_.at(name); module) { + module->To(dtype); + } + } } void Module::Apply(std::function fn) { - for (auto &[_, module] : modules_) { module->Apply(fn); } + for (const auto &name : module_order_) { + if (const auto &module = modules_.at(name); module) { + module->Apply(fn); + } + } fn(this); } diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc index 388b04de5..a97d467e6 100644 --- a/infini_train/src/nn/modules/normalization.cc +++ b/infini_train/src/nn/modules/normalization.cc @@ -14,10 +14,10 @@ LayerNorm::LayerNorm(const std::vector &normalized_shape, float eps, De : CloneableModule(kType), eps_(eps) { device_ = device; - parameters_[kParamWeightName] - = std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad(); - parameters_[kParamBiasName] - = std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad()); + RegisterParameter(kParamBiasName, + std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad()); ResetParameters(); } @@ -33,8 +33,8 @@ void LayerNorm::ResetParameters() { } RMSNorm::RMSNorm(int64_t dim, float eps, Device device) : CloneableModule(kType), eps_(eps) { - parameters_[kParamWeightName] - = std::make_shared(std::vector{dim}, DataType::kFLOAT32, device)->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(std::vector{dim}, DataType::kFLOAT32, device)->RequiresGrad()); nn::init::Ones(parameters_[kParamWeightName]); } diff --git a/infini_train/src/nn/modules/sparse.cc b/infini_train/src/nn/modules/sparse.cc index 9314fe6d3..33fef1e8f 100644 --- a/infini_train/src/nn/modules/sparse.cc +++ b/infini_train/src/nn/modules/sparse.cc @@ -12,9 +12,9 @@ namespace infini_train::nn { Embedding::Embedding(int num_embeddings, int embedding_dim, Device device) : CloneableModule(kType) { device_ = device; - parameters_[kParamWeightName] - = std::make_shared(std::vector{num_embeddings, embedding_dim}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, std::make_shared(std::vector{num_embeddings, embedding_dim}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); ResetParameters(); } diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc index bc4c6bd46..469e75f5e 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc @@ -25,30 +25,30 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon int64_t qkv_dim = (config.n_head + 2 * n_kv_head_) * head_dim_; // qkv: ColumnParallel (do not gather output) - modules_[kCAttnLayerName] = std::make_shared( - /*in_features=*/n_embd_, - /*out_features=*/qkv_dim, - /*bias=*/config_.add_bias_linear, - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kCAttnLayerName, std::make_shared( + /*in_features=*/n_embd_, + /*out_features=*/qkv_dim, + /*bias=*/config_.add_bias_linear, + /*gather_output=*/false, + /*input_is_parallel=*/false, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled())); // proj: RowParallel (input is parallel and output is full) - modules_[kCProjLayerName] = std::make_shared( - /*in_features=*/n_embd_, - /*out_features=*/n_embd_, - /*bias=*/config_.add_bias_linear, - /*reduce_output=*/true, - /*input_is_parallel=*/true, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kCProjLayerName, std::make_shared( + /*in_features=*/n_embd_, + /*out_features=*/n_embd_, + /*bias=*/config_.add_bias_linear, + /*reduce_output=*/true, + /*input_is_parallel=*/true, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled())); // FIXME(zbl): Decouple causal-mask ownership from position embedding. For now, only learned-absolute models use // this precomputed buffer; RoPE callers provide a runtime-sized mask. if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { - buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) - ->View({1, 1, config_.block_size, config_.block_size}); + RegisterBuffer(kParamBiasName, function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) + ->View({1, 1, config_.block_size, config_.block_size})); } } diff --git a/infini_train/src/nn/modules/transformer/mlp.cc b/infini_train/src/nn/modules/transformer/mlp.cc index 115d77e6f..5f926ab99 100644 --- a/infini_train/src/nn/modules/transformer/mlp.cc +++ b/infini_train/src/nn/modules/transformer/mlp.cc @@ -47,40 +47,40 @@ MLP::MLP(const TransformerConfig &config) : CloneableModule(kType) { CHECK_GT(ffn_hidden, 0); // c_fc: ColumnParallel (input full, output parallel) - modules_[kCFcLayerName] = std::make_shared( - /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, - /*bias=*/config.add_bias_linear, - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kCFcLayerName, std::make_shared( + /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, + /*bias=*/config.add_bias_linear, + /*gather_output=*/false, + /*input_is_parallel=*/false, + /*skip_bias_add=*/false, + /*sequence_parallel=*/parallel::global::GetSequenceParallelEnabled())); // For SwiGLU, add second projection if (config.activation_type == MLPType::kSwiGLU) { - modules_[kCFc2LayerName] = std::make_shared( - /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, - /*bias=*/config.add_bias_linear, - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kCFc2LayerName, std::make_shared( + /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, + /*bias=*/config.add_bias_linear, + /*gather_output=*/false, + /*input_is_parallel=*/false, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled())); } // Activation: check for GELU or SwiGLU if (config.activation_type == MLPType::kGELU) { - modules_[kGeluLayerName] = std::make_shared(); + RegisterModule(kGeluLayerName, std::make_shared()); } else if (config.activation_type == MLPType::kSwiGLU) { - modules_[kSiluLayerName] = std::make_shared(); + RegisterModule(kSiluLayerName, std::make_shared()); } // c_proj: RowParallel (input parallel, output full) - modules_[kCProjLayerName] = std::make_shared( - /*in_features=*/ffn_hidden, /*out_features=*/config.n_embd, - /*bias=*/config.add_bias_linear, - /*reduce_output=*/true, - /*input_is_parallel=*/true, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kCProjLayerName, std::make_shared( + /*in_features=*/ffn_hidden, /*out_features=*/config.n_embd, + /*bias=*/config.add_bias_linear, + /*reduce_output=*/true, + /*input_is_parallel=*/true, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled())); } std::vector> diff --git a/infini_train/src/nn/modules/transformer/moe/experts.cc b/infini_train/src/nn/modules/transformer/moe/experts.cc index ede7bc999..047e6da7d 100644 --- a/infini_train/src/nn/modules/transformer/moe/experts.cc +++ b/infini_train/src/nn/modules/transformer/moe/experts.cc @@ -26,7 +26,7 @@ SequentialMLP::SequentialMLP(const TransformerConfig &config) : CloneableModule( CHECK_GT(num_local_experts_, 0); for (int64_t expert_idx = 0; expert_idx < num_local_experts_; ++expert_idx) { - modules_[std::string(kExpertNamePrefix) + std::to_string(expert_idx)] = std::make_shared(config_); + RegisterModule(std::string(kExpertNamePrefix) + std::to_string(expert_idx), std::make_shared(config_)); } } diff --git a/infini_train/src/nn/modules/transformer/moe/moe_layer.cc b/infini_train/src/nn/modules/transformer/moe/moe_layer.cc index 1e15fe814..12536fbf3 100644 --- a/infini_train/src/nn/modules/transformer/moe/moe_layer.cc +++ b/infini_train/src/nn/modules/transformer/moe/moe_layer.cc @@ -18,8 +18,8 @@ MoELayer::MoELayer(const TransformerConfig &config) : CloneableModule(kType), co CHECK(moe_config.token_dispatcher_type == MoEConfig::TokenDispatcherType::kAllGather) << "Current InfiniTrain MoE implementation supports AllGather dispatcher only"; - modules_[kRouterLayerName] = std::make_shared(config_); - modules_[kExpertsLayerName] = std::make_shared(config_); + RegisterModule(kRouterLayerName, std::make_shared(config_)); + RegisterModule(kExpertsLayerName, std::make_shared(config_)); } std::vector> MoELayer::Forward(const std::vector> &input_tensors) { diff --git a/infini_train/src/nn/modules/transformer/moe/router.cc b/infini_train/src/nn/modules/transformer/moe/router.cc index 252086846..4543e0a07 100644 --- a/infini_train/src/nn/modules/transformer/moe/router.cc +++ b/infini_train/src/nn/modules/transformer/moe/router.cc @@ -20,16 +20,16 @@ TopKRouter::TopKRouter(const TransformerConfig &config) : CloneableModule(kType) CHECK_GT(moe_config.num_experts, 0); CHECK_GT(moe_config.router_topk, 0); CHECK_LE(moe_config.router_topk, moe_config.num_experts); - parameters_[kParamWeightName] - = std::make_shared(std::vector{moe_config.num_experts, config_.n_embd}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(std::vector{moe_config.num_experts, config_.n_embd}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); init::KaimingUniform(parameters_[kParamWeightName]); if (config_.add_bias_linear) { - parameters_[kParamBiasName] - = std::make_shared(std::vector{moe_config.num_experts}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamBiasName, std::make_shared(std::vector{moe_config.num_experts}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); parameters_[kParamBiasName]->Fill(0.0f); } } diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc index 99a739d2d..ca14fe23e 100644 --- a/infini_train/src/nn/modules/transformer/transformer.cc +++ b/infini_train/src/nn/modules/transformer/transformer.cc @@ -26,12 +26,13 @@ namespace infini_train::nn { TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config) : CloneableModule(kType), config_(config) { - modules_[kWTELayerName] = std::make_shared( - config_.vocab_size, config_.n_embd, parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kWTELayerName, + std::make_shared(config_.vocab_size, config_.n_embd, + parallel::global::GetSequenceParallelEnabled())); // Only learned absolute position embedding uses a trainable WPE table. if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { - modules_[kWPELayerName] = std::make_shared(config_.block_size, config_.n_embd); + RegisterModule(kWPELayerName, std::make_shared(config_.block_size, config_.n_embd)); } else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) { LOG(FATAL) << "Unsupported position embedding type"; } @@ -77,22 +78,32 @@ std::vector> TransformerFirstStage::Forward(const std::v TransformerLayer::TransformerLayer(const nn::TransformerConfig &config) : CloneableModule(kType) { switch (config.norm_type) { case NormType::kLayerNorm: - modules_[kLn1LayerName] = std::make_shared(std::vector{config.n_embd}); - modules_[kLn2LayerName] = std::make_shared(std::vector{config.n_embd}); + RegisterModule(kLn1LayerName, std::make_shared(std::vector{config.n_embd})); break; case NormType::kRMSNorm: - modules_[kLn1LayerName] = std::make_shared(config.n_embd, config.norm_eps); - modules_[kLn2LayerName] = std::make_shared(config.n_embd, config.norm_eps); + RegisterModule(kLn1LayerName, std::make_shared(config.n_embd, config.norm_eps)); + break; + default: + LOG(FATAL) << "Unsupported norm type"; + } + + RegisterModule(kAttnLayerName, std::make_shared(config)); + + switch (config.norm_type) { + case NormType::kLayerNorm: + RegisterModule(kLn2LayerName, std::make_shared(std::vector{config.n_embd})); + break; + case NormType::kRMSNorm: + RegisterModule(kLn2LayerName, std::make_shared(config.n_embd, config.norm_eps)); break; default: LOG(FATAL) << "Unsupported norm type"; } - modules_[kAttnLayerName] = std::make_shared(config); if (config.ffn_type == FFNType::kMoE) { - modules_[kMlpLayerName] = std::make_shared(config); + RegisterModule(kMlpLayerName, std::make_shared(config)); } else { - modules_[kMlpLayerName] = std::make_shared(config); + RegisterModule(kMlpLayerName, std::make_shared(config)); } } @@ -129,7 +140,7 @@ TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_la auto layer = std::make_shared(config); h.push_back(layer); } - modules_[kHLayerName] = std::make_shared(std::move(h)); + RegisterModule(kHLayerName, std::make_shared(std::move(h))); } std::vector> TransformerChunk::Forward(const std::vector> &x) { @@ -141,10 +152,10 @@ std::vector> TransformerChunk::Forward(const std::vector const auto device = x1->GetDevice(); // Init freqs_cis on device only once - if (buffers_[kFreqsCisName] == nullptr) { + if (!buffers_.contains(kFreqsCisName) || !buffers_.at(kFreqsCisName)) { int64_t head_dim = config_.n_embd / config_.n_head; - buffers_[kFreqsCisName] = PrecomputeFreqsCis(head_dim, config_.block_size * 2, config_.rope_theta, - config_.use_scaled_rope, device); + RegisterBuffer(kFreqsCisName, PrecomputeFreqsCis(head_dim, config_.block_size * 2, config_.rope_theta, + config_.use_scaled_rope, device)); } const auto t = x1->Dims()[1] * nn::parallel::global::GetSequenceParallelSize(); // full_seq_len @@ -176,23 +187,23 @@ std::vector> TransformerChunk::Forward(const std::vector TransformerLastStage::TransformerLastStage(const TransformerConfig &config) : CloneableModule(kType), config_(config) { switch (config.norm_type) { case NormType::kLayerNorm: - modules_[kLnFLayerName] = std::make_shared(std::vector{config_.n_embd}); + RegisterModule(kLnFLayerName, std::make_shared(std::vector{config_.n_embd})); break; case NormType::kRMSNorm: - modules_[kLnFLayerName] = std::make_shared(config.n_embd, config.norm_eps); + RegisterModule(kLnFLayerName, std::make_shared(config.n_embd, config.norm_eps)); break; default: LOG(FATAL) << "Unsupported norm type"; } // NOTE(zbl): weight-tying is possible but torch script did not do so - modules_[kLMHeadLayerName] = std::make_shared( - /*in_features=*/config_.n_embd, /*out_features=*/config_.vocab_size, - /*bias=*/config_.add_bias_lm_head, - // NOTE(zbl): each rank would get sharded [B, T, V_local] as logits - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + RegisterModule(kLMHeadLayerName, std::make_shared( + /*in_features=*/config_.n_embd, /*out_features=*/config_.vocab_size, + /*bias=*/config_.add_bias_lm_head, + // NOTE(zbl): each rank would get sharded [B, T, V_local] as logits + /*gather_output=*/false, + /*input_is_parallel=*/false, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled())); } std::vector> TransformerLastStage::Forward(const std::vector> &x) { @@ -216,14 +227,14 @@ TransformerModel::TransformerModel(const TransformerConfig config) // Here we introduce padding by default, might need modify Tokenizer correspondingly later CHECK_EQ(config.vocab_size % tp_world_size, 0) << "Vocab size should be divisible by TP world size"; - std::unordered_map> transformer; + std::vector transformer; if (stage_info_.is_first_stage) { - modules_[kPPFirstStageName] = std::make_shared(config_); - transformer[TransformerFirstStage::kWTELayerName] - = modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName); + RegisterModule(kPPFirstStageName, std::make_shared(config_)); + transformer.emplace_back(TransformerFirstStage::kWTELayerName, + modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName)); if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { - transformer[TransformerFirstStage::kWPELayerName] - = modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWPELayerName); + transformer.emplace_back(TransformerFirstStage::kWPELayerName, + modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWPELayerName)); } } @@ -241,20 +252,20 @@ TransformerModel::TransformerModel(const TransformerConfig config) for (int idx = 0; idx < layer_size; ++idx) { h.push_back(chunk->mutable_module(TransformerChunk::kHLayerName)->mutable_module(std::to_string(idx))); } - modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)] = std::move(chunk); + RegisterModule(kPPChunkNamePrefix + std::to_string(chunk_idx), std::move(chunk)); ++chunk_idx; } - transformer[TransformerChunk::kHLayerName] = std::make_shared(std::move(h)); + transformer.emplace_back(TransformerChunk::kHLayerName, std::make_shared(std::move(h))); } if (stage_info_.is_last_stage) { - modules_[kPPLastStageName] = std::make_shared(config_); - transformer[TransformerLastStage::kLnFLayerName] - = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLnFLayerName); - modules_[TransformerLastStage::kLMHeadLayerName] - = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLMHeadLayerName); + RegisterModule(kPPLastStageName, std::make_shared(config_)); + transformer.emplace_back(TransformerLastStage::kLnFLayerName, + modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLnFLayerName)); + RegisterModule(TransformerLastStage::kLMHeadLayerName, + modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLMHeadLayerName)); } - modules_[kTransformerModelName] = std::make_shared(std::move(transformer)); + RegisterModule(kTransformerModelName, std::make_shared(std::move(transformer))); // FIXME(jym): Assigning the parameter values of wte to LMHead, which is not real tying operation // TODO: Implement real GPT-2 weight tying: make lm_head.weight share the exact same Parameter/Tensor (same diff --git a/infini_train/src/nn/parallel/data_parallel.cc b/infini_train/src/nn/parallel/data_parallel.cc index cf16f6c39..b773c6f98 100644 --- a/infini_train/src/nn/parallel/data_parallel.cc +++ b/infini_train/src/nn/parallel/data_parallel.cc @@ -71,7 +71,7 @@ DataParallel::DataParallel(const std::shared_ptr &module, int dim, Devic module->To(src_device_); - modules_[kModuleName] = std::move(module); + RegisterModule(kModuleName, std::move(module)); // TODO(dcj): implement check_balance for cuda devices later. diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index 57460ba09..4ea7d06f1 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -45,11 +46,14 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr mod CHECK_EQ(buffer->GetDevice().index(), global::GetDeviceIndex(rank.thread_rank())) << "All buffers must be on the same device as the module"; } - modules_[kModuleName] = std::move(module); + RegisterModule(kModuleName, std::move(module)); if (ddp_config.zero_stage >= 1) { BuildParamAndGradBuffers(); RegisterBackwardHooks(); + if (ddp_config_.overlap_param_gather) { + RegisterForwardPreHooks(); + } } else if (ddp_config.gradient_bucketing_enabled) { // Bucket Assignment auto params = modules_[kModuleName]->Parameters(); @@ -63,6 +67,38 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr mod } } +void DistributedDataParallel::RegisterForwardPreHooks() { + auto &model = modules_.at(kModuleName); + for (auto &module : model->modules()) { + std::unordered_set required_groups; + for (const auto ¶m : module->Parameters(/*recurse=*/false)) { + auto it = param_to_bucket_group_.find(param.get()); + if (it != param_to_bucket_group_.end()) { + required_groups.insert(it->second.get()); + } + } + if (required_groups.empty()) { + continue; + } + + std::vector> ordered_groups; + for (auto it = bucket_groups_.rbegin(); it != bucket_groups_.rend(); ++it) { + if (required_groups.contains(it->get())) { + ordered_groups.emplace_back(*it); + } + } + + module->RegisterForwardPreHook( + [groups = std::move(ordered_groups)](nn::Module *, const std::vector> &) { + for (const auto &weak_group : groups) { + if (auto group = weak_group.lock()) { + group->FinishParamSync(); + } + } + }); + } +} + void DistributedDataParallel::BuildParamAndGradBuffers() { // (param_dtype, grad_dtype) using DTypePair = std::pair; diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 523bcf2d7..b10d2e9ac 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -61,6 +61,9 @@ void DistributedOptimizer::InitializeModelChunks(const std::vectorparam_grad_buffers().end()); bucket_groups_.insert(bucket_groups_.end(), ddp_chunk->bucket_groups().begin(), ddp_chunk->bucket_groups().end()); + if (!ddp_chunk->bucket_groups().empty()) { + first_param_sync_bucket_groups_.push_back(ddp_chunk->bucket_groups().back()); + } } } @@ -174,14 +177,21 @@ void DistributedOptimizer::Step() { // 1. Ensure grads are synced FinishGradSync(); + // Parameter gathers from the previous update must finish before the optimizer writes into the same buffers. + for (auto &group : bucket_groups_) { group->PrepareParamSyncForNextStep(); } + // 2. Base optimizer step on owned param pieces CHECK(base_optimizer_) << "DistributedOptimizer: base optimizer is null."; base_optimizer_->Step(); - // 3. Gather updated param shards back to full params - StartParamSync(/*force_sync=*/false); - // TODO(zbl): Delay sync call until param is actually used in next step - FinishParamSync(/*skip_next_bucket_dispatch=*/true); + // 3. Publish updated parameter shards. With overlap enabled, only launch the first gather in each model + // chunk. Forward pre-hooks wait on it at first use and dispatch subsequent bucket gathers. + CHECK(!bucket_groups_.empty()); + if (bucket_groups_.front()->config().overlap_param_gather) { + for (auto &group : first_param_sync_bucket_groups_) { group->StartParamSync(); } + } else { + StartParamSync(/*force_sync=*/false); + } } std::unordered_map> DistributedOptimizer::StateDict() const { diff --git a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc index ab3a80024..c00046192 100644 --- a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc +++ b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc @@ -130,16 +130,17 @@ ParamAndGradBucketGroup::ParamAndGradBucketGroup(const std::vectorgrad_dtype(), param->GetDevice()); } } + + // Every rank starts with a complete parameter replica. The first gather is needed only after an optimizer update. + param_gather_dispatched_ = ddp_config_.zero_stage >= 1; } void ParamAndGradBucketGroup::Reset() { params_with_grad_.clear(); grad_reduce_work_list_.clear(); grad_reduce_bucket_indices_.clear(); - param_gather_work_list_.clear(); is_last_microbatch_ = true; grad_reduce_dispatched_ = false; - param_gather_dispatched_ = false; if (ddp_config_.zero_stage >= 2) { std::fill(temp_full_grad_buffer_list_.begin(), temp_full_grad_buffer_list_.end(), nullptr); @@ -353,9 +354,13 @@ void ParamAndGradBucketGroup::StartParamSync(bool force_sync) { // force synchronous collective regardless of other settings for (auto work : param_gather_work_list_) { work->WaitNonBlocking(); } param_gather_work_list_.clear(); - return; + if (param_gather_dispatched_) { + return; + } } else { CHECK(param_gather_work_list_.empty()); + CHECK(!param_gather_dispatched_) + << "ParamAndGradBucketGroup: parameter all-gather has already been dispatched for this update."; } auto async_op = ddp_config_.overlap_param_gather && (!force_sync); @@ -371,7 +376,10 @@ void ParamAndGradBucketGroup::StartParamSync(bool force_sync) { param_buffer_shard_list_[i] = ShardBuffer(param_buffer, collective_pg_size_); } auto local_data_view = param_buffer_shard_list_[i][rank_in_collective_pg_]; - param_gather_work_list_.push_back(collective_pg_->AllGather(param_buffer, local_data_view, async_op)); + auto work = collective_pg_->AllGather(param_buffer, local_data_view, async_op); + if (work) { + param_gather_work_list_.push_back(std::move(work)); + } } param_gather_dispatched_ = true; @@ -389,7 +397,6 @@ void ParamAndGradBucketGroup::FinishParamSync(bool skip_next_bucket_dispatch) { if (!param_gather_work_list_.empty()) { for (auto work : param_gather_work_list_) { work->WaitNonBlocking(); } param_gather_work_list_.clear(); - param_gather_dispatched_ = false; if (next_param_gather_bucket_group_ && !skip_next_bucket_dispatch) { if (next_param_gather_bucket_group_->param_gather_dispatched_) { @@ -404,6 +411,12 @@ void ParamAndGradBucketGroup::FinishParamSync(bool skip_next_bucket_dispatch) { } } +void ParamAndGradBucketGroup::PrepareParamSyncForNextStep() { + for (auto &work : param_gather_work_list_) { work->WaitNonBlocking(); } + param_gather_work_list_.clear(); + param_gather_dispatched_ = false; +} + void ParamAndGradBucketGroup::SetNextParamGatherBucketGroup(std::shared_ptr next_group) { next_param_gather_bucket_group_ = next_group; } diff --git a/infini_train/src/nn/parallel/parallel_functional.cc b/infini_train/src/nn/parallel/parallel_functional.cc index ffd218d71..b8d78c127 100644 --- a/infini_train/src/nn/parallel/parallel_functional.cc +++ b/infini_train/src/nn/parallel/parallel_functional.cc @@ -114,25 +114,29 @@ std::vector> Replicate(const std::shared_ptr &ne for (int idx = 0; idx < modules.size(); ++idx) { auto &module = modules[idx]; - for (auto &[name, child] : module->modules_) { - const auto module_idx = module_indices[child.get()]; + for (const auto &name : module->module_order_) { + const auto &child = module->modules_.at(name); for (int replica_idx = 0; replica_idx < num_replicas; ++replica_idx) { auto &replica = module_copies[replica_idx][idx]; - replica->modules_[name] = module_copies[replica_idx][module_idx]; + replica->RegisterModule(name, + child ? module_copies[replica_idx][module_indices.at(child.get())] : nullptr); } } - for (auto &[name, param] : module->parameters_) { - const auto param_idx = param_indices[param.get()]; + for (const auto &name : module->parameter_order_) { + const auto ¶m = module->parameters_.at(name); for (int replica_idx = 0; replica_idx < num_replicas; ++replica_idx) { auto &replica = module_copies[replica_idx][idx]; - replica->parameters_[name] = param_copies[replica_idx][param_idx]; + replica->RegisterParameter(name, + param ? param_copies[replica_idx][param_indices.at(param.get())] : nullptr); } } - for (auto &[name, buffer] : module->buffers_) { - const auto buffer_idx = buffer_indices[buffer.get()]; + for (const auto &name : module->buffer_order_) { + const auto &buffer = module->buffers_.at(name); + const bool persistent = !module->non_persistent_buffers_.contains(name); for (int replica_idx = 0; replica_idx < num_replicas; ++replica_idx) { auto &replica = module_copies[replica_idx][idx]; - replica->buffers_[name] = buffer_copies[replica_idx][buffer_idx]; + replica->RegisterBuffer( + name, buffer ? buffer_copies[replica_idx][buffer_indices.at(buffer.get())] : nullptr, persistent); } } } diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc index c0369cdeb..163756ce8 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -80,7 +80,7 @@ PipelineParallel::PipelineParallel(const std::shared_ptr module, int num const std::vector> &recv_shape, int pp_rank, Device device, int chunk_size) : num_stages_(num_stages), rank_(pp_rank) { - modules_[kModuleName] = std::move(module); + RegisterModule(kModuleName, std::move(module)); int stage_id = pp_rank; int stage_size = num_stages; diff --git a/infini_train/src/nn/parallel/tensor_parallel.cc b/infini_train/src/nn/parallel/tensor_parallel.cc index b16c526e6..2e6446f23 100644 --- a/infini_train/src/nn/parallel/tensor_parallel.cc +++ b/infini_train/src/nn/parallel/tensor_parallel.cc @@ -241,14 +241,14 @@ ColumnParallelLinear::ColumnParallelLinear(int64_t in_features, int64_t out_feat output_size_per_partition_ = out_features / tp_size; // init params shards on local rank - parameters_[kParamWeightName] - = std::make_shared(std::vector{output_size_per_partition_, in_features}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(std::vector{output_size_per_partition_, in_features}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (bias) { - parameters_[kParamBiasName] - = std::make_shared(std::vector{output_size_per_partition_}, DataType::kFLOAT32, device_) - ->RequiresGrad(); + RegisterParameter(kParamBiasName, std::make_shared(std::vector{output_size_per_partition_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); } LinearResetParameters(parameters_[kParamWeightName], bias ? parameters_[kParamBiasName] : nullptr); @@ -298,13 +298,14 @@ RowParallelLinear::RowParallelLinear(int64_t in_features, int64_t out_features, } // init params shards on local rank - parameters_[kParamWeightName] - = std::make_shared(std::vector{out_features, input_size_per_partition_}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(std::vector{out_features, input_size_per_partition_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); if (bias) { - parameters_[kParamBiasName] - = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device_)->RequiresGrad(); + RegisterParameter( + kParamBiasName, + std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device_)->RequiresGrad()); } LinearResetParameters(parameters_[kParamWeightName], bias ? parameters_[kParamBiasName] : nullptr); @@ -354,10 +355,10 @@ VocabParallelEmbedding::VocabParallelEmbedding(int64_t num_embeddings, int64_t e vocab_start_index_ = static_cast(tp_rank) * vocab_size_per_partition_; vocab_end_index_ = vocab_start_index_ + vocab_size_per_partition_; - parameters_[kParamWeightName] - = std::make_shared(std::vector{vocab_size_per_partition_, embedding_dim_}, DataType::kFLOAT32, - device_) - ->RequiresGrad(); + RegisterParameter(kParamWeightName, + std::make_shared(std::vector{vocab_size_per_partition_, embedding_dim_}, + DataType::kFLOAT32, device_) + ->RequiresGrad()); } std::vector> diff --git a/tests/hook/test_hook.cc b/tests/hook/test_hook.cc index 3e4ed6922..f319e9f35 100644 --- a/tests/hook/test_hook.cc +++ b/tests/hook/test_hook.cc @@ -75,4 +75,84 @@ TEST_P(HookTest, HookRemove) { EXPECT_EQ(hook3_count, 3); } +TEST_P(HookTest, ModuleRegistriesPreserveInsertionOrder) { + auto module = std::make_shared(); + auto first = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto second = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto replacement = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + + module->RegisterParameter("first", first); + module->RegisterParameter("second", second); + module->RegisterParameter("first", replacement); + module->RegisterParameter("deferred", nullptr); + + auto direct_params = module->Parameters(/*recurse=*/false); + ASSERT_EQ(direct_params.size(), 2); + EXPECT_EQ(direct_params[0], replacement); + EXPECT_EQ(direct_params[1], second); + + auto child_10 = std::make_shared(); + auto child_2 = std::make_shared(); + auto child_1 = std::make_shared(); + auto child_10_param = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto child_2_param = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto child_1_param = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + child_10->RegisterParameter("weight", child_10_param); + child_2->RegisterParameter("weight", child_2_param); + child_1->RegisterParameter("weight", child_1_param); + module->RegisterModule("10", child_10); + module->RegisterModule("2", child_2); + module->RegisterModule("1", child_1); + + auto modules = module->modules(); + ASSERT_EQ(modules.size(), 4); + EXPECT_EQ(modules[0], module); + EXPECT_EQ(modules[1], child_10); + EXPECT_EQ(modules[2], child_2); + EXPECT_EQ(modules[3], child_1); + + auto recursive_params = module->Parameters(); + ASSERT_EQ(recursive_params.size(), 5); + EXPECT_EQ(recursive_params[0], replacement); + EXPECT_EQ(recursive_params[1], second); + EXPECT_EQ(recursive_params[2], child_10_param); + EXPECT_EQ(recursive_params[3], child_2_param); + EXPECT_EQ(recursive_params[4], child_1_param); +} + +TEST_P(HookTest, ModuleBuffersMatchPersistentStateDictSemantics) { + auto module = std::make_shared(); + auto persistent = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto transient = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + + module->RegisterBuffer("persistent", persistent); + module->RegisterBuffer("transient", transient, false); + module->RegisterBuffer("deferred", nullptr); + + auto buffers = module->Buffers(); + ASSERT_EQ(buffers.size(), 2); + EXPECT_EQ(buffers[0], persistent); + EXPECT_EQ(buffers[1], transient); + + auto state = module->StateDict(); + EXPECT_TRUE(state.contains("persistent")); + EXPECT_FALSE(state.contains("transient")); + EXPECT_FALSE(state.contains("deferred")); + + module->RegisterBuffer("transient", transient, true); + EXPECT_TRUE(module->StateDict().contains("transient")); +} + +TEST_P(HookTest, ModuleRegistriesRejectInvalidAndConflictingNames) { + auto module = std::make_shared(); + auto tensor = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + + EXPECT_DEATH(module->RegisterParameter("", tensor), "cannot be empty"); + EXPECT_DEATH(module->RegisterParameter("nested.weight", tensor), "cannot contain"); + + module->RegisterParameter("weight", tensor); + EXPECT_DEATH(module->RegisterBuffer("weight", tensor), "already used"); + EXPECT_DEATH(module->RegisterModule("weight", std::make_shared()), "already used"); +} + INFINI_TRAIN_REGISTER_TEST(HookTest); diff --git a/tests/module/test_named_parameters.cc b/tests/module/test_named_parameters.cc index c1a4a6216..732775f9f 100644 --- a/tests/module/test_named_parameters.cc +++ b/tests/module/test_named_parameters.cc @@ -33,8 +33,8 @@ TEST_P(ModuleNamedParametersTest, SupportsRecursionAndSharedParameterDeduplicati const auto deduplicated = root->NamedParameters("model"); ASSERT_EQ(deduplicated.size(), 2); - EXPECT_EQ(deduplicated[0].first, "model.0.bias"); - EXPECT_EQ(deduplicated[1].first, "model.0.weight"); + EXPECT_EQ(deduplicated[0].first, "model.0.weight"); + EXPECT_EQ(deduplicated[1].first, "model.0.bias"); std::unordered_set tensors; for (const auto &[name, parameter] : deduplicated) { tensors.insert(parameter.get()); } EXPECT_TRUE(tensors.contains(shared->parameter(nn::Linear::kParamWeightName).get()));