stage-3c: unsloth QLoRA training scaffold for Qwen3.5

Dockerfile.train builds a CUDA 12.4 + unsloth container that consumes the
Trainline JSONL datasets and emits a LoRA adapter at data/adapters/<run>/final.
Defaults target a 24 GB GPU (Qwen3.5-4B-Instruct-bnb-4bit, r=16, bf16, 3 epochs,
effective batch 8). README documents the build + run workflow.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
m17hr1l
2026-05-14 14:17:14 +02:00
parent b8ea4ead02
commit f1ab11f89d
4 changed files with 237 additions and 3 deletions

138
scripts/train_qlora.py Normal file
View File

@@ -0,0 +1,138 @@
"""Train a psyc QLoRA adapter on JSONL Trainline datasets using unsloth + Qwen3.5.
Run inside the psyc training container:
docker run --gpus all -v $(pwd)/data:/data psyc-trainer \
--dataset /data/datasets/ioc_extraction-v1.jsonl \
--dataset /data/datasets/severity_classification-v1.jsonl \
--output /data/adapters/psyc-v1
Defaults target a 24 GB consumer GPU (3090/4090) with Qwen3.5-4B-Instruct at
4-bit + LoRA r=16. For an A100-40/80 bump --base-model to 9B and raise
--batch-size + --max-seq-length.
"""
from __future__ import annotations
# unsloth must be imported BEFORE transformers per their setup notes.
from unsloth import FastLanguageModel # noqa: I001
import argparse
import json
from pathlib import Path
from typing import Dict, List
from datasets import Dataset
from transformers import TrainingArguments
from trl import SFTTrainer
def load_examples(paths: List[Path]) -> List[Dict[str, str]]:
out: List[Dict[str, str]] = []
for p in paths:
with p.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if not line:
continue
ex = json.loads(line)
if not all(k in ex for k in ("instruction", "input", "output")):
continue
out.append(ex)
return out
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("--output", default="/data/adapters/psyc-v1")
parser.add_argument("--epochs", type=int, default=3)
parser.add_argument("--lr", type=float, default=2e-4)
parser.add_argument("--max-seq-length", type=int, default=4096)
parser.add_argument("--batch-size", type=int, default=2)
parser.add_argument("--grad-accum", type=int, default=4)
parser.add_argument("--lora-r", type=int, default=16)
parser.add_argument("--lora-alpha", type=int, default=16)
parser.add_argument("--seed", type=int, default=3407)
args = parser.parse_args()
paths = [Path(p) for p in args.dataset]
examples = load_examples(paths)
if not examples:
raise SystemExit("no examples loaded — check --dataset paths")
print(f"[psyc-train] loaded {len(examples)} example(s) from {len(paths)} dataset(s)")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=args.base_model,
max_seq_length=args.max_seq_length,
dtype=None,
load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
model,
r=args.lora_r,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
lora_alpha=args.lora_alpha,
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=args.seed,
)
def format_one(ex: Dict[str, str]) -> Dict[str, str]:
messages = [
{"role": "user", "content": f"{ex['instruction']}\n\n{ex['input']}"},
{"role": "assistant", "content": ex["output"]},
]
return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}
dataset = Dataset.from_list([format_one(e) for e in examples]).shuffle(seed=args.seed)
output_dir = Path(args.output)
output_dir.mkdir(parents=True, exist_ok=True)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=args.max_seq_length,
args=TrainingArguments(
per_device_train_batch_size=args.batch_size,
gradient_accumulation_steps=args.grad_accum,
warmup_steps=5,
num_train_epochs=args.epochs,
learning_rate=args.lr,
bf16=True,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=args.seed,
output_dir=str(output_dir / "checkpoints"),
save_strategy="epoch",
logging_steps=10,
report_to="none",
),
)
trainer.train()
final_dir = output_dir / "final"
final_dir.mkdir(parents=True, exist_ok=True)
model.save_pretrained(str(final_dir))
tokenizer.save_pretrained(str(final_dir))
(output_dir / "training_meta.json").write_text(json.dumps({
"base_model": args.base_model,
"lora_r": args.lora_r,
"lora_alpha": args.lora_alpha,
"epochs": args.epochs,
"lr": args.lr,
"datasets": [str(p) for p in paths],
"examples": len(examples),
"seed": args.seed,
}, indent=2))
print(f"[psyc-train] adapter saved → {final_dir}")
if __name__ == "__main__":
main()