sft should use pretrain model

pull/376/head
Wenyu Zhang 2025-04-26 11:50:23 -04:00
parent 765c79c7a2
commit c8972f7e89
1 changed files with 1 additions and 1 deletions

View File

@ -97,7 +97,7 @@ def init_model(lm_config):
tokenizer = AutoTokenizer.from_pretrained('../model') tokenizer = AutoTokenizer.from_pretrained('../model')
model = MiniMindForCausalLM(lm_config) model = MiniMindForCausalLM(lm_config)
moe_path = '_moe' if lm_config.use_moe else '' moe_path = '_moe' if lm_config.use_moe else ''
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth' ckp = f'./out/pretrain_{lm_config.hidden_size}{moe_path}.pth'
state_dict = torch.load(ckp, map_location=args.device) state_dict = torch.load(ckp, map_location=args.device)
model.load_state_dict(state_dict, strict=False) model.load_state_dict(state_dict, strict=False)