import argparse import json import math import random import time from pathlib import Path import torch import torch.nn.functional as F from peft import LoraConfig, get_peft_model from torch import nn from torch.utils.checkpoint import checkpoint from transformers import get_cosine_schedule_with_warmup from predict import LABELS, label_ids, load_model, render, thinking_text LOSS_CHUNK = 256 WARMUP_FRACTION = 0.03 MAX_GRAD_NORM = 1.0 MAX_SKIPPED_IN_A_ROW = 10 def target_modules(model): layers = model.model.layers.named_modules(prefix="model.layers") return [name for name, module in layers if isinstance(module, nn.Linear)] def encode_example(tokenizer, labels, row): context = tokenizer.encode(render(tokenizer, row["prompt"], True), add_special_tokens=False) text = thinking_text(tokenizer, row["prompt"], row["thought"]) + LABELS[row["answer"] - 1] ids = tokenizer.encode(text, add_special_tokens=False) if ids[: len(context)] != context or ids[-1] != labels[row["answer"] - 1]: return None return ids, len(context) def collate(batch, pad_id, device): width = max(len(ids) for ids, _ in batch) input_ids = torch.full((len(batch), width), pad_id) mask = torch.zeros((len(batch), width), dtype=torch.long) targets = torch.full((len(batch), width), -100) for row, (ids, context) in enumerate(batch): start = width - len(ids) input_ids[row, start:] = torch.tensor(ids) mask[row, start:] = 1 targets[row, start + context :] = torch.tensor(ids[context:]) return input_ids.to(device), mask.to(device), targets def chunk_loss(base, hidden, targets): logits = base.lm_head(hidden).float() cap = base.config.final_logit_softcapping if cap is not None: logits = torch.tanh(logits / cap) * cap return F.cross_entropy(logits, targets.to(logits.device), reduction="sum") def loss_sum(model, input_ids, mask, targets): base = model.get_base_model() position_ids = (mask.cumsum(-1) - 1).clamp(min=0) hidden = base.model( input_ids=input_ids, attention_mask=mask, position_ids=position_ids, use_cache=False ).last_hidden_state shifted = targets[:, 1:].to(hidden.device) select = shifted != -100 hidden, shifted = hidden[:, :-1][select], shifted[select] chunks = range(0, len(shifted), LOSS_CHUNK) return sum( checkpoint(chunk_loss, base, hidden[i : i + LOSS_CHUNK], shifted[i : i + LOSS_CHUNK], use_reentrant=False) for i in chunks ) def peak_memory(): if not torch.cuda.is_available(): return [] return [round(torch.cuda.max_memory_allocated(i) / 2**30, 2) for i in range(torch.cuda.device_count())] def main(): parser = argparse.ArgumentParser() parser.add_argument("--model", default="google/gemma-4-E4B-it") parser.add_argument("--revision", help="model commit to load; the latest one when not given") parser.add_argument("--data", default="star/star.jsonl") parser.add_argument("--out_dir", default="adapter") parser.add_argument("--epochs", type=int, default=2) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--grad_accum", type=int, default=4) parser.add_argument("--max_length", type=int, default=1024) parser.add_argument("--rank", type=int, default=16) parser.add_argument("--alpha", type=int, default=32) parser.add_argument("--dropout", type=float, default=0.05) parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() tokenizer, model = load_model(args.model, revision=args.revision) base_params = sum(p.numel() for p in model.parameters()) labels = label_ids(tokenizer) rows = [json.loads(line) for line in Path(args.data).read_text(encoding="utf-8").splitlines()] encoded = [encode_example(tokenizer, labels, row) for row in rows] misaligned = sum(e is None for e in encoded) examples = [e for e in encoded if e is not None and len(e[0]) <= args.max_length] too_long = len(encoded) - misaligned - len(examples) print(f"Training on {len(examples)} of {len(rows)} examples: {too_long} longer than {args.max_length} tokens" f" and {misaligned} that do not tokenize as context + target were dropped") if not examples: raise SystemExit("error: no training examples left") dtype = model.dtype torch.manual_seed(args.seed) modules = target_modules(model) config = LoraConfig( r=args.rank, lora_alpha=args.alpha, lora_dropout=args.dropout, bias="none", target_modules=modules ) model = get_peft_model(model, config) model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model.enable_input_require_grads() model.train() params = [p for p in model.parameters() if p.requires_grad] adapter_params = sum(p.numel() for p in params) print(f"LoRA r={args.rank} on {len(modules)} linear modules: {adapter_params:,} adapter parameters," f" {base_params + adapter_params:,} in total") longest = sorted(examples, key=lambda e: -len(e[0]))[: args.batch_size] loss_sum(model, *collate(longest, tokenizer.pad_token_id, model.device)).backward() model.zero_grad(set_to_none=True) print(f"Longest micro-batch ({len(longest[0][0])} tokens) fits: peak GPU memory {peak_memory()} GB") step_size = args.batch_size * args.grad_accum total_steps = args.epochs * math.ceil(len(examples) / step_size) optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0) scheduler = get_cosine_schedule_with_warmup(optimizer, round(WARMUP_FRACTION * total_steps), total_steps) scaler = torch.amp.GradScaler(model.device.type, init_scale=2.0**10, enabled=dtype == torch.float16) rng = random.Random(args.seed) steps, tokens, skipped_in_a_row = [], 0, 0 start_time = time.time() for epoch in range(args.epochs): order = examples[:] rng.shuffle(order) for start in range(0, len(order), step_size): step_examples = sorted(order[start : start + step_size], key=lambda e: len(e[0])) n_targets = sum(len(ids) - context for ids, context in step_examples) step_loss = 0.0 for i in range(0, len(step_examples), args.batch_size): micro = step_examples[i : i + args.batch_size] input_ids, mask, targets = collate(micro, tokenizer.pad_token_id, model.device) loss = loss_sum(model, input_ids, mask, targets) if not math.isfinite(loss.item()): raise RuntimeError(f"non-finite loss {loss.item()} at step {len(steps) + 1}") scaler.scale(loss / n_targets).backward() step_loss += loss.item() tokens += int(mask.sum()) scaler.unscale_(optimizer) grad_norm = float(nn.utils.clip_grad_norm_(params, MAX_GRAD_NORM)) scale = scaler.get_scale() scaler.step(optimizer) scaler.update() skipped = scaler.get_scale() < scale if not math.isfinite(grad_norm) and not skipped: raise RuntimeError(f"non-finite gradients at step {len(steps) + 1}") skipped_in_a_row = skipped_in_a_row + 1 if skipped else 0 if skipped_in_a_row == MAX_SKIPPED_IN_A_ROW: raise RuntimeError(f"{MAX_SKIPPED_IN_A_ROW} steps in a row overflowed in fp16") row = { "step": len(steps) + 1, "epoch": epoch, "loss": step_loss / n_targets, "lr": scheduler.get_last_lr()[0], "grad_norm": grad_norm, "skipped": skipped, "seconds": round(time.time() - start_time, 1), } scheduler.step() optimizer.zero_grad(set_to_none=True) steps.append(row) print(json.dumps(row), flush=True) seconds = time.time() - start_time out_dir = Path(args.out_dir) model.save_pretrained(out_dir) log = { "config": vars(args), "dtype": str(dtype), "grad_scaling": scaler.is_enabled(), "target_modules": len(modules), "base_params": base_params, "adapter_params": adapter_params, "total_params": base_params + adapter_params, "examples": len(rows), "trained_examples": len(examples), "dropped_too_long": too_long, "dropped_misaligned": misaligned, "skipped_steps": sum(row["skipped"] for row in steps), "seconds": round(seconds, 1), "tokens_per_s": round(tokens / seconds, 1), "peak_memory_gb": peak_memory(), "steps": steps, } (out_dir / "train_log.json").write_text(json.dumps(log, indent=1) + "\n", encoding="utf-8") print(f"Saved the adapter to {out_dir}: final loss {steps[-1]['loss']:.4f}, {log['tokens_per_s']} tokens/s") if __name__ == "__main__": main()