Fix parallelism_config being overwritten in TP-only training #42671
+10
−7
Conversation
This file contains hidden or 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
Fixes huggingface#42661 When using TP with a user-provided parallelism_config, the Trainer was incorrectly overwriting the entire config object with a new ParallelismConfig(tp_size=model.tp_size), discarding all user-provided settings (dp_size, pp_size, cp_backend, etc.). Changes: - If user provides parallelism_config, update only the tp_size attribute - If no config is provided, create a new ParallelismConfig with tp_size - Removed redundant nested condition check - Fixed logic flow to check accelerate version first This ensures user-provided parallelism configurations are preserved during TP-only training.
There was a problem hiding this comment.
Pull request overview
This PR fixes a bug where the Trainer was discarding user-provided parallelism_config settings when Tensor Parallelism (TP) was enabled. Previously, the entire config object was being overwritten; now only the tp_size attribute is updated while preserving other user settings.
Key changes:
- Fixed conditional logic to update
tp_sizein existing user configs instead of creating a new config - Restructured version checks to be more clear and avoid redundant nesting
Comment on lines
+5078
to
+5085
| if self.args.parallelism_config is not None: | ||
| # Update tp_size in user-provided config instead of overwriting it | ||
| self.args.parallelism_config.tp_size = self.model.tp_size | ||
| else: | ||
| raise ValueError("Requires accelerate>1.10.1 to use Tensor Parallelism.") | ||
| # Only create new config if user didn't provide one | ||
| from accelerate import ParallelismConfig | ||
|
|
||
| args["parallelism_config"] = ParallelismConfig(tp_size=self.model.tp_size) |
There was a problem hiding this comment.
This bug fix should have test coverage to ensure that user-provided parallelism_config objects are preserved when TP is enabled. Consider adding a test that:
- Creates a
ParallelismConfigwith custom settings (e.g., dp_size, cp_size) - Passes it to TrainingArguments along with a model that has tp_size > 1
- Verifies that after trainer initialization, the parallelism_config still has the original custom settings AND the updated tp_size
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Fixes #42661
Problem
When using TP with user-provided parallelism_config, the Trainer overwrites the entire config object, discarding user settings.
Solution
Changes