enable phi2 autotp (#5436)

This PR aims to enable phi2 model autotp. --------- Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]>
microsoft · May 13, 2024 · 3a7f3aa · 3a7f3aa
1 parent 4696afd
commit 3a7f3aa
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
@@ -307,6 +307,8 @@ def tp_parser(model):
                 # Mixtral-7x8b used w2*act(w1*w3) linear. need to replace w2 to linearallreduce.
                 elif 'w2' in layer and 'Mixtral' in str(type(module)):
                     gem_list = gem_list + [layer]
+                elif "self_attn.dense" in layer and "Phi" in str(type(module)):
+                    gem_list = gem_list + [layer]
 
             layer_list = []
             if gem_list != []: