fix local kv head num (#806)

InternLM · Dec 6, 2023 · 5b9e454 · 5b9e454
1 parent bd7c4e3
commit 5b9e454
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/turbomind/models/llama/LlamaV2.cc b/src/turbomind/models/llama/LlamaV2.cc
@@ -76,7 +76,7 @@ LlamaV2<T>::LlamaV2(size_t                       head_num,
     end_id_(end_id),
     hidden_units_(head_num * size_per_head),
     local_head_num_(head_num / tensor_para.world_size_),
-    local_kv_head_num_(head_num / tensor_para.world_size_),
+    local_kv_head_num_(kv_head_num / tensor_para.world_size_),
     weights_(weights),
     tensor_para_(tensor_para),
     stream_(stream),