From 192b323d8ecdd6bcfae3160e076ffa6a276bde8d Mon Sep 17 00:00:00 2001 From: "Chen, Zhenhuan" Date: Mon, 24 Jul 2023 00:20:29 -0700 Subject: [PATCH] add llama2 autoTP support in replace_module --- deepspeed/module_inject/replace_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 4c488e3bf2cd..6fb4745b2a13 100644 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -428,7 +428,7 @@ def update_mp_params(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" + "all_head_size", "embed_dim", "hidden_size", "num_key_value_heads" ]: if hasattr(child, param): param_val = getattr(child, param)