【AutoParallelism】Add refined_ops_patterns in refined-reompute

[heavyrain-lzy]

PR types

New features

PR changes

Others

Description

PCard-71568

1. Add refined_ops_patterns to support detect the special operators sequence to be excluded in reompute segments

The format is as follow:

[{main_ops:[matmul_v2, elementwise_add], num: 4, pre_ops:[]; suf_ops:[]},
 {main_ops: [flash_attn], num: -1, pre_ops: [], suf_ops: []}]

In the yaml file, you can set as follow:

  refined_ops_patterns:
    - {main_ops: [matmul_v2, elementwise_add], num: 4, pre_ops: [], suf_ops: []}
    - {main_ops: [flash_attn], num: -1, pre_ops: [], suf_ops: []}

In the shell scripts, you can use as follow:

    -o Model.refined_ops_patterns="[{\"main_ops\": [\"matmul_v2\", \"elementwise_add\"], \"num\": 4, \"pre_ops\": [], \"suf_ops\": []},\
                                    {\"main_ops\": [\"flash_attn\"], \"num\": -1, \"pre_ops\": [], \"suf_ops\": []}]"

2. Add "recompute.sr" to specify the first "recompute.sr" chunks in every device not to recompute.
   In the yaml file, you can set as follow:

sr: 0

In the shell scripts, you can use as follow:

 -o Model.sr=1

[zhiqiu]

LGTM