build-small-hackathon/puck

1"""Train Puck's character LoRA on Holotron-12B, on Modal.2 3TRL SFT + PEFT LoRA over the 162-example chat curriculum (build_dataset.py).4Text-only LoRA on the language side of the VLM — character/voice, not vision.5Adapter is saved to a Modal volume (no HF token in env); publish later.6 7  cd molt && uv run build_dataset.py8  modal run --detach train_modal.py              # .spawn() inside → truly detached9  modal volume get puck-lora /puck-holotron-12b-lora ./out10 11⚠️ BLOCKED (2026-06-07): Hcompany/Holotron-12B's published trust_remote_code12modeling.py imports a `_fully_shard.py` that isn't in the repo — an H Company13packaging bug in their *training* path (vLLM inference is unaffected, which is14why the vision endpoint works). transformers can't load it for TRL.15Options when revisiting: (a) wait for H Company to publish the missing file;16(b) train the character LoRA on the BASE nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL17(still Nemotron, still ≤32B, dataset is base-agnostic — voice doesn't need18Holotron's CUA tuning); (c) stub _fully_shard.py if it's a no-op FSDP helper.19Character already lands well via the enriched prompt, so the LoRA (Well-Tuned20badge) is lower priority than vision.21 22Hybrid-Mamba caveat once unblocked: target_modules='all-linear', gradient23checkpointing off (Mamba layers dislike it)."""24 25import json26from pathlib import Path27 28import modal29 30MODEL = "Hcompany/Holotron-12B"31HERE = Path(__file__).resolve().parent32 33image = (34    modal.Image.debian_slim(python_version="3.12")35    .pip_install(36        "torch",37        "transformers>=4.48",38        "trl>=0.12",39        "peft>=0.14",40        "datasets",41        "accelerate",42        "huggingface_hub[hf_transfer]",43        "sentencepiece",44        "einops",  # hybrid-Mamba modeling code often needs it45    )46    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})47)48 49vol = modal.Volume.from_name("puck-lora", create_if_missing=True)50hf_cache = modal.Volume.from_name("puck-hf-cache", create_if_missing=True)51 52app = modal.App("puck-train")53 54 55@app.function(56    image=image,57    gpu="A100-80GB",58    timeout=60 * 60,59    volumes={"/adapter": vol, "/root/.cache/huggingface": hf_cache},60)61def train(records: list[dict]):62    import torch63    from datasets import Dataset64    from peft import LoraConfig65    from transformers import AutoModelForCausalLM, AutoTokenizer66    from trl import SFTConfig, SFTTrainer67 68    # conversational dataset → TRL applies the chat template itself69    ds = Dataset.from_list([{"messages": r["messages"]} for r in records])70 71    tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)72    model = AutoModelForCausalLM.from_pretrained(73        MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"74    )75 76    peft_cfg = LoraConfig(77        r=16,78        lora_alpha=32,79        lora_dropout=0.05,80        bias="none",81        task_type="CAUSAL_LM",82        target_modules="all-linear",  # robust across the hybrid's linear layers83    )84 85    cfg = SFTConfig(86        output_dir="/adapter/run",87        num_train_epochs=4,88        per_device_train_batch_size=1,89        gradient_accumulation_steps=8,90        learning_rate=2e-4,91        warmup_ratio=0.05,92        logging_steps=5,93        save_strategy="epoch",94        bf16=True,95        gradient_checkpointing=False,  # Mamba layers + checkpointing don't mix96        max_length=1024,97        report_to="none",98    )99 100    trainer = SFTTrainer(model=model, args=cfg, train_dataset=ds, peft_config=peft_cfg)101    trainer.train()102 103    out = "/adapter/puck-holotron-12b-lora"104    trainer.save_model(out)105    tok.save_pretrained(out)106    vol.commit()107    print(f"saved adapter → {out}")108    return out109 110 111@app.local_entrypoint()112def main():113    records = [114        json.loads(line) for line in (HERE / "data" / "sft.jsonl").read_text().splitlines()115    ]116    print(f"training on {len(records)} examples")117    # .spawn() so a detached run survives the local caller disconnecting118    # (.remote() is synchronous and Modal cancels it when the CLI exits).119    call = train.spawn(records)120    print(f"spawned: {call.object_id}  —  modal volume get puck-lora /puck-holotron-12b-lora ./out")121