diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index c4f1100..7b361cc 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -1,48 +1,39 @@ -import argparse -import time -import math -import warnings - -warnings.filterwarnings('ignore') import os import sys -import torch -import torch.distributed as dist __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +import argparse +import time +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, AutoModel +from transformers import AutoTokenizer from model.model_vlm import MiniMindVLM, VLMConfig from dataset.lm_dataset import VLMDataset +from trainer.trainer_utils import * + +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): +def train_epoch(epoch, loader, iters, start_step=0, wandb=None): loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() - for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader): + for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): X = X.to(args.device) Y = Y.to(args.device) loss_mask = loss_mask.to(args.device) pixel_values = pixel_values.to(args.device) - lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate) + lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr - with ctx: + with autocast_ctx: res = model(X, pixel_values=pixel_values) loss = loss_fct( res.logits.view(-1, res.logits.size(-1)), @@ -64,27 +55,20 @@ def train_epoch(epoch, wandb): optimizer.zero_grad(set_to_none=True) - if step % args.log_interval == 0: + if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time - Logger( - 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format( - epoch + 1, - args.epochs, - step, - iter_per_epoch, - loss.item(), - optimizer.param_groups[-1]['lr'], - spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) + current_loss = loss.item() * args.accumulation_steps + current_lr = optimizer.param_groups[-1]['lr'] + eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 + + Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + + if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min}) - if (wandb is not None) and (not ddp or dist.get_rank() == 0): - wandb.log({"loss": loss, - "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): + if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() - moe_path = '_moe' if model_config.use_moe else '' - ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth' + moe_suffix = '_moe' if vlm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' if isinstance(model, torch.nn.parallel.DistributedDataParallel): state_dict = model.module.state_dict() else: @@ -94,115 +78,95 @@ def train_epoch(epoch, wandb): } clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存 torch.save(clean_state_dict, ckp) + vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() -def init_model(model_config: VLMConfig): - tokenizer = AutoTokenizer.from_pretrained('../model', use_fast=True) - moe_path = '_moe' if model_config.use_moe else '' - # 加载纯语言模型权重 - ckp = f'{args.save_dir}/llm_{model_config.hidden_size}{moe_path}.pth' - model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16") - state_dict = torch.load(ckp, map_location=args.device) - model.load_state_dict(state_dict, strict=False) - - # 冻结除 vision_proj 外的所有参数 - for name, param in model.named_parameters(): - if 'vision_proj' not in name: - param.requires_grad = False - - Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') - - _, preprocess = model.vision_encoder, model.processor - return model.to(args.device), tokenizer, preprocess - - -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-V Pretrain") - parser.add_argument("--out_dir", type=str, default="../out") - parser.add_argument("--epochs", type=int, default=4) - parser.add_argument("--batch_size", type=int, default=16) - parser.add_argument("--learning_rate", type=float, default=4e-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", default=False, action="store_true") - parser.add_argument("--wandb_project", type=str, default="MiniMind-V") - parser.add_argument("--num_workers", type=int, default=8) - parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl") - parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images") - 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=640, type=int) - parser.add_argument('--use_moe', default=False, type=bool) + parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录") + parser.add_argument('--save_weight', default='pretrain_vlm', type=str, help="保存权重的前缀名") + parser.add_argument("--epochs", type=int, default=4, help="训练轮数") + parser.add_argument("--batch_size", type=int, default=16, help="batch size") + parser.add_argument("--learning_rate", type=float, default=4e-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("--num_workers", type=int, default=8, 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("--log_interval", type=int, default=100, help="日志打印间隔") + parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔") + parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度") + parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") + parser.add_argument('--max_seq_len', default=640, type=int, help="训练的最大截断长度") + parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE") + parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl", help="训练数据路径") + parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images", help="训练图像路径") + parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") + parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是") + parser.add_argument('--freeze_llm', default=True, type=bool, help="是否冻结LLM参数（仅训练vision_proj）") + parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") + parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名") args = parser.parse_args() - model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - max_seq_len=args.max_seq_len) - max_seq_len = model_config.max_seq_len - args.save_dir = os.path.join(args.out_dir) + # ========== 1. 初始化环境和随机种子 ========== + local_rank = init_distributed_mode() + if dist.is_initialized(): args.device = f"cuda:{local_rank}" + setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + + # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) - os.makedirs(args.out_dir, exist_ok=True) - tokens_per_iter = args.batch_size * max_seq_len - torch.manual_seed(1337) + vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + max_seq_len=args.max_seq_len, use_moe=args.use_moe) + ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None + + # ========== 3. 设置混合精度 ========== device_type = "cuda" if "cuda" in args.device else "cpu" - - args.wandb_run_name = f"MiniMind-V 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" - if ddp: - init_distributed_mode() - args.device = torch.device(DEVICE) - - if args.use_wandb and (not ddp or ddp_local_rank == 0): + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 + autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + + # ========== 4. 配wandb ========== + wandb = None + if args.use_wandb and is_main_process(): import swanlab as wandb - wandb.init(project=args.wandb_project, name=args.wandb_run_name) - else: - wandb = None - - model, tokenizer, preprocess = init_model(model_config) - + wandb_id = ckp_data.get('wandb_id') if ckp_data else None + resume = 'must' if wandb_id else None + wandb_run_name = f"MiniMind-V-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" + wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) + + # ========== 5. 定义模型、数据、优化器 ========== + model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, + device=args.device, freeze_llm=args.freeze_llm) train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess, - image_special_token=model_config.image_special_token, - max_length=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'])) + image_special_token=vlm_config.image_special_token, + max_length=vlm_config.max_seq_len) + train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None + scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate) - - if ddp: + + # ========== 6. 从ckp恢复状态 ========== + start_epoch, start_step = 0, 0 + if ckp_data: + model.load_state_dict(ckp_data['model'], strict=False) + optimizer.load_state_dict(ckp_data['optimizer']) + scaler.load_state_dict(ckp_data['scaler']) + start_epoch = ckp_data['epoch'] + start_step = ckp_data.get('step', 0) + + # ========== 7. DDP包模型 ========== + if dist.is_initialized(): model._ddp_params_and_buffers_to_ignore = {"pos_cis"} - model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) - - iter_per_epoch = len(train_loader) - for epoch in range(args.epochs): - train_epoch(epoch, wandb) + model = DistributedDataParallel(model, device_ids=[local_rank]) + + # ========== 8. 开始训练 ========== + for epoch in range(start_epoch, args.epochs): + train_sampler and train_sampler.set_epoch(epoch) + if epoch == start_epoch and start_step > 0: # 第一个epoch且存在检查点 + batch_sampler = SkipBatchSampler(train_sampler or range(len(train_ds)), args.batch_size, start_step + 1) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') + train_epoch(epoch, loader, len(loader) + start_step + 1, start_step, wandb) + else: # 默认从头开始 + loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True) + train_epoch(epoch, loader, len(loader), 0, wandb) diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 204d8ad..98f8d1f 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -1,48 +1,39 @@ -import argparse -import time -import math -import warnings import os import sys -import torch -import torch.distributed as dist __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +import argparse +import time +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, AutoModel +from transformers import AutoTokenizer from model.model_vlm import MiniMindVLM, VLMConfig from dataset.lm_dataset import VLMDataset +from trainer.trainer_utils import * 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): +def train_epoch(epoch, loader, iters, start_step=0, wandb=None): loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() - for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader): + for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): X = X.to(args.device) Y = Y.to(args.device) loss_mask = loss_mask.to(args.device) pixel_values = pixel_values.to(args.device) - lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate) + lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr - with ctx: + with autocast_ctx: res = model(X, pixel_values=pixel_values) loss = loss_fct( res.logits.view(-1, res.logits.size(-1)), @@ -64,27 +55,20 @@ def train_epoch(epoch, wandb): optimizer.zero_grad(set_to_none=True) - if step % args.log_interval == 0: + if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time - Logger( - 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format( - epoch + 1, - args.epochs, - step, - iter_per_epoch, - loss.item(), - optimizer.param_groups[-1]['lr'], - spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) + current_loss = loss.item() * args.accumulation_steps + current_lr = optimizer.param_groups[-1]['lr'] + eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 + + Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') + + if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min}) - if (wandb is not None) and (not ddp or dist.get_rank() == 0): - wandb.log({"loss": loss, - "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): + if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() - moe_path = '_moe' if model_config.use_moe else '' - ckp = f'{args.save_dir}/sft_vlm_{model_config.hidden_size}{moe_path}.pth' + moe_suffix = '_moe' if vlm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' if isinstance(model, torch.nn.parallel.DistributedDataParallel): state_dict = model.module.state_dict() else: @@ -94,116 +78,94 @@ def train_epoch(epoch, wandb): } clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存 torch.save(clean_state_dict, ckp) + vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() -def init_model(model_config: VLMConfig): - tokenizer = AutoTokenizer.from_pretrained('../model') - moe_path = '_moe' if model_config.use_moe else '' - ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth' - - model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16") - state_dict = torch.load(ckp, map_location=args.device) - model.load_state_dict(state_dict, strict=False) - - # 默认全参训练 - # # 只解冻注意力机制中的投影层参数 - # for name, param in model.model.named_parameters(): - # if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']): - # param.requires_grad = True - - Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') - - _, preprocess = model.vision_encoder, model.processor - return model.to(args.device), tokenizer, preprocess - - -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-V Pretrain") - 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=4) - parser.add_argument("--learning_rate", type=float, default=1e-6) - 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", default=False, action="store_true") - parser.add_argument("--wandb_project", type=str, default="MiniMind-V") - parser.add_argument("--num_workers", type=int, default=8) - parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl") - parser.add_argument("--images_path", type=str, default="../dataset/sft_images") - 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=10) - parser.add_argument("--save_interval", type=int, default=10) - 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=1536, type=int) - parser.add_argument('--use_moe', default=False, type=bool) + parser = argparse.ArgumentParser(description="MiniMind-V SFT") + parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录") + parser.add_argument('--save_weight', default='sft_vlm', type=str, help="保存权重的前缀名") + parser.add_argument("--epochs", type=int, default=2, help="训练轮数") + parser.add_argument("--batch_size", type=int, default=4, help="batch size") + parser.add_argument("--learning_rate", type=float, default=1e-6, 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("--num_workers", type=int, default=8, 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("--log_interval", type=int, default=10, help="日志打印间隔") + parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔") + parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度") + parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") + parser.add_argument('--max_seq_len', default=1536, type=int, help="训练的最大截断长度") + parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE") + parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl", help="训练数据路径") + parser.add_argument("--images_path", type=str, default="../dataset/sft_images", help="训练图像路径") + parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") + parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是") + parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") + parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名") args = parser.parse_args() - model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - max_seq_len=args.max_seq_len) - max_seq_len = model_config.max_seq_len - args.save_dir = os.path.join(args.out_dir) + # ========== 1. 初始化环境和随机种子 ========== + local_rank = init_distributed_mode() + if dist.is_initialized(): args.device = f"cuda:{local_rank}" + setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + + # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) - os.makedirs(args.out_dir, exist_ok=True) - tokens_per_iter = args.batch_size * max_seq_len - torch.manual_seed(1337) + vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + max_seq_len=args.max_seq_len, use_moe=args.use_moe) + ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None + + # ========== 3. 设置混合精度 ========== device_type = "cuda" if "cuda" in args.device else "cpu" - - args.wandb_run_name = f"MiniMind-V 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" - if ddp: - init_distributed_mode() - args.device = torch.device(DEVICE) - - if args.use_wandb and (not ddp or ddp_local_rank == 0): + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 + autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + + # ========== 4. 配wandb ========== + wandb = None + if args.use_wandb and is_main_process(): import swanlab as wandb - wandb.init(project=args.wandb_project, name=args.wandb_run_name) - else: - wandb = None - - model, tokenizer, preprocess = init_model(model_config) - + wandb_id = ckp_data.get('wandb_id') if ckp_data else None + resume = 'must' if wandb_id else None + wandb_run_name = f"MiniMind-V-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" + wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) + + # ========== 5. 定义模型、数据、优化器 ========== + model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, + device=args.device) train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess, - image_special_token=model_config.image_special_token, - max_length=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'])) + image_special_token=vlm_config.image_special_token, + max_length=vlm_config.max_seq_len) + train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None + scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) - - if ddp: + + # ========== 6. 从ckp恢复状态 ========== + start_epoch, start_step = 0, 0 + if ckp_data: + model.load_state_dict(ckp_data['model'], strict=False) + optimizer.load_state_dict(ckp_data['optimizer']) + scaler.load_state_dict(ckp_data['scaler']) + start_epoch = ckp_data['epoch'] + start_step = ckp_data.get('step', 0) + + # ========== 7. DDP包模型 ========== + if dist.is_initialized(): model._ddp_params_and_buffers_to_ignore = {"pos_cis"} - model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) - - iter_per_epoch = len(train_loader) - for epoch in range(args.epochs): - train_epoch(epoch, wandb) + model = DistributedDataParallel(model, device_ids=[local_rank]) + + # ========== 8. 开始训练 ========== + for epoch in range(start_epoch, args.epochs): + train_sampler and train_sampler.set_epoch(epoch) + if epoch == start_epoch and start_step > 0: # 第一个epoch且存在检查点 + batch_sampler = SkipBatchSampler(train_sampler or range(len(train_ds)), args.batch_size, start_step + 1) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') + train_epoch(epoch, loader, len(loader) + start_step + 1, start_step, wandb) + else: # 默认从头开始 + loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True) + train_epoch(epoch, loader, len(loader), 0, wandb) diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py new file mode 100644 index 0000000..9d89475 --- /dev/null +++ b/trainer/trainer_utils.py @@ -0,0 +1,158 @@ +""" +训练工具函数集合 +""" +import os +import random +import math +import numpy as np +import torch +import torch.distributed as dist +from torch.utils.data import Sampler +from transformers import AutoTokenizer +from model.model_vlm import MiniMindVLM + + + +def is_main_process(): + return not dist.is_initialized() or dist.get_rank() == 0 + + +def Logger(content): + if is_main_process(): + 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 init_distributed_mode(): + if int(os.environ.get("RANK", -1)) == -1: + return 0 # 非DDP模式 + + dist.init_process_group(backend="nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + return local_rank + + +def setup_seed(seed: int): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../model', + vision_model_path='../model/vision_model/clip-vit-base-patch16', + save_dir='../out', device='cuda', freeze_llm=False): + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + model = MiniMindVLM(vlm_config, vision_model_path=vision_model_path) + + if from_weight != 'none': + moe_suffix = '_moe' if vlm_config.use_moe else '' + weight_path = f'{save_dir}/{from_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' + weights = torch.load(weight_path, map_location=device) + model.load_state_dict(weights, strict=False) + + # Pretrain阶段：冻结除 vision_proj 外的所有参数 + if freeze_llm: + for name, param in model.named_parameters(): + if 'vision_proj' not in name: + param.requires_grad = False + + # 默认全参训练时的可选配置（已注释） + # # 只解冻注意力机制中的投影层参数 + # for name, param in model.model.named_parameters(): + # if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']): + # param.requires_grad = True + + Logger(f'所加载VLM Model可训练参数：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') + preprocess = model.processor + return model.to(device), tokenizer, preprocess + + +def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs): + os.makedirs(save_dir, exist_ok=True) + moe_path = '_moe' if vlm_config.use_moe else '' + ckp_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}.pth' + resume_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}_resume.pth' + + if model is not None: + from torch.nn.parallel import DistributedDataParallel + state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() + # 移除vision_encoder参数（不需要保存，因为是预训练的） + clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('vision_encoder.')} + ckp_tmp = ckp_path + '.tmp' + torch.save({k: v.half() for k, v in clean_state_dict.items()}, ckp_tmp) + os.replace(ckp_tmp, ckp_path) + + wandb_id = None + if wandb: + if hasattr(wandb, 'get_run'): + run = wandb.get_run() + wandb_id = getattr(run, 'id', None) if run else None + else: + wandb_id = getattr(wandb, 'id', None) + + resume_data = { + 'model': state_dict, + 'optimizer': optimizer.state_dict(), + 'epoch': epoch, + 'step': step, + 'world_size': dist.get_world_size() if dist.is_initialized() else 1, + 'wandb_id': wandb_id + } + for key, value in kwargs.items(): + if value is not None: + if hasattr(value, 'state_dict'): + if isinstance(value, DistributedDataParallel): + resume_data[key] = value.module.state_dict() + else: + resume_data[key] = value.state_dict() + else: + resume_data[key] = value + + resume_tmp = resume_path + '.tmp' + torch.save(resume_data, resume_tmp) + os.replace(resume_tmp, resume_path) + else: # 加载模式 + if os.path.exists(resume_path): + ckp_data = torch.load(resume_path, map_location='cpu') + saved_ws = ckp_data.get('world_size', 1) + current_ws = dist.get_world_size() if dist.is_initialized() else 1 + if saved_ws != current_ws: + ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws + Logger(f'GPU数量变化({saved_ws}→{current_ws})，step已自动转换为{ckp_data["step"]}') + return ckp_data + return None + + +class SkipBatchSampler(Sampler): + def __init__(self, sampler, batch_size, skip_batches=0): + self.sampler = sampler + self.batch_size = batch_size + self.skip_batches = skip_batches + + def __iter__(self): + batch = [] + skipped = 0 + for idx in self.sampler: + batch.append(idx) + if len(batch) == self.batch_size: + if skipped < self.skip_batches: + skipped += 1 + batch = [] + continue + yield batch + batch = [] + if len(batch) > 0 and skipped >= self.skip_batches: + yield batch + + def __len__(self): + total_batches = (len(self.sampler) + self.batch_size - 1) // self.batch_size + return max(0, total_batches - self.skip_batches) +