loss: add --use-adaptive-entropy (Skywork-OR1)

docs/user-guide/cli-reference.md

@@ -69,12 +69,23 @@ then push up until you OOM.
 | `--use-kl-loss` | off | Compute KL against the reference model. |
 | `--kl-loss-coef` | `0.0` | Weight of KL in the loss (0 means monitor only). |
 | `--kl-loss-type` | `k1` | `k1`, `k2`, `k3`, `low_var_kl`. |
-| `--entropy-coef` | `0.0` | Entropy bonus weight. |
+| `--entropy-coef` | `0.0` | Entropy bonus weight; starting coefficient in adaptive mode. |
 | `--observe-training-entropy` | off | Log training entropy even when `--entropy-coef` is `0.0`; detached from backward when the coefficient is zero. |
+| `--use-adaptive-entropy` | off | Skywork-OR1 adaptive entropy bonus: applied only while entropy is at or below `--entropy-target`, with the coefficient moving by `--entropy-coef-delta` per step within `[--entropy-coef-min, --entropy-coef-max]`. |
+| `--entropy-target` | `0.2` | Target training entropy in nats, using the `entropy_loss` reduction. |
+| `--entropy-coef-delta` | `0.005` | Adaptive coefficient change per optimizer step. |
+| `--entropy-coef-min` | `0.0` | Lower clamp for the adaptive coefficient. |
+| `--entropy-coef-max` | `1.0` | Upper clamp for the adaptive coefficient. |
 | `--eps-clip` | `0.2` | PPO/GRPO low clip. |
 | `--eps-clip-high` | `–` | Asymmetric high clip (DAPO-style). |
 | `--use-tis` | off | Truncated Importance Sampling for train/inference precision mismatch. |
 
+Adaptive mode always reports `train/entropy_loss`, including steps with no entropy
+bonus. It also reports `train/entropy_coef`: the coefficient that will apply on
+the next optimizer step, including zero when the gate is closed. Entropy is
+detached from backward whenever the applied coefficient is zero. On process
+restart, the controller starts again from `--entropy-coef`.
+
 ### Sampling
 
 | Flag | Default | What |
@@ -243,7 +254,12 @@ Sections mirror the launch-script argument groups.
 | `--use-kl-loss` | flag | off | Compute KL vs. reference. |
 | `--kl-loss-coef` | float | `0.0` | KL weight in loss (0 means monitor). |
 | `--kl-loss-type` | enum | `k1` | `k1`, `k2`, `k3`, `low_var_kl`. |
-| `--entropy-coef` | float | `0.0` | Entropy bonus weight. |
+| `--entropy-coef` | float | `0.0` | Entropy bonus weight; starting coefficient in adaptive mode. |
+| `--use-adaptive-entropy` | flag | off | Adaptive entropy bonus, see `--entropy-target`, `--entropy-coef-delta`, `--entropy-coef-min`, `--entropy-coef-max`. |
+| `--entropy-target` | float | `0.2` | Target training entropy in nats, using the `entropy_loss` reduction. |
+| `--entropy-coef-delta` | float | `0.005` | Adaptive coefficient change per optimizer step. |
+| `--entropy-coef-min` | float | `0.0` | Lower clamp for the adaptive coefficient. |
+| `--entropy-coef-max` | float | `1.0` | Upper clamp for the adaptive coefficient. |
 | `--observe-training-entropy` | flag | off | Log detached training entropy when entropy bonus weight is zero. |
 | `--eps-clip` | float | `0.2` | PPO/GRPO low clip. |
 | `--eps-clip-high` | float | – | Asymmetric high clip. |

miles/backends/megatron_utils/model.py

@@ -30,6 +30,7 @@
 from miles.backends.megatron_utils.optimizer_state_reset import reset_optimizer_states
 from miles.backends.training_utils.data.rollout import DataIterator, get_batch
 from miles.backends.training_utils.data.sampling_mask import get_rollout_sampling_masks
+from miles.backends.training_utils.loss.hub.entropy_control import update_adaptive_entropy
 from miles.backends.training_utils.loss.objective import loss_function
 from miles.backends.training_utils.metrics.checks import check_grad_norm, check_kl
 from miles.backends.training_utils.metrics.log_utils import (
@@ -704,6 +705,8 @@ def train_one_step(
                 if parallel_state.indep_dp.size > 1
                 else aggregate_train_losses(losses_reduced, metric_num_rollouts)
             )
+            if not disable_optimizer and valid_step:
+                update_adaptive_entropy(args, loss_reduced)
             return loss_reduced, grad_norm, outcome
 
     return {}, grad_norm, outcome

miles/backends/training_utils/loss/hub/entropy_control.py

@@ -0,0 +1,45 @@
+"""Adaptive entropy regularization (Skywork-OR1).
+
+With ``--use-adaptive-entropy`` the entropy bonus only acts while the policy's
+entropy is at or below ``--entropy-target``, and its coefficient moves by
+``--entropy-coef-delta`` per optimizer step, up when entropy is at or below the
+target and down otherwise, clamped to ``[--entropy-coef-min, --entropy-coef-max]``.
+The coefficient and the gate use the entropy of the previous step, so every
+micro-batch of a step applies the same bonus; ``update_adaptive_entropy`` runs
+once per step on the reduced ``entropy_loss`` (torch-native backends, and the last
+pipeline stage in Megatron). ``aggregate_train_losses`` reduces over ``effective_dp_cp`` including the
+outer ``indep_dp`` group, so every loss rank sees the same entropy and stays in step.
+
+Runtime state lives on ``args`` as ``adaptive_entropy_coef`` and ``adaptive_entropy_last``;
+``--entropy-coef`` itself is never mutated. Megatron pickles ``args`` into checkpoints but
+does not compare or restore these fields, so a resumed run restarts the controller from
+``--entropy-coef``. Reported as ``entropy_coef`` in the train metrics.
+"""
+
+from typing import Any
+
+ENTROPY_KEY = "entropy_loss"
+
+
+def entropy_coef_to_apply(args) -> float:
+    """The coefficient for this step's policy loss."""
+    if not getattr(args, "use_adaptive_entropy", False):
+        return args.entropy_coef
+    last = getattr(args, "adaptive_entropy_last", None)
+    if last is not None and last > args.entropy_target:
+        return 0.0
+    return getattr(args, "adaptive_entropy_coef", args.entropy_coef)
+
+
+def update_adaptive_entropy(args, loss_dict: dict[str, Any]) -> None:
+    """Update runtime state and log the gated coefficient for the next step in loss_dict."""
+    if not getattr(args, "use_adaptive_entropy", False) or ENTROPY_KEY not in loss_dict:
+        return
+    entropy = float(loss_dict[ENTROPY_KEY])
+    coefficient = getattr(args, "adaptive_entropy_coef", args.entropy_coef)
+    if entropy <= args.entropy_target:
+        args.adaptive_entropy_coef = min(coefficient + args.entropy_coef_delta, args.entropy_coef_max)
+    else:
+        args.adaptive_entropy_coef = max(coefficient - args.entropy_coef_delta, args.entropy_coef_min)
+    args.adaptive_entropy_last = entropy
+    loss_dict["entropy_coef"] = entropy_coef_to_apply(args)

miles/backends/training_utils/loss/hub/losses.py

@@ -11,6 +11,7 @@
 )
 from miles.backends.training_utils.data.sampling_mask import get_rollout_sampling_masks
 from miles.backends.training_utils.loss.hub.corrections import vanilla_tis_function
+from miles.backends.training_utils.loss.hub.entropy_control import ENTROPY_KEY, entropy_coef_to_apply
 from miles.backends.training_utils.loss.hub.logit_processors import get_log_probs_and_entropy, get_values
 from miles.backends.training_utils.loss.hub.math_utils import (
     compute_approx_kl,
@@ -112,7 +113,10 @@ def policy_loss_function(
     response_lengths = batch["response_lengths"]
     total_lengths = batch["total_lengths"]
     max_seq_lens = batch.get("max_seq_lens", None)
-    calculate_entropy = args.entropy_coef != 0 or args.observe_training_entropy
+    use_adaptive_entropy = getattr(args, "use_adaptive_entropy", False)
+    # Adaptive mode gates and sizes the bonus from the previous step's entropy.
+    entropy_coef = entropy_coef_to_apply(args)
+    calculate_entropy = entropy_coef != 0 or args.observe_training_entropy or use_adaptive_entropy
     rollout_sampling_mask = get_rollout_sampling_masks(batch) if args.use_sampling_support_replay else None
 
     log_probs_and_entropy = get_log_probs_and_entropy(
@@ -122,7 +126,7 @@ def policy_loss_function(
         total_lengths=total_lengths,
         response_lengths=response_lengths,
         with_entropy=calculate_entropy,
-        entropy_requires_grad=args.entropy_coef != 0,
+        entropy_requires_grad=entropy_coef != 0,
         max_seq_lens=max_seq_lens,
         rollout_sampling_mask=rollout_sampling_mask,
     )
@@ -310,8 +314,8 @@ def policy_loss_function(
         entropy = log_probs_and_entropy["entropy"]
         entropy = torch.cat(entropy, dim=0)
         entropy_loss = sum_of_sample_mean(entropy)
-        if args.entropy_coef != 0:
-            loss = pg_loss - args.entropy_coef * entropy_loss
+        if entropy_coef != 0:
+            loss = pg_loss - entropy_coef * entropy_loss
         else:
             entropy_loss = entropy_loss.detach()
 
@@ -370,7 +374,7 @@ def policy_loss_function(
     reported_loss = {
         "loss": loss.clone().detach(),
         "pg_loss": pg_loss.clone().detach(),
-        "entropy_loss": entropy_loss.clone().detach(),
+        ENTROPY_KEY: entropy_loss.clone().detach(),
         "pg_clipfrac": pg_clipfrac.clone().detach(),
         "ppo_kl": ppo_kl.clone().detach(),
         "ess_ratio": ess_ratio_sum.squeeze(),

miles/backends/training_utils/torch_native/actor.py

@@ -12,6 +12,7 @@
 
 from miles.backends.training_utils.data.rollout import DataIterator, get_batch, get_data_iterator, get_rollout_data
 from miles.backends.training_utils.data.sampling_mask import get_rollout_sampling_masks
+from miles.backends.training_utils.loss.hub.entropy_control import update_adaptive_entropy
 from miles.backends.training_utils.loss.objective import (
     compute_advantages_and_returns,
     get_log_probs_and_entropy,
@@ -260,9 +261,11 @@ def _optimizer_steps(
                     role="actor",
                     rank=state.intra_dp_cp.rank,
                 )
+            loss_dict = aggregate_train_losses(losses_reduced)
+            update_adaptive_entropy(args, loss_dict)
             log_train_step(
                 args=args,
-                loss_dict=aggregate_train_losses(losses_reduced),
+                loss_dict=loss_dict,
                 grad_norm=metrics.grad_norm,
                 rollout_id=rollout_id,
                 step_id=step_id,

miles/utils/arguments.py

@@ -1752,6 +1752,21 @@ def add_algo_arguments(parser):
                 help="Interval (in rollout steps) to update ref model from actor. If None, ref model is not updated.",
             )
             parser.add_argument("--entropy-coef", type=float, default=0.0, help="Entropy loss coef")
+            parser.add_argument(
+                "--use-adaptive-entropy",
+                action="store_true",
+                default=False,
+                help=(
+                    "Adaptive entropy regularization (Skywork-OR1): the entropy bonus acts only while the training "
+                    "entropy is at or below --entropy-target, and --entropy-coef moves by --entropy-coef-delta per "
+                    "optimizer step, up at or below the target and down above it, within "
+                    "[--entropy-coef-min, --entropy-coef-max]. --entropy-coef is the starting value."
+                ),
+            )
+            parser.add_argument("--entropy-target", type=float, default=0.2, help="Target per-token entropy (nats).")
+            parser.add_argument("--entropy-coef-delta", type=float, default=0.005, help="Per-step coefficient change.")
+            parser.add_argument("--entropy-coef-min", type=float, default=0.0, help="Lower clamp of the coefficient.")
+            parser.add_argument("--entropy-coef-max", type=float, default=1.0, help="Upper clamp of the coefficient.")
             parser.add_argument("--gamma", type=float, default=1.0, help="PPO GAE gamma")
             parser.add_argument("--lambd", type=float, default=1.0, help="PPO GAE lambd")
             parser.add_argument("--normalize-advantages", action="store_true", default=False)

tests/fast/backends/megatron_utils/test_model.py

@@ -9,6 +9,7 @@
 import torch
 
 from miles.backends.training_utils.data.rollout import DataIterator
+from miles.backends.training_utils.loss.hub.entropy_control import entropy_coef_to_apply
 from miles.backends.training_utils.model_companion import ModelCompanion
 
 
@@ -422,3 +423,76 @@ def reject_optimizer_step() -> None:
     )
 
     assert outcome is model_module.TrainStepOutcome.DISCARDED_SHOULD_RETRY
+
+
+@pytest.mark.parametrize(
+    "grad_norm,found_inf,disable_optimizer,has_optimizer,expected_steps",
+    [
+        pytest.param(1.0, False, False, True, 1, id="applied"),
+        pytest.param(float("nan"), False, False, True, 0, id="nan-grad-norm"),
+        pytest.param(float("inf"), False, False, True, 0, id="inf-grad-norm"),
+        pytest.param(torch.tensor(float("nan")), False, False, True, 0, id="tensor-nan-grad-norm"),
+        pytest.param(torch.tensor(float("inf")), False, False, True, 0, id="tensor-inf-grad-norm"),
+        pytest.param(1.0, True, False, True, 0, id="found-inf"),
+        pytest.param(1.0, False, True, True, 0, id="disabled-optimizer"),
+        pytest.param(1.0, False, False, False, 0, id="missing-optimizer"),
+    ],
+)
+def test_adaptive_entropy_updates_only_after_an_applied_optimizer_step(
+    train_one_step_env: TrainOneStepEnv,
+    monkeypatch: pytest.MonkeyPatch,
+    grad_norm: float | torch.Tensor,
+    found_inf: bool,
+    disable_optimizer: bool,
+    has_optimizer: bool,
+    expected_steps: int,
+) -> None:
+    """A skipped step's NaN entropy must not move the coefficient or reopen the bonus gate."""
+    from miles.backends.megatron_utils import model as model_module
+
+    env = train_one_step_env
+    env.args.enable_sample_ownership_checker = False
+    env.args.check_for_nan_in_loss_and_grad = False
+    env.args.debug_disable_optimizer = disable_optimizer
+    env.args.use_adaptive_entropy = True
+    env.args.entropy_coef = 0.01
+    env.args.entropy_target = 0.5
+    env.args.entropy_coef_delta = 0.005
+    env.args.entropy_coef_min = 0.0
+    env.args.entropy_coef_max = 0.03
+    env.args.adaptive_entropy_coef = 0.015
+    env.args.adaptive_entropy_last = 0.9
+    loss_reduced = {"entropy_loss": 0.4 if expected_steps else float("nan"), "pg_loss": 1.0}
+    optimizer = FakeFiniteOptimizer()
+    monkeypatch.setattr(optimizer, "prepare_grads", lambda: found_inf)
+    monkeypatch.setattr(optimizer, "get_grad_norm", lambda: grad_norm)
+    monkeypatch.setattr(model_module, "mpu", FakeMpu(is_last_pipeline_stage=True))
+    monkeypatch.setattr(model_module, "aggregate_train_losses", lambda *_args: loss_reduced)
+
+    losses, _, outcome = model_module.train_one_step(
+        args=env.args,
+        rollout_id=7,
+        step_id=0,
+        data_iterator=env.data_iterator,
+        model=env.model,
+        optimizer=optimizer if has_optimizer else None,
+        opt_param_scheduler=FakeScheduler(),
+        num_microbatches=1,
+        num_rollouts=1,
+        witness_info=None,
+        attempt=2,
+    )
+
+    assert outcome is model_module.TrainStepOutcome.NORMAL
+    assert optimizer.step_count == expected_steps
+    assert losses["pg_loss"] == 1.0
+    if expected_steps:
+        assert env.args.adaptive_entropy_coef == pytest.approx(0.02)
+        assert env.args.adaptive_entropy_last == 0.4
+        assert entropy_coef_to_apply(env.args) == pytest.approx(0.02)
+        assert losses["entropy_coef"] == pytest.approx(0.02)
+    else:
+        assert env.args.adaptive_entropy_coef == 0.015
+        assert env.args.adaptive_entropy_last == 0.9
+        assert entropy_coef_to_apply(env.args) == 0.0
+        assert "entropy_coef" not in losses

tests/fast/backends/training_utils/loss/loss_test_utils.py

@@ -94,6 +94,11 @@ def _trivial_group() -> GroupInfo:
     eps_clip=0.2,
     eps_clip_high=0.2,
     entropy_coef=0.01,
+    use_adaptive_entropy=False,
+    entropy_target=0.2,
+    entropy_coef_delta=0.005,
+    entropy_coef_min=0.0,
+    entropy_coef_max=1.0,
     use_kl_loss=False,
     kl_loss_coef=0.0,
     use_unbiased_kl=False,

tests/fast/backends/training_utils/loss/test_adaptive_entropy.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from miles.backends.training_utils.data.context_parallel import get_sum_of_sample_mean
+from miles.backends.training_utils.loss.hub import losses
+from miles.backends.training_utils.loss.hub.losses import policy_loss_function
+
+from .loss_test_utils import deep_clone, make_args, make_batch, make_inputs, make_parallel_state
+
+
+@pytest.fixture(scope="module")
+def process_group(tmp_path_factory):
+    if dist.is_initialized():
+        yield
+        return
+    rendezvous = tmp_path_factory.mktemp("adaptive-entropy") / "process-group"
+    dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+
+
+def _run(args, inputs):
+    batch = make_batch(inputs, "policy_loss")
+    logits = deep_clone(inputs["policy_logits"])
+    logits.requires_grad_(True)
+    reducer = get_sum_of_sample_mean(
+        batch["total_lengths"],
+        batch["response_lengths"],
+        batch["loss_masks"],
+        args.calculate_per_token_loss,
+        args.qkv_format,
+        batch.get("max_seq_lens"),
+    )
+    loss, metrics = policy_loss_function(args, batch, logits, reducer)
+    return loss.detach(), metrics
+
+
+def _make(**overrides):
+    make_parallel_state()
+    args = make_args(advantage_estimator="grpo", observe_training_entropy=False, **overrides)
+    inputs = make_inputs(seed=11, batch_size=2, prompt_lens=[16, 24], response_lens=[12, 20], vocab_size=64, args=args)
+    return args, inputs
+
+
+def test_bonus_is_gated_off_when_the_last_entropy_was_above_target(process_group):
+    args, inputs = _make(entropy_coef=0.05, use_adaptive_entropy=True, entropy_target=0.1)
+    args.adaptive_entropy_last = 5.0
+    gated_loss, gated_metrics = _run(args, inputs)
+
+    plain_args, _ = _make(entropy_coef=0.0)
+    plain_loss, _ = _run(plain_args, inputs)
+
+    assert torch.equal(gated_loss, plain_loss)
+    # Entropy is still computed and reported so the controller keeps its signal.
+    assert gated_metrics["entropy_loss"].item() > 0
+
+
+def test_bonus_applies_when_the_last_entropy_was_at_or_below_target(process_group):
+    args, inputs = _make(entropy_coef=0.05, use_adaptive_entropy=True, entropy_target=10.0)
+    args.adaptive_entropy_last = 1.0
+    adaptive_loss, _ = _run(args, inputs)
+
+    static_args, _ = _make(entropy_coef=0.05)
+    static_loss, _ = _run(static_args, inputs)
+
+    assert torch.equal(adaptive_loss, static_loss)
+
+
+@pytest.mark.parametrize(
+    "coefficient, last_entropy, requires_grad",
+    [(0.05, 5.0, False), (0.0, 0.1, False), (0.05, 0.1, True), (0.05, None, True)],
+)
+def test_entropy_tracks_gradients_only_when_the_bonus_applies(
+    process_group, monkeypatch, coefficient, last_entropy, requires_grad
+):
+    args, inputs = _make(entropy_coef=coefficient, use_adaptive_entropy=True, entropy_target=0.1)
+    args.adaptive_entropy_last = last_entropy
+    entropy_tensors = []
+    get_log_probs_and_entropy = losses.get_log_probs_and_entropy
+
+    def capture_entropy(*call_args, **call_kwargs):
+        result = get_log_probs_and_entropy(*call_args, **call_kwargs)
+        entropy_tensors.extend(result["entropy"])
+        return result
+
+    monkeypatch.setattr(losses, "get_log_probs_and_entropy", capture_entropy)
+
+    _, metrics = _run(args, inputs)
+
+    assert entropy_tensors
+    assert all(entropy.requires_grad == requires_grad for entropy in entropy_tensors)
+    assert metrics["entropy_loss"].item() > 0