# SFT stage (section-11). Dataclass-backed: every key here must exist in
# align/config.py, and struct mode makes a typo raise at launch rather than
# silently training with the default it failed to override.
#
# CLI contract:
#   torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml --optimizer.lr=1e-5
#
# Prose lives in the config field docstrings, not here. Only values that this
# STAGE chooses (as opposed to the repo-wide defaults) are set below.

run_name = "sft-stage-11"

[model]
precision = "bf16"
# null = auto-select, which resolves to sdpa on this box. triton_fa2 is
# opt-in only and its selector raises until section-06's gate has passed.
# attn_implementation = null  (TOML has no null: omitted, the dataclass default (None) applies)

[optimizer]
# adamw8bit per the 2026-08-17 pre-registration amendment: offload_adamw
# projects ~105GB host RSS (~150GB MemAvailable required) and the shared
# box holds ~85GB — the RSS preflight refuses it, correctly. 8-bit moments
# live on-GPU; the accum>=4 offload-amortization rule does not apply.
name = "adamw8bit"
lr = 1e-05
betas = [0.9, 0.95]
weight_decay = 0.0
# Global-norm clip, ON (G0; OLMoE / OLMo 2 use 1.0). The recorded grad norm
# is the PRE-clip value, which is what the grad_norm tripwire below needs.
max_grad_norm = 1.0
# 2 x micro_batch 3 x world 4 = the pre-registered 24-row global batch
# (train.global_batch_rows), keeping the scheduler horizon arithmetic
# self-consistent. The accum>=4 rule is OFFLOAD economics (2.56s CPU step
# amortization) and is enforced only for offload_adamw.
grad_accum = 2

[train]
full_ft = true
world_size = 4
micro_batch_size = 3
seq_len = 1024
epochs = 1
# ~24 rows per optimizer step; the horizon is derived from the manifest's
# measured corpus size, not from the plan's "~11.8K" prose.
global_batch_rows = 24
warmup_frac = 0.03
# Cadence in MICRO-steps, so changing accumulation does not change how much
# work a host seizure costs. ~400 micro-steps ~ 100 optimizer steps.
save_every_micro_steps = 400
keep_checkpoints = 2
log_every_micro_steps = 10
pinned_staging_shards_per_rank = 2
# Each rank writes its own optimizer_rank{r}.pt instead of gathering every
# rank's state to rank 0 (the 2026-08-12 OOM shape; the gather peak is what
# preflight_rss refuses at ~85GB MemAvailable). Validated on 2 CPU/gloo ranks
# over real FSDP2 shards, bit-exact resume, for adamw8bit and offload_adamw
# (tests/test_trainer_dist.py). Resume must use the SAME world_size.
per_rank_optimizer_state = true
probe_steps = 20
seed = 0

[tripwires]
# Armed (G0). grad_norm is the only tripwire an SFT stage can feed (no
# preference margin / pairs for logps_chosen / length_drift). It HALTS the run
# on a gradient explosion or a non-finite norm instead of training past it.
armed = ["grad_norm"]

[data]
num_workers = 2  # host cap; 16 across ranks SIGBUS'd a live run
sft_split = "sft-train"

[eval]
math500_n = 500
gsm8k_n = 1319  # FULL GSM8K test set, not a subsample

[gates]
# checkpoint-200 measured on this box at 0.35 epochs. A full epoch landing
# below these is near-certain evidence of a new-stack defect.
golden_math500 = 0.09
golden_gsm8k = 0.258
rss_headroom_gb = 45.0
epoch_budget_s = 14400.0  # ~4 GPU-hours budgeted for this stage
schedule_budget_factor = 1.25
# 2 = first corpus build after the doubled-`# Answer`-marker fix.
min_build_version = 2
require_golden_receipt = true