Skip to content

configuration.md

Latest commit

 

History

History
277 lines (195 loc) · 14.2 KB

configuration.md

File metadata and controls

277 lines (195 loc) · 14.2 KB
Raw

Configuration

Every prime-rl entrypoint uses pydantic-config: TOML files for reproducible base configs, CLI flags for one-off overrides.

AI agents working in this repo: the equivalent runbook is at skills/configs/SKILL.md, with extra runtime hints (where config classes live, validator conventions, the trainer-side enable_token_export flag) that aren't surfaced here.

Table of Contents

Sources and Precedence

Field values come from three sources — Pydantic defaults, TOML files (passed with @), and CLI flags. They're layered in this order, with later sources winning:

  1. Defaults declared on the Pydantic model.
  2. TOML files passed with @, left to right — later files override earlier ones.
  3. CLI flags in dotted, kebab-case form (--model.name).

TOML Composition

The @ token introduces a TOML file. Multiple @ arguments compose left-to-right, deep-merged — unset fields in an overlay keep the base value:

uv run rl @ examples/reverse_text/rl.toml                      # one file
uv run rl @ base.toml @ overlay.toml                           # left to right
uv run rl --trainer @ trainer.toml --orchestrator @ orch.toml  # per-section
uv run rl @ base.toml --trainer @ trainer.toml                 # mixed

Mind the space: @ path/to/x.toml, not @path/to/x.toml.

CLI Overrides

CLI flags mirror the TOML tree using dots:

--max-steps 50                              # top-level
--model.name Qwen/Qwen3-4B                  # nested
--trainer.optim.lr 1e-5                     # double-nested
--inference.parallel.tp 4

Field names are snake_case in TOML (max_model_len) and kebab-case on the CLI (--max-model-len).

Renamed fields keep their old name as a validation alias — e.g. rollouts_per_example is still accepted in TOML and CLI after being renamed to group_size. Mixing the two names across sources is safe.

Inspecting and Validating

uv run rl --help                                       # full schema
uv run rl @ rl.toml --dry-run --output-dir /tmp/check  # write resolved configs

Syntax

Booleans

CLI uses paired flags: bare --flag sets True, --no-flag sets False. TOML must be explicit:

uv run rl @ rl.toml --clean-output-dir       # True
uv run rl @ rl.toml --no-clean-output-dir    # False
clean_output_dir = true

Lists

CLI accepts space-separated values or a JSON literal. TOML uses an array literal. Both forms target the same field:

uv run rl @ rl.toml --trainer.model.lora.target-modules q_proj k_proj v_proj
uv run rl @ rl.toml --trainer.model.lora.target-modules '["q_proj", "k_proj", "v_proj"]'
[trainer.model.lora]
target_modules = ["q_proj", "k_proj", "v_proj"]

Overlay TOMLs replace lists wholesale — an overlay that wants to add one item must still spell out the full list. For arrays of tables (e.g. environments), see Environments.

Dicts

CLI takes a JSON literal. TOML uses a table or inline-table. CLI dicts deep-merge with TOML dicts — CLI keys win on conflict but don't wipe the file's keys:

uv run rl @ rl.toml --orchestrator.train.env.0.args \
  '{"dataset_name": "openai/gsm8k", "dataset_subset": "main"}'
[[orchestrator.train.env]]
args = { dataset_name = "openai/gsm8k", dataset_subset = "main" }

Optional Sub-Configs

Many sub-configs are typed SomeConfig | None. Two patterns enable them:

  • Bare flag with defaults: --model.compile or, in TOML, an empty section [model.compile]. The sub-config materializes with all-default values.
  • Enable and set fields together: --model.compile.fullgraph (CLI) or any populated [model.compile] table (TOML).

To disable a sub-config that's on by default, use --no-<name> on the CLI or assign the string "None" in TOML (see None). This is how [ckpt], [model.lora], [model.compile], [trainer.wandb], etc. are turned on and off.

None

TOML has no null. Use the string "None", which the loader coerces:

[inference.model]
max_model_len = "None"

On the CLI: --inference.model.max-model-len None.

Discriminated Unions

Loss, advantage, optimizer, scheduler, weight broadcast transport, and several others are discriminated unions. Set the type field to pick a variant:

[trainer.optim]
type = "muon"
lr = 1e-5
mu = 0.95

Omit type to keep the default variant.

Environments ([[orchestrator.train.env]])

Training environments are an array of tables — set one per env, optionally with sampling weights:

[[orchestrator.train.env]]
id = "math-env"
name = "gsm8k"
args = { dataset_name = "openai/gsm8k", dataset_subset = "main" }
ratio = 3  # 75% of batches

[[orchestrator.train.env]]
id = "reverse-text"
ratio = 1  # default — 25% of batches

[[orchestrator.eval.env]]
id = "math-env"
name = "gsm8k-eval"
args = { dataset_name = "openai/gsm8k", dataset_subset = "main" }

ratio defaults to 1 (equal weight per env); values are relative weights normalized to probabilities across envs.

args is forwarded verbatim to the environment's load_environment(**args).

The same id can appear multiple times across train and eval (or with different args) — useful for evaluating on a held-out split of the env you're training on, or comparing two configurations of the same env side by side. When id is reused, set a distinct name on each entry; name defaults to id and must be unique across all envs in the same group.

Replaying Past Rollouts (Continue / Recheck)

The built-in replay taskset (verifiers v1) re-enters saved rollouts as fresh training tasks. Point it at rollout record files (<output_dir>/rollouts/step_*/*_rollouts_*.jsonl) and at the source taskset the records came from — the source provides tools, setup/finalize, and scoring, so new completions are judged by the original env's verifier:

# Continue: resume each recorded compaction from its handoff summary. The seed is a plain
# string prompt, so it runs under any harness — use the same harness config as the source run.
[[orchestrator.train.env]]
name = "my-env-continue"
taskset = { id = "replay", records = "/data/run1/rollouts/step_*/train_rollouts_<env>.jsonl", mode = "continue", source = { id = "my-env-v1" } }
harness = { id = "rlm" }

# Recheck: replay each finished attempt with a "check your work" turn appended. The seed is a
# full conversation, which needs a message-seeding harness (`default`/`null`).
[[orchestrator.train.env]]
name = "my-env-recheck"
taskset = { id = "replay", records = "/data/run1/rollouts/step_*/train_rollouts_<env>.jsonl", mode = "recheck", source = { id = "my-env-v1" } }
harness = { id = "null" }

Key knobs on the taskset table:

  • mode = "continue" resumes a rollout mid-way; anchor picks the resume point: "compaction" (default; one task per recorded context restart, detected structurally so records from different harnesses mix freely) or "tool-call" (deterministically-drawn resume points right after complete tool-result runs; needs a message-seeding harness). For "tool-call", max_anchors sets how many resume points each source rollout seeds (default 1; None seeds every valid one, in trajectory order).
  • mode = "recheck" replays the final branch of each attempt (truncation artifacts stripped) plus a verification turn (recheck_prompt overrides the wording).
  • max_seed_tokens skips seeds whose context exceeds the budget — set it so seeds leave room to sample under the trainer's seq_len.
  • records globs are followed: new matching files are picked up continuously (append-only, so pool workers stay index-aligned). Point it at a finished run's records, or at the current run's own rollouts/ dir for online self-replay — the env starts empty and grows as steps ship; the replay env's own rollouts are never re-mined as seeds. Avoid sources that grow out of glob order (e.g. several runs writing one tree).
  • Lines that don't validate as the source taskset's task type (e.g. other envs' rollouts in a mixed train_rollouts_<env>.jsonl) are skipped and counted. The filter is structural: for a source whose task type has no distinguishing required fields, use single-env record files instead.

The orchestrator writes a derived index_<env>.jsonl beside each step's train_rollouts_<env>.jsonl: one row per record with selection fields (task name, source reward, branch count, ...) and the record's byte span. When an index is present, replay filters records index-side — already-replayed tasks, out-of-range source rewards, and (for continue/compaction) un-compacted rollouts — and parses only the selected spans; a record file without an index (e.g. hand-assembled) is fully parsed instead. The index is optional and derived; the record files themselves are unchanged.

Group rollouts of one seeded task form a regular GRPO group, so a group-relative algorithm gets contrastive signal at exactly the resumed state. The sandbox is fresh on re-entry: setup runs anew, and no filesystem state from the source rollout is replayed. Records carrying sandbox snapshot refs (trace.info["snapshots"]) restrict resume points to snapshotted anchors and restore the ref during setup — no runtime implements snapshot capture yet, so restore fails loudly if refs ever appear before support lands. Other recycling schemes can subclass ReplayTaskset and override seeds() — see deps/verifiers/verifiers/v1/tasksets/replay/.

Environment Variables

OS environment variables exported into launched component process(es). In rl configs, top-level [env_vars] applies to trainer, inference, and orchestrator:

[env_vars]
HF_HUB_OFFLINE = "1"
TOKENIZERS_PARALLELISM = "false"

Component-specific tables layer on top:

[trainer.env_vars]
NCCL_DEBUG = "INFO"
PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:False"

[inference.env_vars]
VLLM_USE_DEEP_GEMM = "1"

[orchestrator.env_vars]
PI_USAGE_BASE_URL = "https://..."

The rl launcher applies these the same way in both single-node and multi-node (SLURM) runs. Precedence, low to high:

  1. The launcher's own defaults — your env_vars override these.
  2. Your top-level [env_vars].
  3. Your [component.env_vars].
  4. Orchestration-critical vars the launcher always sets last — CUDA_VISIBLE_DEVICES (GPU partitioning) and WANDB_SHARED_* (the single shared W&B run) — these cannot be overridden from env_vars.

For standalone sft and inference configs, [env_vars] applies to that entrypoint's process(es). For disaggregated P/D inference, the role-specific deployment.{prefill,decode}_env_vars layer on top of any shared inference env vars.

Examples

The shipped end-to-end examples in examples/ are the canonical, kept-up-to-date references — the rest of the repo's TOMLs (under configs/) are CI- and debug-internal and may drift. Each example directory has its own README with the full launch story.

Basic (1–8 GPUs):

  • Reverse Text — Qwen3-0.6B reversing a chunk of text. Tiny single-turn SFT + RL; runs on a single consumer GPU in minutes.
  • Wordle — Qwen3-1.7B playing Wordle. Multi-turn SFT + RL; 2–4 H100s.
  • Alphabet Sort — Qwen3-4B-Instruct-2507 sorting names alphabetically. Multi-turn LoRA RL without SFT warmup; one H100.
  • Wiki Search — Qwen3-4B-Instruct-2507 answering trivia by web-searching Wikipedia. Multi-turn with tool use.
  • Hendrycks Sanity — DeepSeek-R1-Distill-Qwen-1.5B on a filtered MATH subset. Useful for algorithm ablations.

Advanced (32–2048 GPUs, SLURM):

Worked Example: Compose, Override, Dry-Run

Start from a shipped base config, override two fields on the CLI, and dry-run:

uv run rl @ examples/reverse_text/rl.toml \
  --wandb.name my-experiment \
  --trainer.optim.lr 5e-6 \
  --output-dir /tmp/reverse-dry \
  --dry-run

Then inspect the resolved config:

ls /tmp/reverse-dry/configs/
# rl.toml  trainer.toml  orchestrator.toml  inference.toml

Each per-process TOML reflects the final, validated configuration that the actual run would consume — exactly what each process sees when started standalone (uv run trainer @ /tmp/reverse-dry/configs/trainer.toml, etc.). This is the easiest way to bisect a misbehaving config: dry-run a known-good base, dry-run your overlay, diff the two.