DynamicKV-LLM Pretrain v1.1.0

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,

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") #数据路径，用于控制数据集的路径。