Set pipeline_stage_id is supported!!

[CPFLAME]

In most cases, the transformer layers of common models have the same computational overhead, so there is no need to set custom_pipeline_stage_id.

But when transformer layers have unbalanced computational overhead, you can set custom_pipeline_stage_id for manually balance the compuation between stages in pipeline_parallelism

For example:

train.dist.pipeline_parallel_size = 4
train.dist.pipeline_num_layers = 24
train.dist.custom_pipeline_stage_id = [0]*6 + [1]*7 + [2]*7 + [3]*4

It means you have [6, 7, 7, 4] layers separately located in stage0~stage3.
Modify custom_pipeline_stage_id according to your own needs.