diff --git a/Dockerfile.train b/Dockerfile.train index 2528f70..115a51b 100644 --- a/Dockerfile.train +++ b/Dockerfile.train @@ -12,36 +12,19 @@ # --dataset /data/datasets/routing_decision-v1.jsonl \ # --dataset /data/datasets/tlp_assignment-v1.jsonl \ # --output /data/adapters/psyc-v1 +# +# Base image already ships Python 3.11 + torch 2.6 + CUDA 12.4 + cuDNN9, so +# there is no apt step and no torch download. Qwen3.5 needs transformers v5 — +# unsloth pulls it automatically. -FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 +FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel -ENV DEBIAN_FRONTEND=noninteractive \ - PYTHONUNBUFFERED=1 \ +ENV PYTHONUNBUFFERED=1 \ PIP_NO_CACHE_DIR=1 \ HF_HOME=/data/.hf-cache -RUN apt-get update && apt-get install -y --no-install-recommends \ - python3.11 python3.11-venv python3-pip \ - git curl ca-certificates \ - && rm -rf /var/lib/apt/lists/* \ - && ln -sf /usr/bin/python3.11 /usr/local/bin/python \ - && ln -sf /usr/bin/python3.11 /usr/local/bin/python3 - -RUN python -m pip install --upgrade pip wheel setuptools && \ - python -m pip install \ - torch==2.5.1 \ - --index-url https://download.pytorch.org/whl/cu124 - -RUN python -m pip install \ - "unsloth @ git+https://github.com/unslothai/unsloth.git" \ - transformers>=4.46 \ - datasets>=3.0 \ - peft>=0.13 \ - trl>=0.12 \ - accelerate>=1.1 \ - bitsandbytes>=0.44 \ - sentencepiece \ - protobuf +RUN pip install --upgrade pip && \ + pip install unsloth unsloth_zoo trl datasets WORKDIR /workspace COPY scripts/train_qlora.py /workspace/train_qlora.py diff --git a/README.md b/README.md index c651e69..28b7ac1 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ To fine-tune Qwen3.5-4B with QLoRA in an NVIDIA Docker container: # 1. build datasets (one-off; re-run after ingestion changes) .venv/bin/psyc train-build-all -# 2. build the training image (CUDA 12.4 + unsloth + Qwen3.5) +# 2. build the training image (pytorch 2.6/CUDA 12.4 base + unsloth + Qwen3.5) docker build -t psyc-trainer -f Dockerfile.train . # 3. fine-tune (mount host data/ so adapters land there) @@ -135,14 +135,25 @@ docker run --gpus all --rm \ --output /data/adapters/psyc-v1 ``` -Defaults target a 24 GB consumer GPU (3090/4090): Qwen3.5-4B-Instruct at 4-bit, +Defaults target a 24 GB consumer GPU (3090/4090): `unsloth/Qwen3.5-4B` at 4-bit, LoRA `r=16`/`alpha=16`, bf16, 3 epochs, effective batch size 8. For A100-40/80 -bump `--base-model unsloth/Qwen3.5-9B-Instruct-bnb-4bit` and raise -`--batch-size` + `--max-seq-length`. +bump `--base-model unsloth/Qwen3.5-9B` and raise `--batch-size` + +`--max-seq-length`. Output: `data/adapters/psyc-v1/final/` (adapter weights) + `training_meta.json` (base model, hyperparameters, dataset list). +Evaluate the adapter against held-out dataset rows: + +```bash +docker run --gpus all --rm \ + --entrypoint python \ + -v $(pwd)/data:/data -v $(pwd)/scripts:/scripts \ + psyc-trainer /scripts/eval_adapter.py \ + --adapter /data/adapters/psyc-v1/final \ + --dataset /data/datasets/ioc_extraction-v1.jsonl --n 5 +``` + ## Status Day 2 of a 48h build. Shipped: Scoutline (URLhaus) → Classifyline → Mapline diff --git a/scripts/eval_adapter.py b/scripts/eval_adapter.py new file mode 100644 index 0000000..bdcad92 --- /dev/null +++ b/scripts/eval_adapter.py @@ -0,0 +1,93 @@ +"""Evaluate a psyc QLoRA adapter — run held-out dataset rows through the model. + +Run inside the psyc training container (override the entrypoint): + docker run --gpus all --rm --entrypoint python \ + -v $(pwd)/data:/data -v $(pwd)/scripts:/scripts \ + psyc-trainer /scripts/eval_adapter.py \ + --adapter /data/adapters/psyc-v1/final \ + --dataset /data/datasets/ioc_extraction-v1.jsonl --n 5 + +Sanity check, not a benchmark: for `--n` rows it prints the prompt, the model's +generation, and the dataset's reference output side by side. With a tiny +dataset the model has seen these rows, so this verifies the adapter learned the +output FORMAT and task shape — not generalization. +""" + +from __future__ import annotations + +# unsloth must be imported BEFORE transformers. +from unsloth import FastLanguageModel # noqa: I001 + +import argparse +import json +import re +from pathlib import Path +from typing import Dict, List + + +def strip_think(text: str) -> str: + """Drop Qwen3.5 thinking-mode blocks so exact-match compares the answer only.""" + return re.sub(r".*?\s*", "", text, flags=re.DOTALL).strip() + + +def load_examples(path: Path, n: int) -> List[Dict[str, str]]: + out: List[Dict[str, str]] = [] + with path.open("r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + out.append(json.loads(line)) + if len(out) >= n: + break + return out + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--adapter", required=True, help="path to adapter final/ dir") + parser.add_argument("--base-model", default="unsloth/Qwen3.5-4B") + parser.add_argument("--dataset", required=True, help="JSONL to sample test rows from") + parser.add_argument("--n", type=int, default=5) + parser.add_argument("--max-seq-length", type=int, default=4096) + parser.add_argument("--max-new-tokens", type=int, default=256) + args = parser.parse_args() + + examples = load_examples(Path(args.dataset), args.n) + if not examples: + raise SystemExit("no examples loaded") + + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=args.adapter, + max_seq_length=args.max_seq_length, + dtype=None, + load_in_4bit=True, + ) + FastLanguageModel.for_inference(model) + + correct = 0 + for i, ex in enumerate(examples, 1): + prompt = f"{ex['instruction']}\n\n{ex['input']}" + messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] + inputs = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + enable_thinking=False, + ).to(model.device) + out = model.generate(input_ids=inputs, max_new_tokens=args.max_new_tokens, do_sample=False) + generated = strip_think(tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)) + expected = ex["output"].strip() + match = generated == expected + correct += int(match) + print(f"\n===== example {i}/{len(examples)} [{ex.get('task', '?')}] {'MATCH' if match else 'DIFF'} =====") + print(f"-- prompt --\n{prompt[:600]}") + print(f"-- expected --\n{expected[:600]}") + print(f"-- generated --\n{generated[:600]}") + + print(f"\n[psyc-eval] exact-match {correct}/{len(examples)}") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_qlora.py b/scripts/train_qlora.py index e6d5d4a..8ae0748 100644 --- a/scripts/train_qlora.py +++ b/scripts/train_qlora.py @@ -22,8 +22,7 @@ from pathlib import Path from typing import Dict, List from datasets import Dataset -from transformers import TrainingArguments -from trl import SFTTrainer +from trl import SFTConfig, SFTTrainer def load_examples(paths: List[Path]) -> List[Dict[str, str]]: @@ -44,7 +43,7 @@ def load_examples(paths: List[Path]) -> List[Dict[str, str]]: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset", action="append", required=True, help="JSONL path (repeatable)") - parser.add_argument("--base-model", default="unsloth/Qwen3.5-4B-Instruct-bnb-4bit") + parser.add_argument("--base-model", default="unsloth/Qwen3.5-4B") parser.add_argument("--output", default="/data/adapters/psyc-v1") parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--lr", type=float, default=2e-4) @@ -96,9 +95,9 @@ def main() -> None: model=model, tokenizer=tokenizer, train_dataset=dataset, - dataset_text_field="text", - max_seq_length=args.max_seq_length, - args=TrainingArguments( + args=SFTConfig( + dataset_text_field="text", + max_seq_length=args.max_seq_length, per_device_train_batch_size=args.batch_size, gradient_accumulation_steps=args.grad_accum, warmup_steps=5,