fix: correct LoRA initialization and forward pass under tensor parallelism

infini_train/include/nn/lora/lora_utils.h

@@ -18,6 +18,18 @@ class Module;
 
 namespace infini_train::nn::lora {
 
+namespace detail {
+
+// Internal helper for packed QKV tensors stored as [Q | K | V] along dim 0.
+std::shared_ptr<Tensor> SlicePackedQKVRowsForTensorParallel(const std::shared_ptr<Tensor> &full_tensor, int64_t q_rows,
+                                                            int tp_rank, int tp_size);
+
+// Internal helper for TP-gathered packed QKV shards stored rank-major as [Q_i | K_i | V_i].
+std::shared_ptr<Tensor> RestorePackedQKVRowsFromTensorParallel(const std::shared_ptr<Tensor> &gathered_tensor,
+                                                               int64_t q_rows, int tp_size);
+
+} // namespace detail
+
 /**
  * Apply LoRA to a model (PEFT-style injection).
  *
@@ -88,23 +100,29 @@ void UnmergeLoRAWeights(std::shared_ptr<Module> model);
 std::shared_ptr<Module> MergeAndUnload(std::shared_ptr<Module> model);
 
 /**
- * Return a state dict containing only LoRA parameters.
+ * Return a state dict containing only LoRA parameters from the current model.
+ *
+ * Under TP this returns the current rank's local shards for sharded LoRA tensors.
+ * Use SaveLoRAWeights() for a portable adapter file with full/unsharded tensors.
  */
 std::unordered_map<std::string, std::shared_ptr<Tensor>> LoRAStateDict(const std::shared_ptr<Module> &model);
 
 /**
- * Load LoRA parameters from a state dict.
+ * Load LoRA parameters from a state dict. Full tensors are sliced for TP-local
+ * parameters when the current model stores shards.
  */
 void LoadLoRAStateDict(std::shared_ptr<Module> model,
                        const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
 
 /**
- * Save only LoRA parameters to file.
+ * Save only LoRA parameters to file. TP-sharded LoRA tensors are exported as
+ * full/unsharded tensors so the file can be loaded with a different TP rank.
  */
 void SaveLoRAWeights(const std::shared_ptr<Module> &model, const std::string &filepath);
 
 /**
- * Load LoRA parameters from file.
+ * Load LoRA parameters from file. Packed QKV LoRA-B tensors are split as
+ * [Qi | Ki | Vi] for the current TP rank.
  */
 void LoadLoRAWeights(std::shared_ptr<Module> model, const std::string &filepath);
 

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

@@ -17,13 +17,20 @@ class Rank;
 
 namespace infini_train::nn::parallel {
 
+namespace {
+constexpr char kModuleName[] = "module";
+
+} // namespace
+
 class DistributedDataParallel : public nn::Module {
 public:
     DistributedDataParallel(std::shared_ptr<nn::Module> module, const Rank &rank,
                             DistributedDataParallelConfig ddp_config);
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
+    std::shared_ptr<nn::Module> module() const;
+
     DistributedDataParallelConfig ddp_config() const { return ddp_config_; }
 
     const std::vector<std::shared_ptr<ParamAndGradBuffer>> &param_grad_buffers() const { return param_grad_buffers_; }

infini_train/include/nn/parallel/process_group.h

@@ -52,21 +52,32 @@ class ProcessGroup {
                                                 function::ReduceOpType reduce_op = function::ReduceOpType::kSum,
                                                 bool async_op = false) const;
 
+    // root_rank_in_group is ProcessGroup-local rank. Broadcast updates tensors in place.
+    virtual std::shared_ptr<Work> Broadcast(const std::vector<std::shared_ptr<Tensor>> &tensors, int root_rank_in_group,
+                                            bool async_op = false) const;
+
+    // Root provides rank-major input_tensors: rank * output_tensors.size() + tensor_index.
+    virtual std::shared_ptr<Work> Scatter(const std::vector<std::shared_ptr<Tensor>> &output_tensors,
+                                          const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                                          int root_rank_in_group, bool async_op = false) const;
+
     virtual std::shared_ptr<Work> Send(std::vector<std::shared_ptr<Tensor>> tensors, int dest_rank,
                                        bool async_op = false) const;
 
     virtual std::shared_ptr<Work> Recv(std::vector<std::shared_ptr<Tensor>> tensors, int src_rank,
                                        bool async_op = false) const;
 
     // Legacy communication APIs (Single-stream)
+    // FIXME(dcj): BroadCast_ and Scatter_ are temporarily retained with trailing underscores for existing DP callers.
+    // Replace direct DP usage with a higher-level communication abstraction.
     virtual std::vector<std::shared_ptr<Tensor>>
-    BroadCast(const std::vector<std::shared_ptr<Tensor>> &input_tensors) const;
+    BroadCast_(const std::vector<std::shared_ptr<Tensor>> &input_tensors) const;
 
     virtual std::vector<std::shared_ptr<Tensor>>
     ReduceAddCoalesced(const std::vector<std::vector<std::shared_ptr<Tensor>>> &grads, Device destination) const;
 
-    virtual std::vector<std::shared_ptr<Tensor>> Scatter(const std::shared_ptr<Tensor> &tensor,
-                                                         std::vector<Device> devices, int64_t dim) const;
+    virtual std::vector<std::shared_ptr<Tensor>> Scatter_(const std::shared_ptr<Tensor> &tensor,
+                                                          std::vector<Device> devices, int64_t dim) const;
 
     virtual std::shared_ptr<Tensor> Gather(const std::vector<std::shared_ptr<Tensor>> &tensors, Device destination,
                                            int64_t dim) const;

infini_train/include/nn/parallel/utils.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <vector>
@@ -22,6 +23,8 @@ std::vector<int> GetTensorParallelGroupRanks(int global_rank);
 std::vector<int> GetPipelineParallelGroupRanks(int global_rank);
 
 // TP/SP Communication Helper Functions
+std::shared_ptr<Tensor> GatherTensorParallelShard(const std::shared_ptr<Tensor> &tensor, int64_t dim);
+
 std::vector<std::shared_ptr<Tensor>> GatherFromTPRegionFunc(const std::shared_ptr<Tensor> &input);
 std::vector<std::shared_ptr<Tensor>> ReduceScatterToSPRegionFunc(const std::shared_ptr<Tensor> &input);
 std::vector<std::shared_ptr<Tensor>> GatherFromSPRegionFunc(const std::shared_ptr<Tensor> &input);

infini_train/src/autograd/comm.cc

@@ -18,7 +18,7 @@ Scatter::Scatter(const std::vector<Device> &target_gpus, int64_t dim,
 std::vector<std::shared_ptr<Tensor>> Scatter::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
     const auto &input = input_tensors[0];
     std::vector<std::shared_ptr<Tensor>> output_tensors;
-    output_tensors = pg_->Scatter(input, target_gpus_, dim_);
+    output_tensors = pg_->Scatter_(input, target_gpus_, dim_);
     return output_tensors;
 }
 
@@ -82,7 +82,7 @@ std::vector<std::shared_ptr<Tensor>> Broadcast::Forward(const std::vector<std::s
     }
 
     // TODO(dcj): mark non differentiable
-    return pg_->BroadCast(input_tensors);
+    return pg_->BroadCast_(input_tensors);
 }
 
 void Broadcast::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,

infini_train/src/nn/lora/lora_parallel_linear.cc

@@ -11,6 +11,7 @@
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/linear.h"
 #include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/process_group.h"
 #include "infini_train/include/nn/parallel/tensor_parallel.h"
 #include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/tensor.h"
@@ -89,22 +90,38 @@ LoRAColumnParallelLinear::LoRAColumnParallelLinear(std::shared_ptr<parallel::Col
 }
 
 void LoRAColumnParallelLinear::InitLoRAWeights() {
-    // LoRA weights stored directly in parameters_
-    // Following PEFT pattern conceptually:
-    // lora_A: [rank, in_features] - replicated
+    // lora_A: [rank, in_features] - replicated across TP ranks
     // lora_B: [out_features_per_partition, rank] - sharded like base weight
-
-    // lora_A: [rank, in_features]
     parameters_[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);
+
+    if (parallel::global::GetTensorParallelSize() > 1) {
+        const auto global_rank = device_.Rank().GlobalRank();
+        auto *tp_group = parallel::ProcessGroupFactory::Instance(device_.type())
+                             ->Get(parallel::GetTensorParallelProcessGroupName(global_rank));
+        const int tp_rank = tp_group->GetGroupRank(global_rank);
+
+        // FIXME: Each TP group's root initializes independently, so corresponding DP replicas are not guaranteed to
+        // start with identical parameters. They may currently match only because every rank uses the same default RNG
+        // seed. DDP should perform a generic parameter broadcast after model construction and before the first forward,
+        // rather than adding a LoRA-specific DP broadcast here.
+        if (tp_rank == 0) {
+            if (config_.use_kaiming_a) {
+                init::KaimingUniform(parameters_[kParamLoraAName], config_.kaiming_a_param);
+            } else {
+                init::Normal(parameters_[kParamLoraAName], 0.0f, 0.02f);
+            }
+        }
+        tp_group->Broadcast({parameters_[kParamLoraAName]}, 0);
     } else {
-        init::Normal(parameters_[kParamLoraAName], 0.0f, 0.02f);
+        if (config_.use_kaiming_a) {
+            init::KaimingUniform(parameters_[kParamLoraAName], config_.kaiming_a_param);
+        } else {
+            init::Normal(parameters_[kParamLoraAName], 0.0f, 0.02f);
+        }
     }
 
-    // lora_B: [out_per_partition, rank] - sharded like base weight
     parameters_[kParamLoraBName]
         = std::make_shared<Tensor>(std::vector<int64_t>{out_features_per_partition_, config_.rank}, DataType::kFLOAT32,
                                    device_)
@@ -126,39 +143,35 @@ LoRAColumnParallelLinear::Forward(const std::vector<std::shared_ptr<Tensor>> &in
         << "Forward() on merged LoRA with requires_grad=true. Call UnmergeWeights() before training.";
 
     if (!merged_) {
-        // 1. Compute base output via parent class
-        auto base_result = ColumnParallelLinear::Forward(input_tensors);
-        auto base_output = base_result[0];
-
-        // 2. Compute LoRA output using the SAME input that base module uses
-        // Match base input path exactly: use direct input if input_is_parallel_ or sequence_parallel_,
-        // otherwise copy to TP region
-        auto lora_input = (input_is_parallel_ || sequence_parallel_)
-                            ? input_tensors[0]
-                            : parallel::CopyToTPRegionFunc(input_tensors[0])[0];
+        // Inline base + LoRA matmuls, add locally, then single collective op.
+        // This avoids 2 separate AllGather ops which cause floating-point divergence.
+        auto input = (input_is_parallel_ || sequence_parallel_) ? input_tensors[0]
+                                                                : parallel::CopyToTPRegionFunc(input_tensors[0])[0];
         if (sequence_parallel_) {
-            // Base uses GatherFromSPRegionFunc to gather sequence dimension
-            lora_input = parallel::GatherFromSPRegionFunc(lora_input)[0];
+            input = parallel::GatherFromSPRegionFunc(input)[0];
         }
 
-        // Compute LoRA: lora_A: [rank, in_features], lora_B: [out_per_partition, rank]
-        auto lora_proj = std::make_shared<autograd::Linear>()->Apply({lora_input, parameters_[kParamLoraAName]})[0];
+        // Base matmul (bias folded in when applicable, matching ColumnParallelLinear::Forward)
+        auto base_shard = std::make_shared<autograd::Linear>()->Apply(
+            (bias_ && !skip_bias_add_)
+                ? std::vector<std::shared_ptr<Tensor>>{input, parameters_.at(kParamWeightName),
+                                                       parameters_[kParamBiasName]}
+                : std::vector<std::shared_ptr<Tensor>>{input, parameters_.at(kParamWeightName)})[0];
+
+        // LoRA matmul (local)
+        // Wrap replicated lora_A through CopyToTPRegion so its gradient gets AllReduced in backward
+        auto lora_A = parallel::CopyToTPRegionFunc(parameters_[kParamLoraAName])[0];
+        auto lora_proj = std::make_shared<autograd::Linear>()->Apply({input, lora_A})[0];
         auto lora_output = std::make_shared<autograd::Linear>()->Apply({lora_proj, parameters_[kParamLoraBName]})[0];
 
-        // Match base output layout (gather if base gathers)
-        if (gather_output_) {
-            lora_output = parallel::GatherFromTPRegionFunc(lora_output)[0];
-        }
-
-        auto scaled_lora = lora_output->Mul(config_.Scaling());
+        // Local add before collective
+        auto combined = base_shard->Add(lora_output->Mul(config_.Scaling()));
 
-        // 3. Add LoRA contribution to base output
-        // Both should now have the same sequence dimension
-        auto output = base_output->Add(scaled_lora);
+        // Single collective op
+        auto output = gather_output_ ? parallel::GatherFromTPRegionFunc(combined)[0] : combined;
 
-        // Return in same format as base module
         return skip_bias_add_
-                 ? std::vector<std::shared_ptr<Tensor>>{output, bias_ ? parameters_[kParamBiasName] : nullptr}
+                 ? std::vector<std::shared_ptr<Tensor>>{output, bias_ ? parameters_.at(kParamBiasName) : nullptr}
                  : std::vector<std::shared_ptr<Tensor>>{output};
     }
 
@@ -294,10 +307,30 @@ void LoRARowParallelLinear::InitLoRAWeights() {
         = std::make_shared<Tensor>(std::vector<int64_t>{config_.rank, in_features_per_partition_}, DataType::kFLOAT32,
                                    device_)
               ->RequiresGrad();
-    if (config_.use_kaiming_a) {
-        init::KaimingUniform(parameters_[kParamLoraAName], config_.kaiming_a_param);
+    if (parallel::global::GetTensorParallelSize() > 1) {
+        const auto global_rank = device_.Rank().GlobalRank();
+        auto *tp_group = parallel::ProcessGroupFactory::Instance(device_.type())
+                             ->Get(parallel::GetTensorParallelProcessGroupName(global_rank));
+        const int tp_rank = tp_group->GetGroupRank(global_rank);
+
+        std::vector<std::shared_ptr<Tensor>> scatter_inputs;
+        if (tp_rank == 0) {
+            auto full_lora_A = std::make_shared<Tensor>(std::vector<int64_t>{config_.rank, in_features_},
+                                                        DataType::kFLOAT32, device_);
+            if (config_.use_kaiming_a) {
+                init::KaimingUniform(full_lora_A, config_.kaiming_a_param);
+            } else {
+                init::Normal(full_lora_A, 0.0f, 0.02f);
+            }
+            scatter_inputs = full_lora_A->Split(in_features_per_partition_, 1);
+        }
+        tp_group->Scatter({parameters_[kParamLoraAName]}, scatter_inputs, 0);
     } else {
-        init::Normal(parameters_[kParamLoraAName], 0.0f, 0.02f);
+        if (config_.use_kaiming_a) {
+            init::KaimingUniform(parameters_[kParamLoraAName], config_.kaiming_a_param);
+        } else {
+            init::Normal(parameters_[kParamLoraAName], 0.0f, 0.02f);
+        }
     }
 
     // lora_B: [out_features, rank]
@@ -321,42 +354,32 @@ LoRARowParallelLinear::Forward(const std::vector<std::shared_ptr<Tensor>> &input
         << "Forward() on merged LoRA with requires_grad=true. Call UnmergeWeights() before training.";
 
     if (!merged_) {
-        // Get effective input - match what base module uses
-        auto effective_input = input_tensors[0];
-        const int64_t in_dim = effective_input->Dims().back();
-
-        if (!input_is_parallel_) {
-            // base would scatter; lora must match
-            effective_input = parallel::ScatterToTPRegionFunc(effective_input)[0];
-            CHECK_EQ(effective_input->Dims().back(), in_features_per_partition_);
-        } else {
-            // input_is_parallel_=true means caller promised shard input
-            CHECK_EQ(in_dim, in_features_per_partition_)
-                << "RowParallel expects sharded input when input_is_parallel_=true. "
-                << "Got full in_dim=" << in_dim << " (likely upstream gathered TP output).";
+        // Inline base + LoRA matmuls, add locally, then single collective op.
+        // This avoids 2 separate AllReduce ops which cause floating-point divergence.
+        auto input = input_is_parallel_ ? input_tensors[0] : parallel::ScatterToTPRegionFunc(input_tensors[0])[0];
+
+        // Base matmul (no bias — RowParallel adds bias AFTER collective)
+        auto base_shard = std::make_shared<autograd::Linear>()->Apply({input, parameters_.at(kParamWeightName)})[0];
+
+        // LoRA matmul (local)
+        // Wrap replicated lora_B through CopyToTPRegion so its gradient gets AllReduced in backward
+        auto lora_proj = std::make_shared<autograd::Linear>()->Apply({input, parameters_[kParamLoraAName]})[0];
+        auto lora_B = parallel::CopyToTPRegionFunc(parameters_[kParamLoraBName])[0];
+        auto lora_output = std::make_shared<autograd::Linear>()->Apply({lora_proj, lora_B})[0];
+
+        // Local add before collective
+        auto combined = base_shard->Add(lora_output->Mul(config_.Scaling()));
+
+        // Single collective op
+        auto output = reduce_output_ ? (sequence_parallel_ ? parallel::ReduceScatterToSPRegionFunc(combined)[0]
+                                                           : parallel::ReduceFromTPRegionFunc(combined)[0])
+                                     : combined;
+
+        // Bias after collective (matching RowParallelLinear::Forward)
+        if (bias_ && !skip_bias_add_) {
+            output = output->Add(parameters_[kParamBiasName]);
         }
 
-        // 1) base output - use effective_input
-        auto base_result = RowParallelLinear::Forward({effective_input});
-        auto base_output = base_result[0];
-
-        // 2) lora branch uses the SAME effective_input
-        auto lora_proj
-            = std::make_shared<autograd::Linear>()->Apply({effective_input, parameters_[kParamLoraAName]})[0];
-        auto lora_output = std::make_shared<autograd::Linear>()->Apply({lora_proj, parameters_[kParamLoraBName]})[0];
-
-        // 3) apply same reduction as base
-        auto lora_out = lora_output;
-        if (reduce_output_) {
-            lora_out = sequence_parallel_ ? parallel::ReduceScatterToSPRegionFunc(lora_out)[0]
-                                          : parallel::ReduceFromTPRegionFunc(lora_out)[0];
-        }
-
-        auto scaled_lora = lora_out->Mul(config_.Scaling());
-        CHECK_EQ(base_output->NumElements(), scaled_lora->NumElements());
-        auto output = base_output->Add(scaled_lora);
-
-        // Return in same format as base module
         return skip_bias_add_
                  ? std::vector<std::shared_ptr<Tensor>>{output, bias_ ? parameters_[kParamBiasName] : nullptr}
                  : std::vector<std::shared_ptr<Tensor>>{output};