# Pretraining stage (section-14). Dataclass-backed: every key here must exist
# in align/config.py (struct mode raises on unknown keys), so this stage rides
# the EXISTING config surface -- reuses are documented inline as TODOs.
#
# CLI contract:
#   torchrun --nproc_per_node=4 scripts/pretrain.py configs/pretrain.toml --optimizer.lr=3e-5
#
# HONEST TOY SCALE: continued pretraining of the base model on a ~50M-token
# streamed slice of open-web-math (see align/data/pretrain_corpus.py -- the cap
# is enforced there, never assumed here). This demonstrates the pretraining
# pillar on 4x A10G; it is not an OLMo-scale budget and does not pretend to be.

run_name = "pretrain-stage-14"

[model]
precision = "bf16"
# null = auto-select, which resolves to sdpa on this box.
# attn_implementation = null  (TOML has no null: omitted, the dataclass default (None) applies)

[optimizer]
# adamw8bit, same reasoning as the SFT stage: offload_adamw projects ~105GB
# host RSS and the shared box holds ~85GB MemAvailable -- the RSS preflight
# refuses it, correctly. 8-bit moments live on-GPU.
name = "adamw8bit"
# Continued-pretraining LR: above the SFT 1e-5 (raw text, one pass, cosine
# to near zero) but far below a from-scratch 3e-4 -- this starts from
# converged v0.3 weights and a from-scratch LR would destroy them.
lr = 3e-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
grad_accum = 2

[train]
full_ft = true
world_size = 4
micro_batch_size = 3
# Matches the corpus builder's packed block length; every row is EXACTLY
# this long (the tail is dropped, not padded), so batches carry no padding.
seq_len = 1024
epochs = 1
# 2 x 3 x 4 = 24: assert_batch_arithmetic_consistent refuses a mismatch.
global_batch_rows = 24
warmup_frac = 0.03
# Cadence in MICRO-steps; the host can seize at any moment (2026-08-12).
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 pretraining 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]
# Pretrain corpus artifacts live beside (not inside) the SFT build so
# neither manifest can shadow the other.
shard_dir = "/apps/development/rlhf/nano-mistral-o1/data/pretrain/shards"
manifest_path = "/apps/development/rlhf/nano-mistral-o1/data/pretrain/manifest.json"
# REUSE, not a typo: DataConfig has no pretrain_split field and config.py
# cannot be edited under the live run. sft_split is "the split this stage
# trains on" everywhere it is read (stage_smoke.open_split).
# TODO(post-run): add DataConfig.pretrain_split and rename.
sft_split = "pretrain"
num_workers = 2  # host cap; 16 across ranks SIGBUS'd a live run
pad_token_id = 2

[gates]
# NO golden anchor for a base-style checkpoint: the golden regression
# certifies the SFT eval stack against known MATH-500/GSM8K scores, which do
# not exist for raw-text continued pretraining. require_golden_receipt
# false makes run_entry_gates log a loud warning instead of refusing.
require_golden_receipt = false
# Pretrain manifests version independently of the SFT corpus (the
# doubled-answer-marker history is meaningless for raw web text);
# pretrain_corpus.PRETRAIN_BUILD_VERSION starts at 1.
min_build_version = 1
rss_headroom_gb = 45.0
# ~50M tokens / (24 rows x 1024 tokens) ~ 2K optimizer steps; budget ~6
# GPU-hours wall-clock with the same 1.25x probe-projection slack as SFT.
epoch_budget_s = 21600.0
schedule_budget_factor = 1.25