import os import sys __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import argparse import time import math import warnings import torch from torch import optim, nn import torch.distributed as dist from contextlib import nullcontext from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import SFTDataset from model.model_lora import load_lora, save_lora, apply_lora warnings.filterwarnings('ignore') # 日志打印函数，在分布式训练时只在主进程上打印 def Logger(content): if not ddp or dist.get_rank() == 0: print(content) # 余弦学习率调度器 def get_lr(current_step, total_steps, lr): # 在训练过程中逐渐降低学习率，最终降到初始值的1/10 return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps)) # 训练一个epoch，使用LoRA方法进行参数高效微调 def train_epoch(epoch, wandb): loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() for step, (X, Y, loss_mask) in enumerate(train_loader): X = X.to(args.device) Y = Y.to(args.device) loss_mask = loss_mask.to(args.device) lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr with ctx: res = model(X) loss = loss_fct( res.logits.view(-1, res.logits.size(-1)), Y.view(-1) ).view(Y.size()) loss = (loss * loss_mask).sum() / loss_mask.sum() loss += res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() if (step + 1) % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) if step % args.log_interval == 0: spend_time = time.time() - start_time Logger( 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format( epoch + 1, args.epochs, step, iter_per_epoch, loss.item() * args.accumulation_steps, optimizer.param_groups[-1]['lr'], spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) if (wandb is not None) and (not ddp or dist.get_rank() == 0): wandb.log({"loss": loss * args.accumulation_steps, "lr": optimizer.param_groups[-1]['lr'], "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0): model.eval() lora_save_path = f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth' os.makedirs(os.path.dirname(lora_save_path), exist_ok=True) # 【区别1】只保存lora权重即可 save_lora(model, lora_save_path) model.train() # 初始化模型和分词器 def init_model(lm_config): # 加载预训练的分词器 tokenizer = AutoTokenizer.from_pretrained('../model/') # 初始化因果语言模型 model = MiniMindForCausalLM(lm_config) # 根据是否使用MoE设置权重路径 moe_path = '_moe' if lm_config.use_moe else '' ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth' # 加载预训练权重 state_dict = torch.load(ckp, map_location=args.device) model.load_state_dict(state_dict, strict=False) return model.to(args.device), tokenizer # 初始化分布式训练环境 def init_distributed_mode(): if not ddp: return global ddp_local_rank, DEVICE # 初始化分布式进程组 dist.init_process_group(backend="nccl") # 获取分布式训练的相关参数 ddp_rank = int(os.environ["RANK"]) # 全局进程编号 ddp_local_rank = int(os.environ["LOCAL_RANK"]) # 本地进程编号 ddp_world_size = int(os.environ["WORLD_SIZE"]) # 总进程数 # 设置当前进程使用的GPU DEVICE = f"cuda:{ddp_local_rank}" torch.cuda.set_device(DEVICE) if __name__ == "__main__": # 配置命令行参数 parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA") # 基础训练参数 parser.add_argument("--out_dir", type=str, default="../out", help="输出目录") parser.add_argument("--epochs", type=int, default=10, help="训练轮数") parser.add_argument("--batch_size", type=int, default=32, help="批次大小") parser.add_argument("--learning_rate", type=float, default=1e-4, help="学习率") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="训练精度") # 日志和监控参数 parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb记录训练过程") parser.add_argument("--wandb_project", type=str, default="MiniMind-LoRA-SFT", help="wandb项目名称") parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔") parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔") # 分布式训练参数 parser.add_argument("--num_workers", type=int, default=1, help="数据加载进程数") parser.add_argument("--ddp", action="store_true", help="是否使用分布式训练") parser.add_argument('--local_rank', type=int, default=-1, help="分布式训练的本地进程编号") # 优化器参数 parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数") # 模型参数 parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度") parser.add_argument('--num_hidden_layers', default=8, type=int, help="Transformer层数") parser.add_argument('--max_seq_len', default=512, type=int, help="最大序列长度") parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE") # 数据和LoRA参数 parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl", help="训练数据路径") parser.add_argument("--lora_name", type=str, default="lora_medical", help="根据任务保存成lora_(英文/医学/心理...)") parser.add_argument("--wandb_api_key", type=str, default=None, help="WandB API Key，用于无交互环境自动登录") args = parser.parse_args() lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe) args.save_dir = os.path.join(args.out_dir) os.makedirs(args.save_dir, exist_ok=True) os.makedirs(args.out_dir, exist_ok=True) tokens_per_iter = args.batch_size * args.max_seq_len device_type = "cuda" if "cuda" in args.device else "cpu" ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast() ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run? ddp_local_rank, DEVICE = 0, "cuda:0" base_seed = 1337 torch.manual_seed(base_seed) torch.cuda.manual_seed(base_seed) if ddp: init_distributed_mode() args.device = torch.device(DEVICE) rank = dist.get_rank() torch.manual_seed(base_seed + rank) # 同时设置 CUDA 的随机种子 torch.cuda.manual_seed(base_seed + rank) args.wandb_run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" if args.use_wandb and (not ddp or ddp_local_rank == 0): import wandb # 优先使用命令行参数，其次环境变量 api_key = args.wandb_api_key or os.environ.get("WANDB_API_KEY", None) if api_key is not None: # 自动登录，适用于无交互环境 wandb.login(key=api_key) # 初始化 wandb 任务 wandb.init(project=args.wandb_project, name=args.wandb_run_name) else: wandb = None # 初始化模型和分词器 model, tokenizer = init_model(lm_config) # 应用LoRA参数高效微调方法 apply_lora(model) # 统计模型参数量 total_params = sum(p.numel() for p in model.parameters()) # 总参数数量 lora_params_count = sum(p.numel() for name, p in model.named_parameters() if 'lora' in name) # LoRA 参数数量 if not ddp or dist.get_rank() == 0: print(f"LLM 总参数量: {total_params}") print(f"LoRA 参数量: {lora_params_count}") print(f"LoRA 参数占比: {lora_params_count / total_params * 100:.2f}%") # 冻结非LoRA参数 for name, param in model.named_parameters(): if 'lora' not in name: param.requires_grad = False # 收集需要优化的LoRA参数 lora_params = [] for name, param in model.named_parameters(): if 'lora' in name: lora_params.append(param) # 初始化优化器，只更新LoRA参数 optimizer = optim.AdamW(lora_params, lr=args.learning_rate) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if ddp else None train_loader = DataLoader( train_ds, batch_size=args.batch_size, pin_memory=True, drop_last=False, shuffle=False, num_workers=args.num_workers, sampler=train_sampler ) scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16'])) iter_per_epoch = len(train_loader) for epoch in range(args.epochs): train_epoch(epoch, wandb)