Merge 03b593621b into 295d681b9c

model/model_minimind.py

@@ -125,10 +125,10 @@ class Attention(nn.Module):
         if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
             output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
         else:
-            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            scores = (xq @ xk.transpose(-2, -1)).float() / math.sqrt(self.head_dim)
             if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
             if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
-            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
+            output = self.attn_dropout(F.softmax(scores, dim=-1).type_as(xq)) @ xv
         output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
         output = self.resid_dropout(self.o_proj(output))
         return output, past_kv

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,