feat: add QAE and adaptive PPO clipping

docs/en/examples/qae-adaptive-clip.md

@@ -0,0 +1,46 @@
+# QAE with adaptive clipping
+
+This recipe combines Quantile Advantage Estimation (QAE) with the adaptive
+upper clipping controller from MAI-Thinking-1.
+
+```bash
+python train.py \
+  ... \
+  --advantage-estimator grpo \
+  --qae-quantile 0.4 \
+  --eps-clip 0.6 \
+  --use-adaptive-clip \
+  --adaptive-clip-target-entropy 0.3 \
+  --adaptive-clip-step-size 0.25 \
+  --adaptive-clip-max-relaxation 2.5 \
+  --entropy-coef 0
+```
+
+`--qae-quantile 0.4` changes only the group reward baseline. The implementation
+uses the paper's right-continuous empirical quantile. For binary rewards this
+produces the two QAE regimes exactly: hard groups update successes and easy
+groups update failures.
+
+With adaptive clipping, the PPO ratio interval at step `t` is
+
+```text
+[1 - eps_clip, 1 / (1 - eps_clip) + k_t]
+```
+
+and the controller applies
+
+```text
+k_{t+1} = clip(k_t + delta * sign(target_entropy - estimated_entropy), 0, k_max).
+```
+
+The entropy estimate is the importance-weighted all-token average from the MAI
+report. The controller updates once after each optimizer step and its state is
+saved beside checkpoints. Monitor these training metrics:
+
+- `adaptive_clip_entropy`: controller entropy estimate
+- `adaptive_clip_relaxation`: current `k`
+- `adaptive_clip_high`: MILES upper clip offset (the upper ratio is `1 + value`)
+
+The MAI values above are a starting point, not universal defaults for every
+model and rollout temperature. A sign controller commonly oscillates in a
+narrow band around the target; it does not guarantee a smooth sine wave.

miles/backends/experimental/fsdp_utils/actor.py

@@ -29,7 +29,12 @@
     log_rollout_data,
     log_train_step,
 )
-from ...training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function
+from ...training_utils.loss import (
+    compute_advantages_and_returns,
+    get_log_probs_and_entropy,
+    loss_function,
+    update_adaptive_clip,
+)
 from ...training_utils.parallel import get_parallel_state, set_parallel_state
 from . import checkpoint
 from .lr_scheduler import get_lr_scheduler
@@ -492,6 +497,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
                     )
 
                 loss_dict = aggregate_train_losses(losses_reduced)
+                update_adaptive_clip(self.args, loss_dict)
 
                 extra_metrics = {}
                 for param_group_id, param_group in enumerate(self.optimizer.param_groups):

miles/backends/experimental/fsdp_utils/checkpoint.py

@@ -13,6 +13,7 @@
 from torch.distributed.checkpoint.stateful import Stateful
 
 from miles.backends.training_utils.log_utils import init_train_step_counter, save_train_step_counter
+from miles.utils.entropy_control import load_adaptive_clip_state, save_adaptive_clip_state
 
 logger = logging.getLogger(__name__)
 
@@ -185,6 +186,7 @@ def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None
         if getattr(actor.args, "start_rollout_id", None) is None:
             actor.args.start_rollout_id = iteration
 
+    load_adaptive_clip_state(actor.args, actor.args.load, iteration)
     if dist.get_rank() == 0:
         init_train_step_counter(actor.args.load, iteration)
 
@@ -251,6 +253,7 @@ def save(actor: Any, iteration: int) -> None:
         tracker_file = base_dir / "latest_checkpointed_iteration.txt"
         tracker_file.write_text(str(step_id))
         save_train_step_counter(actor.args.save, step_id)
+        save_adaptive_clip_state(actor.args, actor.args.save, step_id)
         logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}")
 
     dist.barrier()

miles/backends/megatron_utils/model.py

@@ -22,6 +22,7 @@
 from megatron.training.training import get_model, preprocess_common_state_dict
 
 from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
+from miles.utils.entropy_control import load_adaptive_clip_state, save_adaptive_clip_state
 from miles.utils.memory_utils import clear_memory
 
 from ..training_utils.ci_utils import check_grad_norm, check_kl
@@ -33,7 +34,7 @@
     log_train_step,
     save_train_step_counter,
 )
-from ..training_utils.loss import loss_function
+from ..training_utils.loss import loss_function, update_adaptive_clip
 from ..training_utils.parallel import get_parallel_state
 from .checkpoint import load_checkpoint, save_checkpoint, save_checkpoint_with_lora
 from .ci_utils import (
@@ -626,6 +627,15 @@ def train(
             num_microbatches[step_id],
         )
 
+        role = getattr(model[0], "role", "actor")
+        if role == "actor":
+            update_adaptive_clip(
+                args,
+                loss_dict,
+                pipeline_group=mpu.get_pipeline_model_parallel_group(),
+                pipeline_source_rank=mpu.get_pipeline_model_parallel_last_rank(),
+            )
+
         if step_id == 0:
             # Enable forward pre-hook after training step has successfully run. All subsequent
             # forward passes will use the forward pre-hook / `param_sync_func` in
@@ -659,7 +669,6 @@ def train(
         # per train step log.
         if is_megatron_main_rank():
             accumulated_step_id = rollout_id * num_steps_per_rollout + step_id
-            role = getattr(model[0], "role", "actor")
             role_tag = "" if role == "actor" else f"{role}-"
 
             extra_metrics = {}
@@ -741,6 +750,7 @@ def save(
     clear_memory()
     if is_megatron_main_rank():
         save_train_step_counter(args.save, iteration)
+        save_adaptive_clip_state(args, args.save, iteration)
 
     if hashes is not None:
         save_model_hashes(args, model, iteration, hashes)
@@ -840,6 +850,7 @@ def initialize_model_and_optimizer(
     clear_memory()
 
     check_model_hashes(args, model, iteration)
+    load_adaptive_clip_state(args, args.load, iteration)
 
     # A full resume restores the scheduler in load_checkpoint(). Advancing it
     # again here would double-count the completed iterations. Only reconstruct

miles/backends/training_utils/loss.py

@@ -8,11 +8,13 @@
 from torch.utils.checkpoint import checkpoint
 
 from miles.utils.distributed_utils import distributed_masked_whiten
+from miles.utils.entropy_control import adaptive_clip_high, update_adaptive_clip_relaxation
 from miles.utils.misc import load_function
 from miles.utils.ppo_utils import (
     calculate_log_probs_and_entropy,
     compute_approx_kl,
     compute_gspo_kl,
+    compute_importance_weighted_entropy,
     compute_opd_reward,
     compute_opsm_mask,
     compute_policy_loss,
@@ -34,6 +36,58 @@
 logger = logging.getLogger(__name__)
 
 
+def update_adaptive_clip(
+    args: Namespace,
+    loss_dict: dict[str, float],
+    *,
+    pipeline_group: torch.distributed.ProcessGroup | None = None,
+    pipeline_source_rank: int | None = None,
+) -> None:
+    """Update and synchronize MAI's clipping controller after an optimizer step.
+
+    The policy loss emits an entropy numerator and token count for each
+    microbatch. ``aggregate_train_losses`` sums both globally; their ratio is
+    therefore the paper's all-token importance-weighted entropy estimate.
+    """
+    if not getattr(args, "use_adaptive_clip", False):
+        return
+
+    parallel_state = get_parallel_state()
+    is_last_stage = parallel_state.is_pp_last_stage
+    state = torch.tensor(
+        [float(args.adaptive_clip_relaxation), float("nan")],
+        dtype=torch.float64,
+        device=torch.cuda.current_device(),
+    )
+    if is_last_stage:
+        entropy_sum = loss_dict.pop("adaptive_clip_entropy_sum")
+        token_count = loss_dict.pop("adaptive_clip_token_count")
+        if token_count <= 0:
+            raise RuntimeError("adaptive clipping received a training step with no valid response tokens")
+        estimated_entropy = entropy_sum / token_count
+        relaxation = update_adaptive_clip_relaxation(
+            args.adaptive_clip_relaxation,
+            estimated_entropy,
+            args.adaptive_clip_target_entropy,
+            args.adaptive_clip_step_size,
+            args.adaptive_clip_max_relaxation,
+        )
+        state[0] = relaxation
+        state[1] = estimated_entropy
+
+    if pipeline_group is not None and torch.distributed.get_world_size(pipeline_group) > 1:
+        if pipeline_source_rank is None:
+            raise ValueError("pipeline_source_rank is required when synchronizing adaptive clipping")
+        torch.distributed.broadcast(state, src=pipeline_source_rank, group=pipeline_group)
+
+    args.adaptive_clip_relaxation = state[0].item()
+    args.eps_clip_high = adaptive_clip_high(args.eps_clip, args.adaptive_clip_relaxation)
+    if is_last_stage:
+        loss_dict["adaptive_clip_entropy"] = state[1].item()
+        loss_dict["adaptive_clip_relaxation"] = args.adaptive_clip_relaxation
+        loss_dict["adaptive_clip_high"] = args.eps_clip_high
+
+
 def _nan_dbg_flush_logs() -> None:
     """Flush logger/stdout/stderr so crash diagnostics survive abrupt failures."""
     for handler in logging.getLogger().handlers + logger.handlers:
@@ -763,6 +817,26 @@ def policy_loss_function(
         extra=f"ratio_delta_min={ratio_delta_min} ratio_delta_max={ratio_delta_max}",
     )
 
+    adaptive_clip_metrics = {}
+    if getattr(args, "use_adaptive_clip", False):
+        # MAI Eq. (7): mean over every valid response token of
+        # -log pi_theta(y_t | x, y_<t) * pi_theta(y_t) / pi_old(y_t).
+        # Use token-level ratios even when the optimization objective is GSPO.
+        importance_weighted_entropy = compute_importance_weighted_entropy(log_probs, old_log_probs)
+        token_sum = get_sum_of_sample_mean(
+            total_lengths,
+            response_lengths,
+            batch["loss_masks"],
+            args.calculate_per_token_loss,
+            args.qkv_format,
+            max_seq_lens,
+            loss_agg_mode="token-sum",
+        )
+        adaptive_clip_metrics = {
+            "adaptive_clip_entropy_sum": token_sum(importance_weighted_entropy),
+            "adaptive_clip_token_count": token_sum(torch.ones_like(importance_weighted_entropy)),
+        }
+
     pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high)
 
     if args.use_opsm:
@@ -930,6 +1004,7 @@ def policy_loss_function(
         "nan_dbg/tis_delta_max": _nan_dbg_scalar(tis_delta_max, loss.device),
         "nan_dbg/tis_nonfinite_count": _nan_dbg_scalar(tis_nonfinite_count, loss.device),
     }
+    reported_loss.update(adaptive_clip_metrics)
 
     if train_rollout_logprob_abs_diff is not None:
         reported_loss["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff

miles/ray/rollout.py

@@ -48,6 +48,7 @@
     dict_add_prefix,
 )
 from miles.utils.misc import load_function
+from miles.utils.ppo_utils import compute_group_advantages
 from miles.utils.ray_utils import Box
 from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions
 from miles.utils.tracking_utils import init_tracking
@@ -762,12 +763,12 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
             else:
                 # when samples count are not equal in each group
                 rewards = rewards.view(-1, rewards.shape[-1])
-            mean = rewards.mean(dim=-1, keepdim=True)
-            rewards = rewards - mean
-
-            if self.args.advantage_estimator in ["grpo", "gspo"] and self.args.grpo_std_normalization:
-                std = rewards.std(dim=-1, keepdim=True)
-                rewards = rewards / (std + 1e-6)
+            is_group_policy_estimator = self.args.advantage_estimator in ["grpo", "gspo"]
+            rewards, _ = compute_group_advantages(
+                rewards,
+                quantile=getattr(self.args, "qae_quantile", None) if is_group_policy_estimator else None,
+                normalize_by_std=is_group_policy_estimator and self.args.grpo_std_normalization,
+            )
 
             return raw_rewards, rewards.flatten().tolist()