diff --git a/train_pretrain.py b/train_pretrain.py deleted file mode 100644 index a614933..0000000 --- a/train_pretrain.py +++ /dev/null @@ -1,234 +0,0 @@ -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 -import torch.distributed as dist -from torch import optim, nn -from torch.nn.parallel import DistributedDataParallel -from torch.utils.data import DataLoader, DistributedSampler -from contextlib import nullcontext -from transformers import AutoTokenizer -from model.model_minimind import MiniMindConfig, MiniMindForCausalLM -from dataset.lm_dataset import PretrainDataset - -warnings.filterwarnings('ignore') - - -def Logger(content): - if not ddp or dist.get_rank() == 0: - print(content) - - -def get_lr(current_step, total_steps, lr): - return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps)) - -loss_fct = nn.CrossEntropyLoss(reduction='none') - -def train_epoch(epoch, start_step, wandb): - start_time = time.time() - if epoch == start_epoch: - train_loader_iter = iter(train_loader) - for _ in range(start_step + 1): - next(train_loader_iter) - for step, (X, Y, loss_mask) in enumerate(train_loader_iter, start=start_step + 1): - train_step(epoch, step, X, Y, loss_mask, wandb, start_time) - else: - for step, (X, Y, loss_mask) in enumerate(train_loader): - train_step(epoch, step, X, Y, loss_mask, wandb, start_time) - - -def train_step(epoch, step, X, Y, loss_mask, wandb, start_time): - 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_(model.parameters(), 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.item() * 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() - moe_path = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth' - - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() - - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 - torch.save(state_dict, ckp) - - # 保存训练状态 - training_state = { - 'epoch': epoch, - 'step': step, - 'optimizer_state_dict': optimizer.state_dict(), - 'scaler_state_dict': scaler.state_dict() if scaler is not None else None - } - training_state_path = f'{args.save_dir}/training_state.pth' - torch.save(training_state, training_state_path) - - model.train() - - -def init_model(lm_config): - tokenizer = AutoTokenizer.from_pretrained('../model/') - model = MiniMindForCausalLM(lm_config).to(args.device) - Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') - return model, 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"]) - DEVICE = f"cuda:{ddp_local_rank}" - torch.cuda.set_device(DEVICE) - - -# torchrun --nproc_per_node 2 1-pretrain.py -if __name__ == "__main__": - # 命令行参数解析 - parser = argparse.ArgumentParser(description="MiniMind Pretraining") - parser.add_argument("--out_dir", type=str, default="../out") # 输出目录 - parser.add_argument("--epochs", type=int, default=8) # 训练轮次 - parser.add_argument("--batch_size", type=int, default=8) # 物理批次大小 - parser.add_argument("--learning_rate", type=float, default=5e-4) - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") - parser.add_argument("--dtype", type=str, default="bfloat16") # 混合精度类型 - parser.add_argument("--use_wandb", action="store_true") # 是否使用WandB - parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain") - parser.add_argument("--num_workers", type=int, default=1) # 数据加载线程数 - parser.add_argument("--ddp", action="store_true") # 是否启用DDP - parser.add_argument("--accumulation_steps", type=int, default=4) # 梯度累积步数 - parser.add_argument("--grad_clip", type=float, default=1.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) # DDP自动传入参数 - # 模型架构参数 - parser.add_argument('--hidden_size', default=512, type=int) # 隐藏层维度 - parser.add_argument('--num_hidden_layers', default=8, type=int) # 层数 - parser.add_argument('--max_seq_len', default=4096, type=int) # 序列长度 - parser.add_argument('--use_moe', default=False, type=bool) # 是否使用MoE - parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl") - 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" - - args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" - - 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) - - if args.use_wandb and (not ddp or ddp_local_rank == 0): - import wandb - - wandb.init(project=args.wandb_project, name=args.wandb_run_name) - else: - wandb = None - - model, tokenizer = init_model(lm_config) - train_ds = PretrainDataset(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'])) - optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) - - if ddp: - model._ddp_params_and_buffers_to_ignore = {"pos_cis"} - model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) - - iter_per_epoch = len(train_loader) - - # 检查是否存在保存的训练状态文件 - training_state_path = f'{args.save_dir}/training_state.pth' - if os.path.exists(training_state_path): - training_state = torch.load(training_state_path) - start_epoch = training_state['epoch'] - start_step = training_state['step'] - optimizer.load_state_dict(training_state['optimizer_state_dict']) - if scaler is not None and training_state['scaler_state_dict'] is not None: - scaler.load_state_dict(training_state['scaler_state_dict']) - else: - start_epoch = 0 - start_step = 0 - - for epoch in range(start_epoch, args.epochs): - train_epoch(epoch, start_step if epoch == start_epoch else 0, wandb) \ No newline at end of file