[Possible PR discuss] Automatic Pipeline Parallel Partitioning Considering Embedding Layers

[McmillanTAC]

Currently, pipeline parallel partitioning requires manual configuration of pipeline_parallel_first_stage_less_layers and pipeline_parallel_last_stage_less_layers to address the computation imbalance introduced by embedding layers. This manual setup can be challenging for non-expert users, as these hyperparameters are not intuitive and often require specific adjustments for different model configurations.

To improve usability, we have developed a prototype for an automatic partitioning method that considers the computation cost of embedding layers. This implementation aims to dynamically balance the pipeline stages without needing manual intervention for the aforementioned parameters. We wonder whether a PR of integrating it to torchtitan as an experimental feature is welcomed.

[wwwjn]

cc @H-Huang

[tianyu-l]

Thanks for asking!

IMO this is not the highest priority right now, but if you would like to start sharing a design or submitting a draft PR and see if the community finds it useful, feel free to do so.

[H-Huang]

we have developed a prototype for an automatic partitioning method that considers the computation cost of embedding layers.

Curious how this is implemented!

[McmillanTAC]

Thanks for asking!

IMO this is not the highest priority right now, but if you would like to start sharing a design or submitting a draft PR and see if the community finds it useful, feel free to do so.

Thanks for your response! We'll prepare a draft PR as soon as possible to share with the community for further discussion.

[McmillanTAC]

we have developed a prototype for an automatic partitioning method that considers the computation cost of embedding layers.

Curious how this is implemented!

Thank you for your interest! We would share a more detailed explanation of the approach in the draft PR.