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 <[email protected]>
Co-authored-by: Lev Kurilenko <[email protected]>

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)