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

[fix] residual step order

master
jingyaogong 2026-09-23 00:07:31 +08:00
parent 295d681b9c
commit f900448c60
1 changed files with 1 additions and 10 deletions
Show Stats Download Patch File Download Diff File 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): def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
start_time = time.time() 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): 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) input_ids = input_ids.to(args.device)
labels = labels.to(args.device) labels = labels.to(args.device)
@ -65,7 +64,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
else: else:
pixel_values = pixel_values.to(args.device) pixel_values = pixel_values.to(args.device)
spk_emb = spk_emb.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) lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups: for param_group in optimizer.param_groups:
param_group['lr'] = lr 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 loss = (text_loss + audio_loss + res.aux_loss) / args.accumulation_steps
scaler.scale(loss).backward() scaler.scale(loss).backward()
if step % args.accumulation_steps == 0: if step % args.accumulation_steps == 0 or step == iters:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer) 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 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__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MiniMind-O SFT") parser = argparse.ArgumentParser(description="MiniMind-O SFT")