From c7f58c899f6f099a35d968bdad973f24b842c8c6 Mon Sep 17 00:00:00 2001 From: Yejing-Lai Date: Thu, 31 Oct 2024 08:48:52 +0800 Subject: [PATCH] Add attribute check to support git-base autotp (#6688) Git-base model is an image-text model. After supporting the llama3.2 vision model, we set num_kv_heads dynamically. Git-base only includes vision_config, so we need to add an attribute check for vision_config/text_config when setting num_kv_heads. Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> --- deepspeed/module_inject/replace_module.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index dece5a2c4ecf..8b1455f20c69 100644 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -277,8 +277,10 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None): if hasattr(model_config, "vision_config"): if "MllamaVisionEncoderLayer" in str(module): num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config) - else: + elif hasattr(model_config, "text_config"): num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config) + else: + num_kv_heads = _autotp.get_model_num_kv_heads(model_config) else: num_kv_heads = _autotp.get_model_num_kv_heads(model_config)