Gemma 4 12B SFT-uplift PoC (+ MTP) on NVIDIA RTX 4000 Ada

.env.example

@@ -35,11 +35,10 @@ MAX_NEW_TOKENS=512
 #   x-api-key: <key>
 TEACHER_SERVER_API_KEY=
 
-# Weights & Biases — experiment tracking for the state-classifier trainers
-# (scripts/train_state_classifier_34way*.py, run with --wandb). Grab the key
-# from https://wandb.ai/authorize. The training launch sources .env, so the key
-# is picked up automatically for online logging; leave blank to log offline.
+# Weights & Biases — experiment tracking for SFT trainers (train_sft.py) and
+# state-classifier trainers (train_state_classifier_34way*.py, run with --wandb).
+# Get your key at https://wandb.ai/authorize, then run `wandb login` or set here.
 WANDB_API_KEY=
-# Optional overrides (defaults: project=csen346-state-classifier, run=out-dir basename)
-# WANDB_PROJECT=csen346-state-classifier
-# WANDB_RUN_NAME=
+# Optional overrides (training scripts set sensible defaults if omitted)
+# WANDB_PROJECT=csen346-sft
+# WANDB_RUN_NAME=gemma4-31b-stage2

.gitignore

@@ -69,6 +69,8 @@ references/KELE-EN/.cache/
 # Railguard
 .railguard/*
 
+# W&B local run artifacts
+wandb/
 # macOS
 .DS_Store
 **/.DS_Store

configs/gemma4-12b-local.env

@@ -0,0 +1,30 @@
+# Experiment: Gemma 4 12B-it BASE teacher — fully local, dual-role on llama.cpp.
+#
+# Both teacher and consultant hit the same llama.cpp server on port 8080, mirroring
+# configs/gemma4-31b-local.env. For the SFT-uplift PoC the eval RE-points the
+# consultant role at the Qwen3.5-LoRA state classifier via `--bert-consultant
+# <path>` — but CONSULTANT_* below must still be set, because config.py:load_config
+# require()s them unconditionally at load time.
+#
+# Serve the base weights with:   make serve-gemma4-12b        (UD-Q8_K_XL)
+#   or, to A/B the MTP speculative drafter:  make serve-gemma4-12b-mtp
+# Eval:  make eval-gemma4-12b-base-smoke  /  make eval-gemma4-12b-base-full
+#
+# Shares port 8080 with the other local llama.cpp configs — only one model fits at
+# a time on the RTX 4000 Ada's 20 GB. The teacher alias "Gemma 4 12B" must match
+# the serve wrapper's -a so the eval can't run against the wrong loaded weights.
+
+CONSULTANT_API_KEY=not-needed
+CONSULTANT_BASE_URL=http://localhost:8080/v1
+CONSULTANT_MODEL_NAME=Gemma 4 12B
+# 8192 leaves ample room for structured JSON output. Gemma 4 has no thinking-mode
+# equivalent (Qwen-specific) so no CoT budget needed.
+CONSULTANT_MAX_TOKENS=8192
+
+TEACHER_API_KEY=not-needed
+TEACHER_BASE_URL=http://localhost:8080/v1
+TEACHER_MODEL_NAME=Gemma 4 12B
+
+DEBUG_MODE=false
+MAX_TEACHING_ROUNDS=8
+KELE_TEACHER_LANG=auto

configs/gemma4-12b-sft-local.env

@@ -0,0 +1,26 @@
+# Experiment: Gemma 4 12B-it SFT teacher (Socratic QLoRA, merged → Q6/Q8 GGUF).
+#
+# Same dual-role / fully-local pattern as gemma4-12b-local.env. The ONLY
+# difference is the alias "Gemma 4 12B SFT" — distinct so the eval orchestrator
+# refuses to run against the BASE weights by mistake (it requests this teacher
+# name from the server's /v1/models and fails fast if absent). The Qwen3.5-LoRA
+# consultant is held identical to the baseline run so the metric delta isolates
+# the SFT adapter.
+#
+# Serve the merged SFT weights with:   make serve-gemma4-12b-sft
+# Eval:  make eval-gemma4-12b-sft-smoke  /  make eval-gemma4-12b-sft-full
+#
+# Shares port 8080 — stop the base server before serving the SFT model.
+
+CONSULTANT_API_KEY=not-needed
+CONSULTANT_BASE_URL=http://localhost:8080/v1
+CONSULTANT_MODEL_NAME=Gemma 4 12B SFT
+CONSULTANT_MAX_TOKENS=8192
+
+TEACHER_API_KEY=not-needed
+TEACHER_BASE_URL=http://localhost:8080/v1
+TEACHER_MODEL_NAME=Gemma 4 12B SFT
+
+DEBUG_MODE=false
+MAX_TEACHING_ROUNDS=8
+KELE_TEACHER_LANG=auto

configs/train-sft-gemma4-12b-qlora.env

@@ -0,0 +1,96 @@
+# SFT fine-tuning — Socratic teaching QLoRA on Gemma 4 12B-IT (PoC, 1 epoch)
+#
+# Target hardware: NVIDIA RTX 4000 Ada (20 GB), CUDA. This box is KNOWN-UNSTABLE
+# (see memory: training-host-hardware-fault) — checkpoint frequently and rely on
+# train_sft.py auto-resume rather than expecting a clean single run.
+# Gemma 4 uses standard softmax attention (attn_implementation="sdpa" in
+# train_sft.py) — no flash-linear-attention dependency, no patch-fla step, and
+# none of the gfx1201/ROCm env (TORCH_USE_HIPBLASLT / PYTORCH_HIP_ALLOC_CONF)
+# the 31B AMD config carries.
+#
+# This is the 12B PoC sibling of configs/train-sft-stage2-gemma4-31b.env. Goal:
+# measure whether 1-epoch Socratic SFT gives eval uplift over the BASE Gemma 4
+# 12B teacher (consultant = Qwen3.5-LoRA state classifier, held fixed).
+#
+# Launch (recommended — train from the community bnb-4bit base, no BF16 staging):
+#   make train-gemma4-12b
+# which overrides TRAIN_BASE_MODEL=unsloth/gemma-4-12b-it and TRAIN_PREQ=true.
+#
+# Dry run (no weights):  make train-gemma4-12b-dry-run
+
+# ── Base model ────────────────────────────────────────────────────────────────
+# Default is the BF16 hub model (live-quant path). The make target overrides this
+# to the unsloth bnb-4bit checkpoint with TRAIN_PREQ=true to skip BF16 CPU staging.
+TRAIN_BASE_MODEL=google/gemma-4-12b-it
+
+# ── Training method ───────────────────────────────────────────────────────────
+TRAIN_METHOD=qlora
+
+# ── LoRA hyperparameters ──────────────────────────────────────────────────────
+LORA_RANK=16
+LORA_ALPHA=32
+LORA_DROPOUT=0.05
+# google/gemma-4-12b-it is the same multimodal Gemma4ForConditionalGeneration
+# family as the 31B (text decoder + vision_tower + audio_tower). The anchored
+# regex (build_lora_config detects the leading ^) scopes LoRA to text-decoder
+# nn.Linear paths only; the vision_tower's Gemma4ClippableLinear wrappers crash
+# get_peft_model otherwise. It is layer-count-agnostic (\d+), so it ports from
+# the 31B's 60 decoder layers to the 12B's 48 unchanged.
+LORA_TARGET_MODULES=^model\.language_model\.layers\.\d+\.(self_attn|mlp)\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$
+
+# ── Training hyperparameters ─────────────────────────────────────────────────
+# 1 epoch over the ~77k per-turn records (~4826 steps @ effective-batch 16) is
+# ample SFT signal for a PoC. Bump to 3 only if the downstream eval shows the
+# single epoch underfits.
+# NB: train_sft.py auto-resumes from any checkpoint-* in TRAIN_OUTPUT_DIR built
+# for the OLD schedule. If you re-run after changing TRAIN_EPOCHS, clear them:
+#   rm -rf outputs/sft-gemma4-12b-qlora/checkpoint-*
+TRAIN_EPOCHS=1
+TRAIN_LR=5e-5
+TRAIN_BATCH_SIZE=1
+TRAIN_GRAD_ACCUM=16
+# Measured on socrat-zh-sft + socrat-en-sft (scripts/determine_max_sequence.py):
+# real token max = 909 (train) / 892 (eval), p99 = 673. 1280 truncates NOTHING;
+# with batch_size=1 + dynamic padding the activation peak is set by the actual
+# longest sequence (909), not by this cap, so lowering it frees no VRAM.
+TRAIN_MAX_SEQ_LEN=1280
+TRAIN_BF16=true
+TRAIN_GRAD_CKPT=true
+# One-time tokenization fan-out across CPU cores (procs exit before training).
+TRAIN_DATASET_NUM_PROC=12
+# Step is GPU/dequant-bound; data is pre-tokenized; async prefetch buys nothing.
+TRAIN_DATALOADER_WORKERS=0
+
+# ── Dataset sources ───────────────────────────────────────────────────────────
+# Per-turn, state/action-conditioned format that matches the inference schema
+# socrates_teacher() sends (历史对话记录 / 当前学生输入). Using socrat-*-sft
+# (NOT plain socrat-zh,socrat-en) avoids confounding the SFT uplift measurement
+# with train/serve schema drift.
+TRAIN_SOURCES=socrat-zh-sft,socrat-en-sft
+
+# ── Output & logging ─────────────────────────────────────────────────────────
+TRAIN_OUTPUT_DIR=outputs/sft-gemma4-12b-qlora
+TRAIN_LOGGING_STEPS=10
+# Save every 50 steps so the auto-resume monitor can crawl the run forward across
+# a crash on this known-unstable box instead of restarting at step 0. NOTE: this
+# is NOT the AMD gfx1201 backward page fault (that is ROCm-specific and does not
+# occur on CUDA) — the rationale here is purely the host hardware instability.
+# Disk is bounded by save_total_limit (a couple of ~50 MB LoRA adapters). Drop to
+# 10 only if a crash ever lands inside the 0–50 window and stalls.
+TRAIN_SAVE_STEPS=50
+TRAIN_EVAL_STEPS=200
+# Disable in-training eval — the eval forward's logits.float() cast doubles the
+# 256K-vocab logits tensor (transient multi-GB), which can OOM the 20 GB card.
+# The downstream KELE eval pipeline produces the paper numbers anyway. Also
+# forces load_best_model_at_end=False (which needs eval).
+TRAIN_EVAL_STRATEGY=no
+
+# ── Weights & Biases ──────────────────────────────────────────────────────────
+# Requires `wandb login` or WANDB_API_KEY; train_sft._check_wandb disables
+# reporting if auth fails. Pinned run id + resume=allow so every crash-resume
+# APPENDS to one run with a continuous step axis. Bump WANDB_RUN_ID for a
+# brand-new crawl (else it resumes this run's step history).
+WANDB_PROJECT=csen346-sft
+WANDB_RUN_ID=gemma4-12b-qlora-poc
+WANDB_RESUME=allow
+WANDB_RUN_NAME=gemma4-12b-qlora-poc

configs/train-sft-stage2-gemma4-31b.env

@@ -18,7 +18,14 @@
 #   # 2a → 2b sequence: run 2a first, then point TRAIN_BASE_MODEL at the 2a
 #   # checkpoint dir and re-run with the default TRAIN_SOURCES.
 #
-# Launch (R9700):
+# Launch (R9700) — recommended: skip the 62 GB BF16 CPU staging at load by
+# training from the community pre-quantized unsloth bnb-4bit checkpoint:
+#   make train-gemma4-31b-stage2-unsloth
+# which overrides TRAIN_BASE_MODEL=unsloth/gemma-4-31B-it-unsloth-bnb-4bit and
+# TRAIN_PREQ=true. The default TRAIN_BASE_MODEL below is the live-quant path
+# (make train-gemma4-31b-stage2), which stages the full BF16 model in CPU RAM.
+#
+# Launch (live quant, BF16 base):
 #   mkdir -p outputs/sft-stage2-gemma4-31b
 #   nohup env TORCH_USE_HIPBLASLT=0 PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8 \
 #     uv run --no-sync python scripts/train_sft.py \
@@ -48,31 +55,79 @@ LORA_DROPOUT=0.05
 LORA_TARGET_MODULES=^model\.language_model\.layers\.\d+\.(self_attn|mlp)\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$
 
 # ── Training hyperparameters ─────────────────────────────────────────────────
-TRAIN_EPOCHS=3
+# Stage 2b is a proof-of-concept for the larger SFT plan; 1 epoch over the ~77k
+# per-turn records is ample SFT signal and is a clean 3× wall-time cut
+# (~281h → ~94h) against the gfx1201 ~70s/step ceiling. Bump back to 3 only if
+# the downstream eval shows the single epoch underfits.
+# NB: train_sft.py auto-resumes from any checkpoint-* in TRAIN_OUTPUT_DIR. An
+# epochs=3 run already wrote checkpoints there, built for a ~14.5k-step schedule.
+# Clear them before relaunching at epochs=1, or the trainer resumes the old
+# 3-epoch schedule (wrong LR decay, or exits immediately if past 1-epoch steps):
+#   rm -rf outputs/sft-stage2-gemma4-31b/checkpoint-*
+TRAIN_EPOCHS=1
 TRAIN_LR=5e-5
 TRAIN_BATCH_SIZE=1
 TRAIN_GRAD_ACCUM=16
+# Measured, not assumed (scripts/determine_max_sequence.py, 2026-06-01 on
+# socrat-zh-sft + socrat-en-sft): real token max = 909 (train) / 892 (eval),
+# p99 = 673. Nothing exceeds 1024 → 1280 truncates NOTHING; it is non-binding.
+# NB: with batch_size=1 + dynamic padding the activation peak is set by the
+# actual longest sequence (909), not by this cap, so lowering it frees no VRAM.
 TRAIN_MAX_SEQ_LEN=1280
 TRAIN_BF16=true
 TRAIN_GRAD_CKPT=true
+# CPU parallelism (12C/24T box). dataset_num_proc fans the one-time tokenization
+# of the 77k-record split across processes (single-process by default → pegs one
+# core for minutes at startup); 12 ≈ physical cores. Those procs exit before
+# training.
+TRAIN_DATASET_NUM_PROC=12
+# Dataloader workers are 0 on purpose: the step is GPU/dequant-bound (~70s), data
+# is pre-tokenized, and collation is trivial — async prefetch buys nothing here.
+# Each worker also forks a HIP context that shows up in nvtop and costs a little
+# VRAM on the tight 32GB/31B budget. Raise only if profiling shows data-feed stall.
+TRAIN_DATALOADER_WORKERS=0
 
 # ── Dataset sources ───────────────────────────────────────────────────────────
-# Default: Stage 2b — structural fine-tuning on state-annotated SocratDataset.
-# socrat-synthetic is the contamination-control eval set; never include it here.
-TRAIN_SOURCES=socrat-zh,socrat-en
+# Stage 2b — per-turn inference-matching format (fixes train/serve schema drift).
+# socrat-zh-sft / socrat-en-sft produce one (system, user, assistant) triple per
+# dialogue turn with the exact 历史对话记录 / 当前学生输入 structure that
+# socrates_teacher() sends at inference. socrat-synthetic is eval-only; never here.
+TRAIN_SOURCES=socrat-zh-sft,socrat-en-sft
 
 # ── Output & logging ─────────────────────────────────────────────────────────
 TRAIN_OUTPUT_DIR=outputs/sft-stage2-gemma4-31b
 TRAIN_LOGGING_STEPS=10
-# Save every 100 steps (~32 min at 19s/step on Gemma 4 31B QLoRA). Tighter
-# than the 200-step default — halves the worst-case rollback after a crash
-# at the cost of ~2× I/O for adapter checkpoints (~500 MB each, ~11 GB total
-# across the 23 checkpoints expected for the full 2298-step run).
-TRAIN_SAVE_STEPS=100
+# Save every 10 steps. The gfx1201 GPU page fault (Memory access fault / page
+# not present) is non-deterministic and hits during the QLoRA backward as early
+# as step 14 — BEFORE any 100-step checkpoint lands. With save_steps=100 the
+# auto-resume monitor had nothing to resume from, so every restart began at
+# step 0 and re-crashed before 100: a permanent loop with zero forward progress.
+# save_steps=10 lands a checkpoint before the earliest observed crash, so the
+# resume monitor crawls the run forward across faults instead of looping. Disk
+# is bounded by save_total_limit=2 (~1 GB of LoRA adapters). Drop to 5 if a
+# crash ever lands inside the 10–20 window and stalls.
+TRAIN_SAVE_STEPS=10
 TRAIN_EVAL_STEPS=200
 # Disable in-training eval — the eval forward's logits.float() cast doubles
 # the logits tensor (256K vocab × 1280 seq × 4 bytes = 8.6 GB transient on
 # Gemma 4 31B), which OOMs the 5090 at the first eval step. The downstream
 # eval pipeline (eval_bert_gemma_fewshot10_full.sh) produces paper numbers
 # anyway. Also forces load_best_model_at_end=False (which needs eval).
 TRAIN_EVAL_STRATEGY=no
+
+# ── Weights & Biases ──────────────────────────────────────────────────────────
+# Enabled for the gfx1201 crawl. This lives in the CONFIG, not monitor_stage2.sh,
+# on purpose: the auto-resume monitor is a long-running process with start_stage2()
+# already parsed into memory, so editing the monitor script does NOT change its
+# next relaunch. train_sft.py re-reads this config on every `make` relaunch, so a
+# change here is picked up on the NEXT CRASH without restarting the monitor.
+# (load_env_file setdefaults these into os.environ; a CLI/env override still wins.)
+# Tracking still requires W&B auth — `wandb login` or WANDB_API_KEY. If login fails,
+# train_sft._check_wandb disables reporting regardless of these vars.
+WANDB_PROJECT=csen346-sft
+# Pin a stable run id + resume=allow so every crawl resume APPENDS to one run with a
+# continuous step axis, instead of spawning a fresh W&B run per crash cycle. Bump
+# this id when starting a brand-new crawl (else it resumes this run's step history).
+WANDB_RUN_ID=stage2-gemma4-31b-crawl
+WANDB_RESUME=allow
+WANDB_RUN_NAME=gemma4-31b-stage2-crawl