Repository navigation
loss: add --use-adaptive-entropy (Skywork-OR1) - #3220
Open
behroozazarkhalili wants to merge 2 commits into
Open
behroozazarkhalili wants to merge 2 commits into
behroozazarkhalili wants to merge 2 commits into
Conversation
behroozazarkhalili
requested review from
Rockdu,
Shi-Dong,
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
September 12, 2026 06:52
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
A fixed --entropy-coef either does nothing useful or keeps pushing entropy up after the policy has recovered. Skywork-OR1's adaptive scheme applies the bonus only while entropy is at or below a target and moves the coefficient by a fixed delta per optimizer step, up at or below the target and down above it, within a clamp. entropy_control.py holds the rule: the loss reads the coefficient gated by the previous step's entropy, so every micro-batch of a step applies the same bonus, and both trainers update it once per step from the reduced entropy_loss, which aggregate_train_losses makes identical on every loss rank. Entropy is computed whenever the mode is on so the controller keeps its signal, but it carries a gradient only when the bonus is non-zero. The runtime state lives on args.adaptive_entropy_coef so the user's --entropy-coef is never rewritten, and the coefficient that will apply next is reported as the entropy_coef train metric. Off by default.
…zer step Preserve the adaptive entropy state when a non-finite grad norm skips the optimizer step or --disable-optimizer prevents an update. Also leave it unchanged when no optimizer is present. Previously, NaN entropy decremented the coefficient and became adaptive_entropy_last. Since NaN > target is false, this left the entropy bonus gate open for the next step. Guard the Megatron update with not disable_optimizer and valid_step so it now matches the adaptive KL controller. Cover skipped and applied steps through the existing fast train_one_step fixture.
behroozazarkhalili
force-pushed
the
feat/adaptive-entropy
branch
from
October 15, 2026 19:41
4123847 to
75d6bba
Compare
Contributor
Author
|
@fzyzcjy @yueming-yuan friendly ping. This is rebased onto current |
This branch has not been deployed
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
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.
What
Adds
--use-adaptive-entropy, the Skywork-OR1 adaptive entropy bonus, to the policy loss. With it on, the entropy bonus applies only while the previous step's entropy is at or below--entropy-target, and the coefficient moves by--entropy-coef-deltaper optimizer step (up at or below the target, down above it) within[--entropy-coef-min, --entropy-coef-max], starting from--entropy-coef.loss_hub/entropy_control.pyholds the rule:entropy_coef_to_apply(args)(gated by the previous step's entropy, so every micro-batch of a step applies the same bonus) andupdate_adaptive_entropy(args, loss_dict), called once per optimizer step in both the FSDP actor and the Megatrontrain_one_step(last pipeline stage) on the reducedentropy_loss, whichaggregate_train_lossesmakes identical on every loss rank.args.adaptive_entropy_coef/args.adaptive_entropy_last;--entropy-coefitself is never rewritten. The coefficient that will apply next is reported as theentropy_coeftrain metric.--entropy-coefpath is unchanged.Test
tests/fast/backends/training_utils/loss_hub/test_entropy_control.py: gate on/off, up/down moves, clamps, no-op when off or when the metric is missing,entropy_coefreporting.tests/fast/backends/training_utils/loss/test_adaptive_entropy.py: throughpolicy_loss_functionwith the shared fixtures: gated-off step equals the zero-coefficient loss while still reporting entropy; gated-on step equals the static-coefficient loss; entropy requires grad only when the bonus applies.Verification
On a CPU SLURM node: the 16 new tests pass;
tests/fast/backends/training_utils,tests/fast/backends/test_fsdp_train.py,tests/fast/backends/megatron_utilsandtests/fast/utils/test_arguments.pypass (756 passed, 11 skipped); removing the gate inentropy_coef_to_applymakes 4 of the new tests fail;ruff,black -l 119,isort --profile blackclean. Not run on GPU.Scope
Loss-side only. No change to rollout, checkpoints, or the fixed-coefficient path. Megatron pickles
argsinto checkpoints but does not compare or restore these fields, so a resumed run restarts the controller from--entropy-coef.cc @fzyzcjy @yueming-yuan