feat: Add LoRA alpha scaling factor for improved parameter efficiency and stability

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda",
+    "python-envs.pythonProjects": []
+}

model/model_lora.py

@@ -4,9 +4,11 @@ from torch import optim, nn
 
 # 定义Lora网络结构
 class LoRA(nn.Module):
-    def __init__(self, in_features, out_features, rank):
+    def __init__(self, in_features, out_features, rank, alpha=16):
         super().__init__()
         self.rank = rank  # LoRA的秩（rank），控制低秩矩阵的大小
+        self.alpha = alpha  # LoRA的缩放因子
+        self.scaling = self.alpha / self.rank  # 缩放系数
         self.A = nn.Linear(in_features, rank, bias=False)  # 低秩矩阵A
         self.B = nn.Linear(rank, out_features, bias=False)  # 低秩矩阵B
         # 矩阵A高斯初始化
@@ -15,13 +17,13 @@ class LoRA(nn.Module):
         self.B.weight.data.zero_()
 
     def forward(self, x):
-        return self.B(self.A(x))
+        return self.B(self.A(x)) * self.scaling
 
 
-def apply_lora(model, rank=8):
+def apply_lora(model, rank=8, alpha=16):
     for name, module in model.named_modules():
         if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:
-            lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)
+            lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank, alpha=alpha).to(model.device)
             setattr(module, "lora", lora)
             original_forward = module.forward
 

trainer/train_lora.py

@@ -80,6 +80,8 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
     parser.add_argument("--save_dir", type=str, default="../out/lora", help="模型保存目录")
     parser.add_argument("--lora_name", type=str, default="lora_identity", help="LoRA权重名称(如lora_identity/lora_medical等)")
+    parser.add_argument("--lora_rank", type=int, default=8, help="LoRA秩 rank")
+    parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA缩放系数 alpha")
     parser.add_argument("--epochs", type=int, default=50, help="训练轮数")
     parser.add_argument("--batch_size", type=int, default=32, help="batch size")
     parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
@@ -127,7 +129,7 @@ if __name__ == "__main__":
     
     # ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
-    apply_lora(model)
+    apply_lora(model, rank=args.lora_rank, alpha=args.lora_alpha)
     
     # 统计参数
     total_params = sum(p.numel() for p in model.parameters())