From da9617504739db12f13cd99a46cef5d96c64f9a5 Mon Sep 17 00:00:00 2001
From: bolunz <chamberlain0w0@gmail.com>
Date: Thu, 17 Sep 2026 14:04:53 +0800
Subject: [PATCH 1/4] fix: fix param grad sync under SP

---
 example/gpt2/main.cc                          |  1 +
 example/llama3/main.cc                        |  1 +
 .../nn/parallel/ddp/distributed_optimizer.h   |  3 +-
 infini_train/include/nn/parallel/utils.h      |  2 ++
 infini_train/include/optimizer.h              | 16 +++++++--
 infini_train/include/tensor.h                 |  4 +++
 infini_train/src/nn/modules/normalization.cc  |  3 ++
 .../src/nn/modules/transformer/moe/router.cc  |  7 ++++
 .../src/nn/modules/transformer/transformer.cc |  3 ++
 .../nn/parallel/ddp/distributed_optimizer.cc  | 12 ++++---
 .../src/nn/parallel/tensor_parallel.cc        |  3 ++
 infini_train/src/nn/parallel/utils.cc         | 36 +++++++++++++++++++
 infini_train/src/optimizer.cc                 | 19 ++++++++--
 infini_train/src/tensor.cc                    |  3 ++
 14 files changed, 103 insertions(+), 10 deletions(-)

diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 7697ad213..2017de3ee 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -357,6 +357,7 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
+    optimizer->set_model_grad_finalizer(nn::parallel::FinalizeModelGrads);
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
     TrainingLRSchedulerConfig sched_config;
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index a9f405fa5..6558e2516 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -347,6 +347,7 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
+    optimizer->set_model_grad_finalizer(nn::parallel::FinalizeModelGrads);
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
     TrainingLRSchedulerConfig sched_config;
diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
index d7cea198e..3e8c7fea2 100644
--- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
+++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
@@ -29,7 +29,8 @@ class DistributedOptimizer final : public infini_train::Optimizer {
                          const std::vector<std::shared_ptr<Module>> &model_chunks, size_t ddp_world_size,
                          size_t ddp_rank);
 
-    void Step() override;
+    void FinalizeModelGrads() override;
+    void StepImpl() override;
 
     void ZeroGrad(bool set_to_none = true) override;
 
diff --git a/infini_train/include/nn/parallel/utils.h b/infini_train/include/nn/parallel/utils.h
index 8aa11856a..696b45949 100644
--- a/infini_train/include/nn/parallel/utils.h
+++ b/infini_train/include/nn/parallel/utils.h
@@ -31,4 +31,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<Tensor>> &params);
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h
index d85b1acea..362df76a8 100644
--- a/infini_train/include/optimizer.h
+++ b/infini_train/include/optimizer.h
@@ -18,6 +18,7 @@ using NamedParameter = std::pair<std::string, std::shared_ptr<Tensor>>;
 using NamedParameterList = std::vector<NamedParameter>;
 using OptimizerCreator = std::function<std::shared_ptr<Optimizer>(const std::vector<std::shared_ptr<Tensor>> &params)>;
 using OptimizerCreatorNamed = std::function<std::shared_ptr<Optimizer>(const NamedParameterList &named_params)>;
+using ModelGradFinalizer = std::function<void(const std::vector<std::shared_ptr<Tensor>> &params)>;
 
 class Optimizer {
 public:
@@ -27,7 +28,11 @@ class Optimizer {
 
     virtual void ZeroGrad(bool set_to_none = true);
 
-    virtual void Step() = 0;
+    virtual void Step() final;
+
+    void set_model_grad_finalizer(ModelGradFinalizer finalizer) { model_grad_finalizer_ = std::move(finalizer); }
+
+    const std::vector<std::shared_ptr<Tensor>> &parameters() const { return params_; }
 
     virtual std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const { return {}; };
 
@@ -44,11 +49,16 @@ class Optimizer {
     void set_initial_learning_rate(float lr);
 
 protected:
+    virtual void FinalizeModelGrads();
+    virtual void StepImpl() = 0;
+
     std::vector<std::shared_ptr<Tensor>> params_;
     std::vector<std::string> parameter_names_;
     float learning_rate_ = 0.0f;
     float initial_learning_rate_ = 0.0f;
     bool initial_lr_set_ = false;
+
+    ModelGradFinalizer model_grad_finalizer_;
 };
 
 namespace optimizers {
@@ -57,7 +67,7 @@ class SGD : public Optimizer {
     SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
     SGD(const NamedParameterList &named_params, float learning_rate);
 
-    void Step() override;
+    void StepImpl() override;
 
     static OptimizerCreator Create(float learning_rate);
     static OptimizerCreatorNamed CreateNamed(float learning_rate);
@@ -70,7 +80,7 @@ class Adam : public Optimizer {
     Adam(const NamedParameterList &named_params, float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
          float eps = 1e-8);
 
-    void Step() override;
+    void StepImpl() override;
 
     std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const override;
 
diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h
index dcfd8927f..dbc3860c5 100644
--- 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) { sequence_parallel_ = enabled; }
+    bool sequence_parallel() const { return sequence_parallel_; }
+
     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
index 388b04de5..9791c56a7 100644
--- a/infini_train/src/nn/modules/normalization.cc
+++ b/infini_train/src/nn/modules/normalization.cc
@@ -18,6 +18,8 @@ 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();
+    parameters_[kParamWeightName]->set_sequence_parallel(true);
+    parameters_[kParamBiasName]->set_sequence_parallel(true);
     ResetParameters();
 }
 
@@ -35,6 +37,7 @@ 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();
+    parameters_[kParamWeightName]->set_sequence_parallel(true);
     nn::init::Ones(parameters_[kParamWeightName]);
 }
 
diff --git a/infini_train/src/nn/modules/transformer/moe/router.cc b/infini_train/src/nn/modules/transformer/moe/router.cc
index 252086846..b045400aa 100644
--- 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
index 99a739d2d..f1048058a 100644
--- 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_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
index 523bcf2d7..d1ba03511 100644
--- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
+++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
@@ -107,6 +107,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)});
@@ -170,15 +171,18 @@ float DistributedOptimizer::learning_rate() const {
     return Optimizer::learning_rate();
 }
 
-void DistributedOptimizer::Step() {
-    // 1. Ensure grads are synced
+void DistributedOptimizer::FinalizeModelGrads() {
     FinishGradSync();
+    if (model_grad_finalizer_) {
+        model_grad_finalizer_(base_optimizer_->parameters());
+    }
+}
 
-    // 2. Base optimizer step on owned param pieces
+void DistributedOptimizer::StepImpl() {
     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/tensor_parallel.cc b/infini_train/src/nn/parallel/tensor_parallel.cc
index b16c526e6..65c55bbba 100644
--- 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);
diff --git a/infini_train/src/nn/parallel/utils.cc b/infini_train/src/nn/parallel/utils.cc
index ee28a1694..7f6908ada 100644
--- a/infini_train/src/nn/parallel/utils.cc
+++ b/infini_train/src/nn/parallel/utils.cc
@@ -5,9 +5,38 @@
 #include "infini_train/include/nn/functional.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/process_group.h"
+#include "infini_train/include/nn/parallel/reduce_op_type.h"
 #include "infini_train/include/tensor.h"
 
 namespace infini_train::nn::parallel {
+namespace {
+const ProcessGroup *GetTensorParallelGroup(const Tensor &tensor) {
+    const int global_rank = tensor.GetDevice().Rank().GlobalRank();
+    return ProcessGroupFactory::Instance(tensor.GetDevice().type())
+        ->Get(GetTensorParallelProcessGroupName(global_rank));
+}
+
+void FinalizeSequenceParallelGradients(const std::vector<std::shared_ptr<Tensor>> &params) {
+    // SP replicas see different sequence shards, so replicated parameter grads
+    // must be summed across TP before the optimizer consumes them.
+    if (!global::GetSequenceParallelEnabled() || global::GetTensorParallelSize() <= 1 || params.empty()) {
+        return;
+    }
+
+    const ProcessGroup *tp_group = nullptr;
+    for (const auto &param : params) {
+        if (!param || !param->sequence_parallel() || !param->grad()) {
+            continue;
+        }
+
+        if (tp_group == nullptr) {
+            tp_group = GetTensorParallelGroup(*param);
+            CHECK_NOTNULL(tp_group);
+        }
+        tp_group->AllReduce(param->grad(), function::ReduceOpType::kSum, /*async_op=*/false);
+    }
+}
+} // namespace
 
 std::string GetDataParallelProcessGroupName(int global_rank) {
     return "DP" + std::to_string(global::GetGroupId(global::DP, global_rank));
@@ -60,4 +89,11 @@ std::shared_ptr<Tensor> GatherTensorParallelShard(const std::shared_ptr<Tensor>
     return nn::function::Concat(rank_major_shards, dim)->Contiguous();
 }
 
+void FinalizeModelGrads(const std::vector<std::shared_ptr<Tensor>> &params) {
+    FinalizeSequenceParallelGradients(params);
+
+    // TODO(zbl): Future model-gradient finalization goes here, such as PP tied embeddings,
+    //            MoE shared parameters, and loss normalization.
+}
+
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc
index 39b999c77..2b0322020 100644
--- a/infini_train/src/optimizer.cc
+++ b/infini_train/src/optimizer.cc
@@ -29,6 +29,21 @@ void Optimizer::ZeroGrad(bool set_to_none) {
     for (auto param : params_) { param->ZeroGrad(set_to_none); }
 }
 
+void Optimizer::Step() {
+    FinalizeModelGrads();
+    StepImpl();
+}
+
+// NOTE(zbl): Centralize model-gradient finalization here. This is the boundary
+//            for gradients that a regular DDP reducer cannot express automatically,
+//            such as SP norms, PP tied embeddings, MoE shared parameters, and loss
+//            normalization.
+void Optimizer::FinalizeModelGrads() {
+    if (model_grad_finalizer_) {
+        model_grad_finalizer_(params_);
+    }
+}
+
 void Optimizer::set_learning_rate(float lr) { learning_rate_ = lr; }
 
 float Optimizer::learning_rate() const { return learning_rate_; }
@@ -53,7 +68,7 @@ SGD::SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate
 
 SGD::SGD(const NamedParameterList &named_params, float learning_rate) : Optimizer(named_params, learning_rate) {}
 
-void SGD::Step() {
+void SGD::StepImpl() {
     for (auto param : params_) {
         if (!param->grad()) {
             LOG(INFO) << "Skipping param with null grad.";
@@ -99,7 +114,7 @@ Adam::Adam(const NamedParameterList &named_params, float learning_rate, float be
     }
 }
 
-void Adam::Step() {
+void Adam::StepImpl() {
     ++t_;
 
     for (size_t i = 0; i < params_.size(); ++i) {
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index 4e61e221b..e9a3f4626 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -57,6 +57,7 @@ Tensor::Tensor(const Tensor &tensor, size_t offset, const std::vector<int64_t> &
     : buffer_(tensor.buffer_), offset_(tensor.offset_ + offset), dims_(dims),
       num_elements_(std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int64_t>())), dtype_(tensor.dtype_) {
     CHECK_LE(offset_ + kDataTypeToSize.at(dtype_) * num_elements_, buffer_->Size());
+    sequence_parallel_ = tensor.sequence_parallel_;
 }
 
 Tensor::Tensor(const float *data, const std::vector<int64_t> &dims, DataType dtype, Device device)
@@ -170,6 +171,7 @@ Tensor Tensor::To(Device device) {
     }
 
     new_tensor.requires_grad_ = requires_grad_;
+    new_tensor.sequence_parallel_ = sequence_parallel_;
 
     return new_tensor;
 }
@@ -194,6 +196,7 @@ Tensor Tensor::To(DataType dtype) {
     }
 
     new_tensor.requires_grad_ = requires_grad_;
+    new_tensor.sequence_parallel_ = sequence_parallel_;
 
     return new_tensor;
 }

From 6864ff3941d645340b581061d60672f0dd84c474 Mon Sep 17 00:00:00 2001
From: bolunz <chamberlain0w0@gmail.com>
Date: Tue, 22 Sep 2026 16:22:41 +0800
Subject: [PATCH 2/4] fix: seperate finalize logics from optimzer, add local
 grad map

---
 example/gpt2/main.cc                          |  3 +-
 example/llama3/main.cc                        |  3 +-
 .../parallel/ddp/distributed_data_parallel.h  |  2 +
 .../nn/parallel/ddp/distributed_optimizer.h   |  6 +--
 .../nn/parallel/ddp/param_and_grad_buffer.h   | 10 +++++
 infini_train/include/nn/parallel/utils.h      |  5 ++-
 infini_train/include/optimizer.h              | 16 ++------
 infini_train/include/tensor.h                 |  4 +-
 infini_train/src/nn/modules/normalization.cc  | 11 ++++--
 .../parallel/ddp/distributed_data_parallel.cc |  4 ++
 .../nn/parallel/ddp/distributed_optimizer.cc  | 20 ++--------
 .../nn/parallel/ddp/param_and_grad_buffer.cc  |  6 +++
 .../src/nn/parallel/pp/pipeline_schedule.cc   |  2 +
 infini_train/src/nn/parallel/utils.cc         | 39 ++++++++++++++-----
 infini_train/src/optimizer.cc                 | 19 +--------
 infini_train/src/tensor.cc                    |  4 ++
 16 files changed, 84 insertions(+), 70 deletions(-)

diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 2017de3ee..4af5b8260 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -357,8 +357,6 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
-    optimizer->set_model_grad_finalizer(nn::parallel::FinalizeModelGrads);
-
     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);
@@ -517,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
index 6558e2516..b8342ca90 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -347,8 +347,6 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         optimizer = optimizer_creator(named_parameters);
     }
-    optimizer->set_model_grad_finalizer(nn::parallel::FinalizeModelGrads);
-
     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);
@@ -504,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
index 905816b56..d5d527610 100644
--- 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
index 3e8c7fea2..b27e6f7b0 100644
--- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
+++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h
@@ -29,8 +29,7 @@ class DistributedOptimizer final : public infini_train::Optimizer {
                          const std::vector<std::shared_ptr<Module>> &model_chunks, size_t ddp_world_size,
                          size_t ddp_rank);
 
-    void FinalizeModelGrads() override;
-    void StepImpl() override;
+    void Step() override;
 
     void ZeroGrad(bool set_to_none = true) override;
 
@@ -38,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
index 2c572984d..a7866c4b3 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
@@ -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
index 696b45949..1d3fe975c 100644
--- 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 {
@@ -32,5 +35,5 @@ std::vector<std::shared_ptr<Tensor>> ScatterToTPRegionFunc(const std::shared_ptr
 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<Tensor>> &params);
+void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_chunks);
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h
index 362df76a8..d85b1acea 100644
--- a/infini_train/include/optimizer.h
+++ b/infini_train/include/optimizer.h
@@ -18,7 +18,6 @@ using NamedParameter = std::pair<std::string, std::shared_ptr<Tensor>>;
 using NamedParameterList = std::vector<NamedParameter>;
 using OptimizerCreator = std::function<std::shared_ptr<Optimizer>(const std::vector<std::shared_ptr<Tensor>> &params)>;
 using OptimizerCreatorNamed = std::function<std::shared_ptr<Optimizer>(const NamedParameterList &named_params)>;
-using ModelGradFinalizer = std::function<void(const std::vector<std::shared_ptr<Tensor>> &params)>;
 
 class Optimizer {
 public:
@@ -28,11 +27,7 @@ class Optimizer {
 
     virtual void ZeroGrad(bool set_to_none = true);
 
-    virtual void Step() final;
-
-    void set_model_grad_finalizer(ModelGradFinalizer finalizer) { model_grad_finalizer_ = std::move(finalizer); }
-
-    const std::vector<std::shared_ptr<Tensor>> &parameters() const { return params_; }
+    virtual void Step() = 0;
 
     virtual std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const { return {}; };
 
@@ -49,16 +44,11 @@ class Optimizer {
     void set_initial_learning_rate(float lr);
 
 protected:
-    virtual void FinalizeModelGrads();
-    virtual void StepImpl() = 0;
-
     std::vector<std::shared_ptr<Tensor>> params_;
     std::vector<std::string> parameter_names_;
     float learning_rate_ = 0.0f;
     float initial_learning_rate_ = 0.0f;
     bool initial_lr_set_ = false;
-
-    ModelGradFinalizer model_grad_finalizer_;
 };
 
 namespace optimizers {
@@ -67,7 +57,7 @@ class SGD : public Optimizer {
     SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
     SGD(const NamedParameterList &named_params, float learning_rate);
 
-    void StepImpl() override;
+    void Step() override;
 
     static OptimizerCreator Create(float learning_rate);
     static OptimizerCreatorNamed CreateNamed(float learning_rate);
@@ -80,7 +70,7 @@ class Adam : public Optimizer {
     Adam(const NamedParameterList &named_params, float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
          float eps = 1e-8);
 
-    void StepImpl() override;
+    void Step() override;
 
     std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const override;
 
diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h
index dbc3860c5..3aabe28f1 100644
--- a/infini_train/include/tensor.h
+++ b/infini_train/include/tensor.h
@@ -80,8 +80,8 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     size_t NumElements() const;
     DataType Dtype() const;
 
-    void set_sequence_parallel(bool enabled) { sequence_parallel_ = enabled; }
-    bool sequence_parallel() const { return sequence_parallel_; }
+    void set_sequence_parallel(bool enabled);
+    bool sequence_parallel() const;
 
     std::shared_ptr<Tensor> Detach() const;
 
diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc
index 9791c56a7..4bd0b0dab 100644
--- 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,8 +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();
-    parameters_[kParamWeightName]->set_sequence_parallel(true);
-    parameters_[kParamBiasName]->set_sequence_parallel(true);
+    if (parallel::global::GetSequenceParallelEnabled()) {
+        parameters_[kParamWeightName]->set_sequence_parallel(true);
+        parameters_[kParamBiasName]->set_sequence_parallel(true);
+    }
     ResetParameters();
 }
 
@@ -37,7 +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();
-    parameters_[kParamWeightName]->set_sequence_parallel(true);
+    if (parallel::global::GetSequenceParallelEnabled()) {
+        parameters_[kParamWeightName]->set_sequence_parallel(true);
+    }
     nn::init::Ones(parameters_[kParamWeightName]);
 }
 
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 dd05a8d71..e063fcb6f 100644
--- 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
index d1ba03511..a8fa73e64 100644
--- 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) {
@@ -116,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 +166,7 @@ float DistributedOptimizer::learning_rate() const {
     return Optimizer::learning_rate();
 }
 
-void DistributedOptimizer::FinalizeModelGrads() {
-    FinishGradSync();
-    if (model_grad_finalizer_) {
-        model_grad_finalizer_(base_optimizer_->parameters());
-    }
-}
-
-void DistributedOptimizer::StepImpl() {
+void DistributedOptimizer::Step() {
     CHECK(base_optimizer_) << "DistributedOptimizer: base optimizer is null.";
     base_optimizer_->Step();
 
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 8a3c2d052..970c30f27 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
@@ -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
index 0ea2a4710..40ba16ddf 100644
--- 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/utils.cc b/infini_train/src/nn/parallel/utils.cc
index 7f6908ada..f845ef792 100644
--- a/infini_train/src/nn/parallel/utils.cc
+++ b/infini_train/src/nn/parallel/utils.cc
@@ -3,6 +3,8 @@
 #include "glog/logging.h"
 
 #include "infini_train/include/nn/functional.h"
+#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/process_group.h"
 #include "infini_train/include/nn/parallel/reduce_op_type.h"
@@ -16,16 +18,16 @@ const ProcessGroup *GetTensorParallelGroup(const Tensor &tensor) {
         ->Get(GetTensorParallelProcessGroupName(global_rank));
 }
 
-void FinalizeSequenceParallelGradients(const std::vector<std::shared_ptr<Tensor>> &params) {
+void FinalizeSequenceParallelGradients(const std::vector<LocalGradShard> &param_grads) {
     // SP replicas see different sequence shards, so replicated parameter grads
     // must be summed across TP before the optimizer consumes them.
-    if (!global::GetSequenceParallelEnabled() || global::GetTensorParallelSize() <= 1 || params.empty()) {
+    if (!global::GetSequenceParallelEnabled() || global::GetTensorParallelSize() <= 1 || param_grads.empty()) {
         return;
     }
 
     const ProcessGroup *tp_group = nullptr;
-    for (const auto &param : params) {
-        if (!param || !param->sequence_parallel() || !param->grad()) {
+    for (const auto &[param, grad] : param_grads) {
+        if (!param || !param->sequence_parallel() || !grad) {
             continue;
         }
 
@@ -33,7 +35,7 @@ void FinalizeSequenceParallelGradients(const std::vector<std::shared_ptr<Tensor>
             tp_group = GetTensorParallelGroup(*param);
             CHECK_NOTNULL(tp_group);
         }
-        tp_group->AllReduce(param->grad(), function::ReduceOpType::kSum, /*async_op=*/false);
+        tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false);
     }
 }
 } // namespace
@@ -89,11 +91,30 @@ std::shared_ptr<Tensor> GatherTensorParallelShard(const std::shared_ptr<Tensor>
     return nn::function::Concat(rank_major_shards, dim)->Contiguous();
 }
 
-void FinalizeModelGrads(const std::vector<std::shared_ptr<Tensor>> &params) {
-    FinalizeSequenceParallelGradients(params);
+void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_chunks) {
+    for (const auto &model_chunk : model_chunks) {
+        if (auto ddp = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunk)) {
+            ddp->FinishGradSync();
+        }
+    }
+
+    for (const auto &model_chunk : model_chunks) {
+        std::vector<LocalGradShard> param_grads;
+        auto ddp = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunk);
+        if (ddp && ddp->ddp_config().zero_stage >= 1) {
+            for (const auto &group : ddp->bucket_groups()) {
+                const auto &shards = group->local_grad_shards();
+                param_grads.insert(param_grads.end(), shards.begin(), shards.end());
+            }
+        } else {
+            for (const auto &param : model_chunk->Parameters()) { param_grads.emplace_back(param, param->grad()); }
+        }
+        FinalizeSequenceParallelGradients(param_grads);
+    }
 
-    // TODO(zbl): Future model-gradient finalization goes here, such as PP tied embeddings,
-    //            MoE shared parameters, and loss normalization.
+    // NOTE(zbl): Centralize model-gradient finalization that regular DDP reduction cannot express here,
+    //            including SP replicated parameters, PP tied embeddings, MoE shared parameters, and loss
+    //            normalization.
 }
 
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc
index 2b0322020..39b999c77 100644
--- a/infini_train/src/optimizer.cc
+++ b/infini_train/src/optimizer.cc
@@ -29,21 +29,6 @@ void Optimizer::ZeroGrad(bool set_to_none) {
     for (auto param : params_) { param->ZeroGrad(set_to_none); }
 }
 
-void Optimizer::Step() {
-    FinalizeModelGrads();
-    StepImpl();
-}
-
-// NOTE(zbl): Centralize model-gradient finalization here. This is the boundary
-//            for gradients that a regular DDP reducer cannot express automatically,
-//            such as SP norms, PP tied embeddings, MoE shared parameters, and loss
-//            normalization.
-void Optimizer::FinalizeModelGrads() {
-    if (model_grad_finalizer_) {
-        model_grad_finalizer_(params_);
-    }
-}
-
 void Optimizer::set_learning_rate(float lr) { learning_rate_ = lr; }
 
 float Optimizer::learning_rate() const { return learning_rate_; }
@@ -68,7 +53,7 @@ SGD::SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate
 
 SGD::SGD(const NamedParameterList &named_params, float learning_rate) : Optimizer(named_params, learning_rate) {}
 
-void SGD::StepImpl() {
+void SGD::Step() {
     for (auto param : params_) {
         if (!param->grad()) {
             LOG(INFO) << "Skipping param with null grad.";
@@ -114,7 +99,7 @@ Adam::Adam(const NamedParameterList &named_params, float learning_rate, float be
     }
 }
 
-void Adam::StepImpl() {
+void Adam::Step() {
     ++t_;
 
     for (size_t i = 0; i < params_.size(); ++i) {
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index e9a3f4626..51b60a184 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -105,6 +105,10 @@ size_t Tensor::NumElements() const { return num_elements_; }
 
 DataType Tensor::Dtype() const { return dtype_; }
 
+void Tensor::set_sequence_parallel(bool enabled) { sequence_parallel_ = enabled; }
+
+bool Tensor::sequence_parallel() const { return sequence_parallel_; }
+
 std::shared_ptr<Tensor> Tensor::Detach() const { return std::make_shared<Tensor>(*this, 0, dims_); }
 
 void Tensor::Fill(Scalar value) {

From 7d28b72eb878bd1b597341f96f880e915df798ce Mon Sep 17 00:00:00 2001
From: bolunz <chamberlain0w0@gmail.com>
Date: Thu, 24 Sep 2026 18:12:34 +0800
Subject: [PATCH 3/4] fix: fix qwen3-8b grad sgrad sync behavior under under TP

---
 example/mixtral/main.cc                       |  2 ++
 example/qwen3/main.cc                         |  1 +
 .../transformer/causal_self_attention.cc      |  4 +++
 infini_train/src/nn/parallel/utils.cc         | 34 +++++++++++++++++--
 infini_train/src/tensor.cc                    | 12 ++++---
 tests/tensor/test_tensor_copy.cc              | 30 ++++++++++++++++
 6 files changed, 76 insertions(+), 7 deletions(-)

diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc
index ce4c8cf36..4998675f8 100644
--- 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
index ba969fd21..11249f9fc 100644
--- 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/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc
index 9c32fa23d..31676b1b4 100644
--- 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/parallel/utils.cc b/infini_train/src/nn/parallel/utils.cc
index f845ef792..03d975f5e 100644
--- a/infini_train/src/nn/parallel/utils.cc
+++ b/infini_train/src/nn/parallel/utils.cc
@@ -1,5 +1,7 @@
 #include "infini_train/include/nn/parallel/utils.h"
 
+#include <unordered_set>
+
 #include "glog/logging.h"
 
 #include "infini_train/include/nn/functional.h"
@@ -91,30 +93,56 @@ std::shared_ptr<Tensor> GatherTensorParallelShard(const std::shared_ptr<Tensor>
     return nn::function::Concat(rank_major_shards, dim)->Contiguous();
 }
 
+// Call once after all microbatch backwards, before gradient clipping and optimizer updates.
 void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_chunks) {
+    // Finish DP gradient synchronization for all chunks before TP reduction.
     for (const auto &model_chunk : model_chunks) {
         if (auto ddp = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunk)) {
             ddp->FinishGradSync();
         }
     }
 
+    if (global::GetTensorParallelSize() <= 1) {
+        // DP synchronization is complete; skip TP-specific finalization.
+        return;
+    }
+
     for (const auto &model_chunk : model_chunks) {
+        // Identify replicated Q/K norm weights whose gradients are split across TP-local heads.
+        std::unordered_set<const Tensor *> qk_norm_params;
+        for (const auto &[name, param] : model_chunk->NamedParameters()) {
+            if (name == "q_norm.weight" || name.ends_with(".q_norm.weight") || name == "k_norm.weight"
+                || name.ends_with(".k_norm.weight")) {
+                qk_norm_params.insert(param.get());
+            }
+        }
+        // Pair original parameters with the gradient views consumed by the optimizer.
         std::vector<LocalGradShard> param_grads;
         auto ddp = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunk);
         if (ddp && ddp->ddp_config().zero_stage >= 1) {
+            // ZeRO-1/2 use DP-local gradient shard views registered by DistOpt.
             for (const auto &group : ddp->bucket_groups()) {
                 const auto &shards = group->local_grad_shards();
                 param_grads.insert(param_grads.end(), shards.begin(), shards.end());
             }
         } else {
+            // Ordinary optimizer / ZeRO-0 use full gradients from param.grad().
             for (const auto &param : model_chunk->Parameters()) { param_grads.emplace_back(param, param->grad()); }
         }
+        // TP SUM combines local-token gradients for SP norms, row bias, position embeddings and routers.
         FinalizeSequenceParallelGradients(param_grads);
+
+        // TP SUM combines Q/K norm head-shard gradients regardless of SP; their SP flag is false.
+        for (const auto &[param, grad] : param_grads) {
+            if (param && param->requires_grad() && grad && qk_norm_params.contains(param.get())) {
+                const auto *tp_group = GetTensorParallelGroup(*param);
+                CHECK_NOTNULL(tp_group);
+                tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false);
+            }
+        }
     }
 
-    // NOTE(zbl): Centralize model-gradient finalization that regular DDP reduction cannot express here,
-    //            including SP replicated parameters, PP tied embeddings, MoE shared parameters, and loss
-    //            normalization.
+    // NOTE(zbl): Extend this entry for PP tied embeddings, MoE shared parameters and loss normalization.
 }
 
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index 51b60a184..70d9591af 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -133,9 +133,11 @@ Eigen::Map<Eigen::Matrix<float, 1, Eigen::Dynamic, Eigen::RowMajor>> Tensor::Eig
 Tensor Tensor::To(Device device) {
     const auto buffer_device = buffer_->GetDevice();
     if (device == buffer_device) {
-        auto new_tensor = Tensor(*this, offset_, dims_);
+        auto new_tensor = Tensor(*this, 0, dims_);
+        new_tensor.requires_grad_ = requires_grad_;
+        new_tensor.sequence_parallel_ = sequence_parallel_;
         if (grad_) {
-            new_tensor.grad_ = std::make_unique<Tensor>(*grad_.get(), grad_->offset_, grad_->dims_);
+            new_tensor.grad_ = std::make_unique<Tensor>(*grad_.get(), 0, grad_->dims_);
         }
         return new_tensor;
     }
@@ -182,9 +184,11 @@ Tensor Tensor::To(Device device) {
 
 Tensor Tensor::To(DataType dtype) {
     if (dtype == dtype_) {
-        auto new_tensor = Tensor(*this, offset_, dims_);
+        auto new_tensor = Tensor(*this, 0, dims_);
+        new_tensor.requires_grad_ = requires_grad_;
+        new_tensor.sequence_parallel_ = sequence_parallel_;
         if (grad_) {
-            new_tensor.grad_ = std::make_unique<Tensor>(*grad_.get(), grad_->offset_, grad_->dims_);
+            new_tensor.grad_ = std::make_unique<Tensor>(*grad_.get(), 0, grad_->dims_);
         }
         return new_tensor;
     }
diff --git a/tests/tensor/test_tensor_copy.cc b/tests/tensor/test_tensor_copy.cc
index 88d597b8e..84dcbc434 100644
--- a/tests/tensor/test_tensor_copy.cc
+++ b/tests/tensor/test_tensor_copy.cc
@@ -27,4 +27,34 @@ TEST_P(TensorCopyTest, CopiesPreservesDataType) {
     EXPECT_EQ(target->Dtype(), DataType::kFLOAT32);
 }
 
+TEST_P(TensorCopyTest, NoOpConversionsPreserveParameterMetadataAndViews) {
+    auto storage = std::make_shared<Tensor>(std::vector<int64_t>{12}, DataType::kFLOAT32, GetDevice());
+    auto grad_storage = std::make_shared<Tensor>(std::vector<int64_t>{12}, DataType::kFLOAT32, GetDevice());
+    auto parameter = std::make_shared<Tensor>(*storage, 4 * sizeof(float), std::vector<int64_t>{4});
+    auto grad = std::make_shared<Tensor>(*grad_storage, 4 * sizeof(float), std::vector<int64_t>{4});
+    parameter->RequiresGrad();
+    parameter->set_sequence_parallel(true);
+    parameter->set_grad(grad);
+
+    for (auto converted : {parameter->To(GetDevice()), parameter->To(parameter->Dtype())}) {
+        EXPECT_TRUE(converted.requires_grad());
+        EXPECT_TRUE(converted.sequence_parallel());
+        EXPECT_EQ(converted.DataPtr(), parameter->DataPtr());
+        EXPECT_EQ(converted.Dims(), parameter->Dims());
+        ASSERT_NE(converted.grad(), nullptr);
+        EXPECT_EQ(converted.grad()->DataPtr(), grad->DataPtr());
+        EXPECT_EQ(converted.grad()->Dims(), grad->Dims());
+    }
+}
+
+TEST_P(TensorCopyTest, NoOpConversionsPreserveFrozenTensorMetadata) {
+    auto tensor = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice());
+    for (auto converted : {tensor->To(GetDevice()), tensor->To(tensor->Dtype())}) {
+        EXPECT_FALSE(converted.requires_grad());
+        EXPECT_FALSE(converted.sequence_parallel());
+        EXPECT_EQ(converted.grad(), nullptr);
+        EXPECT_EQ(converted.DataPtr(), tensor->DataPtr());
+    }
+}
+
 INFINI_TRAIN_REGISTER_TEST(TensorCopyTest);

From c15c2d10e3c5dd8134a809e2b05a23e2b4847a33 Mon Sep 17 00:00:00 2001
From: bolunz <chamberlain0w0@gmail.com>
Date: Fri, 9 Oct 2026 10:50:00 +0800
Subject: [PATCH 4/4] fix: add FinalizeQKNormWeightGradients function

---
 infini_train/src/nn/parallel/utils.cc | 38 ++++++++++++++++-----------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/infini_train/src/nn/parallel/utils.cc b/infini_train/src/nn/parallel/utils.cc
index 03d975f5e..9d3e871ec 100644
--- a/infini_train/src/nn/parallel/utils.cc
+++ b/infini_train/src/nn/parallel/utils.cc
@@ -40,6 +40,26 @@ void FinalizeSequenceParallelGradients(const std::vector<LocalGradShard> &param_
         tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false);
     }
 }
+
+void FinalizeQKNormWeightGradients(const nn::Module &model_chunk, const std::vector<LocalGradShard> &param_grads) {
+    // CausalSelfAttention Q/K norm weights (e.g. in Qwen3) are shared across heads and replicated across TP ranks.
+    // Each rank computes gradients for TP-local heads, so sum them across TP regardless of SP.
+    // These weights have sequence_parallel=false and are not reduced by FinalizeSequenceParallelGradients.
+    std::unordered_set<const Tensor *> qk_norm_params;
+    for (const auto &[name, param] : model_chunk.NamedParameters()) {
+        if (name.ends_with("q_norm.weight") || name.ends_with("k_norm.weight")) {
+            qk_norm_params.insert(param.get());
+        }
+    }
+
+    for (const auto &[param, grad] : param_grads) {
+        if (param && param->requires_grad() && grad && qk_norm_params.contains(param.get())) {
+            const auto *tp_group = GetTensorParallelGroup(*param);
+            CHECK_NOTNULL(tp_group);
+            tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false);
+        }
+    }
+}
 } // namespace
 
 std::string GetDataParallelProcessGroupName(int global_rank) {
@@ -108,14 +128,6 @@ void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_ch
     }
 
     for (const auto &model_chunk : model_chunks) {
-        // Identify replicated Q/K norm weights whose gradients are split across TP-local heads.
-        std::unordered_set<const Tensor *> qk_norm_params;
-        for (const auto &[name, param] : model_chunk->NamedParameters()) {
-            if (name == "q_norm.weight" || name.ends_with(".q_norm.weight") || name == "k_norm.weight"
-                || name.ends_with(".k_norm.weight")) {
-                qk_norm_params.insert(param.get());
-            }
-        }
         // Pair original parameters with the gradient views consumed by the optimizer.
         std::vector<LocalGradShard> param_grads;
         auto ddp = std::dynamic_pointer_cast<DistributedDataParallel>(model_chunk);
@@ -132,14 +144,8 @@ void FinalizeModelGrads(const std::vector<std::shared_ptr<nn::Module>> &model_ch
         // TP SUM combines local-token gradients for SP norms, row bias, position embeddings and routers.
         FinalizeSequenceParallelGradients(param_grads);
 
-        // TP SUM combines Q/K norm head-shard gradients regardless of SP; their SP flag is false.
-        for (const auto &[param, grad] : param_grads) {
-            if (param && param->requires_grad() && grad && qk_norm_params.contains(param.get())) {
-                const auto *tp_group = GetTensorParallelGroup(*param);
-                CHECK_NOTNULL(tp_group);
-                tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false);
-            }
-        }
+        // TP SUM combines local-head gradients for replicated Q/K norm weights regardless of SP.
+        FinalizeQKNormWeightGradients(*model_chunk, param_grads);
     }
 
     // NOTE(zbl): Extend this entry for PP tied embeddings, MoE shared parameters and loss normalization.