209 lines
7.8 KiB
Python
209 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Production QLoRA training for BookBuilder bookwriter v1.
|
|
Same recipe as smoke_qlora.py, scaled up via HF Trainer.
|
|
Default target: Qwen 2.5 7B, 1.82B-token corpus, 1 epoch, ctx 2048.
|
|
"""
|
|
import os, sys, json, argparse, time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
from transformers import (
|
|
AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,
|
|
Trainer, TrainingArguments, DataCollatorForLanguageModeling,
|
|
)
|
|
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--base_model", default="Qwen/Qwen2.5-7B")
|
|
p.add_argument("--train_jsonl", required=True)
|
|
p.add_argument("--eval_jsonl", default=None)
|
|
p.add_argument("--output_dir", required=True)
|
|
p.add_argument("--max_seq", type=int, default=2048)
|
|
p.add_argument("--per_device_train_batch_size", type=int, default=8)
|
|
p.add_argument("--per_device_eval_batch_size", type=int, default=8)
|
|
p.add_argument("--grad_accum", type=int, default=4)
|
|
p.add_argument("--lr", type=float, default=2e-4)
|
|
p.add_argument("--num_epochs", type=float, default=1.0)
|
|
p.add_argument("--warmup_ratio", type=float, default=0.03)
|
|
p.add_argument("--lora_r", type=int, default=16)
|
|
p.add_argument("--lora_alpha", type=int, default=32)
|
|
p.add_argument("--lora_dropout", type=float, default=0.05)
|
|
p.add_argument("--save_steps", type=int, default=500)
|
|
p.add_argument("--eval_steps", type=int, default=500)
|
|
p.add_argument("--logging_steps", type=int, default=20)
|
|
p.add_argument("--save_total_limit", type=int, default=3)
|
|
p.add_argument("--hub_repo_id", default=None,
|
|
help="If set, push checkpoints to this HF repo each save.")
|
|
p.add_argument("--hub_token", default=None)
|
|
p.add_argument("--max_train_samples", type=int, default=None,
|
|
help="Cap for sanity runs.")
|
|
p.add_argument("--seed", type=int, default=42)
|
|
p.add_argument("--gradient_checkpointing", action="store_true", default=False)
|
|
p.add_argument("--no_eval", action="store_true")
|
|
return p.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
print(f"=== BookBuilder QLoRA train ===", flush=True)
|
|
print(f"base: {args.base_model}")
|
|
print(f"train: {args.train_jsonl}")
|
|
print(f"eval : {args.eval_jsonl}")
|
|
print(f"out : {args.output_dir}")
|
|
print(f"ctx : {args.max_seq} bs/dev: {args.per_device_train_batch_size} ga: {args.grad_accum} lr: {args.lr}", flush=True)
|
|
|
|
# ---------- tokenizer ----------
|
|
tok = AutoTokenizer.from_pretrained(args.base_model, use_fast=True)
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
tok.padding_side = "right"
|
|
|
|
# ---------- model: 4-bit NF4 ----------
|
|
bnb = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
bnb_4bit_use_double_quant=True,
|
|
)
|
|
print("Loading 4-bit base...", flush=True)
|
|
t0 = time.time()
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
args.base_model,
|
|
quantization_config=bnb,
|
|
torch_dtype=torch.bfloat16,
|
|
attn_implementation="sdpa",
|
|
device_map="cuda:0",
|
|
)
|
|
model.config._attn_implementation = "sdpa"
|
|
model.config.use_cache = False
|
|
print(f" base loaded in {time.time()-t0:.1f}s | VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB", flush=True)
|
|
|
|
model = prepare_model_for_kbit_training(
|
|
model, use_gradient_checkpointing=args.gradient_checkpointing,
|
|
)
|
|
lora = LoraConfig(
|
|
r=args.lora_r,
|
|
lora_alpha=args.lora_alpha,
|
|
lora_dropout=args.lora_dropout,
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
|
|
)
|
|
model = get_peft_model(model, lora)
|
|
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
total = sum(p.numel() for p in model.parameters())
|
|
print(f" trainable: {trainable/1e6:.2f}M / {total/1e6:.1f}M ({100*trainable/total:.3f}%)", flush=True)
|
|
|
|
# ---------- dataset ----------
|
|
data_files = {"train": args.train_jsonl}
|
|
if args.eval_jsonl and not args.no_eval:
|
|
data_files["eval"] = args.eval_jsonl
|
|
ds = load_dataset("json", data_files=data_files)
|
|
if args.max_train_samples:
|
|
ds["train"] = ds["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
|
|
print(f" train rows: {len(ds['train']):,}", flush=True)
|
|
if "eval" in ds:
|
|
print(f" eval rows : {len(ds['eval']):,}", flush=True)
|
|
|
|
def tokenize_fn(batch):
|
|
out = tok(
|
|
batch["text"],
|
|
truncation=True,
|
|
max_length=args.max_seq,
|
|
padding=False,
|
|
)
|
|
return out
|
|
|
|
# Explicit cache paths — HF datasets' default fingerprint isn't stable across
|
|
# Python processes (tokenizer object hash differs), so a plain .map() re-tokenizes
|
|
# every fresh launch. Pin the cache to a deterministic path so subsequent runs hit it.
|
|
cache_root = os.path.dirname(os.path.abspath(args.train_jsonl))
|
|
tag = f"qwen-ctx{args.max_seq}"
|
|
cache_files = {
|
|
split: os.path.join(cache_root, f"tok_cache_{split}_{tag}.arrow")
|
|
for split in ds.keys()
|
|
}
|
|
for split, path in cache_files.items():
|
|
print(f" tokenize cache[{split}] -> {path}", flush=True)
|
|
tokenized = ds.map(
|
|
tokenize_fn,
|
|
batched=True,
|
|
remove_columns=ds["train"].column_names,
|
|
num_proc=8,
|
|
desc="tokenize",
|
|
cache_file_names=cache_files,
|
|
load_from_cache_file=True,
|
|
)
|
|
|
|
collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)
|
|
|
|
# ---------- trainer ----------
|
|
push_to_hub = bool(args.hub_repo_id)
|
|
targs = TrainingArguments(
|
|
output_dir=args.output_dir,
|
|
overwrite_output_dir=False,
|
|
seed=args.seed,
|
|
num_train_epochs=args.num_epochs,
|
|
per_device_train_batch_size=args.per_device_train_batch_size,
|
|
per_device_eval_batch_size=args.per_device_eval_batch_size,
|
|
gradient_accumulation_steps=args.grad_accum,
|
|
learning_rate=args.lr,
|
|
warmup_ratio=args.warmup_ratio,
|
|
lr_scheduler_type="cosine",
|
|
bf16=True,
|
|
fp16=False,
|
|
tf32=True,
|
|
gradient_checkpointing=args.gradient_checkpointing,
|
|
logging_steps=args.logging_steps,
|
|
save_steps=args.save_steps,
|
|
eval_strategy="steps" if "eval" in tokenized else "no",
|
|
eval_steps=args.eval_steps if "eval" in tokenized else None,
|
|
save_total_limit=args.save_total_limit,
|
|
report_to="tensorboard",
|
|
push_to_hub=push_to_hub,
|
|
hub_model_id=args.hub_repo_id if push_to_hub else None,
|
|
hub_token=args.hub_token,
|
|
hub_strategy="every_save" if push_to_hub else "end",
|
|
hub_private_repo=False,
|
|
dataloader_num_workers=4,
|
|
dataloader_pin_memory=True,
|
|
optim="paged_adamw_8bit",
|
|
remove_unused_columns=False,
|
|
)
|
|
|
|
trainer_kwargs = dict(
|
|
model=model,
|
|
args=targs,
|
|
train_dataset=tokenized["train"],
|
|
eval_dataset=tokenized.get("eval"),
|
|
data_collator=collator,
|
|
)
|
|
# transformers 4.46 uses `tokenizer=`; 5.x uses `processing_class=`
|
|
import inspect
|
|
if "processing_class" in inspect.signature(Trainer.__init__).parameters:
|
|
trainer_kwargs["processing_class"] = tok
|
|
else:
|
|
trainer_kwargs["tokenizer"] = tok
|
|
trainer = Trainer(**trainer_kwargs)
|
|
|
|
print(f"=== begin training ===", flush=True)
|
|
trainer.train()
|
|
|
|
print(f"=== saving final adapter ===", flush=True)
|
|
trainer.save_model(args.output_dir)
|
|
tok.save_pretrained(args.output_dir)
|
|
if push_to_hub:
|
|
trainer.push_to_hub(commit_message="final adapter")
|
|
|
|
# Sentinel
|
|
Path(args.output_dir).joinpath(".train_complete").touch()
|
|
print("DONE.", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|