diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 2321c7dc..a6b1c78f 100644 --- a/astrai/trainer/strategy.py +++ b/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"), ) diff --git a/benchmarks/dapo/benchmark_objective.py b/benchmarks/dapo/benchmark_objective.py new file mode 100644 index 00000000..41670c12 --- /dev/null +++ b/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() diff --git a/benchmarks/infraswe/README.md b/benchmarks/infraswe/README.md new file mode 100644 index 00000000..4e31f348 --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,75 @@ +# InfraSWE: configurable GRPO and DAPO objective + +This directory binds AstrAI's configurable GRPO/DAPO objective to the +checked-in NVIDIA L20 evidence. It uses InfraSWE's +`project-fit-system-path-v0.5.1` comparison and scoring models because the +candidate changes the training objective, CLI contract, validation behavior, +and reward-shaping path rather than an isolated kernel. + +InfraSWE v0.5's generic Draft document currently admits kernel and pure-Triton +formula identifiers only. Substituting a kernel formula would misclassify this +change, so the repository stores and validates the native +`ProjectComparisonCell`. The result remains explicitly diagnostic and +unsealed. + +Before the PR was opened, InfraSWE commit +`811bc775ed5b3a6ec853219245f3469f78818020` was used to validate the comparison +cell, run the frozen system-path ProjectFit and BenchmarkTrust functions, run +53 Draft/system-path engine tests, and verify that official scoring remains +unresolved without its required evidence envelope. + +The visible-evidence diagnostic score is **92.15/100** and BenchmarkTrust is +**97.40/100**. The lower load and cold/steady inputs record that this evidence +measures objective math and deterministic training tests, not downstream reward +quality or a long-running training job. Complete inputs and rationales are in +`benchmarks/results/dapo_objective_l20_infraswe_score.json`. + +## Scope + +The candidate adds three independently selectable objective pieces: + +- DAPO Clip-Higher through independent lower and upper ratio bounds; +- token-level DAPO or equal-sequence GRPO loss aggregation; and +- optional linear soft-overlong reward shaping before group advantage + normalization. + +Defaults retain AstrAI's existing symmetric, token-normalized objective and +disable overlong shaping. Dynamic sampling is intentionally excluded because +correct support requires a rollout refill buffer rather than silently dropping +zero-variance groups from an already generated batch. + +## L20 result + +The reproducible objective probe uses FP32 log-probabilities, seed 3407, 20 +warmups, and 100 timed iterations per case. The default candidate loss matches +the baseline exactly. Its median latency ranges from -0.80% to 0.00% relative +to baseline. Opting into Clip-Higher and soft-overlong shaping adds 0.056 to +0.059 ms in the isolated objective microbenchmark (about 22% relative to the +sub-millisecond objective, excluding model forward/backward). + +| Batch×group | Response length | Default parity | Baseline | Default candidate | Full DAPO | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 4×8 | 256 | 0.0 abs | 0.2560 ms | 0.2540 ms | 0.3124 ms | +| 4×8 | 1,024 | 0.0 abs | 0.2570 ms | 0.2570 ms | 0.3164 ms | +| 8×8 | 2,048 | 0.0 abs | 0.2693 ms | 0.2693 ms | 0.3297 ms | +| 4×8 | 4,096 | 0.0 abs | 0.2693 ms | 0.2683 ms | 0.3287 ms | + +Raw timings are checked in at +`benchmarks/results/dapo_objective_l20_sm89.json`. + +## Digest construction + +- target/baseline: SHA-256 of target commit + `88c06db096f197acac2a66953bde445c3d720121`; +- candidate: SHA-256 of the sorted per-file SHA-256 list for implementation, + tests, benchmark tool, and changed project documentation, excluding the + InfraSWE and raw-score artifacts to avoid a recursive digest; +- acceptance: GRPO strategy, online end-to-end, and CLI tests; +- probe/workload: the benchmark tool and checked-in raw result; and +- required deployment cell: the literal + `nvidia-l20-sm89-single-gpu-objective-gpu5`. + +The benchmark source identifier is implementation commit +`bd10b800018794db742cbfe6cf7bcc7c32ea6f44`. The local suite passed 635 tests +with 103 environment-dependent skips, and the focused L20 suite passed 26 +tests. diff --git a/benchmarks/infraswe/astrai-dapo-comparison-cell.json b/benchmarks/infraswe/astrai-dapo-comparison-cell.json new file mode 100644 index 00000000..180e0722 --- /dev/null +++ b/benchmarks/infraswe/astrai-dapo-comparison-cell.json @@ -0,0 +1,16 @@ +{ + "schema_version": "0.5", + "target_project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08", + "target_repository_or_baseline_sha256": "sha256:df1e79ff7151f89d036234564359857fcce3cb4b70e76c72f66d4900b07d17e5", + "change_intent": "integrate", + "semantic_contract_sha256": "sha256:d1a1c9eef4a8d03316f9b805b3f8e9b70c8b9c92f6778c8a8fa750b6562ce119", + "acceptance_contract_sha256": "sha256:9b23206e258c6ee6e3a122fe4925561cc550a28567a67489ada1c830e4ecb2dc", + "probe_set_sha256": "sha256:022b1c0460c9d53e7b2a7bbb0a354af2bfe114f03808d8909bd44386dcae272b", + "workload_portfolio_sha256": "sha256:cfdbe613b9ab7d1872031219783a52e97513db107af68ba2f06173e8338ccdad", + "performance_target_sha256": "sha256:232ef0c0dba913d2333c671203ac41e269e07a18824a8c5200c660d3abc0da0c", + "required_deployment_cell_set_sha256": "sha256:3b17d79e40c95dbc7d9c03c069166a295a08523f867a6fea334bc41a5f7f4448", + "formula_template_id": "project-fit-system-path-v0.5.1", + "evidence_policy_id": "system-path-evidence-v1", + "project_season": "astrai-2026q3", + "cross_project_ranking_allowed": false +} diff --git a/benchmarks/results/dapo_objective_l20_infraswe_score.json b/benchmarks/results/dapo_objective_l20_infraswe_score.json new file mode 100644 index 00000000..31e6a99f --- /dev/null +++ b/benchmarks/results/dapo_objective_l20_infraswe_score.json @@ -0,0 +1,110 @@ +{ + "schema_version": "0.5.1", + "score_kind": "diagnostic-project-fit", + "score_is_official": false, + "candidate_revision": "sha256:bdf90ab8bcde559649302e5721c9db9eb97439210ac2eb491aa4b3a03ac377ad", + "formula_template_id": "project-fit-system-path-v0.5.1", + "diagnostic_project_fit_100": 92.14567151684916, + "component_values": { + "evolutionary_maintainability": 0.8993737273809536, + "project_contract_fit": 0.9592635533087708, + "performance_reuse_utilization": 0.9451290407777138, + "operational_fit": 0.8554459186535495 + }, + "component_floors": { + "evolutionary_maintainability": 0.6, + "project_contract_fit": 0.6, + "performance_reuse_utilization": 0.4, + "operational_fit": 0.6 + }, + "subcomponent_inputs": { + "evolutionary_maintainability": { + "evolution": 0.82, + "locality": 0.9, + "tests": 1.0, + "failure": 0.95, + "contract": 0.95 + }, + "project_contract_fit": { + "integration": 0.95, + "interface": 1.0, + "lifecycle": 0.9, + "buildtest": 1.0, + "policy": 0.95 + }, + "performance_reuse_utilization": { + "attainment": 1.0, + "coverage": 0.85, + "retention": 1.0, + "family": 0.9, + "compile": 1.0 + }, + "operational_fit": { + "replay": 0.85, + "load": 0.75, + "resource": 1.0, + "coldsteady": 0.9 + } + }, + "input_rationale": { + "evolution": "The additive objective controls preserve the existing default but have no upstream maintenance history yet.", + "locality": "Implementation is confined to GRPOStrategy, CLI strategy arguments, focused tests, documentation, and benchmark evidence.", + "tests": "The complete local suite passed 635 tests with 103 environment-dependent skips; 26 focused tests passed on NVIDIA L20 GPU5.", + "failure": "Non-finite or invalid clip bounds, unsupported aggregation modes, incomplete overlong windows, and invalid penalty scales fail explicitly.", + "contract": "Implementation, tests, documentation, benchmark tooling, and raw L20 output are digest-bound, but the artifact is not sealed or maintainer-reviewed.", + "integration": "All options flow through the existing train CLI strategy_kwargs boundary into offline and online GRPO without a parallel strategy implementation.", + "interface": "Unset clip bounds inherit the symmetric epsilon, token aggregation remains the default, and soft-overlong shaping is opt-in.", + "lifecycle": "Reward shaping creates a new tensor before group normalization and does not mutate reward-model output or rollout cache state.", + "buildtest": "The full local suite, focused L20 suite, Ruff checks, and InfraSWE engine tests passed.", + "policy": "No dependency is added and the PR explicitly excludes dynamic sampling rather than approximating its refill semantics.", + "attainment": "Default objective values exactly match baseline in all measured cases and default median latency shows no regression.", + "coverage": "The probe spans batch/group 4x8 and 8x8 with response lengths from 256 to 4096, and tests cover formula, validation, masks, backward, CLI, and online integration.", + "retention": "Default symmetric clipping and token aggregation retain prior behavior while all existing regressions pass.", + "family": "The options reuse GRPOStrategy for offline and online rollout paths and expose both original-GRPO and DAPO reductions.", + "compile": "The objective changes introduce no compilation step and do not compile during timed operations.", + "replay": "Each L20 case uses 20 warmups and 100 timed iterations, but five independent fresh-process runs were not performed.", + "load": "The L20 probe executes the real tensor objective shapes, but excludes model forward/backward and downstream reward-quality evaluation.", + "resource": "Only the explicitly available GPU5 was used; all other users' GPU processes and containers were left untouched.", + "coldsteady": "Warm objective latency is measured; end-to-end cold startup and long training stability are outside this evidence." + }, + "benchmark_trust": { + "formula_version": "benchmark-trust-v0.5", + "status": "scored", + "score_100": 97.40037464252967, + "components": { + "reproducibility": 1.0, + "evidence": 1.0, + "statistics": 0.9, + "environment": 1.0 + }, + "failure_codes": [ + "DRAFT_UNSEALED", + "FRESH_PROCESS_REPLAY_INCOMPLETE", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED", + "END_TO_END_REWARD_QUALITY_UNMEASURED" + ] + }, + "official_project_fit": { + "status": "unresolved", + "score_100": null, + "failure_codes": [ + "DRAFT_SEAL_MISSING", + "FRESH_PROCESS_REPLAYS_BELOW_MINIMUM", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED" + ] + }, + "comparison_cell_path": "benchmarks/infraswe/astrai-dapo-comparison-cell.json", + "execution": { + "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", + "comparison_cell_validation": "pass", + "infraswe_engine_tests": "53 passed", + "astrai_implementation_commit": "bd10b800018794db742cbfe6cf7bcc7c32ea6f44", + "astrai_local_tests": "635 passed, 103 skipped", + "astrai_l20_focused_tests": "26 passed in 2.04s", + "astrai_lint": "passed" + } +} diff --git a/benchmarks/results/dapo_objective_l20_sm89.json b/benchmarks/results/dapo_objective_l20_sm89.json new file mode 100644 index 00000000..4fe193a9 --- /dev/null +++ b/benchmarks/results/dapo_objective_l20_sm89.json @@ -0,0 +1,74 @@ +{ + "schema_version": "1.0", + "candidate_source_commit": "bd10b800018794db742cbfe6cf7bcc7c32ea6f44", + "baseline_commit": "88c06db096f197acac2a66953bde445c3d720121", + "recorded_at": "2026-09-02T22:30:52+08:00", + "gpu": "NVIDIA L20", + "torch": "2.11.0+cu128", + "cuda": "12.8", + "seed": 3407, + "warmup": 20, + "repeats": 100, + "cases": [ + { + "batch": 4, + "group": 8, + "response_len": 256, + "valid_tokens": 6152, + "symmetric_loss_parity_abs": 0.0, + "legacy_median_ms": 0.25600001215934753, + "legacy_p99_ms": 0.32256001234054565, + "candidate_symmetric_median_ms": 0.2539519965648651, + "candidate_symmetric_p99_ms": 0.2672640085220337, + "candidate_symmetric_delta_percent": -0.800006053596447, + "dapo_median_ms": 0.3123840093612671, + "dapo_p99_ms": 0.3246079981327057, + "dapo_delta_percent": 22.024997860868556 + }, + { + "batch": 4, + "group": 8, + "response_len": 1024, + "valid_tokens": 23973, + "symmetric_loss_parity_abs": 0.0, + "legacy_median_ms": 0.25702399015426636, + "legacy_p99_ms": 0.3256640136241913, + "candidate_symmetric_median_ms": 0.25702399015426636, + "candidate_symmetric_p99_ms": 0.27033600211143494, + "candidate_symmetric_delta_percent": 0.0, + "dapo_median_ms": 0.3164159953594208, + "dapo_p99_ms": 0.39215999841690063, + "dapo_delta_percent": 23.107572631452488 + }, + { + "batch": 8, + "group": 8, + "response_len": 2048, + "valid_tokens": 95892, + "symmetric_loss_parity_abs": 0.0, + "legacy_median_ms": 0.2693119943141937, + "legacy_p99_ms": 0.2949120104312897, + "candidate_symmetric_median_ms": 0.2693119943141937, + "candidate_symmetric_p99_ms": 0.2836480140686035, + "candidate_symmetric_delta_percent": 0.0, + "dapo_median_ms": 0.32972800731658936, + "dapo_p99_ms": 0.40857601165771484, + "dapo_delta_percent": 22.433465377673123 + }, + { + "batch": 4, + "group": 8, + "response_len": 4096, + "valid_tokens": 96922, + "symmetric_loss_parity_abs": 0.0, + "legacy_median_ms": 0.2693119943141937, + "legacy_p99_ms": 0.2826240062713623, + "candidate_symmetric_median_ms": 0.2682879865169525, + "candidate_symmetric_p99_ms": 0.27856001257896423, + "candidate_symmetric_delta_percent": -0.38023104015432185, + "dapo_median_ms": 0.32870399951934814, + "dapo_p99_ms": 0.4034560024738312, + "dapo_delta_percent": 22.05323433751878 + } + ] +} diff --git a/docs/developer/internals.md b/docs/developer/internals.md index ccfdd8f3..70e32f29 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -86,7 +86,12 @@ $$ L_{\text{GRPO}} = -\mathbb{E}_t\left[\min\left(\rho_t A,\; \text{clip}\left(\ Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-token importance sampling ratio. Online rollout records $\log \pi_{\text{old}}$ when each token is sampled and reuses those values directly during training; offline batches may fall back to a synchronized `old_model`. Advantages are derived from scalar per-response rewards, group-normalized, and broadcast across all response tokens. Only response tokens contribute to the loss. -Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. +Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. Optional +`clip_eps_low`/`clip_eps_high` values enable DAPO-style asymmetric clipping; +unset values inherit `clip_eps` for backward-compatible symmetric clipping. +The `loss_aggregation` switch selects token-level DAPO weighting or equal +sequence weighting. Optional `overlong_max_len`/`overlong_buffer_len` settings +add the DAPO linear soft-overlong penalty before group advantage normalization. ### MoE Load Balancing diff --git a/docs/guides/params.md b/docs/guides/params.md index 9b23ebd8..08a38104 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -141,6 +141,12 @@ with `--optimizer=muon_adamw`. | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` | | `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo` | | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo`, `online_grpo` | +| `--grpo_clip_eps_low` | Optional lower clip epsilon; defaults to `grpo_clip_eps` | None | `grpo`, `online_grpo` | +| `--grpo_clip_eps_high` | Optional upper clip epsilon for DAPO Clip-Higher | None | `grpo`, `online_grpo` | +| `--grpo_loss_aggregation` | Loss weighting: DAPO-style `token` or equal-weight `sequence` | token | `grpo`, `online_grpo` | +| `--grpo_overlong_max_len` | Optional maximum response length for DAPO soft overlong shaping | None | `grpo`, `online_grpo` | +| `--grpo_overlong_buffer_len` | Linear penalty window before `grpo_overlong_max_len` | 0 | `grpo`, `online_grpo` | +| `--grpo_overlong_penalty_scale` | Scale for the soft overlong reward penalty | 1.0 | `grpo`, `online_grpo` | | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo`, `online_grpo` | | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` | diff --git a/docs/guides/training.md b/docs/guides/training.md index 0d267f12..fc7ef173 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -154,10 +154,25 @@ per-token `logprobs_old` captured by the rollout sampler, avoiding an a compatibility fallback. The KL term regularises $\pi_\theta$ towards a frozen reference model (`ref_model`, typically the SFT checkpoint). -Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. Offline callers that +Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. The optional +`clip_eps_low` and `clip_eps_high` parameters replace the symmetric interval +with $[1-\epsilon_{low}, 1+\epsilon_{high}]$. Leaving both unset preserves the +existing symmetric objective. DAPO Clip-Higher can be selected explicitly, for +example with `clip_eps_low=0.2` and `clip_eps_high=0.28`. Offline callers that do not provide `logprobs_old` must sync `old_model` weights via `sync_old_model()` between data-generation rounds. +`loss_aggregation="token"` (the default) divides by the total number of valid +response tokens, matching DAPO's token-level policy-gradient loss. Set it to +`"sequence"` to first average each response and then weight responses equally, +matching the original GRPO reduction for controlled ablations. + +DAPO soft overlong shaping is enabled by setting `overlong_max_len` and a +positive `overlong_buffer_len`. If $L$ is the valid response length, the added +reward is zero through $L_{max}-L_{buffer}$, falls linearly to -1 at $L_{max}$, +and is multiplied by `overlong_penalty_scale`. It is disabled by default and +does not alter the reward-model output in place. + Keys: `prompts`, `responses`, `masks`, `rewards`, and optional `logprobs_old` (required when `old_model` is not configured). diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 1f9690eb..23654367 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -305,6 +305,48 @@ def _merge_yaml_into_kwargs( group="Algorithm", help="GRPO clip epsilon.", ) +@opt( + "--grpo_clip_eps_low", + type=float, + default=None, + group="Algorithm", + help="Optional lower GRPO clip epsilon; defaults to --grpo_clip_eps.", +) +@opt( + "--grpo_clip_eps_high", + type=float, + default=None, + group="Algorithm", + help="Optional upper GRPO clip epsilon for DAPO Clip-Higher.", +) +@opt( + "--grpo_loss_aggregation", + type=click.Choice(["token", "sequence"]), + default="token", + group="Algorithm", + help="Aggregate GRPO loss by token (DAPO) or equally by sequence.", +) +@opt( + "--grpo_overlong_max_len", + type=int, + default=None, + group="Algorithm", + help="Optional response length limit for DAPO soft overlong shaping.", +) +@opt( + "--grpo_overlong_buffer_len", + type=int, + default=0, + group="Algorithm", + help="Length of the linear DAPO overlong penalty window.", +) +@opt( + "--grpo_overlong_penalty_scale", + type=float, + default=1.0, + group="Algorithm", + help="Scale applied to the DAPO soft overlong penalty.", +) @opt( "--grpo_kl_coef", type=float, @@ -679,6 +721,12 @@ def train( "beta": kwargs.pop("dpo_beta"), "label_smoothing": kwargs.pop("label_smoothing"), "clip_eps": kwargs.pop("grpo_clip_eps"), + "clip_eps_low": kwargs.pop("grpo_clip_eps_low"), + "clip_eps_high": kwargs.pop("grpo_clip_eps_high"), + "loss_aggregation": kwargs.pop("grpo_loss_aggregation"), + "overlong_max_len": kwargs.pop("grpo_overlong_max_len"), + "overlong_buffer_len": kwargs.pop("grpo_overlong_buffer_len"), + "overlong_penalty_scale": kwargs.pop("grpo_overlong_penalty_scale"), "kl_coef": kwargs.pop("grpo_kl_coef"), "group_size": kwargs.pop("group_size"), } diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 473871b2..e9182930 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/tests/trainer/test_grpo_strategy.py @@ -1,6 +1,7 @@ import pytest import torch +import astrai.trainer.strategy as strategy_module from astrai.model.transformer import AutoRegressiveLM from astrai.trainer.strategy import GRPOStrategy from tests.helpers import FakeExecutor, make_frozen, make_model @@ -109,6 +110,147 @@ def test_grpo_rejects_invalid_behavior_logprobs(grpo_strategy, invalid): strategy.compute_loss(batch) +def test_grpo_symmetric_clip_is_backward_compatible(grpo_strategy): + strategy, _device = grpo_strategy + assert strategy.clip_eps == pytest.approx(0.2) + assert strategy.clip_eps_low == pytest.approx(0.2) + assert strategy.clip_eps_high == pytest.approx(0.2) + assert strategy.loss_aggregation == "token" + + +def test_grpo_dapo_clip_higher_changes_positive_advantage_bound( + grpo_strategy, monkeypatch +): + strategy, device = grpo_strategy + batch = { + "prompts": torch.tensor([[5]], device=device), + "responses": torch.tensor([[[6], [7]]], device=device), + "masks": torch.ones(1, 2, 1, device=device), + "rewards": torch.tensor([[1.0, -1.0]], device=device), + } + policy_logprobs = torch.log( + torch.tensor([[1.25], [0.75]], device=device, dtype=torch.float32) + ) + zeros = torch.zeros_like(policy_logprobs) + + def policy_loss(clip_eps_high): + strategy.clip_eps_high = clip_eps_high + outputs = iter( + [ + policy_logprobs, + zeros, + policy_logprobs, + ] + ) + + def fake_get_logprobs(*_args, **_kwargs): + return { + "logprobs": next(outputs), + "aux_loss": None, + "router_stats": None, + } + + monkeypatch.setattr(strategy_module, "get_logprobs", fake_get_logprobs) + return strategy.compute_loss_output(batch)["metrics"]["policy_loss"] + + assert policy_loss(0.2) == pytest.approx(-0.2, abs=1e-6) + assert policy_loss(0.28) == pytest.approx(-0.225, abs=1e-6) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"clip_eps_low": 1.0}, "clip_eps_low"), + ({"clip_eps_high": float("nan")}, "clip_eps_high"), + ( + {"clip_eps_low": 0.3, "clip_eps_high": 0.2}, + "greater than or equal", + ), + ], +) +def test_grpo_rejects_invalid_asymmetric_clip(grpo_strategy, kwargs, match): + strategy, device = grpo_strategy + with pytest.raises(ValueError, match=match): + GRPOStrategy( + model=strategy.model, + device=device, + old_model=strategy.old_model, + ref_model=strategy.ref_model, + clip_eps=0.2, + executor=FakeExecutor(), + **kwargs, + ) + + +def test_grpo_token_and_sequence_aggregation_weight_lengths_differently( + grpo_strategy, +): + strategy, device = grpo_strategy + losses = torch.tensor([[[1.0, 1.0, 0.0, 0.0], [3.0, 3.0, 3.0, 3.0]]], device=device) + masks = torch.tensor([[[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]], device=device) + + strategy.loss_aggregation = "token" + token_loss = strategy._reduce_token_loss(losses, masks) + strategy.loss_aggregation = "sequence" + sequence_loss = strategy._reduce_token_loss(losses, masks) + + assert token_loss.item() == pytest.approx(14.0 / 6.0) + assert sequence_loss.item() == pytest.approx(2.0) + + +def test_grpo_dapo_soft_overlong_reward_shaping(grpo_strategy): + strategy, device = grpo_strategy + strategy.overlong_max_len = 8 + strategy.overlong_buffer_len = 2 + strategy.overlong_penalty_scale = 0.5 + rewards = torch.zeros(1, 4, device=device) + masks = torch.zeros(1, 4, 8, device=device) + for index, length in enumerate((5, 6, 7, 8)): + masks[0, index, :length] = 1 + + shaped, penalty = strategy._shape_overlong_rewards(rewards, masks) + + assert penalty is not None + torch.testing.assert_close( + penalty, torch.tensor([[0.0, 0.0, -0.5, -1.0]], device=device) + ) + torch.testing.assert_close( + shaped, torch.tensor([[0.0, 0.0, -0.25, -0.5]], device=device) + ) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"loss_aggregation": "batch"}, "loss_aggregation"), + ({"overlong_buffer_len": 4}, "requires overlong_max_len"), + ( + {"overlong_max_len": 8, "overlong_buffer_len": 9}, + "no greater than overlong_max_len", + ), + ( + { + "overlong_max_len": 8, + "overlong_buffer_len": 2, + "overlong_penalty_scale": -0.1, + }, + "overlong_penalty_scale", + ), + ], +) +def test_grpo_rejects_invalid_dapo_objective_options(grpo_strategy, kwargs, match): + strategy, device = grpo_strategy + with pytest.raises((TypeError, ValueError), match=match): + GRPOStrategy( + model=strategy.model, + device=device, + old_model=strategy.old_model, + ref_model=strategy.ref_model, + executor=FakeExecutor(), + **kwargs, + ) + + @pytest.mark.parametrize("model_name", ["ref_model", "old_model"]) def test_grpo_frozen_models_not_updated(grpo_strategy, model_name): """Backward should not populate gradients on ref_model or old_model."""