pull/894/merge
liyang 2026-10-25 23:59:50 +08:00 committed by GitHub
commit 465f87c8c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 2 deletions

View File

@ -22,7 +22,7 @@ def init_model(args):
ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True) model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
if args.lora_weight != 'None': if args.lora_weight != 'None':
apply_lora(model) apply_lora(model, rank=args.lora_rank)
load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}{moe_suffix}.pth') load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}{moe_suffix}.pth')
else: else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
@ -35,6 +35,7 @@ def main():
parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录") parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀(pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo)") parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀(pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo)")
parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称(None表示不使用,可选:lora_identity, lora_medical)") parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称(None表示不使用,可选:lora_identity, lora_medical)")
parser.add_argument('--lora_rank', default=16, type=int, help="LoRA rank,需与训练时一致(默认16)")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度") parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)") parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)")

View File

@ -69,6 +69,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning") parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录") parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
parser.add_argument("--lora_rank", default=16, type=int, help="LoRA rank")
parser.add_argument("--lora_name", type=str, default="lora_medical", help="LoRA权重名称(如lora_identity/lora_medical等)") parser.add_argument("--lora_name", type=str, default="lora_medical", help="LoRA权重名称(如lora_identity/lora_medical等)")
parser.add_argument("--epochs", type=int, default=10, help="训练轮数") parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--batch_size", type=int, default=32, help="batch size")
@ -118,7 +119,7 @@ if __name__ == "__main__":
# ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ========== # ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ==========
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
apply_lora(model) apply_lora(model, rank=args.lora_rank)
# 统计参数 # 统计参数
total_params = sum(p.numel() for p in model.parameters()) total_params = sum(p.numel() for p in model.parameters())