whetstone_DSL/tools/mcp/run_qwen7b_persistent.py

#!/usr/bin/env python3
"""
Persistent-weights LoRA training loop for Qwen2.5-Coder-7B.

Loads the model and tokenizer ONCE, then retries training without
reloading on each attempt. Useful for iterating on hyperparameters
or recovering from non-OOM training failures without the 5-7 minute
model load penalty.

Usage:
    .venv/bin/python3 tools/mcp/run_qwen7b_persistent.py [--max-steps N] [--max-retries N]

Reads config from: configs/lora/qwen2.5-coder-7b-quality.env
"""
import argparse
import inspect
import json
import os
import re
import sys
import time
from pathlib import Path


def load_env(env_file):
    cfg = {}
    for line in Path(env_file).read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            k, _, v = line.partition("=")
            cfg[k.strip()] = v.strip().strip('"')
    return cfg


def require(pkg):
    try:
        return __import__(pkg)
    except Exception as e:
        raise SystemExit(f"missing package '{pkg}': {e}")


def build_model_and_tokenizer(cfg, root):
    import torch
    from pathlib import Path
    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    base_model = cfg["BASE_MODEL"]
    local_only = cfg.get("LOCAL_FILES_ONLY", "1") == "1"
    bf16 = cfg.get("BF16", "0") == "1"
    use_4bit = cfg.get("USE_4BIT", "1") == "1"

    model_ref = base_model
    if local_only and "/" in base_model:
        safe = base_model.replace("/", "--")
        cache = Path.home() / ".cache" / "huggingface" / "hub" / f"models--{safe}"
        ref_main = cache / "refs" / "main"
        if ref_main.exists():
            rev = ref_main.read_text().strip()
            snap = cache / "snapshots" / rev
            if snap.exists():
                model_ref = str(snap)

    # Always load in float16 — bitsandbytes 0.49.x only applies 4-bit quantization
    # reliably when the base dtype is float16. BF16 base dtype causes bnb to skip
    # quantization and load the full model (~14 GB), which OOMs on 12 GB cards.
    load_dtype = torch.float16
    bnb_config = None
    model_kwargs = {}

    if use_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=load_dtype,
        )
        model_kwargs["quantization_config"] = bnb_config
        model_kwargs["device_map"] = "auto"
        model_kwargs["low_cpu_mem_usage"] = True

    print(f"[loader] loading tokenizer from {model_ref}")
    tokenizer = AutoTokenizer.from_pretrained(
        model_ref, use_fast=True, local_files_only=local_only
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"[loader] loading model ({load_dtype}, 4bit={use_4bit})")
    model = AutoModelForCausalLM.from_pretrained(
        model_ref,
        local_files_only=local_only,
        torch_dtype=load_dtype,
        **model_kwargs,
    )
    print("[loader] model loaded")
    return model, tokenizer


def build_trainer(model, tokenizer, cfg, root, max_steps_override=None):
    import torch
    from datasets import load_dataset
    from peft import LoraConfig
    from trl import SFTConfig, SFTTrainer

    bf16 = cfg.get("BF16", "0") == "1"
    max_seq_len = int(cfg.get("MAX_SEQ_LEN", "1536"))
    out_dir = Path(root) / cfg["OUTPUT_DIR"]
    out_dir.mkdir(parents=True, exist_ok=True)

    train_path = Path(root) / cfg["SFT_DIR"] / "train.jsonl"
    val_path = Path(root) / cfg["SFT_DIR"] / "val.jsonl"

    ds = load_dataset("json", data_files={"train": str(train_path), "validation": str(val_path)})

    def to_text(ex):
        msgs = ex["messages"]
        if hasattr(tokenizer, "apply_chat_template"):
            txt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        else:
            txt = "\n".join(f"[{m['role']}] {m['content']}" for m in msgs)
        return {"text": txt}

    ds = ds.map(to_text, remove_columns=[c for c in ds["train"].column_names if c != "messages"])

    sft_params = inspect.signature(SFTConfig.__init__).parameters
    eval_key = "evaluation_strategy" if "evaluation_strategy" in sft_params else "eval_strategy"

    max_steps = max_steps_override or int(cfg.get("MAX_STEPS", "0"))
    # For short benchmark runs auto-compute a sensible save cadence.
    save_steps = int(cfg.get("SAVE_STEPS", "200"))
    if max_steps > 0:
        save_steps = max(1, max_steps // 4)

    targs_kwargs = dict(
        output_dir=str(out_dir),
        num_train_epochs=float(cfg.get("NUM_EPOCHS", "2")),
        learning_rate=float(cfg.get("LEARNING_RATE", "0.0002")),
        weight_decay=float(cfg.get("WEIGHT_DECAY", "0.01")),
        per_device_train_batch_size=int(cfg.get("PER_DEVICE_TRAIN_BATCH", "1")),
        per_device_eval_batch_size=int(cfg.get("PER_DEVICE_EVAL_BATCH", "1")),
        gradient_accumulation_steps=int(cfg.get("GRAD_ACCUM", "16")),
        warmup_steps=int(float(cfg.get("WARMUP_RATIO", "0.03")) * max(max_steps, 100)),
        logging_steps=int(cfg.get("LOGGING_STEPS", "20")),
        eval_steps=int(cfg.get("EVAL_STEPS", "200")),
        save_steps=save_steps,
        save_strategy="steps",
        # No AMP — bitsandbytes 4-bit computes internally in fp16 via bnb_4bit_compute_dtype.
        # Using fp16=True or bf16=True creates a GradScaler that crashes on the mixed
        # fp16/bf16 dtypes that PEFT 0.18+ + bnb 0.49.x produce. No AMP = no GradScaler.
        bf16=False,
        fp16=False,
        report_to="none",
        seed=int(cfg.get("SEED", "42")),
        max_steps=max_steps,
        max_length=max_seq_len,
    )
    targs_kwargs[eval_key] = "steps"
    targs = SFTConfig(**targs_kwargs)

    peft_config = LoraConfig(
        r=int(cfg.get("LORA_R", "32")),
        lora_alpha=int(cfg.get("LORA_ALPHA", "64")),
        lora_dropout=float(cfg.get("LORA_DROPOUT", "0.05")),
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=[m.strip() for m in cfg.get("LORA_TARGET_MODULES", "q_proj,v_proj").split(",") if m.strip()],
    )

    trainer_kwargs = dict(
        model=model,
        args=targs,
        train_dataset=ds["train"],
        eval_dataset=ds["validation"],
        peft_config=peft_config,
        processing_class=tokenizer,
        formatting_func=lambda ex: ex["text"],
    )
    if "max_seq_length" in inspect.signature(SFTTrainer.__init__).parameters:
        trainer_kwargs["max_seq_length"] = max_seq_len
    trainer = SFTTrainer(**trainer_kwargs)

    return trainer, out_dir


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/lora/qwen2.5-coder-7b-quality.env")
    ap.add_argument("--max-steps", type=int, default=20)
    ap.add_argument("--max-retries", type=int, default=3)
    args = ap.parse_args()

    require("torch")
    require("transformers")
    require("peft")
    require("trl")
    require("datasets")

    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    root = Path(__file__).resolve().parent.parent.parent
    cfg = load_env(root / args.config)

    print(f"[main] loading model once (max_steps={args.max_steps}, max_retries={args.max_retries})")
    model, tokenizer = build_model_and_tokenizer(cfg, root)

    for attempt in range(1, args.max_retries + 1):
        print(f"\n[attempt {attempt}/{args.max_retries}] building trainer")
        try:
            trainer, out_dir = build_trainer(model, tokenizer, cfg, root, max_steps_override=args.max_steps)
            print(f"[attempt {attempt}] starting training")
            trainer.train()
            trainer.save_model(str(out_dir / "adapter"))
            tokenizer.save_pretrained(str(out_dir / "adapter"))
            print(f"[attempt {attempt}] SUCCESS — adapter saved to {out_dir / 'adapter'}")
            return
        except RuntimeError as e:
            print(f"[attempt {attempt}] RuntimeError: {e}")
            if attempt < args.max_retries:
                print("[retry] rebuilding trainer with same loaded model in 3s")
                time.sleep(3)
            else:
                print("[failed] all retries exhausted")
                sys.exit(1)
        except Exception as e:
            print(f"[attempt {attempt}] unexpected error: {e}")
            raise


if __name__ == "__main__":
    main()