perf(pipeline): implement auto-partition algorithm

[TXacs]

Auto-Partition in torchtitan

Overview

This PR provides an automatic partitioning method that considers the computation cost of embedding layers.
Thsi method involves calculating the floating-point operations (FLOPs) of the embedding layers and constructing an array that incorporates the FLOPs of both the transformer and embedding layers. Subsequently, a heuristic algorithm is employed to identify a balanced pipeline partition.

Solution Architecture

1. Dynamic Cost Analysis
2. Adaptive Partitioning Algorithm
3. Workload Balancing

Performance

Hardware configuration: 4x RTX 3090 24GB, pipeline parallelism dimension is 4.

llama3 配置对比

hidden size	layers	autopipe TPS	default TPS	Speedup
dim=256	6	31,094	29,549	+5.2%
dim=256	12	21,803	21,923	-0.5%
dim=2048	12	3,348	2,616	+28.0%
dim=4096	12	981	761	+28.9%

deepseekv3(without moe) 配置对比

hidden size	layers	autopipe TPS	default TPS	Speedup
dim=256	6	13,373	13,059	+2.4%
dim=256	12	7,714	6,859	+12.5%
dim=2048	12	4,331	3,810	+13.7%
dim=4096	12	2,888	2,561	+12.8%
dim=4096	16	2,207	2,008	+9.9%
dim=8192	16	4,331	3,935	+10.1%

[meta-cla]

Hi @TXacs!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

[tianyu-l]

Thanks. Is it true that the only "real" deltas are

• autopipe.cpp
• pipeline_parallel.py
• profiler.py

[tianyu-l]

This looks interesting -- how much benefit you'd get from having a c++ implementation, compared with a python one?

[TXacs]

Thanks. Is it true that the only "real" deltas are

• autopipe.cpp
• pipeline_parallel.py
• profiler.py

Yes，actually, profile.py also uses the file from DeepSpeed. It would be even better if TorchTitan could provide a more authoritative FLOPs calculation method in the future, so that we could also adapt it for MoE models.

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[H-Huang]

why is it assumed to be 3x?

[McmillanTAC]

Similar to the computation time, in which the backward computation time is twice that of the forward computation time, by default, we assume that the floating-point operations (FLOPs) required for backward computation are twice those for forward computation. During the modeling process, we take the recomputation technique into account. This technique inserts an additional forward computation before the backward computation. Consequently, we set the default FLOPs for backward computation to be three times those of forward computation.

[H-Huang]

I guess the FlopsProfiler does not estimate the backward flops?

[McmillanTAC]

You are correct that the FlopsProfiler does not estimate backward flops. Instead, we use a heuristic rule by default: given the adoption of recomputation techniques in the backward computation, backward flops are set to three times those of the forward computation.

[tianyu-l]

curious why three times? IIUC the convention was two times.