[3D-parallelism] Hybrid Model Parallelism

[JZ-LIANG]

PR types

New features

PR changes

APIs

Describe

• new features

  • Hybrid Model Parallelism：
    • Combine 3 single Model parallelism strategies (megatron, sharding, pipline) together as a Hybrid Parallelism strategy
    • Uniform switch to turn on/off each single parallelism strategy (temporally use Sharding & Pipeline as the uniform API)

• performance optimization：

  • speed
    • remove potential un-necessary sync_calc & sync_calc in hybrid model parallelism
  • Memory usage
    • optimizer offload
    • optimizer temporary vars inplacement [commit withdraw, update in next pr]

• performance-related

  • the order of parallelism from inner to outer is : mp --> sharidng --> pp
  • mp (megatron) and sharding parallelism will introduce large communication and is recommend: to be used within node. (mp_degree * sharding_degree = number of gpu per node)
  • pp parallelism have less communication load compared with above two, which make them more suitable to use across the nodes (pp_degree = number node)

example

• assume we have 4 nodes with 8 gpus per node:

• mp-sharding-pp 3D parallelism

        dist_strategy.sharding = True
        dist_strategy.pipeline = True
        dist_strategy.sharding_configs = {"segment_broadcast_MB": 32,
                                            "pp_degree": 4,
                                            "sharding_degree":4,
                                            "mp_degree": 2,
                                            "optimize_offload": True,
                                            }
        dist_strategy.pipeline_configs = {"schedule_mode": "1F1B",
                                            "micro_batch_size": 1,
                                            "accumulate_steps": 4,
                                            }

• mp-pp 2D parallelism

        dist_strategy.sharding = True
        dist_strategy.pipeline = True
        dist_strategy.sharding_configs = {"segment_broadcast_MB": 32,
                                            "pp_degree": 4,
                                            "sharding_degree":1,
                                            "mp_degree": 8,
                                            "optimize_offload": True,
                                            }
        dist_strategy.pipeline_configs = {"schedule_mode": "1F1B",
                                            "micro_batch_size": 1,
                                            "accumulate_steps": 4,
                                            }

• shardign-pp 2D parallelism

        dist_strategy.sharding = True
        dist_strategy.pipeline = True
        dist_strategy.sharding_configs = {"segment_broadcast_MB": 32,
                                            "pp_degree": 1,
                                            "sharding_degree":4,
                                            "mp_degree": 8,
                                            "optimize_offload": True,
                                            }
        dist_strategy.pipeline_configs = {"schedule_mode": "1F1B",
                                            "micro_batch_size": 1,
                                            "accumulate_steps": 4,
                                            }

• mp-sharding 2D parallelism

        dist_strategy.sharding = True
        dist_strategy.sharding_configs = {"segment_broadcast_MB": 32,
                                            "pp_degree": 1,
                                            "sharding_degree":4,
                                            "mp_degree": 8,
                                            "optimize_offload": False,
                                            "gradient_merge_acc_step": 4,
                                            }

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-bot-old]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[wangxicoding]

Enum comments

[JZ-LIANG]

recorded, document will be added in fluiddoc and fleetx

[wangxicoding]

Also need add comments to this code.

[wangxicoding]

Add some comments, in 3d or 4d parallel, allreduce_in_optimize=True can reduce communication, allreduce_in_optimize=False can reduce memory

[JZ-LIANG]

recorded, document will be added in fluiddoc and fleetx and .py file where the feature is called.

but I think this should be a feature for internal project now, and we should not expose It to users ?

[zhiqiu]

LGTM for backward.py

[wangxicoding]

LGTM