From 3a7f3aa8498582860cfdaca6d1e19a968964ea76 Mon Sep 17 00:00:00 2001 From: Yejing-Lai Date: Tue, 14 May 2024 04:10:53 +0800 Subject: [PATCH] enable phi2 autotp (#5436) This PR aims to enable phi2 model autotp. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Olatunji Ruwase --- deepspeed/module_inject/auto_tp.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 4944e1954..5e0aa0286 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -307,6 +307,8 @@ class AutoTP(): # 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 != []: