[rl] Add the curriculum-RL sampling experiment - #8769
Open
rjpower wants to merge 63 commits into
Open
Conversation
Wire a fixed-budget comparison of environment-sampling curricula for SkyRL GRPO on Qwen3-0.6B. A shared pool step writes six difficulty-graded math bins (GSM8K, MATH levels 1-2/3/4/5, AIME 1983-2024; 10,427 train rows) in the SkyRL parquet schema with per-row env_class and grade, plus a fixed gsm8k/MATH-500 validation set. A mirror step stages the pinned Qwen/Qwen3-0.6B snapshot as the starting policy, and each arm builds a skyrl_step on cw-rno2a (interactive priority) with a terminal eval step. Arms share the pool, policy, and step budget and differ only in prompt sampling. Only the naive arm (the pinned trainer's uniform shuffle) is buildable; the Thompson and graded-curriculum arms fail fast until the MarinSkyRL sampler branch lands. Scale presets cover a 2-node smoke and an 8-node (64 H100) measurement run. Part of #8765
The pinned trainer (b45a6ed2) fails at startup with 'SKYRL_DEBUG_ARTIFACT_DIR is required for a debug process manifest' when debug mode is off; MarinSkyRL main fixed this in the next commit (marin-community/MarinSkyRL#461). This failed every attempt of the curriculum-rl smoke run. Part of #8765
Eight single-GPU vLLM engines on one node can race on the same torch TCPStore port at startup (EADDRINUSE); the failure is transient and the run resumes from the latest checkpoint, so raise max_retries from 3 to 6. Part of #8765
The aime env asserts extras['reward_model']['ground_truth'] (the gsm8k env
reads reward_spec), so every batch with a MATH/AIME row crashed at env
construction; rows now carry both columns. The aime verifier extracts
answers with the Minerva 'Answer: ...' regex, not \boxed{}, so the prompt
instruction now elicits a final Answer line. Representative ground truths
(integers, degrees, fractions, tuples, radicals, polynomials) pass the
skyrl-gym aime contract's two-sided preflight. Fail fast on launch errors:
max_retries drops to 1.
Part of #8765
At d018d672 the terminal policy export dies in the torch-free launcher
environment ('No module named torch': skyrl_train/hf_export.py pulls
skyrl_train.utils, whose package init imports torch). Track
fix-hf-export-torch-free (618f6b2d) until the fix merges to main.
Part of #8765
The launcher's terminal JSON response was corrupted by the export subprocess writing [rl-iris] lines to inherited fd 1; the branch now routes that subprocess's stdout to stderr. Part of #8765
The export subprocess runs from the package root, so the caller-relative cluster config no longer resolved and the nested launch exited 1. The branch now absolutizes config paths before spawning. Part of #8765
The launcher defaults trainer.hf_hub_repo_id to laion/<job_name>, and the export job 403s creating that repo. Arms override it to null; the terminal export stays in object storage. Part of #8765
At 0.6B the Qwen3 thinking mode consumed the whole generation budget in the smoke run (85% truncation at a 1024-token cap, 4% parseable answers), leaving almost no reward signal. Switch the generator to the non-batched path and pass chat_template_kwargs enable_thinking=false; tighten the full preset to a 2048-token window with 1024 new tokens. Part of #8765
The config flattener emits bare generator.* keys and hydra struct mode rejects new children under the empty chat_template_kwargs mapping, which failed every r8 attempt at startup. Route the flag through the caller override list, which is applied with ++. Part of #8765
Pin MarinSkyRL to the curriculum-sampling branch, which adds the data.sampling weighted sampler over graded bins plus the unstepped episode metrics fix, and route each non-naive arm through data.sampling.kind. The naive arm keeps the stock uniform shuffle. Part of #8765
Prompts above the trainer's max_input_length (1024 tokens in both presets) skip generation, and a fully skipped GRPO group fails group admission (r10 smoke: TrainingGroupInvariantError fully_masked on an AIME row). Filter train rows whose non-thinking templated prompt exceeds 1000 tokens when writing the pool; they could never contribute gradient signal and waste a rollout-budget slot in every arm. Part of #8765
apply_chat_template(tokenize=True) returns an encoding mapping in the pinned transformers, so the length filter measured every prompt at 2 tokens and dropped nothing (r11 failed on the same over-length AIME group as r10). Request return_dict and measure input_ids; the filter now drops 8 train rows. Part of #8765
Rename the misnamed BOXED_* constants (the aime env verifies an Answer: line, not \boxed), delete the unused POOL_BINS/POOL_GRADE_COUNT pair, assert at import that the pool's prompt-length cap stays under every preset's input budget, and rewrite the eval reporter on the canonical record.json reader instead of a hand-rolled glob that duplicated arm names and result-layout knowledge. Part of #8765
rjpower
added
agent-generated
Contributor
IaC preview
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 656a1c8384
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
A resumed run starts a new W&B entry whose _step restarts at zero, so a single merged token history collapses entries onto the same keys and undercounts spend. Walk entries chronologically, accumulate tokens across them, and key eval points by (entry, step). This is the accounting used for the round-2 report charts.
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
…export Round 3 rebuilds the pool on a 0-13 difficulty ladder anchored to school grade and contest tier. ASDiv joins with its per-problem school grades (read from the original XML at a pinned commit; the HF mirrors drop the grade attribute), SVAMP at the same band, and the chain_sum bins are re-anchored to round-2 measured pass rates. The non-math spelling and base-conversion bins leave the ladder. Grades g04/g09/g11-g13 are reserved for AMC, olympiad, Putnam-AXIOM, and graduate bins. The Snowball 67B-A2B stage-2 SFT export is registered in experiments/models.py via ArtifactStep.adopt so pipelines can consume the existing S3 export as a typed checkpoint artifact without copying 134GB.
NuminaMath source tiers (cn_k12 at grade 4, synthetic AMC at 9, olympiads at 11), Omni-MATH AoPS-rated bands (7.0+ at grade 11, 8.5+ at 12, with a held-out validation slice), TheoremQA numeric rows and HARDMath single-expression results at grade 13. Putnam-AXIOM stays out: every HF mirror is gated, so the 8.5+ Omni band carries that tier. HARDMath ground truths arrive as statements; the loader keeps the right-hand side the question's rounding instruction targets.
Arms now select a policy model: the existing Qwen mirror or the adopted Snowball 67B-A2B SFT export. Snowball arms run on cw-us-east-02a beside the export, roll out through one node-sized expert-parallel vLLM engine (tp=1, dp=8, ep=8, matching the evaluation serving profile), and train on four FSDP2 nodes; the snowball-smoke preset stays at the 2k window where the #7786 campaign was stable. Evaluation serving branches to the expert-parallel profile for Snowball checkpoints.
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
The learnability and grade-prior samplers now weight bins by 1 - p^n - (1-p)^n (the probability a GRPO group survives dynamic-sampling filtering, n = n_samples_per_prompt) instead of per-sample reward variance p(1-p), which under-samples low-pass-rate bins by ~3x relative to their actual refill cost. Adds the learnability-dapo and grade-prior-dapo arms for the round-3 catalog and repins MarinSkyRL to 44817814, which carries the new data.sampling.weighting knob.
The grade-weighted end metric needs visibility above grade 5; val-amc (128 rows past the train quota) and val-theoremqa (every fifth numeric row) extend held-out coverage to grades 9 and 13 with splits disjoint from training by construction.
The 67B export streams ~134GB of bf16 shards through host buffers on load; Ray OOM-killed the first snowball-smoke at the 128GB request. The evaluation serving profile already sizes this model's nodes at 512GB.
Four policy plus four expert-parallel engine nodes, 60 steps at 128x8 responses, and a 3072-token window with a 2048 response budget: the smoke averaged 884 tokens against a 1024 cap, so the thinking model gets room before truncation while every pool prompt still fits.
Charts now compute the fixed end metric (per-validation-bin scores weighted by 1 + grade), a frontier-grade-over-tokens headline chart (highest grade scoring >= 0.25), and an empirical learning-velocity-vs-pass-rate curve with the p(1-p) and group-informative weighting references overlaid. Token accounting is preset-aware: snowball arms use the SNOWBALL_FULL batch shape (128x8) instead of the Qwen FULL shape (512x8).
Snowball checkpoint names keep the -snowball-full scale label (only the Qwen FULL preset drops its suffix), so run names parse to snowball-naive-snowball-full; strip the suffix before matching ARM_ORDER.
The naive probe measured ~1770s steps with ~1660s in policy_train: at micro_train_batch_size_per_gpu=1 the 1024-response batch runs 32 sequential FSDP micro-steps per step, each re-gathering the full 134GB of shards. micro=4 quarters the gather traffic, projecting ~10-minute steps so 60 steps fit a 24h driver window. Generation (40s), weight sync (10-20s), and the logprob pass (55s) are immaterial.
eval/*/avg_score is a per-environment reward with inconsistent shaping (gsm8k logs 0/1, math bins log +/-1), so the grade-weighted end metric and frontier threshold now read eval/*/pass_at_1. Snowball arms match only the 2026.08.31.1 relaunch so the cancelled micro=1 probe cannot merge into the snowball-naive series.
Sampling 50 history rows for key discovery misses the per-step curriculum rows in snowball runs, where engine metrics outnumber them; summaries carry every logged key.
Replace the flat matplotlib chart set with seaborn-styled figures split by model family: EMA-smoothed grade-weighted headlines with raw evals drawn faintly underneath, 2x3 per-grade validation breakouts over steps and tokens, and a grade-attainment bar chart (tokens to first sustain pass@1 >= 0.25 per grade) replacing the unreadable integer frontier-grade step chart. Curriculum weight charts move to a log scale so floor starvation is visible, and the velocity chart clips its axis to the binned-mean range. seaborn is not a workspace dependency; the docstring documents the uv --with overlay.
…version Round 4 presets (snowball-smoke-r4, snowball-full-r4) switch the optimizer to MuonH (the grug_moe_muonh_v1 recipe the base model trained under) at 1e-5, halve the batch to 64 prompts at 120 steps (all-gather cost per token is fixed by the micro batch; smaller steps buy 2x-fresher rollouts for ~1% wall time), and set data.sampling.reversion_mass=2.0 so starved bins recover re-probeable weight (MarinSkyRL b7eed273). The pool gains a system message carrying the per-bin answer-format contract; round-3 arms showed a trailing user instruction loses to the SFT \boxed style. Charts treat round 4 as its own snowball-r4 family.
A 4-variant sweep against the served Snowball base model (40 gsm8k + 40 MATH-500 + 20 TheoremQA rows, temperature 1.0, rule-graded): the round-3 suffixes scored 0.00 on gsm8k with correct math ending in \boxed; the system message plus a restated exact-form final-line instruction scored 0.55. Both instruction constants now carry the strong form. MuonH validates weight_decay=0, so the round-4 overrides set it explicitly (the base config's AdamW default is 1e-2).
The r4 smoke (MuonH lr 1e-5, weight_decay 0, micro=8, batch 64x8) ran clean on cw-us-east-02a: step 222s vs 534s in round 3, policy_train ~112s, raw grad norm 0.38-0.47 under the 1.0 clip. Round-4 charts read the 2026.09.01.1 pool version, which carries the strengthened format contracts.
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
Aggregates think-token usage, truncation, and answer-line compliance per step bucket and bin, segmented by retention stream: truncated and failed rollouts are retained mandatorily, so the hash-sampled remainder covers successful terminating rollouts only. Replaces the round-3 one-shot inspect_trajectories.py diagnostic.
…prompt Round 4's binding constraint was the 2048-token generation cap (40-48% of rollouts truncated at end of training; AIME/Omni/TheoremQA near-total), so the round-5 presets raise max_new_tokens to 8192 over the same 1024-token prompt budget. micro_train drops to 2 because activation memory tracks tokens per micro batch and 8x3072 OOMed twice in the round-4 naive arm. The system prompt shrinks from 111 to 37 tokens: a served-model A/B (gsm8k/MATH-500/TheoremQA, temp 1.0, 100-250 rows per variant) showed rule-graded compliance flat across 111/66/37-token wordings, with pass rates tying or favoring the shortest.
…hared constants Rewrites round-numbered comment narration as standing rationale, renames SNOWBALL_R4_OVERRIDES to SNOWBALL_MUONH_OVERRIDES (both the round-4 and round-5 presets use it), hoists the trajectory-stats tail window to a constant, and spells the Snowball SFT export URI once (SNOWBALL_SFT_EXPORT_URI) so the RL policy and eval model cannot drift.
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
…rriculum # Conflicts: # config/external/MarinSkyRL/uv.lock # lib/marin/src/marin/external_dependencies.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add
experiments/post_training/curriculum_rl, a SkyRL GRPO experiment measuring what dynamic sampling curricula buy at a fixed budget. Five arms train Qwen3-0.6B on one shared six-bin graded math pool (GSM8K, MATH by level, AIME; 10,419 rows) and differ only in thedata.samplingpolicy: naive shuffle, Thompson sampling on informative-group fraction, and three pre-graded variants (uniform, adaptive advance-on-exhaustion, grade-seeded prior). Each full arm is 120 steps x 512 prompts x 8 samples on 8 H100 nodes (16 training GPUs, 48 vLLM engines), wired as cached ArtifactSteps: pool build, model mirror, RL training with terminal HF export, and Evalchemy evaluation.The pool builder pins dataset revisions, carries both
reward_specandreward_modelground-truth columns (the gsm8k and aime envs read different keys), elicitsAnswer:-line answers for the Minerva-style aime verifier, and drops the 8 rows whose non-thinking templated prompt exceeds 1000 tokens — over-length prompts skip generation and a fully skipped GRPO group fails the trainer's admission invariant. Arms run in non-thinking mode (enable_thinking=falsevia a++override; the config flattener emits bare keys and hydra rejects new children under the emptychat_template_kwargs): thinking mode at the 1024-token generation cap truncated 85% of rollouts and parsed 4%, non-thinking truncates 20% and parses 70%.The MarinSkyRL pin moves to the
curriculum-samplingbranch, which adds the stateful curriculum sampler plus fixes found during bring-up (torch-free export imports, JSON-only launcher stdout, export config path resolution, unstepped-episode metrics aggregation).Results (single seed per arm): at a fixed generated-token budget of 146M, grade-adaptive reached val-math500 0.156 vs naive 0.080 at ~53% of naive's token spend; at fixed steps/GPU-hours, naive led MATH500 (0.460) and thompson led gsm8k-0shot (0.516). Full tables and recommendations are on the tracking issue.
Part of #8765