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:4c9755ea49ad21288c1e0868571df1f6cbff3fed
Branches Tags
github:master
github:website
..
compare: github:db4d1751f7746081c1183e9ab2da7d49f4259a80
Branches Tags
github:master
github:website

1 Commits
4c9755ea49 ... db4d1751f7

Author SHA1 Message Date
Sihua ZHANG db4d1751f7
Merge 03b593621b into 295d681b9c 2026-09-20 23:41:48 +08:00
1 changed files with 10 additions and 1 deletions
Show Stats Expand all files Collapse all files

11
trainer/train_sft_omni.py
View File

@ -51,6 +51,7 @@ 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)
@ -64,6 +65,7 @@ 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
@ -94,7 +96,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 or step == iters: if step % args.accumulation_steps == 0:
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)
@ -127,6 +129,13 @@ 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")