fix

Author: grimoire

lmdeploy/pytorch/models/deepseek_v32.py

@@ -256,11 +256,8 @@ def forward(
     ):
         """Rewrite of LlamaAttention.forward."""
         dist_ctx = get_dist_manager().current_context()
-        if dist_ctx.dp > 1:
-            num_heads = self.num_heads
-        else:
-            world_size = dist_ctx.world_size
-            num_heads = self.num_heads // world_size
+        tp_world_size = dist_ctx.dist_config.attn_tp
+        num_heads = self.num_heads // tp_world_size
         nope_size = self.kv_lora_rank
         q_len = hidden_states.size(1)
 

lmdeploy/pytorch/nn/rotary_embedding.py

@@ -201,12 +201,12 @@ def __init__(self):
     def forward(self, query: Tensor, key: Tensor, cos: Tensor, sin: Tensor, inplace: bool = True):
         """forward."""
 
-        assert query.dim() == key.dim() == 3, 'Expected query key (seq_len, heads, head_dim)'
         assert cos.dim() <= 3 and sin.dim() <= 3
 
         need_reshape = False
         if cos.dim() == 3:
             # for fope
+            assert query.dim() == key.dim() == 3, 'Expected query key (seq_len, heads, head_dim)'
             need_reshape = True
             query_shape = query.shape
             key_shape = key.shape