Merge 138c367ddb into 3f1a7cc25b

model/model_lora.py

@@ -20,8 +20,8 @@ class LoRA(nn.Module):
 
 def apply_lora(model, rank=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)
+        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
+            lora = LoRA(module.in_features, module.out_features, rank=rank).to(model.device)
             setattr(module, "lora", lora)
             original_forward = module.forward
 

trainer/train_agent.py

@@ -487,4 +487,6 @@ if __name__ == "__main__":
         else:
             rl_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
 
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_distillation.py

@@ -243,4 +243,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), teacher_model, lm_config_student, 0, wandb, args.alpha, args.temperature)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_dpo.py

@@ -223,4 +223,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), ref_model, lm_config, 0, wandb, args.beta)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_full_sft.py

@@ -168,4 +168,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), 0, wandb)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_grpo.py

@@ -329,4 +329,6 @@ if __name__ == "__main__":
             grpo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_lora.py

@@ -181,4 +181,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), lora_params, 0, wandb)
     
     # ========== 10. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_ppo.py

@@ -432,4 +432,6 @@ if __name__ == "__main__":
             ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

trainer/train_pretrain.py

@@ -167,4 +167,6 @@ if __name__ == "__main__":
             train_epoch(epoch, loader, len(loader), 0, wandb)
     
     # ========== 9. 清理分布进程 ==========
-    if dist.is_initialized(): dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()