training/modal_train.py

"""Run the QLoRA fine-tune + GGUF export on a Modal GPU and publish to HF.

This is a thin Modal wrapper around the existing pipeline — it does NOT reinvent
the training logic. It ships the repo to a serverless A100/H100, then runs, in order:

    python training/make_dataset.py     # build dataset.jsonl from seeds (if missing)
    python training/train_qlora.py      # Unsloth QLoRA + merge to fp16
    bash   training/export_gguf.sh      # fp16 -> GGUF -> Q4_K_M -> upload to HF

Everything here runs offline before the app serves anything, so the "no cloud AI
API" rule (Off the Grid) is untouched — see PLAN.md.

Setup (once):
    pip install modal
    modal token new
    modal secret create huggingface HF_TOKEN=hf_xxxxxxxx   # write token

Run:
    modal run training/modal_train.py                      # A100-80GB, gemma-4-31B
    modal run training/modal_train.py --gpu H100           # faster
    modal run training/modal_train.py \
        --base-model google/gemma-3-27b-it \
        --hf-repo n8mauer/gemma-cal-gguf
    modal run training/modal_train.py --base-model google/gemma-4-E4B-it  # cheap validation (~$)

Rough cost (A100-80GB @ ~$2.5/hr, per-second billing): a few-hundred-to-2000
example QLoRA run is ~1-3 hr ≈ $5-15. $250 of credit ≈ 15-40 full iterations.
"""
from __future__ import annotations

import os
from pathlib import Path

import modal

REPO_ROOT = Path(__file__).resolve().parent.parent

# llama.cpp is built into the image so export_gguf.sh has convert_hf_to_gguf.py +
# llama-quantize available without a runtime clone.
LLAMA_CPP = "/opt/llama.cpp"

# --- Image: CUDA base (unsloth/bitsandbytes need it) + training deps + llama.cpp ---
image = (
    modal.Image.from_registry(
        "nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.11"
    )
    .apt_install("build-essential", "cmake", "git", "curl")
    # Let unsloth resolve its OWN compatible transformers/trl/peft/accelerate/
    # bitsandbytes set — hand-pinning those (e.g. trl<0.13) drags transformers to a
    # version that breaks unsloth's internals (the auto_docstring / ConstantLengthDataset
    # import errors). Only add what unsloth doesn't pull: the GGUF-convert deps + hf CLI.
    .pip_install(
        "unsloth",
        "huggingface_hub[cli]",
        # convert_hf_to_gguf.py needs these:
        "gguf",
        "sentencepiece",
        "protobuf",
    )
    # Overlay the newest unsloth + unsloth_zoo from git (the documented fix for the
    # auto_docstring / ConstantLengthDataset import errors against current transformers/
    # trl). --no-deps so it patches only unsloth's code, not the resolved dependency set.
    .run_commands(
        "pip install --no-deps --upgrade "
        "git+https://github.com/unslothai/unsloth.git "
        "git+https://github.com/unslothai/unsloth-zoo.git"
    )
    # Build llama.cpp (CPU is enough for convert + quantize) once, into the image.
    .run_commands(
        f"git clone --depth 1 https://github.com/ggml-org/llama.cpp {LLAMA_CPP}",
        f"cmake -S {LLAMA_CPP} -B {LLAMA_CPP}/build -DLLAMA_CURL=OFF",
        f"cmake --build {LLAMA_CPP}/build --target llama-quantize -j",
    )
    # Ship the repo's code (server/, training/) — read at runtime by the scripts.
    .add_local_dir(
        str(REPO_ROOT),
        "/root/repo",
        ignore=[
            ".git",
            "**/__pycache__",
            "training/outputs",
            "training/data/.smcalflow_cache",  # ~70MB raw download; not needed remotely
            "**/*.gguf",
        ],
    )
)

app = modal.App("imessage-cal-train", image=image)

# Persist the HF model cache (base weights) and outputs across runs so re-runs
# skip the multi-GB base-model download and a failed upload isn't catastrophic.
hf_cache = modal.Volume.from_name("imessage-cal-hf-cache", create_if_missing=True)
outputs = modal.Volume.from_name("imessage-cal-outputs", create_if_missing=True)


@app.function(
    gpu="A100-80GB",  # overridable per-run via .with_options(gpu=...) below
    timeout=6 * 60 * 60,  # QLoRA on a 31B can run hours
    secrets=[modal.Secret.from_name("huggingface")],  # injects HF_TOKEN
    volumes={"/cache/hf": hf_cache, "/outputs": outputs},
)
def train(
    base_model: str = "google/gemma-4-31B-it",
    hf_repo: str = "n8mauer/gemma-4-cal-gguf",
    max_seq_len: int = 4096,
    num_epochs: int = 2,
    mmproj_src_repo: str = "unsloth/gemma-4-31B-it-GGUF",
    mmproj_file: str = "mmproj-F16.gguf",
    out_name: str = "gemma-cal",       # produced file: <out_name>-Q4_K_M.gguf
    skip_mmproj: bool = False,         # True for eval-gated staging uploads
    hand_upsample: int = 4,            # x-factor for thread-style hand-authored rows
) -> str:
    """Build dataset -> QLoRA -> merge -> GGUF/Q4_K_M -> upload to `hf_repo`."""
    import shutil
    import subprocess

    # The image layer is read-only; make_dataset.py writes into training/data/,
    # so copy the code to a writable workspace and run there.
    workspace = "/workspace"
    if os.path.exists(workspace):
        shutil.rmtree(workspace)
    shutil.copytree("/root/repo", workspace)

    # Absolute output paths on the persisted volume; both scripts honor these envs.
    env = {
        **os.environ,
        "HF_HOME": "/cache/hf",
        "BASE_MODEL": base_model,
        "MAX_SEQ_LEN": str(max_seq_len),
        "NUM_EPOCHS": str(num_epochs),
        # Per-run dirs (keyed by out_name): different-size runs sharing one merged dir
        # leave stale shards behind, and convert_hf_to_gguf then reads a MIXED model
        # ("Can not map tensor 'model.layers.42...'" when 31B leftovers meet an E4B).
        "OUTPUT_DIR": f"/outputs/{out_name}-lora",
        # export_gguf.sh inputs:
        "MERGED_DIR": f"/outputs/{out_name}-lora-merged",
        "OUT": f"/outputs/{out_name}",
        "LLAMA_CPP": LLAMA_CPP,
        "HF_REPO": hf_repo,
        "MMPROJ_SRC_REPO": mmproj_src_repo,
        "MMPROJ_FILE": mmproj_file,
        "SKIP_MMPROJ": "1" if skip_mmproj else "0",
        "HAND_UPSAMPLE": str(hand_upsample),
    }

    # Start from clean output dirs — save_pretrained does not clear its target, so a
    # re-run with the same out_name would otherwise mix old and new shards.
    for d in (env["OUTPUT_DIR"], env["MERGED_DIR"]):
        if os.path.exists(d):
            shutil.rmtree(d)

    def run(*cmd: str) -> None:
        print(f"\n$ {' '.join(cmd)}", flush=True)
        subprocess.run(cmd, cwd=workspace, env=env, check=True)

    # `modal run` ships the local working tree, which is CRLF on Windows; bash then
    # chokes on the carriage returns ("set: pipefail: invalid option name"). Normalize
    # the shell scripts to LF in the (writable) workspace before running them.
    for rel in ("training/export_gguf.sh", "scripts/start_space.sh"):
        path = f"{workspace}/{rel}"
        if os.path.exists(path):
            with open(path, "rb") as fh:
                data = fh.read()
            with open(path, "wb") as fh:
                fh.write(data.replace(b"\r\n", b"\n"))

    # 1) Dataset (skip if you've already committed a fuller dataset.jsonl).
    if not os.path.exists(f"{workspace}/training/data/dataset.jsonl"):
        run("python", "training/make_dataset.py")

    # 2) QLoRA fine-tune + merge to fp16 (writes to OUTPUT_DIR + -merged).
    run("python", "training/train_qlora.py")

    # 3) Convert -> quantize -> upload GGUF + mmproj to HF (Well-Tuned).
    run("bash", "training/export_gguf.sh")

    hf_cache.commit()  # persist the base-model download so re-runs skip it
    outputs.commit()  # flush the volume so artifacts persist
    print(
        f"\nDone. Point the Space at:\n"
        f"  MODEL_REPO={hf_repo}\n"
        f"  MODEL_FILE={out_name}-Q4_K_M.gguf\n"
        f"  MMPROJ_FILE={mmproj_file}   # enables Gemma vision"
    )
    return hf_repo


@app.local_entrypoint()
def main(
    gpu: str = "A100-80GB",
    base_model: str = "google/gemma-4-31B-it",
    hf_repo: str = "n8mauer/gemma-4-cal-gguf",
    max_seq_len: int = 4096,
    num_epochs: int = 2,
    out_name: str = "gemma-cal",
    skip_mmproj: bool = False,
    hand_upsample: int = 4,
):
    """`modal run training/modal_train.py [--gpu ...] [--base-model ...] ...`

    For eval-gated staging: `--out-name gemma-cal-staging --skip-mmproj` uploads
    `gemma-cal-staging-Q4_K_M.gguf` alongside production without touching it; the
    gate (training/gated_retrain.py) promotes it only if it beats the eval.
    """
    # gpu is fixed at decoration time, so override it for this run via with_options.
    repo = train.with_options(gpu=gpu).remote(
        base_model=base_model,
        hf_repo=hf_repo,
        max_seq_len=max_seq_len,
        num_epochs=num_epochs,
        out_name=out_name,
        skip_mmproj=skip_mmproj,
        hand_upsample=hand_upsample,
    )
    print(f"Published to https://huggingface.co/{repo}")