Update app.py

app.py

@@ -58,10 +58,12 @@ class MultiScaleRetention(nn.Module):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
+        
+        # ✅ 실제 hidden_size 가져오기
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
         
-        # ✅ Head dimension 정확하게 계산
+        # ✅ Head dimension 계산
         self.head_dim = self.hidden_size // self.num_heads
         
         # ✅ 나누어떨어지는지 확인
@@ -71,12 +73,13 @@ class MultiScaleRetention(nn.Module):
                 f"num_attention_heads ({self.num_heads})"
             )
         
-        print(f"  📐 Layer {layer_idx} Retention config:")
+        print(f"  📐 Layer {layer_idx} Retention initialized:")
         print(f"     - hidden_size: {self.hidden_size}")
         print(f"     - num_heads: {self.num_heads}")
         print(f"     - head_dim: {self.head_dim}")
         
-        # Q, K, V projections (hidden_size → hidden_size)
+        # ✅ Projections - input과 output 크기 명시
+        # input: hidden_size -> output: hidden_size
         self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
         self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
         self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
@@ -107,61 +110,47 @@ class MultiScaleRetention(nn.Module):
         """
         O(n) 복잡도 Retention 메커니즘
         """
-        batch_size, seq_len, _ = hidden_states.shape
+        batch_size, seq_len, input_dim = hidden_states.shape
+        
+        # ✅ 입력 차원 확인
+        if input_dim != self.hidden_size:
+            raise ValueError(
+                f"Input hidden_states has dimension {input_dim} "
+                f"but model expects {self.hidden_size}"
+            )
         
         if past_key_values is not None:
             past_key_value = past_key_values
         
         # Q, K, V 계산
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
+        query_states = self.q_proj(hidden_states)  # [B, L, H]
+        key_states = self.k_proj(hidden_states)    # [B, L, H]
+        value_states = self.v_proj(hidden_states)  # [B, L, H]
+        
+        # ✅ Projection 후 크기 확인
+        assert query_states.shape[-1] == self.hidden_size, \
+            f"Q projection output is {query_states.shape[-1]}, expected {self.hidden_size}"
+        
+        # ✅ Multi-head reshape
+        # [B, L, H] -> [B, L, num_heads, head_dim] -> [B, num_heads, L, head_dim]
+        query_states = query_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
         
-        # ✅ Shape 디버깅
-        print(f"\n  🔍 Retention forward shapes:")
-        print(f"     - Input hidden_states: {hidden_states.shape}")
-        print(f"     - After projection Q: {query_states.shape}")
-        print(f"     - Expected reshape: [{batch_size}, {seq_len}, {self.num_heads}, {self.head_dim}]")
+        key_states = key_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
         
-        # ✅ Multi-head reshape - 정확한 차원으로
-        try:
-            query_states = query_states.view(
-                batch_size, seq_len, self.num_heads, self.head_dim
-            ).transpose(1, 2)  # [B, H, L, D]
-            
-            key_states = key_states.view(
-                batch_size, seq_len, self.num_heads, self.head_dim
-            ).transpose(1, 2)
-            
-            value_states = value_states.view(
-                batch_size, seq_len, self.num_heads, self.head_dim
-            ).transpose(1, 2)
-            
-            print(f"     - After reshape Q: {query_states.shape}")
-            print(f"     ✅ Reshape successful!")
-            
-        except RuntimeError as e:
-            print(f"\n  ❌ Reshape failed!")
-            print(f"     - query_states shape: {query_states.shape}")
-            print(f"     - query_states size: {query_states.numel()}")
-            print(f"     - Target shape: [{batch_size}, {seq_len}, {self.num_heads}, {self.head_dim}]")
-            print(f"     - Target size: {batch_size * seq_len * self.num_heads * self.head_dim}")
-            print(f"     - Error: {e}")
-            
-            # ✅ 실제 크기 계산
-            actual_total = query_states.numel()
-            actual_per_token = actual_total // (batch_size * seq_len)
-            print(f"     - Actual hidden per token: {actual_per_token}")
-            
-            raise
+        value_states = value_states.view(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
         
         # Retention 계산
         retention_states = self._compute_retention(
-            query_states, key_states, value_states,
-            past_key_value
+            query_states, key_states, value_states, past_key_value
         )
         
-        # Reshape back
+        # Reshape back: [B, num_heads, L, head_dim] -> [B, L, H]
         retention_states = retention_states.transpose(1, 2).contiguous()
         retention_states = retention_states.reshape(
             batch_size, seq_len, self.hidden_size
@@ -187,11 +176,6 @@ class MultiScaleRetention(nn.Module):
         """O(n) Retention 계산"""
         batch_size, num_heads, seq_len, head_dim = queries.shape
         
-        print(f"  🔄 Computing retention:")
-        print(f"     - queries: {queries.shape}")
-        print(f"     - keys: {keys.shape}")
-        print(f"     - values: {values.shape}")
-        
         # State 초기화
         if past_state is not None:
             state = past_state
@@ -222,12 +206,8 @@ class MultiScaleRetention(nn.Module):
         
         output = torch.stack(outputs, dim=2)  # [B, H, L, D]
         
-        print(f"     - output: {output.shape}")
-        
         return output
 
-
-
 class HierarchicalRetention(nn.Module):
     """
     PHOENIX의 계층적 Retention
@@ -263,16 +243,15 @@ class HierarchicalRetention(nn.Module):
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
-        cache_position: Optional[torch.Tensor] = None,  # ✅ 추가
-        past_key_values: Optional[Tuple[torch.Tensor]] = None,  # ✅ 추가
-        **kwargs  # ✅ 추가 - 기타 모든 인자 받기
+        cache_position: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[torch.Tensor]] = None,
+        **kwargs
     ):
         """
         Granite 모델과 호환되는 forward 메서드
         """
         batch_size, seq_len, hidden_size = hidden_states.shape
         
-        # past_key_values와 past_key_value 통합 처리
         if past_key_values is not None:
             past_key_value = past_key_values
         
@@ -347,62 +326,39 @@ def replace_attention_with_retention(model, use_hierarchical=True):
     
     total_layers = len(layers)
     
+    # ✅ 첫 번째 레이어에서 실제 hidden_size 확인
+    first_layer = layers[0]
+    if hasattr(first_layer, 'self_attn') and hasattr(first_layer.self_attn, 'q_proj'):
+        actual_output_dim = first_layer.self_attn.q_proj.weight.shape[0]
+        actual_input_dim = first_layer.self_attn.q_proj.weight.shape[1]
+        
+        print(f"\n📐 Detected dimensions from first layer:")
+        print(f"   - Input dim: {actual_input_dim}")
+        print(f"   - Output dim: {actual_output_dim}")
+        print(f"   - Config hidden_size: {model.config.hidden_size}")
+        
+        # ✅ Config 업데이트
+        if actual_output_dim != model.config.hidden_size:
+            print(f"   ⚠️ Updating config to match actual dimensions")
+            model.config.hidden_size = actual_output_dim
+    
     for layer_idx, layer in enumerate(layers):
         try:
-            # Attention 레이어 찾기
             if hasattr(layer, 'self_attn'):
                 old_attn = layer.self_attn
-                config = model.config
-                
-                print(f"\n  📐 Layer {layer_idx} - Original Attention:")
-                
-                # ✅ 실제 가중치 shape 확인
-                if hasattr(old_attn, 'q_proj'):
-                    print(f"     - Q weight: {old_attn.q_proj.weight.shape}")
-                    print(f"     - K weight: {old_attn.k_proj.weight.shape}")
-                    print(f"     - V weight: {old_attn.v_proj.weight.shape}")
-                    print(f"     - O weight: {old_attn.o_proj.weight.shape}")
-                    
-                    # ✅ 실제 output 크기 확인
-                    actual_hidden = old_attn.q_proj.weight.shape[0]
-                    actual_input = old_attn.q_proj.weight.shape[1]
-                    
-                    print(f"     - Actual output dim: {actual_hidden}")
-                    print(f"     - Actual input dim: {actual_input}")
-                    print(f"     - Config hidden_size: {config.hidden_size}")
-                    
-                    # ✅ Config가 맞지 않으면 조정
-                    if actual_hidden != config.hidden_size or actual_input != config.hidden_size:
-                        print(f"     ⚠️ Dimension mismatch detected!")
-                        print(f"     Using actual dimensions: {actual_input} → {actual_hidden}")
-                        
-                        # 새로운 config 생성
-                        class CustomConfig:
-                            def __init__(self, hidden, heads):
-                                self.hidden_size = hidden
-                                self.num_attention_heads = heads
-                        
-                        config = CustomConfig(actual_hidden, model.config.num_attention_heads)
                 
                 # PHOENIX Retention 생성
-                print(f"\n  🔄 Creating PHOENIX Retention for layer {layer_idx}...")
-                
                 if use_hierarchical:
-                    new_retention = HierarchicalRetention(config, layer_idx)
+                    new_retention = HierarchicalRetention(model.config, layer_idx)
                 else:
-                    new_retention = MultiScaleRetention(config, layer_idx)
+                    new_retention = MultiScaleRetention(model.config, layer_idx)
                 
-                # ✅ 가중치 복사 (shape 완벽히 확인)
+                # ✅ 가중치 복사
                 if hasattr(old_attn, 'q_proj'):
-                    old_q_shape = old_attn.q_proj.weight.shape
-                    new_q_shape = new_retention.base_retention.q_proj.weight.shape
-                    
-                    print(f"\n  📋 Weight copy:")
-                    print(f"     - Old Q: {old_q_shape}")
-                    print(f"     - New Q: {new_q_shape}")
-                    
-                    if old_q_shape == new_q_shape:
-                        # Shape 일치 - 복사
+                    # Shape 확인
+                    if (old_attn.q_proj.weight.shape == 
+                        new_retention.base_retention.q_proj.weight.shape):
+                        
                         new_retention.base_retention.q_proj.weight.data = \
                             old_attn.q_proj.weight.data.clone()
                         new_retention.base_retention.k_proj.weight.data = \
@@ -412,9 +368,9 @@ def replace_attention_with_retention(model, use_hierarchical=True):
                         new_retention.base_retention.o_proj.weight.data = \
                             old_attn.o_proj.weight.data.clone()
                         
-                        print(f"     ✅ Weights copied successfully")
+                        print(f"  ✅ Layer {layer_idx}: Weights copied")
                     else:
-                        print(f"     ⚠️ Shape mismatch - using random initialization")
+                        print(f"  ⚠️ Layer {layer_idx}: Shape mismatch, random init")
                 
                 # 교체
                 layer.self_attn = new_retention
@@ -423,8 +379,7 @@ def replace_attention_with_retention(model, use_hierarchical=True):
                 print(f"  ✅ Layer {layer_idx}: Attention → Retention")
                 
         except Exception as e:
-            print(f"\n  ❌ Layer {layer_idx}: Conversion failed")
-            print(f"     Error: {e}")
+            print(f"  ❌ Layer {layer_idx}: Failed - {e}")
             import traceback
             traceback.print_exc()
             continue