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 <[email protected]>
microsoft · Oct 31, 2024 · c7f58c8 · c7f58c8
1 parent 9b54731
commit c7f58c8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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)