diff --git a/model/LMConfig.py b/model/LMConfig.py index fedfc1c..dadfa9f 100644 --- a/model/LMConfig.py +++ b/model/LMConfig.py @@ -9,8 +9,8 @@ class LMConfig(PretrainedConfig): self, dim: int = 512, n_layers: int = 8, - n_heads: int = 8, - n_kv_heads: int = 2, + n_heads: int = 32, + n_kv_heads: int = 8, vocab_size: int = 6400, hidden_dim: int = None, multiple_of: int = 64, diff --git a/train_pretrain.py b/train_pretrain.py index ecb8a12..d2a86fa 100644 --- a/train_pretrain.py +++ b/train_pretrain.py @@ -179,23 +179,23 @@ if __name__ == "__main__": parser.add_argument("--out_dir", type=str, default="out") # 若要以最快速度实现zero则epochs设置为1轮;否则应当利用有限的数据训练2~6个epochs。 parser.add_argument("--epochs", type=int, default=3) - parser.add_argument("--batch_size", type=int, default=32) - parser.add_argument("--learning_rate", type=float, default=5e-4) + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--learning_rate", type=float, default=2e-4) parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") #如果GPU可用,则使用GPU,否则使用CPU。 parser.add_argument("--dtype", type=str, default="bfloat16") parser.add_argument("--use_wandb", default=True, action="store_true") parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain") parser.add_argument("--num_workers", type=int, default=8) parser.add_argument("--ddp", action="store_true") - parser.add_argument("--accumulation_steps", type=int, default=8) #梯度累积步数,用于控制梯度更新频率。 + parser.add_argument("--accumulation_steps", type=int, default=64) #梯度累积步数,用于控制梯度更新频率。 parser.add_argument("--grad_clip", type=float, default=1.0) #梯度裁剪阈值,用于防止梯度爆炸。 parser.add_argument("--warmup_iters", type=int, default=0) #预热迭代次数,用于控制学习率预热过程。 parser.add_argument("--log_interval", type=int, default=100) #日志打印间隔,用于控制日志打印的频率。 parser.add_argument("--save_interval", type=int, default=100) #模型保存间隔,用于控制模型保存的频率。 parser.add_argument('--local_rank', type=int, default=-1) #本地进程编号,用于分布式训练。 - parser.add_argument('--dim', default=1024, type=int) #模型维度,用于控制模型的大小。 - parser.add_argument('--n_layers', default=24, type=int) #层数,用于控制模型层数。 - parser.add_argument('--max_seq_len', default=1024, type=int) #最大序列长度,用于控制输入序列的最大长度。 + parser.add_argument('--dim', default=4096, type=int) #模型维度,用于控制模型的大小。 + parser.add_argument('--n_layers', default=32, type=int) #层数,用于控制模型层数。 + parser.add_argument('--max_seq_len', default=2048, type=int) #最大序列长度,用于控制输入序列的最大长度。 parser.add_argument('--use_moe', default=False, type=bool) #是否使用MOE,用于控制是否使用MOE。 parser.add_argument('--disable_db', action='store_true', help="禁用数据库功能,使用固定值1e-4替代") #禁用数据库功能,启用特殊模式 parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl") #数据路径,用于控制数据集的路径。