From 3babaea826adfd19001f266c6a83f2a2cf9c9fc7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=B4=AB=E8=8B=8F?= <148778141+zisu09@users.noreply.github.com> Date: Wed, 28 May 2025 17:48:38 +0800 Subject: [PATCH] Delete train_full_sft.py --- train_full_sft.py | 229 ---------------------------------------------- 1 file changed, 229 deletions(-) delete mode 100644 train_full_sft.py diff --git a/train_full_sft.py b/train_full_sft.py deleted file mode 100644 index cc6e47d..0000000 --- a/train_full_sft.py +++ /dev/null @@ -1,229 +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 contextlib import nullcontext -from torch import optim, nn -from torch.nn.parallel import DistributedDataParallel -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 - -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)) - - -def train_epoch(epoch, wandb, start_step=0): - loss_fct = nn.CrossEntropyLoss(reduction='none') - start_time = time.time() - for step, (X, Y, loss_mask) in enumerate(train_loader): - if step < start_step: - continue - 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 * 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}/full_sft_{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) - - # 保存优化器状态、当前epoch和step - checkpoint = { - 'epoch': epoch, - 'step': step + 1, - 'optimizer_state_dict': optimizer.state_dict(), - 'scaler_state_dict': scaler.state_dict() - } - torch.save(checkpoint, f'{args.save_dir}/checkpoint.pth') - model.train() - - -def init_model(lm_config): - tokenizer = AutoTokenizer.from_pretrained('../model') - model = MiniMindForCausalLM(lm_config) - moe_path = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth' - state_dict = torch.load(ckp, map_location=args.device) - model.load_state_dict(state_dict, strict=False) - - Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') - model = model.to(args.device) - 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) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="MiniMind Full SFT") - parser.add_argument("--out_dir", type=str, default="../out") - parser.add_argument("--epochs", type=int, default=2) - parser.add_argument("--batch_size", type=int, default=16) - parser.add_argument("--learning_rate", type=float, default=5e-7) - 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") - parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT") - parser.add_argument("--num_workers", type=int, default=1) - parser.add_argument("--ddp", action="store_true") - parser.add_argument("--accumulation_steps", type=int, default=1) - 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('--hidden_size', default=512, type=int) - parser.add_argument('--num_hidden_layers', default=8, type=int) - parser.add_argument('--max_seq_len', default=512, type=int) - parser.add_argument('--use_moe', default=False, type=bool) - parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl") - parser.add_argument("--resume", action="store_true", help="Resume training from checkpoint") - - 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-Full-SFT-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 = 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'])) - 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) - - start_epoch = 0 - start_step = 0 - if args.resume: - checkpoint_path = f'{args.save_dir}/checkpoint.pth' - if os.path.exists(checkpoint_path): - checkpoint = torch.load(checkpoint_path) - start_epoch = checkpoint['epoch'] - start_step = checkpoint['step'] - optimizer.load_state_dict(checkpoint['optimizer_state_dict']) - scaler.load_state_dict(checkpoint['scaler_state_dict']) - Logger(f"Resuming training from epoch {start_epoch + 1}, step {start_step}") - else: - Logger("Checkpoint not found, starting training from scratch.") - - for epoch in range(start_epoch, args.epochs): - train_epoch(epoch, wandb, start_step if epoch == start_epoch else 0) \ No newline at end of file