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

Compare commits

base: github:db4d1751f7746081c1183e9ab2da7d49f4259a80
Branches Tags
github:master
github:website
..
compare: github:4c9755ea49ad21288c1e0868571df1f6cbff3fed
Branches Tags
github:master
github:website

2 Commits
db4d1751f7 ... 4c9755ea49

Author SHA1 Message Date
Sihua ZHANG 4c9755ea49
Merge 03b593621b into f900448c60 2026-09-23 18:43:35 +08:00
jingyaogong f900448c60 [fix] residual step order 2026-09-23 00:07:31 +08:00
1 changed files with 1 additions and 10 deletions
Show Stats Expand all files Collapse all files

11
trainer/train_sft_omni.py
View File

@ -51,7 +51,6 @@ def omni_collate_fn(batch):
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
start_time = time.time()
last_step = start_step
for step, (input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb) in enumerate(loader, start=start_step + 1):
input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
@ -65,7 +64,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
else:
pixel_values = pixel_values.to(args.device)
spk_emb = spk_emb.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
@ -96,7 +94,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
loss = (text_loss + audio_loss + res.aux_loss) / args.accumulation_steps
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)
scaler.step(optimizer)
@ -129,13 +127,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
del input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb, 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-O SFT")