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:
138
scripts/train_qlora.py
Normal file
138
scripts/train_qlora.py
Normal 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()
|
||||
Reference in New Issue
Block a user