| # 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 | |