# 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= --paths.checkpoint_dir= # # 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"