From abc4e6d6f97d468f788bde38b540941304782f2c Mon Sep 17 00:00:00 2001 From: Dino Chen Date: Sat, 11 Nov 2023 02:55:18 +0800 Subject: [PATCH] fix num_kv_heads sharding in autoTP for the new in-repo Falcon-40B (#4654) to be compatible with the latest Falcon-40B's `num_kv_heads` in https://huggingface.co/tiiuae/falcon-40b/commit/4a70170c215b36a3cce4b4253f6d0612bb7d4146 ![image](https://github.com/microsoft/DeepSpeed/assets/5948851/d20aa6f2-b9af-4104-b9d3-8ba1ab588a6e) error message like: ![image](https://github.com/microsoft/DeepSpeed/assets/5948851/06ef6dd2-25d5-4b51-8789-36e1b3f94a32) Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Lev Kurilenko <113481193+lekurile@users.noreply.github.com> --- deepspeed/module_inject/auto_tp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 50fab4cced37..c235cc766209 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -385,7 +385,7 @@ def update_mp_params(self, child): return for param in [ "n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", - "all_head_size", "embed_dim", "hidden_size", "num_key_value_heads" + "all_head_size", "embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads" ]: if hasattr(child, param): param_val = getattr(child, param)