feat: add distopt AllGather - next step Forward overlap

example/mnist/net.cc

@@ -18,8 +18,8 @@ MNIST::MNIST() {
     std::vector<std::shared_ptr<nn::Module>> layers;
     layers.push_back(std::make_shared<nn::Linear>(784, 30));
     layers.push_back(std::make_shared<nn::Sigmoid>());
-    modules_["sequential"] = std::make_shared<nn::Sequential>(std::move(layers));
-    modules_["linear2"] = std::make_shared<nn::Linear>(30, 10);
+    RegisterModule("sequential", std::make_shared<nn::Sequential>(std::move(layers)));
+    RegisterModule("linear2", std::make_shared<nn::Linear>(30, 10));
 }
 
 std::vector<std::shared_ptr<infini_train::Tensor>>

infini_train/include/nn/modules/container.h

@@ -1,7 +1,8 @@
 #pragma once
 
 #include <memory>
-#include <unordered_map>
+#include <string>
+#include <utility>
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
@@ -23,8 +24,8 @@ class Sequential : public CloneableModule<Sequential> {
 class ModuleDict : public CloneableModule<ModuleDict> {
 public:
     static constexpr char kType[] = "ModuleDict";
-    // TODO(dcj): in torch, there is a dict with the order of insertion
-    explicit ModuleDict(std::unordered_map<std::string, std::shared_ptr<Module>> modules);
+    using Item = std::pair<std::string, std::shared_ptr<Module>>;
+    explicit ModuleDict(std::vector<Item> modules);
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 };

infini_train/include/nn/modules/module.h

@@ -47,10 +47,16 @@ class Module : public std::enable_shared_from_this<Module> {
 
     const std::string &type() const;
 
-    virtual std::vector<std::shared_ptr<Tensor>> 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<Tensor> RegisterParameter(const std::string &name, std::shared_ptr<Tensor> parameter);
+    std::shared_ptr<Tensor> RegisterBuffer(const std::string &name, std::shared_ptr<Tensor> buffer,
+                                           bool persistent = true);
+    std::shared_ptr<Module> RegisterModule(const std::string &name, std::shared_ptr<Module> module);
 
-    // InfiniTrain's NamedParameters returns results ordered by full parameter name.
-    // TODO: Align with PyTorch's ordering in the future.
+    virtual std::vector<std::shared_ptr<Tensor>> Parameters(bool recurse = true) const;
+
+    // Results follow parameter and module registration order, matching torch.nn.Module.named_parameters().
     std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
     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<Module> {
     std::unordered_map<std::string, std::shared_ptr<Module>> modules_;
     std::unordered_map<std::string, std::shared_ptr<Tensor>> parameters_;
     std::unordered_map<std::string, std::shared_ptr<Tensor>> buffers_;
+    std::vector<std::string> module_order_;
+    std::vector<std::string> parameter_order_;
+    std::vector<std::string> buffer_order_;
+    std::unordered_set<std::string> non_persistent_buffers_;
 
     std::vector<ModulePreHook> forward_pre_hooks_;
     std::vector<ModulePostHook> forward_post_hooks_;

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<Tensor> &param);
 
 private:

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;
 

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<std::shared_ptr<ParamAndGradBuffer>> param_grad_buffers_;
     std::vector<std::shared_ptr<ParamAndGradBucketGroup>> bucket_groups_;
+    std::vector<std::shared_ptr<ParamAndGradBucketGroup>> first_param_sync_bucket_groups_;
 
     // DP info
     size_t ddp_world_size_;

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<ParamAndGradBucketGroup> next_group);
 

infini_train/src/nn/lora/lora_linear.cc

@@ -35,11 +35,11 @@ LoRALinear::LoRALinear(std::shared_ptr<nn::Module> 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<nn::Module> 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<Tensor>(std::vector<int64_t>{config_.rank, in_features_}, DataType::kFLOAT32, device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraAName, std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{out_features_, config_.rank}, DataType::kFLOAT32, device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraBName, std::make_shared<Tensor>(std::vector<int64_t>{out_features_, config_.rank},
+                                                                DataType::kFLOAT32, device_)
+                                           ->RequiresGrad());
     init::Zeros(parameters_[kParamLoraBName]);
 }
 

infini_train/src/nn/lora/lora_parallel_linear.cc

@@ -34,15 +34,15 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptr<parallel::Col
     device_ = base_module->parameter(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();
     out_features_per_partition_ = weight_dims[0];
 
     // Transfer bias if exists
     if (base_module->has_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_ptr<parallel::Col
     device_ = base_module->parameter(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_ptr<parallel::Col
 
     // Transfer bias if exists
     if (base_module->has_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<parallel::Col
 void LoRAColumnParallelLinear::InitLoRAWeights() {
     // lora_A: [rank, in_features] - replicated across TP ranks
     // lora_B: [out_features_per_partition, rank] - sharded like base weight
-    parameters_[kParamLoraAName]
-        = std::make_shared<Tensor>(std::vector<int64_t>{config_.rank, in_features_}, DataType::kFLOAT32, device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraAName, std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{out_features_per_partition_, config_.rank}, DataType::kFLOAT32,
-                                   device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraBName,
+                      std::make_shared<Tensor>(std::vector<int64_t>{out_features_per_partition_, config_.rank},
+                                               DataType::kFLOAT32, device_)
+                          ->RequiresGrad());
     init::Zeros(parameters_[kParamLoraBName]);
 }
 
@@ -243,15 +243,15 @@ LoRARowParallelLinear::LoRARowParallelLinear(std::shared_ptr<parallel::RowParall
     device_ = base_module->parameter(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();
     in_features_per_partition_ = weight_dims[1];
 
     // Transfer bias if exists
     if (base_module->has_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_ptr<parallel::RowParall
     device_ = base_module->parameter(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_ptr<parallel::RowParall
 
     // Transfer bias if exists
     if (base_module->has_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<Tensor>(std::vector<int64_t>{config_.rank, in_features_per_partition_}, DataType::kFLOAT32,
-                                   device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraAName,
+                      std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{out_features_, config_.rank}, DataType::kFLOAT32, device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamLoraBName, std::make_shared<Tensor>(std::vector<int64_t>{out_features_, config_.rank},
+                                                                DataType::kFLOAT32, device_)
+                                           ->RequiresGrad());
     init::Zeros(parameters_[kParamLoraBName]);
 }
 

infini_train/src/nn/modules/container.cc

@@ -10,7 +10,7 @@ namespace infini_train::nn {
 Sequential::Sequential(std::vector<std::shared_ptr<Module>> &&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<std::shared_ptr<Tensor>> Sequential::Forward(const std::vector<std::
     return x;
 }
 
-ModuleDict::ModuleDict(std::unordered_map<std::string, std::shared_ptr<Module>> modules) : CloneableModule(kType) {
-    for (auto &[name, layer] : modules) { modules_[name] = std::move(layer); }
+ModuleDict::ModuleDict(std::vector<Item> modules) : CloneableModule(kType) {
+    for (auto &[name, layer] : modules) { RegisterModule(name, std::move(layer)); }
 }
 
 std::vector<std::shared_ptr<Tensor>> ModuleDict::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
@@ -33,7 +33,7 @@ ModuleList::ModuleList(std::vector<std::shared_ptr<Module>> &&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;
     }
 }

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<Tensor>(std::vector<int64_t>{out_features, in_features}, DataType::kFLOAT32, device_)
-              ->RequiresGrad();
+    RegisterParameter(kParamWeightName, std::make_shared<Tensor>(std::vector<int64_t>{out_features, in_features},
+                                                                 DataType::kFLOAT32, device_)
+                                            ->RequiresGrad());
     if (bias) {
-        parameters_[kParamBiasName]
-            = std::make_shared<Tensor>(std::vector<int64_t>{out_features}, DataType::kFLOAT32, device_)->RequiresGrad();
+        RegisterParameter(
+            kParamBiasName,
+            std::make_shared<Tensor>(std::vector<int64_t>{out_features}, DataType::kFLOAT32, device_)->RequiresGrad());
     }
     ResetParameters();
 }