diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index 393c138..c4c76e0 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -23,10 +23,12 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() + last_step = start_step for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) labels = labels.to(args.device) pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()} if isinstance(pixel_values, dict) else pixel_values.to(args.device) + last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -38,7 +40,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -47,17 +49,17 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): optimizer.zero_grad(set_to_none=True) - if step % args.log_interval == 0 or step == iters - 1: + if step % args.log_interval == 0 or step == iters: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] - eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 + eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) - if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): + if (step % args.save_interval == 0 or step == iters) and is_main_process(): model.eval() moe_suffix = '_moe' if vlm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' @@ -76,6 +78,13 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): del input_ids, labels, pixel_values, res, loss + if last_step > start_step and last_step % 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 __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind-V Pretrain") diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 8a42bd7..109f4d5 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -23,10 +23,12 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() + last_step = start_step for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) labels = labels.to(args.device) pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()} if isinstance(pixel_values, dict) else pixel_values.to(args.device) + last_step = step lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -38,7 +40,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -47,17 +49,17 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): optimizer.zero_grad(set_to_none=True) - if step % args.log_interval == 0 or step == iters - 1: + if step % args.log_interval == 0 or step == iters: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] - eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 + eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) - if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): + if (step % args.save_interval == 0 or step == iters) and is_main_process(): model.eval() moe_suffix = '_moe' if vlm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' @@ -76,6 +78,13 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): del input_ids, labels, pixel_values, res, loss + if last_step > start_step and last_step % 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 __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind-V SFT")