| 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() |
|
|
