# examples/lora_gpt2.py """Fine-tune real GPT-2 attention with LoRA on Tiny Shakespeare. This experiment adapts every attention fused QKV (c_attn) and output (c_proj) projection. MLPs, embeddings, norms and the tied language-model head stay frozen. It is a worked example, not a reproduction of a benchmark in the LoRA paper. Commands (only train/generate download pretrained weights): python examples/lora_gpt2.py train --device cuda --steps 100 python examples/lora_gpt2.py generate --device cuda --merge Dependencies: see requirements-lora.txt. No PEFT or LabML is used. """ import argparse import hashlib import json import math from pathlib import Path import urllib.request import torch from torch import nn from torch.nn import functional as F from transformers import AutoTokenizer, GPT2LMHeadModel from transformers.pytorch_utils import Conv1D MODEL_ID = "openai-community/gpt2" DATA_URL = ("https://raw.githubusercontent.com/karpathy/char-rnn/" "master/data/tinyshakespeare/input.txt") CACHE_DIR = Path.home() / ".cache" / "abhik-lora-gpt2" ARCHITECTURE_KEYS = ("vocab_size", "n_positions", "n_embd", "n_layer", "n_head") def freeze_model(model): """Freezes the entire pretrained model before adding trainable adapters.""" # region: freeze for parameter in model.parameters(): parameter.requires_grad_(False) # endregion: freeze class LoRAConv1D(nn.Module): """Adds (alpha / rank) * B A to GPT-2's transposed linear storage. Conv1D stores weight as (in_features, out_features), unlike nn.Linear. A has shape (rank, in_features); B has shape (out_features, rank). """ def __init__(self, base, rank, alpha): """Initializes A randomly and B at zero for an unchanged first output.""" super().__init__() in_features, out_features = base.weight.shape if (isinstance(rank, bool) or not isinstance(rank, int) or not 1 <= rank <= min(in_features, out_features)): raise ValueError("rank must be an integer between 1 and min(in, out)") if not math.isfinite(alpha) or alpha <= 0: raise ValueError("alpha must be positive and finite") # region: adapter-init self.base = base freeze_model(self.base) self.rank = rank self.alpha = float(alpha) self.scale = self.alpha / rank self.A = nn.Parameter(base.weight.new_empty(rank, in_features)) self.B = nn.Parameter(base.weight.new_zeros(out_features, rank)) nn.init.normal_(self.A, std=0.01) self.merged = False # endregion: adapter-init # region: adapter-forward def forward(self, inputs): """Adds the low-rank branch until its weights have been merged.""" result = self.base(inputs) if not self.merged: result = result + F.linear(F.linear(inputs, self.A), self.B) * self.scale return result # endregion: adapter-forward # region: merge @torch.no_grad() def merge(self): """Merges once for inference, transposing BA into Conv1D storage.""" if not self.merged: self.base.weight.add_((self.B @ self.A).T * self.scale) self.merged = True self.A.requires_grad_(False) self.B.requires_grad_(False) # endregion: merge def adapter_modules(model): """Returns installed adapters keyed by reconstructible module paths.""" return {name: module for name, module in model.named_modules() if isinstance(module, LoRAConv1D)} def install_lora(model, rank=8, alpha=16): """Adapts only GPT-2 attention QKV and output projections in every block.""" if adapter_modules(model): raise ValueError("model already has adapters") # region: install freeze_model(model) for block in model.transformer.h: for name in ("c_attn", "c_proj"): projection = getattr(block.attn, name) if not isinstance(projection, Conv1D): raise ValueError("expected GPT-2 attention Conv1D projections") setattr(block.attn, name, LoRAConv1D(projection, rank, alpha)) # endregion: install return adapter_modules(model) # region: loss def language_model_loss(model, input_ids): """Lets GPT-2 shift labels internally: position t predicts token t + 1.""" return model(input_ids=input_ids, labels=input_ids, use_cache=False).loss # endregion: loss def architecture(model): """Returns dimensions needed to reject an incompatible base model.""" return {key: getattr(model.config, key) for key in ARCHITECTURE_KEYS} def save_adapter(model, directory, model_id, revision, training=None): """Saves only adapters and JSON metadata; never duplicates base weights.""" modules = adapter_modules(model) if not modules or any(module.merged for module in modules.values()): raise ValueError("save requires installed, unmerged adapters") first = next(iter(modules.values())) metadata = { "format_version": 1, "model_id": model_id, "revision": revision, "architecture": architecture(model), "rank": first.rank, "alpha": first.alpha, "targets": list(modules), "training": training or {}, } # region: save state = {f"{name}.{key}": getattr(module, key).detach().cpu() for name, module in modules.items() for key in ("A", "B")} directory = Path(directory) directory.mkdir(parents=True, exist_ok=True) torch.save(state, directory / "adapter.pt") (directory / "metadata.json").write_text( json.dumps(metadata, indent=2) + "\n", encoding="utf-8") # endregion: save def load_adapter(model, directory, model_id=None, revision=None): """Validates identity, dimensions and tensors before installing adapters.""" directory = Path(directory) metadata = json.loads((directory / "metadata.json").read_text(encoding="utf-8")) if metadata.get("format_version") != 1: raise ValueError("unsupported adapter format") if metadata.get("architecture") != architecture(model): raise ValueError("adapter architecture does not match the base model") for key, expected in (("model_id", model_id), ("revision", revision)): if expected is not None and metadata.get(key) != expected: raise ValueError(f"adapter {key} does not match the base model") targets = [f"transformer.h.{index}.attn.{name}" for index in range(model.config.n_layer) for name in ("c_attn", "c_proj")] if metadata.get("targets") != targets or adapter_modules(model): raise ValueError("adapter targets do not match an unmodified GPT-2") state = torch.load(directory / "adapter.pt", map_location="cpu", weights_only=True) expected_keys = {f"{name}.{key}" for name in targets for key in ("A", "B")} if set(state) != expected_keys: raise ValueError("adapter tensor keys do not match targets") rank, alpha = metadata["rank"], metadata["alpha"] for name in targets: in_features, out_features = model.get_submodule(name).weight.shape if (isinstance(rank, bool) or not isinstance(rank, int) or not 1 <= rank <= min(in_features, out_features) or not math.isfinite(alpha) or alpha <= 0): raise ValueError("invalid adapter rank or alpha") for key, shape in (("A", (rank, in_features)), ("B", (out_features, rank))): tensor = state[f"{name}.{key}"] if (not isinstance(tensor, torch.Tensor) or tensor.shape != shape or not tensor.is_floating_point() or not torch.isfinite(tensor).all()): raise ValueError(f"invalid adapter tensor: {name}.{key}") modules = install_lora(model, rank, alpha) with torch.no_grad(): for name, module in modules.items(): for key in ("A", "B"): getattr(module, key).copy_(state[f"{name}.{key}"]) return metadata @torch.no_grad() def evaluate(model, tokens, seq_len, batch_size, batches, device): """Computes next-token loss on fixed, nonoverlapping held-out blocks.""" was_training = model.training model.eval() block_count = min(len(tokens) // seq_len, batches * batch_size) if block_count < 1: raise ValueError("held-out split is too short for evaluation") blocks = tokens[:block_count * seq_len].reshape(block_count, seq_len) total = 0.0 for batch in blocks.split(batch_size): total += language_model_loss(model, batch.to(device)).item() * len(batch) model.train(was_training) return total / block_count def split_tokens(tokens, seq_len): """Splits contiguous tokens 90/10, preventing training/evaluation overlap.""" split = int(len(tokens) * .9) if seq_len < 2 or min(split, len(tokens) - split) < seq_len: raise ValueError("each split must hold at least one sequence of 2+ tokens") return tokens[:split], tokens[split:] def sample_batch(tokens, seq_len, batch_size, generator): """Samples complete training windows using a dedicated seeded CPU stream.""" starts = torch.randint(len(tokens) - seq_len + 1, (batch_size,), generator=generator) return torch.stack([tokens[start:start + seq_len] for start in starts.tolist()]) @torch.no_grad() def generate_tokens(model, input_ids, max_new_tokens=40): """Generates deterministic greedy continuations with dropout disabled.""" model.eval() if input_ids.shape[-1] + max_new_tokens > model.config.n_positions: raise ValueError("prompt plus generated tokens exceeds model context length") return model.generate( input_ids, attention_mask=torch.ones_like(input_ids), do_sample=False, max_new_tokens=max_new_tokens, pad_token_id=model.config.eos_token_id, ) @torch.no_grad() def verify_merge(model, input_ids): """Checks inference logits before/after merging; returns max absolute error.""" model.eval() before = model(input_ids, use_cache=False).logits for module in adapter_modules(model).values(): module.merge() after = model(input_ids, use_cache=False).logits torch.testing.assert_close(after, before, rtol=1e-4, atol=1e-4) return (after - before).abs().max().item() def load_pretrained(model_id, revision, cache_dir, device): """Downloads/caches GPT-2 and its tokenizer only when a CLI command runs.""" tokenizer = AutoTokenizer.from_pretrained( model_id, revision=revision, cache_dir=str(cache_dir / "huggingface")) model = GPT2LMHeadModel.from_pretrained( model_id, revision=revision, cache_dir=str(cache_dir / "huggingface"), attn_implementation="eager", ).to(device) # A resolved immutable Hugging Face commit identifies the frozen base weights. resolved_revision = model.config._commit_hash if not resolved_revision: raise ValueError("pretrained model did not resolve to an immutable revision") return model, tokenizer, resolved_revision def load_data(tokenizer, cache_dir): """Caches the corpus and returns GPT-2 tokens plus source provenance.""" path = cache_dir / "tiny-shakespeare.txt" path.parent.mkdir(parents=True, exist_ok=True) if not path.exists(): with urllib.request.urlopen(DATA_URL, timeout=60) as response: path.write_bytes(response.read()) raw = path.read_bytes() # Suppress the tokenizer's single-sequence context warning: we window later. ids = tokenizer.encode(raw.decode("utf-8"), add_special_tokens=False, verbose=False) return torch.tensor(ids, dtype=torch.long), { "url": DATA_URL, "sha256": hashlib.sha256(raw).hexdigest(), "tokens": len(ids), "split": "contiguous first 90% train, final 10% held out", } def train(args): """Runs adapter training and reports honest held-out and generation results.""" torch.manual_seed(args.seed) generator = torch.Generator().manual_seed(args.seed) model, tokenizer, revision = load_pretrained( args.model, args.revision, args.cache_dir, args.device) if args.seq_len > model.config.n_positions: raise ValueError("sequence length exceeds GPT-2 context length") tokens, data_info = load_data(tokenizer, args.cache_dir) train_tokens, held_out = split_tokens(tokens, args.seq_len) install_lora(model, args.rank, args.alpha) parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] print(json.dumps({"base_model": args.model, "revision": revision, "trainable_parameters": sum(p.numel() for p in parameters), "total_parameters": sum(p.numel() for p in model.parameters()), "data": data_info}), flush=True) prompt = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(args.device) baseline_loss = evaluate(model, held_out, args.seq_len, args.batch_size, args.eval_batches, args.device) print(json.dumps({"stage": "before", "held_out_loss": baseline_loss, "generation": tokenizer.decode(generate_tokens( model, prompt, args.max_new_tokens)[0])}), flush=True) optimizer = torch.optim.AdamW(parameters, lr=args.lr) model.train() for step in range(1, args.steps + 1): batch = sample_batch(train_tokens, args.seq_len, args.batch_size, generator).to(args.device) optimizer.zero_grad(set_to_none=True) loss = language_model_loss(model, batch) loss.backward() torch.nn.utils.clip_grad_norm_(parameters, 1.0) optimizer.step() if step == 1 or step % 10 == 0 or step == args.steps: print(json.dumps({"step": step, "train_loss": loss.item()}), flush=True) final_loss = evaluate(model, held_out, args.seq_len, args.batch_size, args.eval_batches, args.device) continuation = tokenizer.decode(generate_tokens(model, prompt, args.max_new_tokens)[0]) training_info = {"seed": args.seed, "steps": args.steps, "seq_len": args.seq_len, "batch_size": args.batch_size, "learning_rate": args.lr, "eval_batches": args.eval_batches, "data": data_info, "before_held_out_loss": baseline_loss, "after_held_out_loss": final_loss} save_adapter(model, args.output, args.model, revision, training_info) difference = verify_merge(model, prompt) print(json.dumps({"stage": "after", "held_out_loss": final_loss, "generation": continuation, "merge_max_abs_error": difference, "adapter": str(args.output)}), flush=True) def generate(args): """Loads the same frozen pretrained revision, then applies saved adapters.""" torch.manual_seed(args.seed) metadata = json.loads((args.adapter / "metadata.json").read_text(encoding="utf-8")) model, tokenizer, revision = load_pretrained( metadata["model_id"], metadata["revision"], args.cache_dir, args.device) load_adapter(model, args.adapter, metadata["model_id"], revision) prompt = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(args.device) if args.merge: print(json.dumps({"merge_max_abs_error": verify_merge(model, prompt)}), flush=True) print(tokenizer.decode(generate_tokens(model, prompt, args.max_new_tokens)[0]), flush=True) def positive_int(value): """Validates positive integer CLI options before downloads can begin.""" result = int(value) if result < 1: raise argparse.ArgumentTypeError("must be a positive integer") return result def positive_float(value): """Validates positive finite numeric CLI options.""" result = float(value) if not math.isfinite(result) or result <= 0: raise argparse.ArgumentTypeError("must be a positive finite number") return result def build_parser(): """Builds train/generate commands with documented, reproducible defaults.""" parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) commands = parser.add_subparsers(dest="command", required=True) train_parser = commands.add_parser("train", help="train adapters on Tiny Shakespeare") generate_parser = commands.add_parser("generate", help="load adapters for inference") for command in (train_parser, generate_parser): command.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") command.add_argument("--seed", type=int, default=42) command.add_argument("--cache-dir", type=Path, default=CACHE_DIR) command.add_argument("--prompt", default="To be, or not to be,") command.add_argument("--max-new-tokens", type=positive_int, default=40) train_parser.add_argument("--model", default=MODEL_ID) train_parser.add_argument("--revision", default="main") train_parser.add_argument("--steps", type=positive_int, default=100) train_parser.add_argument("--seq-len", type=positive_int, default=128) train_parser.add_argument("--batch-size", type=positive_int, default=2) train_parser.add_argument("--eval-batches", type=positive_int, default=8) train_parser.add_argument("--rank", type=positive_int, default=8) train_parser.add_argument("--alpha", type=positive_float, default=16) train_parser.add_argument("--lr", type=positive_float, default=3e-4) train_parser.add_argument("--output", type=Path, default=CACHE_DIR / "adapter") generate_parser.add_argument("--adapter", type=Path, default=CACHE_DIR / "adapter") generate_parser.add_argument("--merge", action="store_true", help="merge once and verify logits") return parser def main(argv=None): """Dispatches commands; importing this file never downloads weights or data.""" parser = build_parser() args = parser.parse_args(argv) if args.command == "train" and args.seq_len < 2: parser.error("--seq-len must be at least 2") try: (train if args.command == "train" else generate)(args) except ValueError as error: parser.error(str(error)) if __name__ == "__main__": main()