fix: rename dp_* to ddp_*, remove unnecessary comments

Author: Chamberlain0w0

infini_train/include/nn/parallel/ddp/distributed_optimizer.h

@@ -21,8 +21,8 @@ class DistributedOptimizer final : public infini_train::Optimizer {
 public:
     DistributedOptimizer(OptimizerCreator base_optimizer_creator,
                          const std::vector<std::shared_ptr<Tensor>> &full_params,
-                         const std::vector<std::shared_ptr<Module>> &model_chunks, size_t dp_world_size,
-                         size_t dp_rank);
+                         const std::vector<std::shared_ptr<Module>> &model_chunks, size_t ddp_world_size,
+                         size_t ddp_rank);
 
     void Step() override;
 
@@ -43,8 +43,8 @@ class DistributedOptimizer final : public infini_train::Optimizer {
     std::vector<std::shared_ptr<ParamAndGradBucketGroup>> bucket_groups_;
 
     // DP info
-    size_t dp_world_size_;
-    size_t dp_rank_;
+    size_t ddp_world_size_;
+    size_t ddp_rank_;
 
     // shard params
     std::vector<std::shared_ptr<Tensor>> shard_params_;

infini_train/src/nn/parallel/ddp/distributed_optimizer.cc

@@ -9,10 +9,10 @@ namespace infini_train::nn::parallel {
 DistributedOptimizer::DistributedOptimizer(OptimizerCreator creator,
                                            const std::vector<std::shared_ptr<Tensor>> &full_params,
                                            const std::vector<std::shared_ptr<Module>> &model_chunks,
-                                           size_t dp_world_size, size_t dp_rank)
-    : Optimizer(full_params), dp_world_size_(dp_world_size), dp_rank_(dp_rank) {
+                                           size_t ddp_world_size, size_t ddp_rank)
+    : Optimizer(full_params), ddp_world_size_(ddp_world_size), ddp_rank_(ddp_rank) {
 
-    CHECK(dp_world_size_ > 1) << "DistributedOptimizer: dp_world_size must be greater than 1.";
+    CHECK(ddp_world_size_ > 1) << "DistributedOptimizer: ddp_world_size must be greater than 1.";
 
     for (size_t i = 0; i < model_chunks.size(); ++i) {
         auto ddp_chunk = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunks[i]);
@@ -43,9 +43,9 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads() {
             CHECK(bucket_param) << "DistributedOptimizer requires param buffer.";
             CHECK(bucket_grad) << "DistributedOptimizer requires grad buffer.";
 
-            CHECK_EQ(bucket_param->NumElements() % dp_world_size_, 0);
-            const size_t bucket_shard_numel = bucket_param->NumElements() / dp_world_size_;
-            const size_t bucket_shard_start = dp_rank_ * bucket_shard_numel;
+            CHECK_EQ(bucket_param->NumElements() % ddp_world_size_, 0);
+            const size_t bucket_shard_numel = bucket_param->NumElements() / ddp_world_size_;
+            const size_t bucket_shard_start = ddp_rank_ * bucket_shard_numel;
             const size_t bucket_shard_end = bucket_shard_start + bucket_shard_numel;
 
             // Iterate param in bucket, build each param(or param_shard) seperately

infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc

@@ -545,8 +545,6 @@ PartitionBuckets(const std::vector<std::shared_ptr<ParamAndGradBuffer>> &buffers
         auto ddp_world_size = buffers.front()->ddp_world_size();
 
         for (const auto &buffer : buffers) {
-            // TODO(zbl): override == for ddp config
-            // CHECK(buffer->ddp_config() == ddp_config) << "PartitionBuckets: buffers have different ddp_config.";
             CHECK(buffer->ddp_pg() == ddp_pg) << "PartitionBuckets: buffers have different ddp_pg.";
             CHECK(buffer->ddp_world_size() == ddp_world_size)
                 << "PartitionBuckets: buffers have different ddp_world_size.";