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

[fix] residual step order

master
jingyaogong 2026-09-23 00:07:30 +08:00
parent dacc687889
commit 1862b633fc
2 changed files with 2 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
trainer/train_pretrain_vlm.py
View File

@ -24,12 +24,10 @@ 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
@ -41,7 +39,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.scale(loss).backward()
if step % args.accumulation_steps == 0:
if step % args.accumulation_steps == 0 or step == iters:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@ -79,13 +77,6 @@ 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")

11
trainer/train_sft_vlm.py
View File

@ -24,12 +24,10 @@ 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
@ -41,7 +39,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.scale(loss).backward()
if step % args.accumulation_steps == 0:
if step % args.accumulation_steps == 0 or step == iters:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@ -79,13 +77,6 @@ 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")