Merge 03b593621b into f900448c60

model/model_omni.py

@@ -1,5 +1,7 @@
 import os, math, torch, soundfile as sf, librosa, warnings, numpy as np, onnxruntime as ort, logging, contextlib, io
 from types import SimpleNamespace
+from dataclasses import dataclass
+from typing import Optional
 from torch import nn
 from torch.nn import functional as F
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
@@ -28,6 +30,11 @@ class OmniConfig(MiniMindConfig):
         self.image_token_len = kwargs.get("image_token_len", 64)
         self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)
 
+# 避免 DDP + compile 重建 OmniCausalLMOutputWithPast（不支持 audio_logits 参数）
+@dataclass
+class OmniCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
+    audio_logits: Optional[list[torch.FloatTensor]] = None
+
 class MMAudioProjector(nn.Module):
     def __init__(self, in_dim, out_dim):
         super().__init__()
@@ -311,9 +318,12 @@ class MiniMindOmni(MiniMindForCausalLM):
         text_logits = self.thinker.lm_head(h_thinker[:, slice_indices, :])
         audio_logits = self.talker.lm_head(h_talker[:, slice_indices, :])
         
-        out = MoeCausalLMOutputWithPast(aux_loss=aux_loss, logits=text_logits, past_key_values=presents)
-        out.audio_logits = audio_logits
-        return out
+        return OmniCausalLMOutputWithPast(
+            aux_loss=aux_loss,
+            logits=text_logits,
+            past_key_values=presents,
+            audio_logits=audio_logits,
+        )
 
     @torch.inference_mode()
     def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,

trainer/train_sft_omni.py

@@ -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")