n_heads for qwen3

2026-02-23 20:07:10 +00:00
parent c90a0cec78
commit b48f3c530a
1 changed files with 2 additions and 0 deletions
@@ -852,6 +852,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                layer.self_attn.o_proj = self.sharded_to_all_linear(
                    layer.self_attn.o_proj
                )
+                layer.self_attn.n_heads //= self.N
+                layer.self_attn.n_kv_heads //= self.N
            else:
                assert isinstance(layer, Qwen3NextDecoderLayer)
                if hasattr(layer, "linear_attn"):