feat: add configurable DAPO objective controls

astrai/trainer/strategy.py

@@ -1,6 +1,8 @@
 """Training strategy implementations with factory pattern."""
 
+import math
 from abc import ABC
+from numbers import Real
 from typing import Callable, Dict, List, Optional, TypedDict, Union
 
 import torch
@@ -560,17 +562,114 @@ def __init__(
         old_model: Optional[nn.Module],
         ref_model: nn.Module,
         clip_eps: float = 0.2,
+        clip_eps_low: Optional[float] = None,
+        clip_eps_high: Optional[float] = None,
         kl_coef: float = 0.01,
         group_size: int = 4,
+        loss_aggregation: str = "token",
+        overlong_max_len: Optional[int] = None,
+        overlong_buffer_len: int = 0,
+        overlong_penalty_scale: float = 1.0,
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
         self.old_model = old_model
         self.ref_model = ref_model
-        self.clip_eps = clip_eps
+        self.clip_eps = self._validate_clip_epsilon(clip_eps, "clip_eps", upper=True)
+        self.clip_eps_low = self._validate_clip_epsilon(
+            self.clip_eps if clip_eps_low is None else clip_eps_low,
+            "clip_eps_low",
+            upper=True,
+        )
+        self.clip_eps_high = self._validate_clip_epsilon(
+            self.clip_eps if clip_eps_high is None else clip_eps_high,
+            "clip_eps_high",
+        )
+        if self.clip_eps_high < self.clip_eps_low:
+            raise ValueError(
+                "clip_eps_high must be greater than or equal to clip_eps_low"
+            )
+        if loss_aggregation not in {"token", "sequence"}:
+            raise ValueError("loss_aggregation must be 'token' or 'sequence'")
+        self.loss_aggregation = loss_aggregation
+        self.overlong_max_len, self.overlong_buffer_len = (
+            self._validate_overlong_window(overlong_max_len, overlong_buffer_len)
+        )
+        self.overlong_penalty_scale = self._validate_non_negative_real(
+            overlong_penalty_scale, "overlong_penalty_scale"
+        )
         self.kl_coef = kl_coef
         self.group_size = group_size
 
+    @staticmethod
+    def _validate_clip_epsilon(value: float, name: str, upper: bool = False) -> float:
+        if isinstance(value, bool) or not isinstance(value, Real):
+            raise TypeError(f"{name} must be a real number")
+        value = float(value)
+        if not math.isfinite(value) or value < 0 or (upper and value >= 1):
+            interval = "[0, 1)" if upper else "[0, infinity)"
+            raise ValueError(f"{name} must be finite and in {interval}")
+        return value
+
+    @staticmethod
+    def _validate_non_negative_real(value: float, name: str) -> float:
+        if isinstance(value, bool) or not isinstance(value, Real):
+            raise TypeError(f"{name} must be a real number")
+        value = float(value)
+        if not math.isfinite(value) or value < 0:
+            raise ValueError(f"{name} must be finite and non-negative")
+        return value
+
+    @staticmethod
+    def _validate_overlong_window(
+        max_len: Optional[int], buffer_len: int
+    ) -> tuple[Optional[int], int]:
+        if max_len is None:
+            if buffer_len != 0:
+                raise ValueError(
+                    "overlong_buffer_len requires overlong_max_len to be set"
+                )
+            return None, 0
+        if isinstance(max_len, bool) or not isinstance(max_len, int) or max_len <= 0:
+            raise ValueError("overlong_max_len must be a positive integer or None")
+        if (
+            isinstance(buffer_len, bool)
+            or not isinstance(buffer_len, int)
+            or buffer_len <= 0
+            or buffer_len > max_len
+        ):
+            raise ValueError(
+                "overlong_buffer_len must be a positive integer no greater "
+                "than overlong_max_len"
+            )
+        return max_len, buffer_len
+
+    def _reduce_token_loss(self, loss: Tensor, mask: Tensor) -> Tensor:
+        """Reduce response-token losses with GRPO or DAPO weighting."""
+        mask = mask.float()
+        if self.loss_aggregation == "token":
+            return (loss * mask).sum() / mask.sum().clamp(min=1.0)
+
+        lengths = mask.sum(dim=-1)
+        valid_sequences = lengths > 0
+        per_sequence = (loss * mask).sum(dim=-1) / lengths.clamp(min=1.0)
+        return (per_sequence * valid_sequences).sum() / valid_sequences.sum().clamp(
+            min=1
+        )
+
+    def _shape_overlong_rewards(
+        self, rewards: Tensor, token_masks: Tensor
+    ) -> tuple[Tensor, Optional[Tensor]]:
+        if self.overlong_max_len is None:
+            return rewards, None
+
+        lengths = token_masks.sum(dim=-1)
+        penalty_start = self.overlong_max_len - self.overlong_buffer_len
+        penalty = ((penalty_start - lengths) / self.overlong_buffer_len).clamp(
+            min=-1.0, max=0.0
+        )
+        return rewards + self.overlong_penalty_scale * penalty, penalty
+
     def sync_old_model(self):
         """Copy current policy weights to old model."""
         if self.old_model is None:
@@ -667,6 +766,7 @@ def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
 
         # Group-normalized advantages from scalar per-response rewards.
         eps = 1e-8
+        rewards, overlong_penalty = self._shape_overlong_rewards(rewards, token_masks)
         mean = rewards.mean(dim=-1, keepdim=True)
         std = rewards.std(dim=-1, keepdim=True, unbiased=False)
         advantages = (rewards - mean) / (std + eps)
@@ -678,22 +778,35 @@ def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         ratio = torch.exp(log_ratio)
 
         surr1 = ratio * advantages
-        surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        surr2 = (
+            torch.clamp(
+                ratio,
+                1 - self.clip_eps_low,
+                1 + self.clip_eps_high,
+            )
+            * advantages
+        )
         per_token_policy_loss = -torch.min(surr1, surr2)
-        token_count = token_masks.sum().clamp(min=1.0)
-        policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
+        policy_loss = self._reduce_token_loss(per_token_policy_loss, token_masks)
 
         # KL penalty to frozen reference model with k1 estimator (non-negative):
         # k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
         log_ref_ratio = token_log_probs_ref - token_log_probs_policy
         r = torch.exp(log_ref_ratio)
         kl_per_token = r - torch.log(r + eps) - 1.0
-        kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
+        kl_penalty = self.kl_coef * self._reduce_token_loss(kl_per_token, token_masks)
 
         task_loss = policy_loss + kl_penalty
+        metrics = {
+            "policy_loss": policy_loss,
+            "kl_loss": kl_penalty,
+        }
+        if overlong_penalty is not None:
+            metrics["overlong_penalty_mean"] = overlong_penalty.mean()
+            metrics["overlong_fraction"] = (overlong_penalty < 0).float().mean()
         return self._loss_output(
             task_loss,
-            {"policy_loss": policy_loss, "kl_loss": kl_penalty},
+            metrics,
             aux_loss,
             policy_output.get("router_stats"),
         )

benchmarks/dapo/benchmark_objective.py

@@ -0,0 +1,161 @@
+"""Benchmark GRPO and DAPO objective reductions on one CUDA device."""
+
+import json
+
+import torch
+
+
+def reduce_loss(loss, mask, aggregation):
+    if aggregation == "token":
+        return (loss * mask).sum() / mask.sum().clamp(min=1.0)
+    lengths = mask.sum(dim=-1)
+    valid = lengths > 0
+    per_sequence = (loss * mask).sum(dim=-1) / lengths.clamp(min=1.0)
+    return (per_sequence * valid).sum() / valid.sum().clamp(min=1)
+
+
+def objective(
+    log_policy,
+    log_old,
+    log_ref,
+    rewards,
+    mask,
+    *,
+    clip_low,
+    clip_high,
+    aggregation="token",
+    overlong_buffer=0,
+    penalty_scale=1.0,
+):
+    if overlong_buffer:
+        max_len = mask.shape[-1]
+        lengths = mask.sum(dim=-1)
+        penalty_start = max_len - overlong_buffer
+        penalty = ((penalty_start - lengths) / overlong_buffer).clamp(-1.0, 0.0)
+        rewards = rewards + penalty_scale * penalty
+
+    advantages = (rewards - rewards.mean(-1, keepdim=True)) / (
+        rewards.std(-1, keepdim=True, unbiased=False) + 1e-8
+    )
+    advantages = advantages.unsqueeze(-1)
+    ratio = torch.exp(log_policy - log_old)
+    surrogate = torch.minimum(
+        ratio * advantages,
+        torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantages,
+    )
+    policy_loss = reduce_loss(-surrogate, mask, aggregation)
+
+    ref_ratio = torch.exp(log_ref - log_policy)
+    kl = ref_ratio - torch.log(ref_ratio + 1e-8) - 1.0
+    return policy_loss + 0.01 * reduce_loss(kl, mask, aggregation)
+
+
+def median_ms(fn, warmup=20, repeats=100):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end))
+    samples.sort()
+    return samples[len(samples) // 2], samples[int(len(samples) * 0.99) - 1]
+
+
+def main():
+    torch.manual_seed(3407)
+    device = torch.device("cuda")
+    output = {
+        "gpu": torch.cuda.get_device_name(device),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "seed": 3407,
+        "warmup": 20,
+        "repeats": 100,
+        "cases": [],
+    }
+    for batch, group, response_len in (
+        (4, 8, 256),
+        (4, 8, 1024),
+        (8, 8, 2048),
+        (4, 8, 4096),
+    ):
+        shape = (batch, group, response_len)
+        log_policy = torch.randn(shape, device=device) * 0.2
+        log_old = torch.randn(shape, device=device) * 0.2
+        log_ref = torch.randn(shape, device=device) * 0.2
+        rewards = torch.randn((batch, group), device=device)
+        lengths = torch.randint(
+            response_len // 2,
+            response_len + 1,
+            (batch, group),
+            device=device,
+        )
+        mask = (
+            torch.arange(response_len, device=device)[None, None, :]
+            < lengths[..., None]
+        ).float()
+
+        legacy = lambda: objective(  # noqa: E731
+            log_policy,
+            log_old,
+            log_ref,
+            rewards,
+            mask,
+            clip_low=0.2,
+            clip_high=0.2,
+        )
+        symmetric = lambda: objective(  # noqa: E731
+            log_policy,
+            log_old,
+            log_ref,
+            rewards,
+            mask,
+            clip_low=0.2,
+            clip_high=0.2,
+        )
+        dapo = lambda: objective(  # noqa: E731
+            log_policy,
+            log_old,
+            log_ref,
+            rewards,
+            mask,
+            clip_low=0.2,
+            clip_high=0.28,
+            aggregation="token",
+            overlong_buffer=max(1, response_len // 8),
+        )
+
+        legacy_median, legacy_p99 = median_ms(legacy)
+        symmetric_median, symmetric_p99 = median_ms(symmetric)
+        dapo_median, dapo_p99 = median_ms(dapo)
+        output["cases"].append(
+            {
+                "batch": batch,
+                "group": group,
+                "response_len": response_len,
+                "valid_tokens": int(mask.sum().item()),
+                "symmetric_loss_parity_abs": abs(legacy().item() - symmetric().item()),
+                "legacy_median_ms": legacy_median,
+                "legacy_p99_ms": legacy_p99,
+                "candidate_symmetric_median_ms": symmetric_median,
+                "candidate_symmetric_p99_ms": symmetric_p99,
+                "candidate_symmetric_delta_percent": (
+                    symmetric_median / legacy_median - 1.0
+                )
+                * 100.0,
+                "dapo_median_ms": dapo_median,
+                "dapo_p99_ms": dapo_p99,
+                "dapo_delta_percent": (dapo_median / legacy_median - 1.0) * 100.0,
+            }
+        )
+    print(json.dumps(output, indent=2))
+
+
+if __name__ == "__main__":
+    main()