feat(train): flip Path 3 -> Path 1 — production 1.5B GRPO on a100-large

Previous job (69ece94a, a10g-largex2) loaded
unsloth/Qwen2.5-Coder-0.5B-Instruct-bnb-4bit because job_train.sh
hardcoded `model=qwen_0_5b_smoke` and never got switched to the
production 1.5B target documented in HACKATHON_FINAL_PLAN.md and
PROJECT.md ("A 1.5B model trained on our env solves 200K-token QA
tasks via recursion"). Submission was about to ship as 0.5B.

This commit switches both the SFT and GRPO blocks to:
model=qwen_1_5b -> unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit
(LoRA r16, max_seq=16384)
train.max_steps=200 -> enough room for the format-gate-crossing
transition + visible reward climb (50 was a
smoke; previous 50-step run with broken
sys_msg had reward=0 every step).

Memory budget retuned for the bigger model on the bigger GPU:
num_generations 4 -> 8 (back to TRL/STACK default)
max_prompt_length 2048 -> 4096 (full prompt budget)
max_completion_length 512 -> 1024 (long-form answers fit)
vllm_gpu_mem_util 0.35 -> 0.50 (A100-80GB has the headroom)

Hardware: scripts/hf_jobs_helper.py default flavor moves
a10g-largex2 -> a100-large
The CUDA_VISIBLE_DEVICES=0 guardrail stays as a defensive no-op for
single-GPU instances (still trips the multi-rank entropy bug if we
ever fall back to a 2x flavor).

Estimated run: SFT ~12 min + GRPO 200 steps ~75 min = ~90 min total,
~$5-7 on HF Jobs a100-large.

Made-with: Cursor

scripts/hf_jobs_helper.py

@@ -75,10 +75,15 @@ def submit() -> None:
         'git clone -b main https://oauth2:$HF_TOKEN@huggingface.co/Pratham-math/fathom-code /w && '
         'bash /w/scripts/job_train.sh',
     ]
+    # a100-large = 1x A100-80GB. Required for the 1.5B production run:
+    # 1.5B 4-bit weights + LoRA + grad/optim state + vLLM colocate KV cache
+    # for num_generations=8 @ max_completion=1024 needs ~30-40 GB.
+    # The previous a10g-largex2 (24 GB usable on 1 GPU) only held 0.5B safely.
+    flavor = os.environ.get("HF_JOB_FLAVOR", "a100-large")
     job = api.run_job(
         image="pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
         command=cmd,
-        flavor="a10g-largex2",
+        flavor=flavor,
         namespace=USER,
         secrets={
             "HF_TOKEN": os.environ["HF_TOKEN"],

scripts/job_train.sh

@@ -5,11 +5,12 @@ set -e
 cd /w
 export PYTHONPATH="/w:${PYTHONPATH}"
 
-# Force single-GPU training. The 0.5B (and 1.5B) 4-bit + LoRA model fits on
-# one A10G easily, and TRL's SFTTrainer entropy logging path on multi-GPU
-# (a10g-largex2) crashes with `RuntimeError: tensor a (2) must match tensor
-# b (4) at non-singleton dimension 0` because the per-rank logits get mixed
-# with the full-batch attention mask. Hide the second GPU completely.
+# Force single-GPU training. On a100-large (1x A100-80GB) this is a no-op,
+# but kept as a defensive guardrail in case we ever fall back to a 2x flavor:
+# TRL's SFTTrainer entropy logging path on multi-GPU crashes with
+# `RuntimeError: tensor a (2) must match tensor b (4) at non-singleton
+# dimension 0` because per-rank logits get mixed with the full-batch
+# attention mask. Hiding extra GPUs completely sidesteps this.
 export CUDA_VISIBLE_DEVICES=0
 
 # Install huggingface_hub first so we can download data files (parity with job_smoke.sh)
@@ -62,7 +63,8 @@ python -c "import urllib.request, sys; sys.exit(0 if urllib.request.urlopen('htt
 python -m train.smoke_test --env-url http://localhost:8001
 grep -q "VERDICT: GO" outputs/smoke/SMOKE_RESULT.md || (cat outputs/smoke/SMOKE_RESULT.md; exit 1)
 
-# SFT (Path 3 — 0.5B sanity check; switch to qwen_1_5b for Path 1)
+# SFT (Path 1 — production 1.5B). qwen_1_5b config points at
+# unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit per STACK §3.1.
 # NOTE: use initialize_config_dir (absolute path) — heredocs have no caller file,
 # so the relative `../configs` resolves against CWD and breaks. Absolute is robust.
 python <<'PY'
@@ -70,17 +72,19 @@ from hydra import initialize_config_dir, compose
 from train.model_load import load_model_and_tokenizer
 from train.sft import run_sft
 with initialize_config_dir(config_dir="/w/configs", version_base="1.3"):
-    cfg = compose(config_name="config", overrides=["model=qwen_0_5b_smoke","train=sft"])
+    cfg = compose(config_name="config", overrides=["model=qwen_1_5b","train=sft"])
 m, t = load_model_and_tokenizer(cfg)
 print("SFT adapter:", run_sft(cfg, m, t))
 PY
 
-# GRPO 50 steps (Path 3 — 0.5B sanity check; switch to qwen_1_5b + max_steps=400 for Path 1)
-# Memory overrides for single-A10G + vLLM colocate at 0.5B:
-#  - num_generations=4 (was 8)            → halves rollout VRAM
-#  - max_prompt_length=2048 (was 4096)    → halves KV cache per prompt
-#  - max_completion_length=512 (was 2048) → quarters generation memory
-#  - vllm_gpu_memory_utilization=0.35     → leaves ~16GB headroom for trainer
+# GRPO 200 steps (Path 1 — production 1.5B on A100-80GB).
+# Memory budget on a100-large (~80GB):
+#  - 1.5B 4-bit + LoRA-r16 weights + optim state ≈ 6-8 GB
+#  - vLLM colocate KV cache @ num_generations=8, max_prompt=4096, max_completion=1024 ≈ 25-35 GB
+#  - vllm_gpu_memory_utilization=0.50 → leaves ~30 GB for trainer + activations
+# Ramping max_steps from 50 → 200 to give the reward curve room to climb
+# (50 steps was a smoke; 200 is enough to show the gate-crossing transition
+# from format_gate=0 → format_gate=1 → composite reward rising).
 python <<'PY'
 from hydra import initialize_config_dir, compose
 from train.model_load import load_model_and_tokenizer
@@ -88,13 +92,13 @@ from train.grpo import run_grpo
 from rewards.compose import make_reward_fn
 with initialize_config_dir(config_dir="/w/configs", version_base="1.3"):
     cfg = compose(config_name="config", overrides=[
-        "model=qwen_0_5b_smoke",
+        "model=qwen_1_5b",
         "train=grpo",
-        "train.max_steps=50",
-        "train.num_generations=4",
-        "train.max_prompt_length=2048",
-        "train.max_completion_length=512",
-        "train.vllm_gpu_memory_utilization=0.35",
+        "train.max_steps=200",
+        "train.num_generations=8",
+        "train.max_prompt_length=4096",
+        "train.max_completion_length=1024",
+        "train.vllm_gpu_memory_utilization=0.50",
     ])
 m, t = load_model_and_tokenizer(cfg)
 print("GRPO merged:", run_grpo(cfg, m, t, make_reward_fn(cfg.reward), "http://localhost:8001"))