Merge 03b593621b into 295d681b9c

trainer/train_sft_omni.py

@@ -51,6 +51,7 @@ 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)
@@ -64,6 +65,7 @@ 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
@@ -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
 
         scaler.scale(loss).backward()
-        if step % args.accumulation_steps == 0 or step == iters:
+        if step % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
             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
 
+    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")