"""End-to-end training pipeline: SFT warm-up -> GRPO. Designed to run top-to-bottom in Colab or on HF compute. Matches the organizer guidance: verifier-based rewards, multiple independent reward components, SFT warm-start before RL, and per-component logging. Usage (Colab / HF Job): !pip install -q unsloth trl datasets wandb !python train.py Toggle QUICK_MODE=True for a 5-min pipeline validation on Colab free tier. Set to False on onsite HF compute for the real run. """ from __future__ import annotations import os import random import sys import warnings from typing import Optional import numpy as np import torch # Silence transformers generation-config spam during offline eval warnings.filterwarnings("ignore", message="Both `max_new_tokens`.*") warnings.filterwarnings("ignore", category=FutureWarning, module="transformers") os.environ["TRANSFORMERS_VERBOSITY"] = "error" # Module-level globals; real values set in main() after GPU detection. USE_BF16: bool = False USE_FP16: bool = False from drift_env.dataset import build_dataset, dataset_stats from drift_env.prompts import SYSTEM_PROMPT from drift_env.training.rewards import ( reward_compliance, reward_appropriateness, reward_drift_bonus, ) # --------------------------------------------------------------------------- # Config (edit here or via env vars) # --------------------------------------------------------------------------- QUICK_MODE = os.getenv("QUICK_MODE", "true").lower() == "true" # Model / hardware MODEL_NAME = os.getenv("MODEL_NAME", "unsloth/Qwen2.5-0.5B-Instruct" if QUICK_MODE else "unsloth/Qwen2.5-3B-Instruct") MAX_SEQ_LEN = int(os.getenv("MAX_SEQ_LEN", "4096")) LOAD_IN_4BIT = os.getenv("LOAD_IN_4BIT", "true").lower() == "true" # Data N_EPISODES_TRAIN = 50 if QUICK_MODE else 800 N_EPISODES_EVAL = 5 if QUICK_MODE else 40 SEED = 42 # SFT SFT_EPOCHS = 1 SFT_BATCH = 2 if QUICK_MODE else 4 SFT_LR = 2e-4 # GRPO GRPO_NUM_GEN = 4 if QUICK_MODE else 8 # K completions per prompt GRPO_MAX_STEPS = 50 if QUICK_MODE else 600 GRPO_BATCH = 1 if QUICK_MODE else 2 GRPO_GRAD_ACCUM = 4 GRPO_LR = 5e-6 GRPO_MAX_COMPLETION = 128 # LoRA LORA_R = 16 LORA_ALPHA = 16 LORA_DROPOUT = 0.0 # Output OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs") WANDB_PROJECT = os.getenv("WANDB_PROJECT", "drift-env") USE_WANDB = os.getenv("USE_WANDB", "false").lower() == "true" def seed_all(s: int) -> None: random.seed(s); np.random.seed(s); torch.manual_seed(s) if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) def _push_outputs_to_hub(label: str = "training outputs") -> None: """Push the entire OUTPUT_DIR to HF Hub. Safe to call multiple times. Survives ephemeral HF Jobs containers — without this, work is lost when the container terminates.""" hub_repo = os.getenv("HUB_REPO_ID") if not hub_repo: return try: from huggingface_hub import HfApi, create_repo create_repo(hub_repo, repo_type="model", exist_ok=True, private=True, token=os.getenv("HF_TOKEN")) HfApi().upload_folder( folder_path=OUTPUT_DIR, repo_id=hub_repo, repo_type="model", token=os.getenv("HF_TOKEN"), commit_message=label, ) print(f"[hub] {label} pushed -> https://huggingface.co/{hub_repo}") except Exception as e: print(f"[hub] WARN: failed to push '{label}': {e}") def _precision_flags() -> tuple[bool, bool]: """Return (use_bf16, use_fp16). T4 (Turing, CC 7.5) doesn't support bf16 — only Ampere (CC 8.0) and later do. Fall back to fp16 on T4.""" if not torch.cuda.is_available(): return False, False major, _ = torch.cuda.get_device_capability() if major >= 8: return True, False # Ampere+ — use bf16 return False, True # T4 etc — use fp16 # --------------------------------------------------------------------------- # 1. Build train/eval datasets from our environment # --------------------------------------------------------------------------- def build_hf_datasets(): from datasets import Dataset train_rows = build_dataset(n_episodes=N_EPISODES_TRAIN, start_seed=0) eval_rows = build_dataset(n_episodes=N_EPISODES_EVAL, start_seed=10_000) print(f"Train: {dataset_stats(train_rows)}") print(f"Eval: {dataset_stats(eval_rows)}") # SFT wants `prompt` + `completion` fields (chat template applied later) # GRPO wants `prompt` + the extra columns used in the reward functions. def to_chat(row, include_completion=False): out = { "prompt": [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": row["prompt"]}, ], # Kept for TRL to forward into the reward funcs "correct_action_hint": row["correct_action_hint"], "email_kind": row["email_kind"] or "", "can_earn_drift_bonus": bool(row["can_earn_drift_bonus"]), "drift_sensitive_to": row["drift_sensitive_to"] or "", "is_admin_email": bool(row["is_admin_email"]), } if include_completion: out["completion"] = [{"role": "assistant", "content": row["correct_action_json"]}] return out sft_train = Dataset.from_list([to_chat(r, include_completion=True) for r in train_rows]) grpo_train = Dataset.from_list([to_chat(r) for r in train_rows]) grpo_eval = Dataset.from_list([to_chat(r) for r in eval_rows]) return sft_train, grpo_train, grpo_eval # --------------------------------------------------------------------------- # 2. Load Qwen via Unsloth with LoRA adapters # --------------------------------------------------------------------------- def load_model_and_tokenizer(): from unsloth import FastLanguageModel model, tokenizer = FastLanguageModel.from_pretrained( model_name=MODEL_NAME, max_seq_length=MAX_SEQ_LEN, dtype=None, # auto load_in_4bit=LOAD_IN_4BIT, ) model = FastLanguageModel.get_peft_model( model, r=LORA_R, target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", use_gradient_checkpointing="unsloth", random_state=SEED, ) return model, tokenizer # --------------------------------------------------------------------------- # 3. SFT warm-up (1 epoch) # --------------------------------------------------------------------------- def run_sft(model, tokenizer, train_ds): from trl import SFTTrainer, SFTConfig args = SFTConfig( output_dir=f"{OUTPUT_DIR}/sft", num_train_epochs=SFT_EPOCHS, per_device_train_batch_size=SFT_BATCH, gradient_accumulation_steps=4, learning_rate=SFT_LR, logging_steps=10, save_strategy="no", bf16=USE_BF16, fp16=USE_FP16, report_to="wandb" if USE_WANDB else "none", seed=SEED, max_length=MAX_SEQ_LEN, # Cap tokenization parallelism — Unsloth defaults to num_proc=cpu_count() # which spawns 64+ workers on cloud GPUs and OOMs the container. dataset_num_proc=4, ) trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=train_ds, args=args, ) print(f"\n=== SFT warm-up: {len(train_ds)} samples, {SFT_EPOCHS} epoch(s) ===") trainer.train() # Save structured log history for plotting. import json as _json with open(f"{OUTPUT_DIR}/sft_log.json", "w") as f: _json.dump(trainer.state.log_history, f, indent=2) print(f"SFT log saved -> {OUTPUT_DIR}/sft_log.json") # CRITICAL: save adapter immediately after SFT in case GRPO fails later. # Without this, an error during GRPO loses the entire SFT adapter. sft_adapter_path = f"{OUTPUT_DIR}/lora_adapters_sft" trainer.model.save_pretrained(sft_adapter_path) tokenizer.save_pretrained(sft_adapter_path) print(f"SFT-only adapter checkpoint saved -> {sft_adapter_path}") # Push outputs to Hub now too, so the SFT artifact survives any later crash. _push_outputs_to_hub("post-SFT checkpoint") return trainer.model # --------------------------------------------------------------------------- # 4. GRPO training # --------------------------------------------------------------------------- def run_grpo(model, tokenizer, train_ds, eval_ds): from trl import GRPOTrainer, GRPOConfig args = GRPOConfig( output_dir=f"{OUTPUT_DIR}/grpo", num_generations=GRPO_NUM_GEN, max_completion_length=GRPO_MAX_COMPLETION, per_device_train_batch_size=GRPO_BATCH, gradient_accumulation_steps=GRPO_GRAD_ACCUM, learning_rate=GRPO_LR, max_steps=GRPO_MAX_STEPS, logging_steps=5, save_strategy="no", bf16=USE_BF16, fp16=USE_FP16, report_to="wandb" if USE_WANDB else "none", seed=SEED, # Keep rollouts fast for Colab: temperature=0.7, top_p=0.9, ) print(f"\n=== GRPO: {len(train_ds)} prompts, max_steps={GRPO_MAX_STEPS}, K={GRPO_NUM_GEN} ===") trainer = GRPOTrainer( model=model, processing_class=tokenizer, reward_funcs=[ reward_compliance, reward_appropriateness, reward_drift_bonus, ], args=args, train_dataset=train_ds, eval_dataset=eval_ds, ) trainer.train() # Save structured log history for plotting. import json as _json with open(f"{OUTPUT_DIR}/grpo_log.json", "w") as f: _json.dump(trainer.state.log_history, f, indent=2) print(f"GRPO log saved -> {OUTPUT_DIR}/grpo_log.json") return trainer.model # --------------------------------------------------------------------------- # 5. Offline eval (drift-sensitive accuracy before vs after) # --------------------------------------------------------------------------- def offline_eval(model, tokenizer, eval_ds, label: str, max_rows: int = 200): """Greedy-decode each eval prompt, score with total_reward, print summary. Logs drift-sensitive accuracy broken down by drift direction (tightening / loosening / neutral). The *tightening* number is the one that actually matters for the pitch — it measures the leniency bias that our training specifically aims to remove. Loosening accuracy is often high even pre-training because the looser rule matches the base model's internet prior. """ from drift_env.training.rewards import parse_generated_action, total_reward from drift_env.policy import drift_direction model.eval() comp_total = appr_total = bonus_total = 0.0 drift_total = drift_correct = 0 per_dir = {"tightening": [0, 0], "loosening": [0, 0], "neutral": [0, 0]} # [correct, total] n = min(len(eval_ds), max_rows) for i in range(n): row = eval_ds[i] chat = row["prompt"] inputs = tokenizer.apply_chat_template( chat, add_generation_prompt=True, return_tensors="pt", ).to(model.device) with torch.no_grad(): out = model.generate( inputs, max_new_tokens=GRPO_MAX_COMPLETION, do_sample=False, pad_token_id=tokenizer.eos_token_id, use_cache=True, ) text = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True) r = total_reward( completion=text, correct_action_hint=row["correct_action_hint"], email_kind=row["email_kind"] or None, can_earn_drift_bonus=row["can_earn_drift_bonus"], drift_sensitive_to=row["drift_sensitive_to"] or None, is_admin_email=row["is_admin_email"], ) comp_total += r["compliance"] appr_total += r["appropriateness"] bonus_total += r["drift_bonus"] # per-direction tracking: a row counts only if it is drift-sensitive # (i.e. the correct answer is different from what it'd be pre-drift) drift_to = row["drift_sensitive_to"] or None if drift_to: direction = drift_direction(drift_to) if direction is not None: per_dir[direction][1] += 1 if r["compliance"] >= 1.0: per_dir[direction][0] += 1 if row["can_earn_drift_bonus"]: drift_total += 1 if r["compliance"] >= 1.0: drift_correct += 1 drift_acc = drift_correct / drift_total if drift_total else None def _acc(pair): c, t = pair return (c / t) if t else None per_dir_acc = {k: _acc(v) for k, v in per_dir.items()} print(f"\n=== Offline eval [{label}] over {n} samples ===") print(f" compliance avg : {comp_total / n:.3f} / 1.0") print(f" appropriateness avg: {appr_total / n:.3f} / 0.5") print(f" drift_bonus avg : {bonus_total / n:.3f} / 0.5") print(f" total avg : {(comp_total + appr_total + bonus_total) / n:.3f} / 2.0") if drift_acc is not None: print(f" drift-sens acc : {drift_acc:.1%} ({drift_correct}/{drift_total})") for direction in ("tightening", "loosening", "neutral"): c, t = per_dir[direction] acc_str = f"{c/t:.1%}" if t else "n/a" print(f" {direction:<11}: {acc_str} ({c}/{t})") return { "compliance": comp_total / n, "appropriateness": appr_total / n, "drift_bonus": bonus_total / n, "drift_acc": drift_acc, "drift_acc_by_direction": per_dir_acc, "drift_counts_by_direction": {k: {"correct": v[0], "total": v[1]} for k, v in per_dir.items()}, } # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> int: seed_all(SEED) os.makedirs(OUTPUT_DIR, exist_ok=True) global USE_BF16, USE_FP16 USE_BF16, USE_FP16 = _precision_flags() print(f"QUICK_MODE={QUICK_MODE} MODEL={MODEL_NAME}") if torch.cuda.is_available(): name = torch.cuda.get_device_name(0) cap = torch.cuda.get_device_capability(0) print(f"GPU: {name} capability={cap} bf16={USE_BF16} fp16={USE_FP16}") else: print("CUDA not available — training will be extremely slow") sft_train, grpo_train, grpo_eval = build_hf_datasets() model, tokenizer = load_model_and_tokenizer() # Pre-training offline eval — the "before" number for the pitch clip. pre = offline_eval(model, tokenizer, grpo_eval, label="pre-training") # SFT warm-start. model = run_sft(model, tokenizer, sft_train) post_sft = offline_eval(model, tokenizer, grpo_eval, label="post-SFT") # GRPO. If this fails we still want the SFT outputs pushed. try: model = run_grpo(model, tokenizer, grpo_train, grpo_eval) post_grpo = offline_eval(model, tokenizer, grpo_eval, label="post-GRPO") except Exception as exc: print(f"\n[grpo] FAILED: {type(exc).__name__}: {exc}") print("[grpo] Continuing with post-SFT model as the final artifact.") post_grpo = post_sft # fall back so the summary still prints sensibly print("\n=== Improvement summary ===") for k in ("compliance", "appropriateness", "drift_bonus"): print(f" {k:<20} {pre[k]:.3f} -> {post_sft[k]:.3f} -> {post_grpo[k]:.3f}") print(f" drift-sens acc (all) {_fmt_acc(pre['drift_acc'])} -> " f"{_fmt_acc(post_sft['drift_acc'])} -> {_fmt_acc(post_grpo['drift_acc'])}") for direction in ("tightening", "loosening", "neutral"): pre_a = pre["drift_acc_by_direction"].get(direction) sft_a = post_sft["drift_acc_by_direction"].get(direction) grpo_a = post_grpo["drift_acc_by_direction"].get(direction) print(f" {direction:<12} {_fmt_acc(pre_a)} -> {_fmt_acc(sft_a)} -> {_fmt_acc(grpo_a)}") # Save eval snapshots for the plotting script. import json as _json with open(f"{OUTPUT_DIR}/evals.json", "w") as f: _json.dump({ "pre": pre, "post_sft": post_sft, "post_grpo": post_grpo, "model_name": MODEL_NAME, "quick_mode": QUICK_MODE, }, f, indent=2, default=str) print(f"Eval snapshots saved -> {OUTPUT_DIR}/evals.json") # Save LoRA adapters ONLY (organizer warning: do not naively merge 4-bit). adapter_path = f"{OUTPUT_DIR}/lora_adapters" model.save_pretrained(adapter_path) tokenizer.save_pretrained(adapter_path) print(f"\nLoRA adapters saved to {adapter_path}") # Final push of all outputs (SFT + GRPO + plots + adapter). _push_outputs_to_hub("post-GRPO final") if not os.getenv("HUB_REPO_ID"): print("\n[note] HUB_REPO_ID not set; outputs only exist inside this container.") return 0 def _fmt_acc(a: Optional[float]) -> str: return f"{a:.1%}" if a is not None else "n/a" if __name__ == "__main__": sys.exit(main())