any plans to add improved PP in llama-3?

[SeunghyunSEO]

batch constraint removed Pipeline Parallelism (PP) in llama-3 paper seems to added torch 2.4.
do you have any plans to add this feature to torchtitan ?!

[tianyu-l]

@wconstab @H-Huang anything to share?

[H-Huang]

Hi @SeunghyunSEO, the ScheduleFlexibleInterleaved1F1B was not introduced in 2.4, but it is available in the PyTorch nightly build (and will be available in the 2.5 release).

For torchtitan I can quickly add a PR to add a flag to support this schedule.

[H-Huang]

merged in 9cf4b2f