scripts/modal_qlora_train.py

"""
Modal QLoRA fine-tune: Qwen3-8B → ai-sherpa/Qwen3-8B-Kintsugi.

Plan reference: subagent ad6ef461338cb47b6 §3 (Training compute), §4
(Publishing artifacts), §7 (Risks & time estimate).

WHAT IT DOES (in order):
  1. Mount Modal volumes (HF cache + checkpoints).
  2. Load Qwen3-8B base in 4-bit NF4 (BitsAndBytes).
  3. Apply QLoRA adapters (r=16, α=32) to attention + MLP projections.
  4. Build a `datasets.Dataset` from docs/finetune/training-data/{train,eval}.jsonl
     using the existing chat-message format from build_training_data.py.
  5. Train with TRL SFTTrainer, 3 epochs, ~90 min on H100.
  6. Merge LoRA into base, copy tokenizer_config.json from the base repo
     (per plan §7 risk #2 — prevents chat_template drift), push merged
     model to HF Hub as `ai-sherpa/Qwen3-8B-Kintsugi`.

WHAT IT DOES NOT DO:
  - Convert merged model to GGUF Q4_K_M — that's a separate llama.cpp
    convert + quantize step on a CPU machine (no GPU needed).
  - Publish the dataset card — that's a separate `huggingface-cli` step
    or a Python helper, run after the training data is final.
  - Run the QA acceptance harness — that's local-CPU work via
    scripts/qa_acceptance_harness.py once LLAMA_REPO is flipped.

PREREQUISITES (one-time setup):
  1. `pip install modal && modal setup` — set up Modal account locally.
  2. Create the HF Hub repos (do this from a browser to confirm
     namespace ownership; modal can't create them):
         - https://huggingface.co/new (model) → ai-sherpa/Qwen3-8B-Kintsugi
  3. Set the HF token as a Modal Secret:
         modal secret create huggingface HF_TOKEN=hf_xxxxxxxx
  4. Generate training data locally:
         python3.10 scripts/build_training_data.py \\
             --input docs/finetune/seed-examples.jsonl
     (Plus the 150 self-distilled rows when ready, per plan §2.)

USAGE:
  # Dry-run: validate env + show config, no GPU spend.
  modal run scripts/modal_qlora_train.py::main --dry-run

  # Real training run (~$8, ~90 min on H100).
  modal run scripts/modal_qlora_train.py::main

DRAFT STATUS:
  This is a draft. Verify against current library docs before launching
  a paid run:
    - Modal API: https://modal.com/docs
    - TRL SFTTrainer: https://huggingface.co/docs/trl/sft_trainer
    - PEFT LoraConfig: https://huggingface.co/docs/peft/package_reference/lora
  The pinned library versions in IMAGE below were the stable set as of
  2026-Q2; if Modal's pre-built CUDA image diverges, expect to adjust
  bitsandbytes/torch combos. Run `--dry-run` first to surface any
  version mismatch before paying for GPU.
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import modal


# ----------------------------------------------------------------------------
# Constants — from plan §3
# ----------------------------------------------------------------------------

BASE_MODEL_ID = "Qwen/Qwen3-8B"
HUB_MODEL_ID = "ai-sherpa/Qwen3-8B-Kintsugi"

# QLoRA hyperparameters (plan §3)
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05
# Apply LoRA to attention + MLP — covers the projections that matter for
# style transfer without ballooning trainable params.
LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]

# Training hyperparameters
NUM_EPOCHS = 3
PER_DEVICE_BATCH_SIZE = 2
GRAD_ACCUMULATION = 4          # effective batch = 8
LEARNING_RATE = 2e-4
WARMUP_RATIO = 0.03
MAX_SEQ_LENGTH = 4096          # the OUTPUT_FORMAT + lexicon scaffold + assistant
                               # turn together fit comfortably under this
WEIGHT_DECAY = 0.01

# Modal config
GPU_TYPE = "H100"              # ~$6/hr × 90 min ≈ $9; A100 (~$3/hr) is fine
                               # for cost-sensitive runs at slower wall time
TIMEOUT_SECONDS = 60 * 60 * 2  # 2-hour ceiling


# ----------------------------------------------------------------------------
# Modal app + image
# ----------------------------------------------------------------------------

# Qwen3 architecture support was added in transformers 4.51. TRL's
# SFTTrainer API also shifted in this window (tokenizer→processing_class,
# max_seq_length → SFTConfig). These pins are the post-shift compatible
# set as of 2026-Q2 — see the SFTTrainer call below for the new API
# this script relies on.
IMAGE = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "torch==2.5.1",
        "transformers>=4.52.0,<5.0",
        "peft>=0.15.0",
        "accelerate>=1.5.0",
        "bitsandbytes>=0.45.0",
        "trl>=0.18.0,<0.20",
        "datasets>=3.2.0",
        "huggingface_hub>=0.28.0",
        "sentencepiece",
        "protobuf",
    )
    .env({
        # HF_HOME points at the persistent volume so the 5GB base download
        # is paid once across all runs.
        "HF_HOME": "/hf_cache",
        # TRL has been moving the chat-template assembly between
        # tokenizers and the trainer; setting this avoids the deprecation
        # warning and is a no-op when not relevant.
        "TRL_USE_RICH": "0",
    })
)

app = modal.App(name="kintsugi-qlora-train", image=IMAGE)

# Volumes — survive across runs.
hf_cache = modal.Volume.from_name("kintsugi-hf-cache", create_if_missing=True)
checkpoints = modal.Volume.from_name("kintsugi-checkpoints", create_if_missing=True)


# ----------------------------------------------------------------------------
# Remote function: the train run
# ----------------------------------------------------------------------------

@app.function(
    gpu=GPU_TYPE,
    timeout=TIMEOUT_SECONDS,
    volumes={"/hf_cache": hf_cache, "/checkpoints": checkpoints},
    secrets=[modal.Secret.from_name("huggingface")],
)
def train_qlora(
    train_jsonl: bytes,
    eval_jsonl: bytes,
    num_epochs: int = NUM_EPOCHS,
    push_to_hub: bool = True,
    run_tag: str = "v1",
) -> dict:
    """Train Qwen3-8B with QLoRA on the supplied (train, eval) JSONL bytes.

    Returns a dict with hub_model_id (if pushed) and final losses.
    """
    import torch
    from transformers import (
        AutoModelForCausalLM, AutoTokenizer,
        BitsAndBytesConfig,
    )
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
    from trl import SFTConfig, SFTTrainer
    from datasets import Dataset

    hf_token = os.environ["HF_TOKEN"]

    print(f"[train] CUDA available: {torch.cuda.is_available()}")
    print(f"[train] device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu'}")

    # ---- 1. Tokenizer ----
    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL_ID, token=hf_token, trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        # Qwen3 ships an explicit pad token; this is belt-and-braces.
        tokenizer.pad_token = tokenizer.eos_token

    # ---- 2. Base model, 4-bit NF4 ----
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        token=hf_token,
        trust_remote_code=True,
    )
    base = prepare_model_for_kbit_training(base)

    # ---- 3. LoRA config ----
    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        target_modules=LORA_TARGET_MODULES,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(base, lora_config)
    model.print_trainable_parameters()

    # ---- 4. Datasets ----
    def parse_jsonl(blob: bytes) -> Dataset:
        rows = []
        for line in blob.decode("utf-8").splitlines():
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            # SFTTrainer with chat-format expects a 'messages' field.
            rows.append({"messages": row["messages"]})
        return Dataset.from_list(rows)

    train_ds = parse_jsonl(train_jsonl)
    eval_ds = parse_jsonl(eval_jsonl)
    print(f"[train] train rows: {len(train_ds)}  eval rows: {len(eval_ds)}")

    # ---- 5. SFTConfig (TRL ≥0.18 — replaces TrainingArguments for SFT) ----
    output_dir = f"/checkpoints/{run_tag}"
    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=PER_DEVICE_BATCH_SIZE,
        per_device_eval_batch_size=PER_DEVICE_BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUMULATION,
        learning_rate=LEARNING_RATE,
        warmup_ratio=WARMUP_RATIO,
        weight_decay=WEIGHT_DECAY,
        bf16=True,
        optim="paged_adamw_8bit",          # bitsandbytes optimizer — keeps memory low
        logging_steps=2,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,                # keep last 2 checkpoints only
        report_to="none",                  # no wandb/etc unless you set it up
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        load_best_model_at_end=False,      # eval set is small; best-loss is noisy here
        max_seq_length=MAX_SEQ_LENGTH,     # moved from SFTTrainer kwarg into SFTConfig in TRL 0.13+
        # NOTE: assistant_only_loss=True would be ideal for voice transfer
        # (train on assistant turn only, not the static lexicon scaffold in
        # the user turn). But it requires the tokenizer's chat template to
        # mark assistant spans with {% generation %} jinja tags, which
        # Qwen3's stock template does not. Adding a custom template is
        # possible but invasive — for a 30-row dataset over 3 epochs the
        # extra gradient cost from training on the (static) user turn is
        # minimal. Leave full-sequence loss for now.
    )

    # ---- 6. SFTTrainer ----
    # TRL ≥0.13: tokenizer= → processing_class=.
    # With messages-format datasets, SFTTrainer auto-applies the tokenizer's
    # chat_template. Qwen3 ships a Qwen3-formatted template producing
    # <|im_start|>role<|im_end|> markers.
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        processing_class=tokenizer,
    )

    # ---- 7. Train ----
    train_result = trainer.train()
    metrics = train_result.metrics
    trainer.save_model(output_dir)
    checkpoints.commit()

    print(f"[train] final train loss: {metrics.get('train_loss')}")

    if not push_to_hub:
        return {"hub_model_id": None, "metrics": metrics, "output_dir": output_dir}

    # ---- 8. Merge LoRA into base + push ----
    # Reload the base in fp16 (not 4-bit) for merge; merging into a
    # quantized base would lose precision in the adapter direction.
    print("[merge] reloading base in bf16 for merge...")
    del model, base, trainer
    torch.cuda.empty_cache()

    base_fp = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=hf_token,
        trust_remote_code=True,
    )
    peft_model = PeftModel.from_pretrained(base_fp, output_dir, token=hf_token)
    merged = peft_model.merge_and_unload()

    merged_dir = f"/checkpoints/{run_tag}-merged"
    merged.save_pretrained(merged_dir, safe_serialization=True)

    # ---- 9. Copy tokenizer config from base (plan §7 risk #2) ----
    # tokenizer.save_pretrained() captures the chat_template; without
    # this step the merged repo may lack the field that transformers
    # fallback / Gradio inference rely on.
    tokenizer.save_pretrained(merged_dir)
    checkpoints.commit()

    # ---- 10. Push to hub ----
    print(f"[push] pushing merged model to {HUB_MODEL_ID}...")
    merged.push_to_hub(
        HUB_MODEL_ID,
        token=hf_token,
        private=False,
        commit_message=f"QLoRA fine-tune from {BASE_MODEL_ID} ({run_tag})",
    )
    tokenizer.push_to_hub(HUB_MODEL_ID, token=hf_token)
    print(f"[push] done.")

    return {
        "hub_model_id": HUB_MODEL_ID,
        "metrics": metrics,
        "output_dir": merged_dir,
        "run_tag": run_tag,
    }


# ----------------------------------------------------------------------------
# Local entrypoint
# ----------------------------------------------------------------------------

@app.local_entrypoint()
def main(
    train_path: str = "docs/finetune/training-data/train.jsonl",
    eval_path: str = "docs/finetune/training-data/eval.jsonl",
    epochs: int = NUM_EPOCHS,
    push: bool = True,
    run_tag: str = "v1",
    dry_run: bool = False,
):
    """Local entrypoint — reads JSONL from disk and dispatches to Modal.

    Modal's CLI passes args as keyword strings, so booleans accept "true"/"false".
    """
    repo_root = Path(__file__).resolve().parent.parent
    train_file = repo_root / train_path
    eval_file = repo_root / eval_path

    if not train_file.exists():
        print(f"ERROR: train file not found: {train_file}", file=sys.stderr)
        print("  Generate it first with:", file=sys.stderr)
        print("    python3.10 scripts/build_training_data.py "
              "--input docs/finetune/seed-examples.jsonl", file=sys.stderr)
        return 1
    if not eval_file.exists():
        print(f"ERROR: eval file not found: {eval_file}", file=sys.stderr)
        return 1

    train_bytes = train_file.read_bytes()
    eval_bytes = eval_file.read_bytes()
    train_rows = train_bytes.decode("utf-8").count("\n")
    eval_rows = eval_bytes.decode("utf-8").count("\n")

    print(f"Train: {train_rows} rows ({len(train_bytes)} bytes)")
    print(f"Eval:  {eval_rows} rows ({len(eval_bytes)} bytes)")
    print(f"Epochs: {epochs}")
    print(f"GPU: {GPU_TYPE}")
    print(f"Push to hub: {push} ({HUB_MODEL_ID if push else 'skipped'})")
    print(f"Run tag: {run_tag}")

    # Cost estimate
    est_minutes = 30 if GPU_TYPE == "H100" else 75
    est_minutes *= max(1, epochs / NUM_EPOCHS)
    est_cost = (est_minutes / 60) * (6.0 if GPU_TYPE == "H100" else 3.0)
    print(f"\nEstimate: ~{est_minutes:.0f} min wall, ~${est_cost:.2f}")

    if dry_run:
        print("\n--dry-run: not dispatching to Modal.")
        return 0

    print("\nDispatching to Modal...")
    result = train_qlora.remote(
        train_jsonl=train_bytes,
        eval_jsonl=eval_bytes,
        num_epochs=epochs,
        push_to_hub=push,
        run_tag=run_tag,
    )
    print("\nResult:")
    print(json.dumps(result, indent=2, default=str))
    if result.get("hub_model_id"):
        print(f"\nNext steps:")
        print(f"  1. Verify the model card at "
              f"https://huggingface.co/{result['hub_model_id']}")
        print(f"  2. Convert to GGUF Q4_K_M (separate llama.cpp step).")
        print(f"  3. Publish GGUF as ai-sherpa/Qwen3-8B-Kintsugi-GGUF.")
        print(f"  4. Flip LLAMA_REPO in app.py and re-run the QA harness.")
    return 0