Merge 695aa1fdb6 into f659b55761
commit
465f87c8c5
|
|
@ -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=是)")
|
||||||
|
|
|
||||||
|
|
@ -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())
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue