[snimu]

Stacked on #4062 (feat/pretrain).

What

The pretrain trainer stops using the shared torch LR scheduler and gets per-hyperparameter schedules. SFT, RL and trainer/scheduler.py are unchanged.

[schedules.optim.lr]        # default: wsd with warmup 0.1, decay 0.9 (linear decay to 0 after the warmup)
type = "wsd"
warmup = 0.05               # steps, or a fraction of max_steps
decay = 0.2
end_ratio = 0.1
decay_shape = "1-sqrt"      # linear | cosine | 1-sqrt

[schedules.optim.betas2]
type = "cosine"
warmup = 0
end_ratio = 1.04

[schedules.clip_norm]       # multiplier on optim.max_norm
type = "constant"

Design

• A schedule is a pure function f(step) -> multiplier, step = optimizer updates completed so far. Update step uses base_g * f(step), with base_g each param group's own base value, so ratios between groups hold at every step (the shared CosineAnnealingLR decays every group to one absolute eta_min).
• Endpoints are what you configure: f(0) = start_ratio, the warmup reaches 1 at update warmup - 1 (warmup 0 or 1 starts at 1), the decay reaches end_ratio at the last update max_steps - 1. Both ratios default to 0.
• Types: wsd (default), cosine, constant, custom (import_path + kwargs, called as fn(step, total_steps=max_steps, **kwargs)). Plain functions bound with functools.partial.
• schedules.optim keys are optim field names. Each optimizer config declares its schedulable fields as data (schedulable ClassVar: field -> param group key, tuple index, algorithm filter): AdamW lr, weight_decay, betas1/betas2 (-> betas[0/1]); Muon lr, weight_decay, mu (Muon groups), betas1/betas2 (-> beta1/beta2, AdamW groups); SignSGD lr, weight_decay. There is no eps field on these configs on this branch, so eps is not schedulable (Dion's Muon-group epsilon is the Newton-Schulz epsilon anyway). A PretrainConfig validator rejects keys the chosen optimizer can't schedule, clip_norm without max_norm, decays/fractions without max_steps, and phases longer than max_steps.
• Runtime: Schedules is built right after the optimizer and captures each group's base values before any checkpoint load. schedules.apply(step) at the top of each step writes the values (by group index, since loading an optimizer state dict replaces the group dicts) and returns the clip norm plus metrics. Nothing is checkpointed: values are recomputed from the step on resume. ckpt.skip_scheduler restarts the schedules at the resumed step over the remaining steps.
• Logging: optim/<name> (first targeted group's value; optim/lr as before) and schedule/<name> multipliers, plus optim/max_norm / schedule/clip_norm.
• Ladder examples: [scheduler] cosine, warmup_steps = N, min_lr = X -> [schedules.optim.lr] cosine, warmup = N, end_ratio = X / lr.

Validation

• Unit tests on the research cluster (CPU SLURM job): uv run pytest tests/unit/train/pretrain tests/unit/test_configs.py -q -> 260 passed (5 new tests in tests/unit/train/pretrain/test_schedules.py; test_configs.py loads all TOMLs including the updated ladder examples).
• GPU smoke run, 1x H200: examples/pretrain/llama3/17m.toml with 60 steps, batch 32, lr wsd (warmup = 6, decay = 0.3, end_ratio = 0.1), betas2 cosine (warmup = 0, end_ratio = 1.04), clip_norm cosine (warmup = 0, end_ratio = 0.5), checkpoint every 30 steps. It finished and the logged values match the convention exactly: LR 0 at step 1, rising by 6e-4 per step to 3e-3 at step 6, flat through step 42, then linear down to 3e-4 at step 60. betas2 went from 0.95 to 0.988 and max_norm from 1.0 to 0.5 at the last step. Loss 11.82 -> 7.56.
• Resume: a second run with resume.dir set to the step-30 checkpoint logs the same scheduled values at steps 31-60. Step 31 matches bit for bit (loss 7.852322, grad norm 0.323856), and losses later agree to about 1e-4 (normal GPU nondeterminism).

Notes / open questions

• Default wsd is warmup 0.1, decay 0.9, linear, end_ratio = 0: linear decay to 0 over all updates after the warmup. On our ablations (CompleteP + Muon, 537M tokens, 65k-token batches), WSD with 10/20/40/60/80% decay lost +0.069/+0.041/+0.018/+0.007/+0.001 nats to it at each schedule's best LR, and cosine tied. This matches Bergsma et al. 2025, Schaipp et al. 2025 and Wen et al. 2025 (Muon). For cooldown branches, set a shorter decay (e.g. 0.2) with decay_shape = "1-sqrt" (−0.010 vs linear at 20%). Unit tests after the change: 260 passed (CPU job on the research cluster).
• The default WSD uses fractions, so a pretrain config without max_steps must now set a schedule without decay/fractions (e.g. constant). Previously the default scheduler was constant.
• Old ladder warmups started at min_lr (torch LinearLR start factor); the new ones start at start_ratio = 0. Set start_ratio = 0.1 to keep the old start.
• Setting schedules.optim in TOML replaces the default dict, so a config that schedules only weight_decay has a constant LR. Documented.
• CLI: dotted overrides of dict entries work (--schedules.optim.lr.end-ratio 0.2, adding --schedules.optim.weight-decay.type constant). Changing the type of an entry merges with the TOML fields, so switching a cosine entry to constant from the CLI fails on the leftover fields (same as other unions).
• CheckpointConfig.skip_scheduler keeps its shared name; for pretraining it means "restart the schedules at the resumed step".

[snimu]

Folded into #4062: merge commit 2f706b4 brings this branch into feat/pretrain, so GitHub shows this PR as merged. In #4062, 379971f then cuts the schedules down to one WSD learning-rate schedule ([schedule]: warmup, decay, end_ratio, decay_shape). The branch stays on the remote.