【AutoParallelism】add "refined_ops_patterns" in recompute strategy (#7349

) * add refined_ops_patterns in recompute strategy * modify the yaml file * add sr * fix yaml file"
PaddlePaddle · Nov 2, 2023 · c013ccb · c013ccb
1 parent fc2bf9c
commit c013ccb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/model_zoo/gpt-3/ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_6.7B_dp2_mp2_pp2_sharding2.yaml b/model_zoo/gpt-3/ppfleetx/configs/nlp/gpt/auto/pretrain_gpt_6.7B_dp2_mp2_pp2_sharding2.yaml
@@ -20,6 +20,10 @@ Model:
   fuse_attn_qkv: True
   use_recompute: True
   recompute_granularity:
+  sr: 0
+  refined_ops_patterns:
+    - {main_ops: [matmul_v2, elementwise_add], num: 0, pre_ops: [], suf_ops: []}
+    - {main_ops: [flash_attn], num: 0, pre_ops: [], suf_ops: []}
   no_recompute_layers:
 
 

diff --git a/model_zoo/gpt-3/ppfleetx/utils/auto_config.py b/model_zoo/gpt-3/ppfleetx/utils/auto_config.py
@@ -178,6 +178,8 @@ def process_strategy(config):
             config.Model["no_recompute_layers"] = sorted(list(set(config.Model["no_recompute_layers"])))
         recompute = strategy.recompute
         recompute.enable = config.Model.get("use_recompute", False)
+        recompute.sr = config.Model.pop("sr", 0)
+        recompute.refined_ops_patterns = config.Model.pop("refined_ops_patterns", []) # gpt.GPTModelAuto don't need this parameter
         recompute.no_recompute_segments = config.Model.pop("no_recompute_layers", [])
         recompute.enable_tuning = config.get("Tuning", False) and config.Tuning.get("tuning_recompute", False)