[PP] Don't allow for num_microbatches > num_stages for single stage schedules #144702
Conversation
…chedules [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/144702
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 20aba94 with merge base e6b9e67 ( UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@kwen2501, yeah this is just to guard against a bug. It should not be a restriction and if we move to a different runtime then we will allow it |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
I think this is a regression as the bug was only in the 1F1B scheduler, GPipe worked with MB < PP. It's a pretty important feature for people training very big models on long context. |
Stack from ghstack (oldest at bottom):
There is an edge case where
Schedule1F1Bwill hang when num_microbatches=1 (pytorch/torchtitan#775). For validation it makes sense to check that the number of stages should be >= number of microbatches otherwise there will be an even larger bubble.This can be removed when we have the single stage schedules to use an IR and updated to run with schedule runtime (issue tracker #144701)
cc @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o