-
Notifications
You must be signed in to change notification settings - Fork 758
add RAFT #363
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
add RAFT #363
Conversation
|
@DearAJ please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
agentlightning/verl/trainer.py
Outdated
| self._balance_batch(positive_batch, metrics=metrics) | ||
|
|
||
| # Pad batch for distributed training | ||
| positive_batch, pad_size = pad_dataproto_to_divisor(positive_batch, self.actor_rollout_wg.world_size) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why three paddings (here and L514, L577) are needed?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why three paddings (here and L514, L577) are needed?
The first padding operation was a legacy from the previous code and has now been removed. The latter two padding operations are necessary:
compute_log_prob requires distributed training, and the batch size must be divisible by the world_size.
update_actor requires distributed training, and the batch size must be divisible by the world_size.
| raft_batch.batch["returns"] = raft_batch.batch["advantages"].clone() | ||
|
|
||
| # Store labels in batch for potential use | ||
| raft_batch.batch["labels"] = labels |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
even labels are specified here. will these labels be used by the following update_actor method? or will they be used as expected?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
even labels are specified here. will these labels be used by the following update_actor method? or will they be used as expected?
These labels will not be used, but the update_actor method requires these fields. Therefore, I have retained the original GRPO processing logic to ensure the data format meets the requirements.
| # Update actor with pure SFT loss | ||
| # With advantages=1.0 and clip_ratio=1.0, this becomes standard cross-entropy | ||
| # This mimics SFTTrainer.compute_loss() behavior | ||
| actor_output = self.actor_rollout_wg.update_actor(raft_batch) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
the called update_actor function is still the RL one, not the one in "SFTTrainer"? not sure if verl has SFTTrainer...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Verl does not have a SFTrainer. SFTrainer inherits from transformers.Trainer and requires the complete HuggingFace Trainer infrastructure, whereas VERL uses Ray distributed training and a custom worker group. Directly adopting SFTrainer would disrupt VERL's existing architecture. Additionally, SFTrainer and VERL use different data formats.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Make sense. Then why SFT can be implemented in this way? Will update_actor actually compute SFT loss as we expected?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Make sense. Then why SFT can be implemented in this way? Will update_actor actually compute SFT loss as we expected?
Here is the derivation of SFT from PPO:
The standard PPO objective (with clipping) is defined as:
L = - E(min( rA, clip( r, 1-ϵ,1+ϵ )A ))
To adapt the PPO framework for Supervised Fine-Tuning (SFT) mode, we set the following neutral conditions:
Set advantages: A=1
Disable clipping: ϵ = 1.0
clip(r, 1-1, 1+0) = clip(r, 0, 2) ≈ r (r = exp(log_prob - old_log_prob), which is usually ranges in [0, 2])
When A=1 and clipping is disabled, the loss simplifies to:
L = - E(min( r, clip( r, 1-ϵ,1+ϵ ))) = -E(r)
logr = logπ(a|s)-logπold(a|s)
Therefore:
L = - E( exp(logπ(a|s)-logπold(a|s)))
In SFT, we directly optimize the current policy without relying on importance sampling. When the old policy is equal to (or very close to) the current policy, the PPO objective is replaced with the Negative Log-Likelihood loss, which is what we want to minimize:
L = - E( logπ(a|s))
For language models, this is the standard Cross-Entropy Loss.
Summary, by:
Setting A=1.0$to remove advantage weighting.
Disabling clipping to remove the PPO clipping mechanism.
The final loss effectively degenerates (or is replaced by) the standard Cross-Entropy Loss (Negative Log-Likelihood), which is the SFT loss.
Thus, the PPO framework, under these specific conditions, becomes equivalent to standard supervised learning (SFT).
| original_clip_low = self.config.actor_rollout_ref.actor.get("clip_ratio_low", 0.2) | ||
| original_clip_high = self.config.actor_rollout_ref.actor.get("clip_ratio_high", 0.3) | ||
|
|
||
| # Disable clipping: set both ratios to 1.0 (no clipping in pure SFT) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why clipping can be disabled by setting them to 1.0? This might not be related to RAFT.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here is the derivation of SFT from PPO:
The standard PPO objective (with clipping) is defined as:
L = - E(min( rA, clip( r, 1-ϵ,1+ϵ )A ))
To adapt the PPO framework for Supervised Fine-Tuning (SFT) mode, we set the following neutral conditions:
- Set advantages: A=1
- Disable clipping: ϵ = 1.0
clip(r, 1-1, 1+0) = clip(r, 0, 2) ≈ r (r = exp(log_prob - old_log_prob), which is usually ranges in [0, 2])
When A=1 and clipping is disabled, the loss simplifies to:
L = - E(min( r, clip( r, 1-ϵ,1+ϵ ))) = -E(r)
logr = logπ(a|s)-logπold(a|s)
Therefore:
L = - E( exp(logπ(a|s)-logπold(a|s)))
In SFT, we directly optimize the current policy without relying on importance sampling. When the old policy is equal to (or very close to) the current policy, the PPO objective is replaced with the Negative Log-Likelihood loss, which is what we want to minimize:
L = - E( logπ(a|s))
For language models, this is the standard Cross-Entropy Loss.
Summary, by:
- Setting A=1.0$to remove advantage weighting.
- Disabling clipping to remove the PPO clipping mechanism.
The final loss effectively degenerates (or is replaced by) the standard Cross-Entropy Loss (Negative Log-Likelihood), which is the SFT loss.
Thus, the PPO framework, under these specific conditions, becomes equivalent to standard supervised learning (SFT).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here are some things I am not very sure.
- clipping is applied to importance sampling ratio, which is a ratio between prob and old_prob. So the the numerical range of importance sampling ratio is from 0 to positive infinity. Setting high and low to 1 leads to clip(r, 0, 2), which may not disabled clipping.
- If you want to use A to implement RAFT and SFT loss, what are the "labels" used for?
Set "adv_estimator": "raft" in config.algorithm to enable support for the RAFT algorithm.