github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add ckp_dir and tokenizer path for train_lora.py

pull/189/head
xunuohope1107 2025-02-19 10:18:26 +08:00 committed by GitHub
parent 5d3f3449ed
commit 68762f084a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
train_lora.py
View File

@ -85,10 +85,10 @@ def train_epoch(epoch, wandb):
def init_model(lm_config): def init_model(lm_config):
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
model = MiniMindLM(lm_config) model = MiniMindLM(lm_config)
moe_path = '_moe' if lm_config.use_moe else '' moe_path = '_moe' if lm_config.use_moe else ''
ckp = f'./out/rlhf_{lm_config.dim}{moe_path}.pth' ckp = f'{args.ckp_dir}/rlhf_{lm_config.dim}{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)
return model.to(args.device), tokenizer return model.to(args.device), tokenizer
@ -109,6 +109,8 @@ def init_distributed_mode():
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA") parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA")
parser.add_argument("--out_dir", type=str, default="out") parser.add_argument("--out_dir", type=str, default="out")
parser.add_argument("--ckp_dir", type=str, default="out", help="Directory to save/load checkpoints")
parser.add_argument("--tokenizer_path", type=str, default="./model/minimind_tokenizer", help="Path to the tokenizer")
parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--epochs", type=int, default=50)
parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--learning_rate", type=float, default=5e-5)