github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[fix] sync minimind

pull/115/head
jingyaogong 2026-04-02 16:34:38 +08:00
parent 887a42d9cb
commit 4bc6524ac2
2 changed files with 26 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
trainer/train_pretrain_vlm.py
View File

@ -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")

17
trainer/train_sft_vlm.py
View File

@ -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")