feat(lora): add GetLoRAParameters and MergeAndUnload APIs

- Refactor GetLoRAParameters() to retrieve only LoRA parameters for optimizer
- Add MergeAndUnload() to merge weights and export as standard model
- Update gpt2/llama3 examples to use new GetLoRAParameters API
- Refactor LoRA linear modules and fix dimension mismatch
- Improve LoRA tests and update documentation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: 2 people

docs/lora_usage.md

@@ -41,7 +41,7 @@ config.alpha = 16.0f;  // 缩放因子
 auto lora_model = GetLoRAModel(model, config);
 
 // 3. 获取可训练参数用于优化器
-auto trainable_params = lora_model->TrainableParameters();
+auto trainable_params = nn::lora::GetLoRAParameters(lora_model);
 auto optimizer = std::make_shared<Adam>(trainable_params, lr);
 
 // 4. 训练循环
@@ -94,9 +94,6 @@ public:
     LoRAModel(std::shared_ptr<Module> base_model,
               const LoRAConfig &config);
 
-    // 获取可训练参数
-    std::vector<std::shared_ptr<Tensor>> TrainableParameters() const;
-
     // 获取所有参数
     std::vector<std::shared_ptr<Tensor>> Parameters() const override;
 
@@ -237,11 +234,15 @@ void MergeLoRAWeights(std::shared_ptr<Module> model);
 
 // 恢复原始基础权重
 void UnmergeLoRAWeights(std::shared_ptr<Module> model);
+
+// 合并权重并卸载 LoRA 模块，返回纯基础模型
+std::shared_ptr<Module> MergeAndUnload(std::shared_ptr<Module> model);
 ```
 
 **使用场景：**
 - 推理时合并权重可以消除额外计算开销
 - 导出模型时合并权重得到标准模型格式
+- `MergeAndUnload`: 导出完整的标准模型，替换所有 LoRA 模块为普通 Linear 层
 
 ### 保存/加载函数
 
@@ -351,7 +352,7 @@ int main() {
     // =========================================
 
     // 创建优化器（只优化 LoRA 参数）
-    auto trainable_params = lora_model->TrainableParameters();
+    auto trainable_params = nn::lora::GetLoRAParameters(lora_model);
     auto optimizer = std::make_shared<Adam>(trainable_params, /*lr=*/1e-4);
 
     // 训练循环
@@ -441,7 +442,39 @@ auto output = (*lora_model)({input_ids});
 UnmergeLoRAWeights(lora_model);
 ```
 
-### 示例 4: 自定义目标层
+### 示例 4: 导出标准模型 (MergeAndUnload)
+
+使用 `MergeAndUnload` 将 LoRA 模型转换为标准模型，可以直接保存为普通模型文件：
+
+```cpp
+// 加载基础模型并应用 LoRA
+auto model = std::make_shared<GPT2>(config);
+model->LoadWeights("gpt2_weights.bin");
+
+LoRAConfig lora_config;
+lora_config.rank = 8;
+lora_config.alpha = 16.0f;
+lora_config.target_modules = ParseLoRATargetModules("c_attn,c_proj");
+auto lora_model = GetLoRAModel(model, lora_config);
+
+// 训练...
+// ...
+
+// 加载训练好的 LoRA 权重
+LoadLoRAWeights(lora_model, "gpt2_lora.bin");
+
+// 合并并卸载 LoRA，返回标准模型
+// lora_model 中的所有 LoRALinear 都被替换为普通 Linear
+auto merged_model = MergeAndUnload(lora_model);
+
+// 保存为标准模型（与原始模型格式相同）
+merged_model->SaveWeights("gpt2_finetuned.bin");
+
+// 现在 merged_model 是一个普通模型，无需 LoRA 即可推理
+auto output = (*merged_model)({input_ids});
+```
+
+### 示例 5: 自定义目标层
 
 ```cpp
 // 对所有线性层应用
@@ -599,7 +632,7 @@ int main() {
     auto lora_model = std::make_shared<LoRAModel>(base_model, lora_config);
 
     // 4. 获取可训练参数用于优化器
-    auto trainable_params = lora_model->TrainableParameters();
+    auto trainable_params = nn::lora::GetLoRAParameters(lora_model);
     auto optimizer = std::make_shared<Adam>(trainable_params, 1e-5);
 
     // 5. 打印摘要

example/gpt2/main.cc

@@ -239,7 +239,7 @@ void Train(const nn::parallel::Rank &rank) {
 
     auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size);
 
-    // Create optimizer - use LoRAModel's TrainableParameters() if LoRA is enabled
+    // Create optimizer - use GetLoRAParameters if LoRA is enabled
     std::vector<std::shared_ptr<Tensor>> params_to_optimize;
     if (lora_enabled) {
         params_to_optimize = nn::lora::GetLoRAParameters(model);

example/llama3/main.cc

@@ -263,8 +263,14 @@ void Train(const nn::parallel::Rank &rank) {
     auto optimizer_creator = optimizers::Adam::Create(FLAGS_learning_rate);
     std::shared_ptr<Optimizer> optimizer = nullptr;
 
-    // Create optimizer - use TrainableParameters() as single source of truth
-    std::vector<std::shared_ptr<Tensor>> params_to_optimize = model->TrainableParameters();
+    std::vector<std::shared_ptr<Tensor>> params_to_optimize;
+    if (lora_enabled) {
+        params_to_optimize = nn::lora::GetLoRAParameters(model);
+        LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters";
+    } else {
+        params_to_optimize = model->Parameters();
+        LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters";
+    }
 
     if (FLAGS_use_distributed_optimizer) {
         auto model_chunks = (pp_world_size > 1)

infini_train/include/nn/lora/lora_linear.h

@@ -4,7 +4,13 @@
 #include <vector>
 
 #include "infini_train/include/nn/lora/lora_config.h"
-#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/modules/linear.h"
+
+// Forward declarations for test functions (required for friend declarations)
+void test_lora_linear_init();
+void test_lora_linear_forward();
+void test_lora_linear_merge();
+void test_lora_utils();
 
 namespace infini_train {
 class Tensor;
@@ -16,19 +22,13 @@ namespace infini_train::nn::lora {
 // LoRA wrapper for standard Linear layer
 // Implements: y = Wx + b + (alpha/r) * x @ A^T @ B^T
 // Where W is frozen, A and B are trainable low-rank matrices
-class LoRALinear : public nn::CloneableModule<LoRALinear> {
+class LoRALinear : public nn::Linear {
 public:
     static constexpr char kType[] = "LoRALinear";
 
-    // Parameter names
-    static constexpr char kParamWeightName[] = "weight"; // Frozen base weight
-    static constexpr char kParamBiasName[] = "bias";     // Frozen base bias
-    static constexpr char kParamLoraAName[] = "lora_A";  // Trainable A matrix [rank, in_features]
-    static constexpr char kParamLoraBName[] = "lora_B";  // Trainable B matrix [out_features, rank]
-
-    // Constructor from scratch
-    LoRALinear(int64_t in_features, int64_t out_features, const LoRAConfig &config, bool bias = true,
-               const Device *device = nullptr);
+    // Parameter names for LoRA-specific parameters
+    static constexpr char kParamLoraAName[] = "lora_A"; // Trainable A matrix [rank, in_features]
+    static constexpr char kParamLoraBName[] = "lora_B"; // Trainable B matrix [out_features, rank]
 
     // Constructor wrapping existing Linear module (transfers ownership of parameters)
     LoRALinear(std::shared_ptr<nn::Module> base_linear, const LoRAConfig &config);
@@ -43,33 +43,29 @@ class LoRALinear : public nn::CloneableModule<LoRALinear> {
     // Get only LoRA parameters (for optimizer)
     std::vector<std::shared_ptr<Tensor>> LoRAParameters() const;
 
-    // Override Parameters() to return all parameters (frozen base + trainable LoRA)
-    std::vector<std::shared_ptr<Tensor>> Parameters() const override;
-
-    // Get trainable parameters (requires_grad == true)
-    std::vector<std::shared_ptr<Tensor>> TrainableParameters() const;
-
-    // Get all parameters including frozen base weights (for state dict)
-    std::vector<std::shared_ptr<Tensor>> AllParameters() const;
-
     // Accessors
     int64_t in_features() const;
     int64_t out_features() const;
     int64_t rank() const;
     float scaling() const;
 
 private:
+    // Test-only: Create LoRA module from scratch (normal usage goes through InjectLoRALayers)
+    LoRALinear(int64_t in_features, int64_t out_features, const LoRAConfig &config, bool bias, const Device *device);
+
+    // Test access
+    friend void ::test_lora_linear_init();
+    friend void ::test_lora_linear_forward();
+    friend void ::test_lora_linear_merge();
+    friend void ::test_lora_utils();
+
     void InitLoRAWeights();
     void FreezeBaseWeights();
 
     LoRAConfig config_;
-    int64_t in_features_;
-    int64_t out_features_;
-    bool bias_;
+    int64_t in_features_ = 0;
+    int64_t out_features_ = 0;
     bool merged_ = false;
-
-    // Store original weight for unmerge
-    std::shared_ptr<Tensor> original_weight_;
 };
 
 } // namespace infini_train::nn::lora

infini_train/include/nn/lora/lora_parallel_linear.h

@@ -4,7 +4,7 @@
 #include <vector>
 
 #include "infini_train/include/nn/lora/lora_config.h"
-#include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/parallel/tensor_parallel.h"
 
 namespace infini_train {
 class Tensor;
@@ -18,21 +18,19 @@ namespace infini_train::nn::lora {
 // LoRA A: [rank, in_features] - replicated across TP ranks (implemented as Linear)
 // LoRA B: [out_features_per_partition, rank] - sharded like base weight (implemented as ColumnParallelLinear with
 // gather_output)
-class LoRAColumnParallelLinear : public nn::CloneableModule<LoRAColumnParallelLinear> {
+class LoRAColumnParallelLinear : public nn::parallel::ColumnParallelLinear {
 public:
     static constexpr char kType[] = "LoRAColumnParallelLinear";
 
-    static constexpr char kParamWeightName[] = "weight";
-    static constexpr char kParamBiasName[] = "bias";
     static constexpr char kParamLoraAName[] = "lora_A";
     static constexpr char kParamLoraBName[] = "lora_B";
 
     // Constructor wrapping existing ColumnParallelLinear
-    LoRAColumnParallelLinear(std::shared_ptr<nn::Module> base_module, const LoRAConfig &config, int64_t in_features,
-                             int64_t out_features);
+    LoRAColumnParallelLinear(std::shared_ptr<parallel::ColumnParallelLinear> base_module, const LoRAConfig &config,
+                             int64_t in_features, int64_t out_features);
 
     // Constructor wrapping existing ColumnParallelLinear (auto-infer dimensions from weight)
-    LoRAColumnParallelLinear(std::shared_ptr<nn::Module> base_module, const LoRAConfig &config);
+    LoRAColumnParallelLinear(std::shared_ptr<parallel::ColumnParallelLinear> base_module, const LoRAConfig &config);
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
@@ -41,9 +39,6 @@ class LoRAColumnParallelLinear : public nn::CloneableModule<LoRAColumnParallelLi
     bool IsMerged() const;
 
     std::vector<std::shared_ptr<Tensor>> LoRAParameters() const;
-    std::vector<std::shared_ptr<Tensor>> Parameters() const override;
-    std::vector<std::shared_ptr<Tensor>> TrainableParameters() const;
-    std::vector<std::shared_ptr<Tensor>> AllParameters() const;
 
     int64_t in_features() const;
     int64_t out_features() const;
@@ -54,39 +49,28 @@ class LoRAColumnParallelLinear : public nn::CloneableModule<LoRAColumnParallelLi
     void FreezeBaseWeights();
 
     LoRAConfig config_;
-    int64_t in_features_;
-    int64_t out_features_;
-    int64_t out_features_per_partition_;
-    bool bias_;
-    bool gather_output_;
-    bool input_is_parallel_;
-    bool skip_bias_add_;
-    bool sequence_parallel_;
+    int64_t in_features_ = 0;
+    int64_t out_features_ = 0;
+    int64_t out_features_per_partition_ = 0;
     bool merged_ = false;
-
-    std::shared_ptr<Tensor> original_weight_;
-    std::shared_ptr<nn::Module> base_module_; // Not registered in modules_ to avoid double-counting
 };
 
 // LoRA wrapper for RowParallelLinear
 // Weight shape: [out_features, in_features_per_partition]
 // LoRA A: [rank, in_features_per_partition] - sharded like base weight (implemented as RowParallelLinear with
 // input_is_parallel) LoRA B: [out_features, rank] - replicated (implemented as Linear)
-class LoRARowParallelLinear : public nn::CloneableModule<LoRARowParallelLinear> {
+class LoRARowParallelLinear : public nn::parallel::RowParallelLinear {
 public:
     static constexpr char kType[] = "LoRARowParallelLinear";
-
-    static constexpr char kParamWeightName[] = "weight";
-    static constexpr char kParamBiasName[] = "bias";
     static constexpr char kParamLoraAName[] = "lora_A";
     static constexpr char kParamLoraBName[] = "lora_B";
 
     // Constructor wrapping existing RowParallelLinear
-    LoRARowParallelLinear(std::shared_ptr<nn::Module> base_module, const LoRAConfig &config, int64_t in_features,
-                          int64_t out_features);
+    LoRARowParallelLinear(std::shared_ptr<parallel::RowParallelLinear> base_module, const LoRAConfig &config,
+                          int64_t in_features, int64_t out_features);
 
     // Constructor wrapping existing RowParallelLinear (auto-infer dimensions from weight)
-    LoRARowParallelLinear(std::shared_ptr<nn::Module> base_module, const LoRAConfig &config);
+    LoRARowParallelLinear(std::shared_ptr<parallel::RowParallelLinear> base_module, const LoRAConfig &config);
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
@@ -95,9 +79,6 @@ class LoRARowParallelLinear : public nn::CloneableModule<LoRARowParallelLinear>
     bool IsMerged() const;
 
     std::vector<std::shared_ptr<Tensor>> LoRAParameters() const;
-    std::vector<std::shared_ptr<Tensor>> Parameters() const override;
-    std::vector<std::shared_ptr<Tensor>> TrainableParameters() const;
-    std::vector<std::shared_ptr<Tensor>> AllParameters() const;
 
     int64_t in_features() const;
     int64_t out_features() const;
@@ -108,18 +89,10 @@ class LoRARowParallelLinear : public nn::CloneableModule<LoRARowParallelLinear>
     void FreezeBaseWeights();
 
     LoRAConfig config_;
-    int64_t in_features_;
-    int64_t out_features_;
-    int64_t in_features_per_partition_;
-    bool bias_;
-    bool reduce_output_;
-    bool input_is_parallel_;
-    bool skip_bias_add_;
-    bool sequence_parallel_;
+    int64_t in_features_ = 0;
+    int64_t out_features_ = 0;
+    int64_t in_features_per_partition_ = 0;
     bool merged_ = false;
-
-    std::shared_ptr<Tensor> original_weight_;
-    std::shared_ptr<nn::Module> base_module_; // Not registered in modules_ to avoid double-counting
 };
 
 } // namespace infini_train::nn::lora

infini_train/include/nn/lora/lora_utils.h

@@ -74,6 +74,19 @@ void MergeLoRAWeights(std::shared_ptr<Module> model);
  */
 void UnmergeLoRAWeights(std::shared_ptr<Module> model);
 
+/**
+ * Merge LoRA weights and remove LoRA modules, returning a clean base model.
+ * Similar to PEFT's merge_and_unload().
+ *
+ * For each LoRA module:
+ * 1. Merge weights: W += (alpha/r) * B @ A
+ * 2. Replace LoRA module with a base module sharing the merged weight/bias
+ *
+ * After this call, the model contains no LoRA parameters.
+ * Root module may be replaced (same pattern as InjectLoRALayers).
+ */
+std::shared_ptr<Module> MergeAndUnload(std::shared_ptr<Module> model);
+
 /**
  * Return a state dict containing only LoRA parameters.
  */

infini_train/include/nn/modules/linear.h

@@ -22,6 +22,8 @@ class Linear : public CloneableModule<Linear> {
     Linear(int64_t in_features, int64_t out_features, bool bias = true, Device device = Device());
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
+    bool has_bias() const { return bias_; }
+
 private:
     void ResetParameters();
     bool bias_ = true;

infini_train/include/nn/modules/module.h

@@ -47,19 +47,17 @@ class Module : public std::enable_shared_from_this<Module> {
 
     const std::string &type() const;
 
+    // TODO: Change return type to filterable iterator (like PyTorch's named_parameters with prefix matching)
     virtual std::vector<std::shared_ptr<Tensor>> Parameters() const;
-    // Get parameters with requires_grad == true (trainable parameters)
-    std::vector<std::shared_ptr<Tensor>> TrainableParameters() const;
     bool has_parameter(const std::string &name) const;
     std::shared_ptr<Tensor> *mutable_parameter(const std::string &name);
     const std::shared_ptr<Tensor> &parameter(const std::string &name) const;
 
     virtual std::vector<std::shared_ptr<Tensor>> Buffers() const;
 
     std::vector<std::shared_ptr<Module>> modules();
-    std::shared_ptr<Module> mutable_module(const std::string &name);
+    std::shared_ptr<Module> &mutable_module(const std::string &name);
     const Module &module(const std::string &name) const;
-    void replace_module(const std::string &name, std::shared_ptr<Module> new_module);
 
     std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const;
 

infini_train/include/nn/parallel/tensor_parallel.h

@@ -37,7 +37,7 @@ class ColumnParallelLinear : public nn::CloneableModule<ColumnParallelLinear> {
     bool skip_bias_add() const;
     bool sequence_parallel() const;
 
-private:
+protected:
     bool bias_ = true;
     bool gather_output_ = false;     // whether to return full local output tensor after forward (need gather)
     bool input_is_parallel_ = false; // will perform an autograd-aware copy when false
@@ -66,7 +66,7 @@ class RowParallelLinear : public nn::CloneableModule<RowParallelLinear> {
     bool skip_bias_add() const;
     bool sequence_parallel() const;
 
-private:
+protected:
     bool bias_ = true;
     bool reduce_output_ = false;     // whether to return full local output tensor after forward (need reduce)
     bool input_is_parallel_ = false; // will perform an autograd-aware copy when false