From d881b3dcbd6d841c482392f1fcf1487cda56e9a7 Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Mon, 10 Apr 2023 11:35:40 +0800 Subject: [PATCH 1/2] Enable auto TP policy for llama model --- deepspeed/module_inject/auto_tp.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index bf49df9781f5..396fe7db2447 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -104,6 +104,10 @@ def tp_parser(model): gem_list = gem_list + [layer_list[i - 1]] elif 'out_proj' in layer: gem_list = gem_list + [layer] + elif 'o_proj' in layer: + gem_list = gem_list + [layer] + elif 'down_proj' in layer: + gem_list = gem_list + [layer] layer_list = [] if gem_list != []: gem_list = list(set(gem_list)) From fc885d0003017710a420b9effd48160b3380df22 Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Wed, 12 Apr 2023 15:31:25 +0800 Subject: [PATCH 2/2] Update automatic-tensor-parallelism.md --- docs/_tutorials/automatic-tensor-parallelism.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/_tutorials/automatic-tensor-parallelism.md b/docs/_tutorials/automatic-tensor-parallelism.md index 6991d5caf925..c4654af180f9 100644 --- a/docs/_tutorials/automatic-tensor-parallelism.md +++ b/docs/_tutorials/automatic-tensor-parallelism.md @@ -120,6 +120,7 @@ The following model families have been successfully tested with automatic tensor - gpt-neox - longt5 - luke +- llama - m2m_100 - marian - mvp