From 8baa9c5f714694a280e6fc752ca6e88fbdd916db Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Thu, 14 May 2026 14:28:11 +0800 Subject: [PATCH] fix: add missing --use_mla to train_full_sft.py --- trainer/train_full_sft.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 475e8a4..a24d3e5 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -14,6 +14,7 @@ from torch import optim, nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig +from model.model_minimind_mla import MiniMindMLAConfig from dataset.lm_dataset import SFTDataset from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, \ init_model, SkipBatchSampler @@ -101,6 +102,8 @@ if __name__ == "__main__": parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）") parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") + parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1], help="是否使用MLA注意力架构（0=否，1=是）") + parser.add_argument('--kv_lora_rank', default=128, type=int, help="MLA的KV压缩秩（仅use_mla=1时生效）") parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径") parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") @@ -118,8 +121,12 @@ if __name__ == "__main__": # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) - lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - use_moe=bool(args.use_moe)) + if args.use_mla: + lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank) + else: + lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe)) ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None