#!/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()