| # P2 SFT: LoRA fine-tune of Goedel-Prover-V2-8B (Qwen3ForCausalLM) on verifier-accepted proofs. | |
| # | |
| # torchrun --standalone --nproc_per_node=4 scripts/sft.py configs/p2_sft.toml \ | |
| # --data.shard_dir=<P2-B shard root> --paths.checkpoint_dir=<run root> | |
| # | |
| # Struct mode: every key must exist in align/config.py. Prose lives in the config docstrings; this | |
| # file states the values this STAGE chooses and the MEASUREMENT behind each. Every [measured] number | |
| # was produced on 4x A10G (23,028 MiB) with scripts/qwen3_sft_probe.py through the SAME | |
| # scripts/sft.py functions this file drives, on real Lean-statement tokens, 10 optimizer steps each: | |
| # | |
| # micro-batch tokens | alloc peak | reserved peak | s / micro-step | tokens/s (4 GPUs) | |
| # 4,096 | 11.72 GiB | (see note) | 11.75 | 1,394 (comm-bound) | |
| # 8,192 | 13.04 GiB | (see note) | 16.86 | 1,943 | |
| # 16,384 | 15.84 GiB | 20.05 GiB | 30.45 | 2,152 | |
| # 18,432 | 16.91 GiB | 20.78 GiB | 34.36 | 2,146 (compute-bound) | |
| # | |
| # note: the 4,096 / 8,192 runs predate the probe's empty_cache() after the model build, so their | |
| # reserved column reads the load phase (18.88 GiB), not the step; allocated peak is the valid column. | |
| # | |
| # A "micro-batch" is a TOKEN BUDGET (train.seq_len x micro_batch_size padded tokens, packed by the | |
| # length-grouped sampler), not a row count: short proofs are packed several to a micro-step, so the | |
| # per-step cost is flat in the data's length mix. Throughput saturates at ~2.15k tokens/s (537 per | |
| # GPU) once the 16.4 GB FSDP all-gathers (4 GB/s over host-shared-memory NCCL on this PCIe box, | |
| # ~8.5 s per micro-step, un-overlappable at short length) are amortised, i.e. from ~16k tokens. | |
| run_name = "p2-goedel8b-lora-sft" | |
| [model] | |
| # arch routes scripts/sft.py to align.train.qwen3_stage; the architecture itself is read from the | |
| # snapshot's config.json. The numbers below must AGREE with that snapshot (the stage refuses a | |
| # mismatch): n_params is count_params() of the released config, not the Mistral default. | |
| arch = "qwen3" | |
| name = "goedel-prover-v2-8b" | |
| n_params = 8_190_735_360 | |
| n_layers = 36 | |
| precision = "bf16" | |
| vocab_size = 151936 | |
| hidden_size = 4096 | |
| intermediate_size = 12288 | |
| n_heads = 32 | |
| n_kv_heads = 8 | |
| head_dim = 128 | |
| rms_norm_eps = 1e-6 | |
| rope_theta = 1e6 | |
| max_position_embeddings = 40960 | |
| # sdpa = torch's flash kernel (enable_gqa) on sm86. triton_fa2 has no QK-norm seam and is refused. | |
| attn_implementation = "sdpa" | |
| [optimizer] | |
| # adamw8bit: the in-repo torch-only block-wise 8-bit AdamW. The adapters are fp32 (175M params), | |
| # so there is no bf16 write-back rounding; the 8-bit moments are ~0.35 GB total, sharded 4 ways. | |
| name = "adamw8bit" | |
| # LR 2e-5 (the allowed band is 1e-5..5e-5). Why this and not the LoRA folk value 1e-4: | |
| # * the base is ALREADY a post-trained (SFT + RL) prover; this stage moves it a little toward its own | |
| # verifier-accepted outputs, it does not teach it Lean. Large steps on a few thousand self-generated | |
| # rows is how a policy loses the capability the verifier loop depends on; | |
| # * rank 64 / alpha 128 scales the adapter by 2, and Adam moves each B entry ~lr per step, so the | |
| # relative weight change grows ~linearly in the summed LR. MEASURED (scripts/qwen3_export_check.py, | |
| # 252 adapted matrices, this config after 12 optimizer steps, summed LR ~1.2e-4): | |
| # ||W' - W||_F / ||W||_F mean 3.1e-4, max 1.0e-3. The planned 100-150 steps at mean LR ~1e-5 | |
| # (cosine) sum to ~1-1.5e-3, i.e. ~10x: ~0.3% mean / ~1% max of the weight norm -- the scale of | |
| # a gentle full-FT at 1e-5, which is what a post-trained prover should take. 5e-5 would be 2.5x | |
| # that; 1e-5 would leave a 2-epoch run on a few thousand rows barely above bf16 rounding (the | |
| # merged checkpoint is rounded to bf16 once, ~0.2% per element); | |
| # * measured grad norms are 0.09-0.14 (pre-clip) at every length, so max_grad_norm = 1.0 never | |
| # binds: the update size is set by the LR alone, which is what we want to be able to reason about. | |
| lr = 2e-5 | |
| betas = [0.9, 0.95] | |
| weight_decay = 0.0 | |
| max_grad_norm = 1.0 | |
| # 2 micro-batches per rank per optimizer step x 4 ranks = 8 micro-batches of up to 18,432 padded tokens | |
| # (~100-150k real tokens) per step: with the few-thousand-row datasets this experiment produces that | |
| # is ~50-150 steps per epoch, enough for a warmup + cosine to mean something. | |
| grad_accum = 2 | |
| [lora] | |
| enabled = true | |
| r = 64 | |
| alpha = 128.0 | |
| dropout = 0.0 | |
| adapter = "sft" | |
| target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] | |
| [train] | |
| full_ft = false | |
| world_size = 4 | |
| micro_batch_size = 1 | |
| seq_len = 18432 | |
| # 2 epochs: rejection-sampling SFT on a small verified set; the adapter is low-capacity (175M params, | |
| # 2.1% of the model) and the LR decays to 0, so a second pass sharpens rather than memorises. The | |
| # held-out Lean pass-rate (P2-C) is the judge; a single epoch is `--train.epochs=1`. | |
| epochs = 2 | |
| # accum 2 x micro 1 x world 4 = 8 (the stage refuses a mismatch). "Rows" here are micro-batches. | |
| global_batch_rows = 8 | |
| warmup_frac = 0.1 | |
| # Micro-steps (a multiple of grad_accum, or the Trainer refuses); 40 micro = 20 optimizer steps = | |
| # ~23 min at 34 s/micro. A save costs ~3.3 min [measured: 19 GiB = merged 16.4 GiB bf16 + four | |
| # 0.7 GiB adapter/optimizer files; rank-0 FSDP gather + write], i.e. ~13% overhead, and bounds the | |
| # work a host seizure can destroy to ~23 minutes. keep_checkpoints = 2 => 38 GiB on disk. | |
| save_every_micro_steps = 40 | |
| keep_checkpoints = 2 | |
| log_every_micro_steps = 1 | |
| pinned_staging_shards_per_rank = 2 | |
| # U03-2: every target token weighs the same across the accumulation window and the 4 ranks (summed CE / | |
| # window target-token count all-reduced), instead of the mean of 8 micro-batch means. Token-budget batches | |
| # differ in target count (280 vs 6 seen), so under "micro_mean" a short batch's tokens weigh ~47x more. | |
| # NEW P2 runs only: runs recorded under the default are not comparable to this objective. | |
| loss_normalization = "global_token_mean" | |
| per_rank_optimizer_state = true | |
| probe_steps = 20 | |
| seed = 0 | |
| [tripwires] | |
| armed = ["grad_norm"] | |
| [data] | |
| num_workers = 2 # host cap (assert_worker_count); shards are memory-mapped and the loop is GPU-bound | |
| sft_split = "sft-train" | |
| pad_token_id = 151643 # <|endoftext|>; right-padding only, masked out of the loss (labels -100) | |
| [paths] | |
| model_snapshot = "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots" | |
| checkpoint_dir = "/apps/development/rlhf/align/.tmp/p2/sft" | |