From bafbd06744a8e6ddccefcd1fe2a9dee5ac6b5d96 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Fri, 24 Jan 2025 15:06:50 +0100 Subject: [PATCH] Update transformers_flash_causal_lm.py --- .../models/transformers_flash_causal_lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/models/transformers_flash_causal_lm.py b/server/text_generation_server/models/transformers_flash_causal_lm.py index be035417..7513ce01 100644 --- a/server/text_generation_server/models/transformers_flash_causal_lm.py +++ b/server/text_generation_server/models/transformers_flash_causal_lm.py @@ -110,7 +110,7 @@ transformers.models.cohere.configuration_cohere.CohereConfig.base_model_tp_plan "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } -transformers.models.cohere.configuration_cohere2.Cohere2Config.base_model_tp_plan = { +transformers.models.cohere2.configuration_cohere2.Cohere2Config.base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise",