kilinchange · Oct 9, 2026 · Sep 17, 2026 · Sep 22, 2026 · Sep 24, 2026 · Oct 9, 2026
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
@@ -357,7 +357,6 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
-
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
     TrainingLRSchedulerConfig sched_config;
     sched_config.lr = static_cast<float>(FLAGS_learning_rate);
@@ -516,6 +515,7 @@ void Train(const nn::parallel::Rank &rank) {
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 
+            nn::parallel::FinalizeModelGrads({model});
             optimizer->Step();
             if (scheduler) {
                 scheduler->Step();

diff --git a/example/llama3/main.cc b/example/llama3/main.cc
@@ -347,7 +347,6 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
-
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
     TrainingLRSchedulerConfig sched_config;
     sched_config.lr = static_cast<float>(FLAGS_learning_rate);
@@ -503,6 +502,7 @@ void Train(const nn::parallel::Rank &rank) {
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 
+            nn::parallel::FinalizeModelGrads({model});
             optimizer->Step();
             if (scheduler) {
                 scheduler->Step();

diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc
@@ -22,6 +22,7 @@
 #include "infini_train/include/nn/modules/loss.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
 #include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/optimizer.h"
 #include "infini_train/include/tensor.h"
 #ifdef PROFILE_MODE
@@ -150,6 +151,7 @@ int main(int argc, char *argv[]) {
             auto loss_cpu = loss->To(Device());
             lossf += static_cast<const float *>(loss_cpu.DataPtr())[0];
         }
+        infini_train::nn::parallel::FinalizeModelGrads({model});
         optimizer->Step();
 
         device_impl->SynchronizeDevice(train_device);

diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc
@@ -502,6 +502,7 @@ void Train(const nn::parallel::Rank &rank) {
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 
+            nn::parallel::FinalizeModelGrads({model});
             optimizer->Step();
             if (scheduler) {
                 scheduler->Step();

diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h
@@ -34,6 +34,8 @@ class DistributedDataParallel : public nn::Module {
 
     std::unique_ptr<nn::NoSyncGuard> no_sync() override;
 
+    void FinishGradSync();
+
     DistributedDataParallelConfig ddp_config() const { return ddp_config_; }
 
     const std::vector<std::shared_ptr<ParamAndGradBuffer>> &param_grad_buffers() const { return param_grad_buffers_; }

diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
@@ -37,9 +37,6 @@ class DistributedOptimizer final : public infini_train::Optimizer {
 
     void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict) override;
 
-    void StartGradSync();
-    void FinishGradSync();
-
     void StartParamSync(bool force_sync = false);
     void FinishParamSync(bool skip_next_bucket_dispatch = false);
 

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
@@ -5,6 +5,7 @@
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "infini_train/include/datatype.h"
@@ -19,6 +20,9 @@ class Work;
 } // namespace infini_train
 
 namespace infini_train::nn::parallel {
+// Original parameter and its local gradient view, registered by DistributedOptimizer.
+using LocalGradShard = std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>;
+
 class ParamAndGradBucket {
 public:
     /**
@@ -123,6 +127,10 @@ class ParamAndGradBucketGroup {
     // ZeRO-2: Get a bucket's local grad shard buffer
     std::shared_ptr<Tensor> GetLocalGradShardBuffer(size_t bucket_idx) const;
 
+    const std::vector<LocalGradShard> &local_grad_shards() const;
+
+    void set_local_grad_shards(std::vector<LocalGradShard> shards);
+
     const DistributedDataParallelConfig &config() const { return ddp_config_; }
 
 private:
@@ -146,6 +154,8 @@ class ParamAndGradBucketGroup {
     // ZeRO-2: persistent grad shard buffers and temporary full grad buffers
     std::vector<std::shared_ptr<Tensor>> grad_shard_buffer_list_;
     std::vector<std::shared_ptr<Tensor>> temp_full_grad_buffer_list_;
+    // These views share optimizer grad storage and persist across iteration resets.
+    std::vector<LocalGradShard> local_grad_shards_;
 
     std::shared_ptr<ParamAndGradBucketGroup> next_param_gather_bucket_group_ = nullptr;
 

diff --git a/infini_train/include/nn/parallel/utils.h b/infini_train/include/nn/parallel/utils.h
@@ -7,6 +7,9 @@
 
 namespace infini_train {
 class Tensor;
+namespace nn {
+class Module;
+}
 } // namespace infini_train
 
 namespace infini_train::nn::parallel {
@@ -31,4 +34,6 @@ std::vector<std::shared_ptr<Tensor>> GatherFromSPRegionFunc(const std::shared_pt
 std::vector<std::shared_ptr<Tensor>> ScatterToTPRegionFunc(const std::shared_ptr<Tensor> &input);
 std::vector<std::shared_ptr<Tensor>> ReduceFromTPRegionFunc(const std::shared_ptr<Tensor> &input);
 std::vector<std::shared_ptr<Tensor>> CopyToTPRegionFunc(const std::shared_ptr<Tensor> &input);
+
+void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_chunks);
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h
@@ -80,6 +80,9 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     size_t NumElements() const;
     DataType Dtype() const;
 
+    void set_sequence_parallel(bool enabled);
+    bool sequence_parallel() const;
+
     std::shared_ptr<Tensor> Detach() const;
 
     void Fill(Scalar value);
@@ -242,6 +245,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
 private:
     std::shared_ptr<Tensor> grad_ = nullptr;
     bool requires_grad_ = false;
+    bool sequence_parallel_ = false;
     bool is_leaf_ = true;
     std::shared_ptr<autograd::Function> grad_fn_ = nullptr;
     int output_idx_ = 0;

diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc
@@ -7,6 +7,7 @@
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/functional.h"
 #include "infini_train/include/nn/init.h"
+#include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
 namespace infini_train::nn {
@@ -18,6 +19,10 @@ LayerNorm::LayerNorm(const std::vector<int64_t> &normalized_shape, float eps, De
         = std::make_shared<Tensor>(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad();
     parameters_[kParamBiasName]
         = std::make_shared<Tensor>(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad();
+    if (parallel::global::GetSequenceParallelEnabled()) {
+        parameters_[kParamWeightName]->set_sequence_parallel(true);
+        parameters_[kParamBiasName]->set_sequence_parallel(true);
+    }
     ResetParameters();
 }
 
@@ -35,6 +40,9 @@ void LayerNorm::ResetParameters() {
 RMSNorm::RMSNorm(int64_t dim, float eps, Device device) : CloneableModule(kType), eps_(eps) {
     parameters_[kParamWeightName]
         = std::make_shared<Tensor>(std::vector<int64_t>{dim}, DataType::kFLOAT32, device)->RequiresGrad();
+    if (parallel::global::GetSequenceParallelEnabled()) {
+        parameters_[kParamWeightName]->set_sequence_parallel(true);
+    }
     nn::init::Ones(parameters_[kParamWeightName]);
 }
 

diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc
@@ -26,6 +26,10 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon
     if (config_.qk_layernorm) {
         modules_[kQNormLayerName] = std::make_shared<nn::RMSNorm>(head_dim_, config_.norm_eps);
         modules_[kKNormLayerName] = std::make_shared<nn::RMSNorm>(head_dim_, config_.norm_eps);
+        // NOTE(zbl): In Qwen3-8B, Q/K norm sees full sequences and TP-local heads.
+        //            So we only need to finalize its replicated weights across TP regardless of SP.
+        modules_[kQNormLayerName]->parameter(RMSNorm::kParamWeightName)->set_sequence_parallel(false);
+        modules_[kKNormLayerName]->parameter(RMSNorm::kParamWeightName)->set_sequence_parallel(false);
     }
 
     int64_t qkv_dim = (config.n_head + 2 * n_kv_head_) * head_dim_;

diff --git a/infini_train/src/nn/modules/transformer/moe/router.cc b/infini_train/src/nn/modules/transformer/moe/router.cc
@@ -11,6 +11,7 @@
 #include "infini_train/include/nn/functional.h"
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/transformer/moe/moe_utils.h"
+#include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
 namespace infini_train::nn::moe {
@@ -24,12 +25,18 @@ TopKRouter::TopKRouter(const TransformerConfig &config) : CloneableModule(kType)
         = std::make_shared<Tensor>(std::vector<int64_t>{moe_config.num_experts, config_.n_embd}, DataType::kFLOAT32,
                                    device_)
               ->RequiresGrad();
+    if (parallel::global::GetSequenceParallelEnabled()) {
+        parameters_[kParamWeightName]->set_sequence_parallel(true);
+    }
     init::KaimingUniform(parameters_[kParamWeightName]);
 
     if (config_.add_bias_linear) {
         parameters_[kParamBiasName]
             = std::make_shared<Tensor>(std::vector<int64_t>{moe_config.num_experts}, DataType::kFLOAT32, device_)
                   ->RequiresGrad();
+        if (parallel::global::GetSequenceParallelEnabled()) {
+            parameters_[kParamBiasName]->set_sequence_parallel(true);
+        }
         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
@@ -32,6 +32,9 @@ TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config)
     // Only learned absolute position embedding uses a trainable WPE table.
     if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
         modules_[kWPELayerName] = std::make_shared<Embedding>(config_.block_size, config_.n_embd);
+        if (parallel::global::GetSequenceParallelEnabled()) {
+            modules_[kWPELayerName]->parameter(Embedding::kParamWeightName)->set_sequence_parallel(true);
+        }
     } else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) {
         LOG(FATAL) << "Unsupported position embedding type";
     }

diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc
@@ -231,6 +231,10 @@ std::unique_ptr<nn::NoSyncGuard> DistributedDataParallel::no_sync() {
     return std::make_unique<nn::NoSyncGuard>([this, previous] { SetIsLastMicrobatch(previous); });
 }
 
+void DistributedDataParallel::FinishGradSync() {
+    for (auto &group : bucket_groups_) { group->FinishGradSync(); }
+}
+
 void DistributedDataParallel::SetIsLastMicrobatch(bool is_last_microbatch) {
     is_last_microbatch_->store(is_last_microbatch, std::memory_order_relaxed);
     if (reducer_) {

diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
@@ -68,6 +68,7 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add
     size_t num_shard_params = 0;
 
     for (const auto &group : bucket_groups_) {
+        std::vector<LocalGradShard> local_grad_shards;
         const bool use_grad_shard = group->config().zero_stage >= 2;
         const auto &buckets = group->buckets();
         for (size_t bucket_idx = 0; bucket_idx < buckets.size(); ++bucket_idx) {
@@ -107,6 +108,7 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add
 
                 auto param_piece = std::make_shared<Tensor>(*bucket_param, param_piece_offset_bytes,
                                                             std::vector<int64_t>{static_cast<int64_t>(piece_numel)});
+                param_piece->set_sequence_parallel(param->sequence_parallel());
 
                 auto grad_piece = std::make_shared<Tensor>(*bucket_grad, grad_piece_offset_bytes,
                                                            std::vector<int64_t>{static_cast<int64_t>(piece_numel)});
@@ -115,24 +117,18 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add
                 // NOTE(zbl): Do not call `param->set_grad(grad_piece);` under ZeRO-2.
                 //            The base optimizer updates param_piece views only; original param->grad()
                 //            would be a partial flattened shard and does not represent the full parameter grad.
+                local_grad_shards.emplace_back(param, grad_piece);
                 add_shard_param(param, param_piece);
                 ++num_shard_params;
             }
         }
+        group->set_local_grad_shards(std::move(local_grad_shards));
     }
 
     CHECK_GT(num_shard_params, 0) << "DistributedOptimizer: this DP rank owns no param pieces. "
                                   << "Check bucket padding/divisibility and param bucketing order.";
 }
 
-void DistributedOptimizer::StartGradSync() {
-    for (auto &group : bucket_groups_) { group->StartGradSync(); }
-}
-
-void DistributedOptimizer::FinishGradSync() {
-    for (auto &group : bucket_groups_) { group->FinishGradSync(); }
-}
-
 void DistributedOptimizer::StartParamSync(bool force_sync) {
     for (auto &group : bucket_groups_) { group->StartParamSync(force_sync); }
 }
@@ -171,14 +167,10 @@ float DistributedOptimizer::learning_rate() const {
 }
 
 void DistributedOptimizer::Step() {
-    // 1. Ensure grads are synced
-    FinishGradSync();
-
-    // 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
+    // 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);

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
@@ -240,6 +240,12 @@ std::shared_ptr<Tensor> ParamAndGradBucketGroup::GetLocalGradShardBuffer(size_t
     return grad_shard_buffer_list_[bucket_idx];
 }
 
+const std::vector<LocalGradShard> &ParamAndGradBucketGroup::local_grad_shards() const { return local_grad_shards_; }
+
+void ParamAndGradBucketGroup::set_local_grad_shards(std::vector<LocalGradShard> shards) {
+    local_grad_shards_ = std::move(shards);
+}
+
 void ParamAndGradBucketGroup::StartGradSync() {
     if (!collective_pg_) {
         LOG(FATAL) << "ParamAndGradBucketGroup: StartGradSync() called with null collective_pg_.";

diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc
@@ -15,6 +15,7 @@
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_stage.h"
 #include "infini_train/include/nn/parallel/pp/send_recv.h"
+#include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/optimizer.h"
 #include "infini_train/include/tensor.h"
 
@@ -299,6 +300,7 @@ float PipelineSchedule::Step(std::shared_ptr<Tensor> input, std::shared_ptr<Tens
 
     float lossf = StepMicroBatches(micro_batches, target_mbs, loss_fn, dtype);
 
+    FinalizeModelGrads(stage_->chunks());
     optimizer->Step();
 
     return lossf;

diff --git a/infini_train/src/nn/parallel/tensor_parallel.cc b/infini_train/src/nn/parallel/tensor_parallel.cc
@@ -305,6 +305,9 @@ RowParallelLinear::RowParallelLinear(int64_t in_features, int64_t out_features,
     if (bias) {
         parameters_[kParamBiasName]
             = std::make_shared<Tensor>(std::vector<int64_t>{out_features}, DataType::kFLOAT32, device_)->RequiresGrad();
+        if (sequence_parallel_) {
+            parameters_[kParamBiasName]->set_sequence_parallel(true);
+        }
     }
 
     LinearResetParameters(parameters_[kParamWeightName], bias ? parameters_[kParamBiasName] : nullptr);