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

fix: add missing --use_mla to train_full_sft.py

pull/791/head
baihaochen 2026-05-14 14:28:11 +08:00
parent 49a05a1307
commit 8baa9c5f71
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
trainer/train_full_sft.py
View File

@ -14,6 +14,7 @@ from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from model.model_minimind_mla import MiniMindMLAConfig
from dataset.lm_dataset import SFTDataset from dataset.lm_dataset import SFTDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, \ from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, \
init_model, SkipBatchSampler init_model, SkipBatchSampler
@ -101,6 +102,8 @@ if __name__ == "__main__":
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") 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('--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_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("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
parser.add_argument('--from_weight', default='pretrain', type=str, parser.add_argument('--from_weight', default='pretrain', type=str,
help="基于哪个权重训练，为none则不基于任何权重训练") help="基于哪个权重训练，为none则不基于任何权重训练")
@ -118,8 +121,12 @@ if __name__ == "__main__":
# ========== 2. 配置目录、模型参数、检查ckp ========== # ========== 2. 配置目录、模型参数、检查ckp ==========
os.makedirs(args.save_dir, exist_ok=True) os.makedirs(args.save_dir, exist_ok=True)
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, if args.use_mla:
use_moe=bool(args.use_moe)) 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, ckp_data = lm_checkpoint(lm_config, weight=args.save_weight,
save_dir='../checkpoints') if args.from_resume == 1 else None save_dir='../checkpoints') if args.from_resume == 1 else None