[SFT] Support context parallelism for SFT #132
Merged
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
…and make FSDP wrapping policy conditional
|
ok the training script is working now! 
|
|
This PR should be ready for review! I just added a CI that checks for loss match
|
xingyaoww
changed the title
verl/trainer/fsdp_sft_trainer.py
Outdated
| @@ -165,6 +193,14 @@ def _build_model_optimizer(self): | |||
| trust_remote_code = self.config.model.trust_remote_code | |||
| # load config first | |||
| config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=trust_remote_code) | |||
| if self.use_remove_padding: | |||
| assert self.config.ulysses_sequence_parallel_size > 1, "Remove padding is only supported with sequence parallel" | |||
There was a problem hiding this comment.
I assume it should be the opposite? Sequence parallel is only support when remove_padding is enabled?
verl/trainer/fsdp_sft_trainer.py
Outdated
| loss.backward() | ||
| if self.use_remove_padding and self.config.ulysses_sequence_parallel_size > 1: | ||
| # micro_batch = micro_batch.to('cuda') | ||
| loss = self._compute_loss_and_backward_sp(batch=micro_batch) / n_micro_batches |
There was a problem hiding this comment.
Can we combine the two functions as there are plenty replicated code?
|
Nice work! |
1. Fix assertion logic - sequence parallel requires remove_padding 2. Combine loss computation functions to reduce code duplication
|
Could you run the formatting script? |
|
done! @vermouth1992 |
|
There are conflicts with newly merge MR.. |
vermouth1992
approved these changes
Jan 27, 2025
Chendong98
pushed a commit
to Chendong98/verl
that referenced
this pull request
Feb 4, 2025
# Add Sequence Parallelism and Padding Removal to SFT Trainer This PR adds sequence parallelism (SP) and padding removal optimizations to the SFT trainer, which can help improve training efficiency for large language models. ## Key Changes ### Core Features 1. **Sequence Parallelism**: Added support for sequence parallelism through the Ulysses framework - Configurable via `ulysses_sequence_parallel_size` parameter - Properly handles data distribution across SP ranks - Maintains consistent loss computation across distributed setup 2. **Padding Removal**: Added support for efficient handling of variable-length sequences - Enabled via `use_remove_padding` flag (requires SP to be enabled) - Uses flash-attention's padding removal utilities - Handles proper re-padding and loss computation 3. **Training Improvements**: - Added label smoothing support to loss computation - Added progress bar with epoch information - Added RoPE scaling configuration support - Improved error messages for batch size validation ### Testing - Added comprehensive test suite (`test_trainer.py`) to verify: - Forward pass consistency between original and SP+rmpad implementations - Loss computation correctness across distributed setup - Proper handling of micro-batches ### Example Usage Added example script `examples/sft/gsm8k/run_qwen_05_sp2.sh` demonstrating how to use the new features with Qwen-2.5B model. ## Implementation Details - Uses device mesh for proper distributed training setup - Handles data distribution ensuring same sequences within SP groups but different across DP groups - Carefully manages backward pass timing with gradient checkpointing - Maintains compatibility with existing FSDP features ## Testing Instructions 1. Run the example script with sequence parallelism: ```bash bash examples/sft/gsm8k/run_qwen_05_sp2.sh <nproc_per_node> <save_path> ``` 2. Run the test suite: ```bash tests/sft/run_sft_sp_loss_match.sh``` ^^ These are PR description generated by [OpenHands](https://github.com/All-Hands-AI/OpenHands) --------- Co-authored-by: Jiayi Pan <[email protected]> Co-authored-by: openhands <[email protected]>
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.
Add Sequence Parallelism and Padding Removal to SFT Trainer
This PR adds sequence parallelism (SP) and padding removal optimizations to the SFT trainer, which can help improve training efficiency for large language models.
Key Changes
Core Features
Sequence Parallelism: Added support for sequence parallelism through the Ulysses framework
ulysses_sequence_parallel_sizeparameterPadding Removal: Added support for efficient handling of variable-length sequences
use_remove_paddingflag (requires SP to be enabled)Training Improvements:
Testing
test_trainer.py) to verify:Example Usage
Added example script
examples/sft/gsm8k/run_qwen_05_sp2.shdemonstrating how to use the new features with Qwen-2.5B model.Implementation Details
Testing Instructions
bash tests/sft/run_sft_sp_loss_match.sh^^ These are PR description generated by OpenHands