From 44f30197ac6e94084f1b5ee65a60825db484c0bf Mon Sep 17 00:00:00 2001 From: xunuohope1107 <148425765+xunuohope1107@users.noreply.github.com> Date: Tue, 18 Feb 2025 16:36:29 +0800 Subject: [PATCH] Add ckp_dir and tokenizer path --- train_full_sft.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/train_full_sft.py b/train_full_sft.py index 5185861..a0b99b9 100644 --- a/train_full_sft.py +++ b/train_full_sft.py @@ -96,10 +96,10 @@ def train_epoch(epoch, wandb): def init_model(lm_config): - tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) model = MiniMindLM(lm_config) moe_path = '_moe' if lm_config.use_moe else '' - ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth' + ckp = f'{args.ckp_dir}/pretrain_{lm_config.dim}{moe_path}.pth' state_dict = torch.load(ckp, map_location=args.device) model.load_state_dict(state_dict, strict=False) Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') @@ -122,6 +122,8 @@ def init_distributed_mode(): if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind Full SFT") 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=1) parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--learning_rate", type=float, default=5e-5)