Update app.py

app.py

@@ -687,6 +687,7 @@ def generate_modeling_phoenix_code():
     return '''"""
 PHOENIX Retention Model v1.4.3
 ✅ v1.4.3 CRITICAL FIX: forward() 시그니처 Transformers 호환
+✅ v1.4.3 HOTFIX: dtype 불일치 수정 (bfloat16 지원)
 ✅ PhoenixPreTrainedModel 베이스 클래스 포함
 ✅ 모든 Retention 클래스 완전 구현
 """
@@ -748,7 +749,8 @@ class MultiScaleRetention(nn.Module):
         b, s, _ = hidden_states.shape
         device, dtype = hidden_states.device, hidden_states.dtype
         
-        if self.q_proj.weight.device != device:
+        # ✅ FIX: dtype과 device 모두 일치시킴
+        if self.q_proj.weight.device != device or self.q_proj.weight.dtype != dtype:
             self.to(device=device, dtype=dtype)
         
         q = self.q_proj(hidden_states).view(b, s, self.num_heads, self.head_dim).transpose(1, 2)
@@ -801,7 +803,9 @@ class HierarchicalRetention(nn.Module):
     ):
         b, s, h = hidden_states.shape
         device, dtype = hidden_states.device, hidden_states.dtype
-        if next(self.short_proj.parameters()).device != device:
+        
+        # ✅ FIX: dtype과 device 모두 일치시킴
+        if next(self.short_proj.parameters()).device != device or next(self.short_proj.parameters()).dtype != dtype:
             self.to(device=device, dtype=dtype)
         
         ret_out = self.base_retention(hidden_states)[0]
@@ -824,10 +828,23 @@ def replace_attention_with_retention_for_loading(model, use_hierarchical=True):
     layers = getattr(layers, 'layers', getattr(layers, 'h', getattr(layers, 'layers', None)))
     if layers is None: return model, 0, 0
     
+    # ✅ FIX: 원본 모델의 dtype 감지
+    original_dtype = None
+    for param in model.parameters():
+        original_dtype = param.dtype
+        break
+    
     cnt = 0
     for i, layer in enumerate(layers):
         if hasattr(layer, 'self_attn'):
-            layer.self_attn = HierarchicalRetention(model.config, i) if use_hierarchical else MultiScaleRetention(model.config, i)
+            # 새 Retention 생성
+            new_retention = HierarchicalRetention(model.config, i) if use_hierarchical else MultiScaleRetention(model.config, i)
+            
+            # ✅ FIX: 원본 dtype으로 변환
+            if original_dtype is not None:
+                new_retention = new_retention.to(dtype=original_dtype)
+            
+            layer.self_attn = new_retention
             cnt += 1
     return model, cnt, len(layers)
 
@@ -1871,6 +1888,7 @@ with gr.Blocks(
     **Complete Integrated Version with All Fixes**
     
     ✅ **NEW v1.4.3!** forward() 시그니처 Transformers 호환 - 완벽 수정!
+    ✅ **NEW v1.4.3!** dtype 불일치 수정 - bfloat16 완벽 지원!
     ✅ Embedding Tying 저장 시점 처리
     ✅ State Dict 직접 로드로 Retention 보존
     ✅ Model Structure Pre-Analysis
@@ -2003,6 +2021,7 @@ with gr.Blocks(
     
     ### What's New in v1.4.3 (Complete Integrated Version)
     - ✅ **CRITICAL FIX: forward() Signature** - Transformers 호환성 완벽 수정
+    - ✅ **HOTFIX: dtype 불일치** - bfloat16 완벽 지원
     - ✅ **Embedding Tying** - 저장 시점에 자동 처리
     - ✅ **Qwen3-0.6B Generation Fixed** - 정상적인 텍스트 생성
     - ✅ **완전 통합** - 모든 수정사항 포함