[PP] microbatch split config #947
Open
+17
−6
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
tianyu-l
reviewed
Mar 9, 2025
| @@ -359,6 +359,14 @@ def __init__(self): | |||
| If using looped schedules, this still specifies the number of physical ranks, not the number | |||
| of stages. Stages per rank are inferred from split points degree, and schedule.""", | |||
| ) | |||
| self.parser.add_argument( | |||
| "--experimental.pipeline_parallel_batch_split_dim", | |||
There was a problem hiding this comment.
It seems PP has been in experimental for a while. Do you think it's time we extract pipeline_parallel into a standalone section and put all configs over there?
It doesn't have to happen in this PR.
| @@ -119,17 +121,18 @@ def build_pipeline_schedule( | |||
| f"of stages ({num_total_stages}) which may result in a bubble in the pipeline." | |||
| ) | |||
|
|
|||
| # validate that the batch size is divisible by the number of microbatches otherwise we'll hang or error during training | |||
There was a problem hiding this comment.
I have several questions here:
- If
pipeline_parallel_batch_split_dim == 0, what would happen ifif job_config.training.batch_size % num_total_stages != 0? - If
pipeline_parallel_batch_split_dimis on the sequence dim or other dims, don't we need similar checks in the extremal cases e.g. seq_len < num_stages - Btw this divisibility requirement seems not exactly the same as "batch_size >= num stages" you mentioned in the PR summary.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Allow the input batch to be split on the sequence dimension in pipeline parallelism (remove the requirement for batch_size >= num stages)
depends on pytorch/pytorch#148458
The new config to set this is
pipeline_parallel_batch_split_dim = 1