From 4ebc2ac497a4a1bb0d0e6a91b4529aeea14b6f43 Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Fri, 11 Sep 2026 18:58:00 -0700 Subject: [PATCH 1/2] loss: add --use-adaptive-entropy (Skywork-OR1) 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. --- docs/user-guide/cli-reference.md | 20 +++- miles/backends/megatron_utils/model.py | 2 + .../loss/hub/entropy_control.py | 45 ++++++++ .../training_utils/loss/hub/losses.py | 14 ++- .../training_utils/torch_native/actor.py | 5 +- miles/utils/arguments.py | 15 +++ .../training_utils/loss/loss_test_utils.py | 5 + .../loss/test_adaptive_entropy.py | 97 ++++++++++++++++ .../loss_hub/test_entropy_control.py | 104 ++++++++++++++++++ 9 files changed, 299 insertions(+), 8 deletions(-) create mode 100644 miles/backends/training_utils/loss/hub/entropy_control.py create mode 100644 tests/fast/backends/training_utils/loss/test_adaptive_entropy.py create mode 100644 tests/fast/backends/training_utils/loss_hub/test_entropy_control.py diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index 2ff73528239..4dd3340258a 100644 --- a/docs/user-guide/cli-reference.md +++ b/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. | diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 79665bd4e7c..d1c6ab18e66 100644 --- a/miles/backends/megatron_utils/model.py +++ b/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,7 @@ def train_one_step( if parallel_state.indep_dp.size > 1 else aggregate_train_losses(losses_reduced, metric_num_rollouts) ) + update_adaptive_entropy(args, loss_reduced) return loss_reduced, grad_norm, outcome return {}, grad_norm, outcome diff --git a/miles/backends/training_utils/loss/hub/entropy_control.py b/miles/backends/training_utils/loss/hub/entropy_control.py new file mode 100644 index 00000000000..1a1d4bd8c92 --- /dev/null +++ b/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) diff --git a/miles/backends/training_utils/loss/hub/losses.py b/miles/backends/training_utils/loss/hub/losses.py index b48b7ab7629..1c0de0d1aa0 100644 --- a/miles/backends/training_utils/loss/hub/losses.py +++ b/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(), diff --git a/miles/backends/training_utils/torch_native/actor.py b/miles/backends/training_utils/torch_native/actor.py index 4c0c1a41513..e6764af0a82 100644 --- a/miles/backends/training_utils/torch_native/actor.py +++ b/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, diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 5eb3230ed12..8d90910f9bd 100644 --- a/miles/utils/arguments.py +++ b/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) diff --git a/tests/fast/backends/training_utils/loss/loss_test_utils.py b/tests/fast/backends/training_utils/loss/loss_test_utils.py index 4254a18ab4d..ede2c083b12 100644 --- a/tests/fast/backends/training_utils/loss/loss_test_utils.py +++ b/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, diff --git a/tests/fast/backends/training_utils/loss/test_adaptive_entropy.py b/tests/fast/backends/training_utils/loss/test_adaptive_entropy.py new file mode 100644 index 00000000000..c4e50e6e371 --- /dev/null +++ b/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 diff --git a/tests/fast/backends/training_utils/loss_hub/test_entropy_control.py b/tests/fast/backends/training_utils/loss_hub/test_entropy_control.py new file mode 100644 index 00000000000..369443f439c --- /dev/null +++ b/tests/fast/backends/training_utils/loss_hub/test_entropy_control.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import pickle +from types import SimpleNamespace + +import pytest + +from miles.backends.training_utils.loss.hub.entropy_control import entropy_coef_to_apply, update_adaptive_entropy + + +def _args(**overrides): + values = dict( + use_adaptive_entropy=True, + entropy_coef=0.01, + entropy_target=0.5, + entropy_coef_delta=0.005, + entropy_coef_min=0.0, + entropy_coef_max=0.02, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def test_static_mode_applies_the_coefficient_as_is(): + args = _args(use_adaptive_entropy=False, entropy_coef=0.3) + loss_dict = {"entropy_loss": 0.0} + + assert entropy_coef_to_apply(args) == 0.3 + assert update_adaptive_entropy(args, loss_dict) is None + assert args.entropy_coef == 0.3 + assert loss_dict == {"entropy_loss": 0.0} + assert not hasattr(args, "adaptive_entropy_coef") + + +def test_first_step_applies_the_starting_coefficient(): + assert entropy_coef_to_apply(_args()) == 0.01 + + +def test_coefficient_rises_while_entropy_is_at_or_below_target_and_is_clamped(): + args = _args() + + update_adaptive_entropy(args, {"entropy_loss": 0.4}) + assert args.adaptive_entropy_coef == 0.015 + assert args.adaptive_entropy_last == 0.4 + assert entropy_coef_to_apply(args) == 0.015 + + update_adaptive_entropy(args, {"entropy_loss": 0.5}) + update_adaptive_entropy(args, {"entropy_loss": 0.1}) + assert args.adaptive_entropy_coef == 0.02 # clamped at max + + +def test_coefficient_falls_and_bonus_is_gated_off_above_target(): + args = _args(entropy_coef=0.004) + + update_adaptive_entropy(args, {"entropy_loss": 0.9}) + + assert args.adaptive_entropy_coef == 0.0 # clamped at min + assert entropy_coef_to_apply(args) == 0.0 + + +@pytest.mark.parametrize("entropy, expected_coef", [(0.4, 0.015), (0.5, 0.015), (0.9, 0.0)]) +def test_update_reports_the_coefficient_to_apply_on_the_next_step(entropy, expected_coef): + args = _args() + loss_dict = {"entropy_loss": entropy, "pg_loss": 1.0} + + assert update_adaptive_entropy(args, loss_dict) is None + + assert loss_dict == {"entropy_loss": entropy, "pg_loss": 1.0, "entropy_coef": expected_coef} + assert entropy_coef_to_apply(args) == expected_coef + + +def test_bonus_uses_the_adapted_coefficient_when_the_gate_reopens(): + args = _args() + + update_adaptive_entropy(args, {"entropy_loss": 0.9}) + assert args.adaptive_entropy_coef == 0.005 + assert entropy_coef_to_apply(args) == 0.0 + + update_adaptive_entropy(args, {"entropy_loss": 0.4}) + assert entropy_coef_to_apply(args) == 0.01 + + +def test_updates_preserve_the_starting_coefficient_in_checkpoint_args(): + args = _args() + + update_adaptive_entropy(args, {"entropy_loss": 0.4}) + # Megatron includes the args namespace in its checkpoint state dict. + checkpoint = pickle.loads(pickle.dumps({"args": args})) + + assert args.entropy_coef == 0.01 + assert checkpoint["args"].entropy_coef == 0.01 + assert checkpoint["args"].adaptive_entropy_coef == 0.015 + + +def test_missing_entropy_metric_leaves_state_alone(): + args = _args() + loss_dict = {"pg_loss": 1.0} + + update_adaptive_entropy(args, loss_dict) + + assert loss_dict == {"pg_loss": 1.0} + assert args.entropy_coef == 0.01 + assert not hasattr(args, "adaptive_entropy_coef") + assert not hasattr(args, "adaptive_entropy_last") From 75d6bba9df84aff7e3338b5b88cf7c7f126e97ab Mon Sep 17 00:00:00 2001 From: Behrooz Azarkhalili Date: Thu, 15 Oct 2026 12:00:27 -0700 Subject: [PATCH 2/2] loss: move the adaptive entropy coefficient only on an applied optimizer 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. --- miles/backends/megatron_utils/model.py | 3 +- .../backends/megatron_utils/test_model.py | 74 +++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index d1c6ab18e66..1149a34aee2 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -705,7 +705,8 @@ def train_one_step( if parallel_state.indep_dp.size > 1 else aggregate_train_losses(losses_reduced, metric_num_rollouts) ) - update_adaptive_entropy(args, loss_reduced) + if not disable_optimizer and valid_step: + update_adaptive_entropy(args, loss_reduced) return loss_reduced, grad_norm, outcome return {}, grad_norm, outcome diff --git a/tests/fast/backends/megatron_utils/test_model.py b/tests/fast/backends/megatron_utils/test_model.py index 9398092d8ee..e8a28a045f1 100644 --- a/tests/fast/backends/megatron_utils/test_model.py +++ b/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