code/repro_experiment.py

#!/usr/bin/env python3
"""Clean-room scaled reproduction of BP-checkpointing vs FmAD vs ZO.

The experiment uses the paper's AG News task but a small, from-scratch residual
text model with frozen random features and trainable rank-1 adapters.  This makes
all three gradient estimators practical on a single 24 GB GPU while preserving
the paper's core comparison: exact reverse-mode gradients with activation
checkpointing, one-direction forward-mode AD, and one-direction two-sided ZO.
"""

from __future__ import annotations

import argparse
import collections
import csv
import gc
import json
import math
import os
import random
import re
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.func import functional_call, jvp
from torch.utils.checkpoint import checkpoint


TOKEN_RE = re.compile(r"[A-Za-z0-9']+")


def seed_everything(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def tokenize(text: str) -> list[str]:
    return TOKEN_RE.findall(text.lower())


def load_ag_news(train_size: int, eval_size: int, vocab_size: int, seq_len: int, seed: int):
    from datasets import load_dataset

    dataset = load_dataset("fancyzhx/ag_news")
    train_rows = dataset["train"].shuffle(seed=seed).select(range(train_size))
    eval_rows = dataset["test"].shuffle(seed=seed).select(range(eval_size))

    counts: collections.Counter[str] = collections.Counter()
    for text in train_rows["text"]:
        counts.update(tokenize(text))
    vocab = {word: idx + 2 for idx, (word, _) in enumerate(counts.most_common(vocab_size - 2))}

    def encode(texts):
        ids = torch.zeros((len(texts), seq_len), dtype=torch.long)
        mask = torch.zeros((len(texts), seq_len), dtype=torch.float32)
        for row, text in enumerate(texts):
            tokens = tokenize(text)[:seq_len]
            if not tokens:
                tokens = ["<unk>"]
            token_ids = [vocab.get(token, 1) for token in tokens]
            ids[row, : len(token_ids)] = torch.tensor(token_ids)
            mask[row, : len(token_ids)] = 1.0
        return ids, mask

    train_x, train_mask = encode(train_rows["text"])
    eval_x, eval_mask = encode(eval_rows["text"])
    return (
        train_x,
        train_mask,
        torch.tensor(train_rows["label"], dtype=torch.long),
        eval_x,
        eval_mask,
        torch.tensor(eval_rows["label"], dtype=torch.long),
        len(vocab) + 2,
    )


def load_synthetic(train_size: int, eval_size: int, vocab_size: int, seq_len: int, seed: int):
    generator = torch.Generator().manual_seed(seed)

    def make(n):
        y = torch.randint(0, 4, (n,), generator=generator)
        x = torch.randint(6, vocab_size, (n, seq_len), generator=generator)
        x[:, :4] = y[:, None] + 2
        mask = torch.ones((n, seq_len), dtype=torch.float32)
        return x, mask, y

    return (*make(train_size), *make(eval_size), vocab_size)


class AdapterBlock(nn.Module):
    def __init__(self, width: int, rank: int, generator: torch.Generator):
        super().__init__()
        base = torch.empty(width, width).normal_(generator=generator)
        base = torch.linalg.qr(base).Q
        self.register_buffer("base", base)
        self.lora_a = nn.Parameter(torch.empty(rank, width).normal_(std=0.02, generator=generator))
        self.lora_b = nn.Parameter(torch.empty(width, rank).normal_(std=0.02, generator=generator))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base = F.linear(x, self.base)
        adapter = F.linear(F.linear(x, self.lora_a), self.lora_b)
        return x + 0.25 * F.gelu(base + adapter)


class AdapterTextModel(nn.Module):
    def __init__(self, vocab_size: int, width: int, depth: int, rank: int, seed: int):
        super().__init__()
        generator = torch.Generator().manual_seed(seed)
        self.embedding = nn.Embedding(vocab_size, width, padding_idx=0)
        with torch.no_grad():
            self.embedding.weight.normal_(std=1.0 / math.sqrt(width), generator=generator)
            self.embedding.weight[0].zero_()
        self.embedding.weight.requires_grad_(False)
        self.blocks = nn.ModuleList([AdapterBlock(width, rank, generator) for _ in range(depth)])
        self.classifier = nn.Linear(width, 4)

    def forward(self, token_ids: torch.Tensor, mask: torch.Tensor, use_checkpoint: bool = False):
        hidden = self.embedding(token_ids)
        for block in self.blocks:
            if use_checkpoint and self.training:
                hidden = checkpoint(block, hidden, use_reentrant=False)
            else:
                hidden = block(hidden)
        pooled = (hidden * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp_min(1)
        return self.classifier(pooled)


def trainable_parameters(model: nn.Module):
    return {name: param for name, param in model.named_parameters() if param.requires_grad}


@torch.inference_mode()
def evaluate(model, x, mask, y, batch_size=256):
    model.eval()
    correct = 0
    loss_sum = 0.0
    for start in range(0, len(y), batch_size):
        stop = min(start + batch_size, len(y))
        logits = model(x[start:stop], mask[start:stop], use_checkpoint=False)
        loss_sum += F.cross_entropy(logits, y[start:stop], reduction="sum").item()
        correct += (logits.argmax(dim=-1) == y[start:stop]).sum().item()
    model.train()
    return loss_sum / len(y), correct / len(y)


def sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize()


def make_batches(n: int, batch_size: int, steps: int, seed: int):
    generator = torch.Generator().manual_seed(seed)
    batches = []
    while len(batches) < steps:
        permutation = torch.randperm(n, generator=generator)
        for start in range(0, n - batch_size + 1, batch_size):
            batches.append(permutation[start : start + batch_size])
            if len(batches) == steps:
                break
    return batches


def one_step(model, optimizer, method, x, mask, y, perturbations, zo_eps, grad_clip):
    optimizer.zero_grad(set_to_none=True)
    if method == "bp_checkpoint":
        logits = model(x, mask, use_checkpoint=True)
        loss = F.cross_entropy(logits, y)
        loss.backward()
    else:
        params = trainable_parameters(model)
        estimated_grads = {name: torch.zeros_like(param) for name, param in params.items()}
        losses = []
        for _ in range(perturbations):
            directions = {name: torch.randn_like(param) for name, param in params.items()}
            if method == "fmad":
                def loss_fn(candidate_params):
                    logits = functional_call(
                        model,
                        candidate_params,
                        (x, mask),
                        {"use_checkpoint": False},
                        strict=False,
                    )
                    return F.cross_entropy(logits, y)

                primal_loss, directional_derivative = jvp(loss_fn, (params,), (directions,))
                losses.append(primal_loss.detach())
                scalar = directional_derivative.detach()
            elif method == "zo":
                positive = {name: param + zo_eps * directions[name] for name, param in params.items()}
                negative = {name: param - zo_eps * directions[name] for name, param in params.items()}
                with torch.no_grad():
                    positive_loss = F.cross_entropy(
                        functional_call(model, positive, (x, mask), {"use_checkpoint": False}, strict=False), y
                    )
                    negative_loss = F.cross_entropy(
                        functional_call(model, negative, (x, mask), {"use_checkpoint": False}, strict=False), y
                    )
                losses.append((positive_loss + negative_loss) / 2)
                scalar = (positive_loss - negative_loss) / (2 * zo_eps)
            else:
                raise ValueError(method)
            for name in estimated_grads:
                estimated_grads[name].add_(directions[name], alpha=float(scalar) / perturbations)
        for name, param in params.items():
            param.grad = estimated_grads[name]
        loss = torch.stack(losses).mean()

    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    optimizer.step()
    return float(loss.detach()), float(grad_norm.detach())


def run_method(method, args, data, device):
    train_x, train_mask, train_y, eval_x, eval_mask, eval_y, vocab_size = data
    seed_everything(args.seed)
    model = AdapterTextModel(vocab_size, args.width, args.depth, args.rank, args.seed).to(device)
    optimizer = torch.optim.AdamW(
        [parameter for parameter in model.parameters() if parameter.requires_grad],
        lr=getattr(args, f"lr_{'bp' if method == 'bp_checkpoint' else method}"),
        weight_decay=0.0,
    )
    batch_override = getattr(args, f"batch_size_{'bp' if method == 'bp_checkpoint' else method}")
    method_batch_size = batch_override if batch_override is not None else args.batch_size
    batches = make_batches(len(train_y), method_batch_size, args.steps, args.seed + 17)
    eval_x = eval_x.to(device)
    eval_mask = eval_mask.to(device)
    eval_y = eval_y.to(device)
    model.train()

    initial_loss, initial_accuracy = evaluate(model, eval_x, eval_mask, eval_y)
    rows = [{
        "method": method,
        "step": 0,
        "train_loss": None,
        "val_loss": initial_loss,
        "val_accuracy": initial_accuracy,
        "elapsed_train_s": 0.0,
        "peak_memory_mb": None,
        "step_flops": None,
        "grad_norm": None,
    }]
    print(json.dumps({"event": "eval", **rows[-1]}), flush=True)

    train_elapsed = 0.0
    peak_memory_mb = 0.0
    step_flops = None
    first_threshold_step = None
    first_threshold_time = None
    last_loss = None
    last_grad_norm = None
    for step, indices in enumerate(batches, start=1):
        bx = train_x[indices].to(device)
        bm = train_mask[indices].to(device)
        by = train_y[indices].to(device)
        if device.type == "cuda":
            torch.cuda.reset_peak_memory_stats(device)
        sync(device)
        start = time.perf_counter()
        if step == 1:
            try:
                from torch.utils.flop_counter import FlopCounterMode
                with FlopCounterMode(display=False) as flop_counter:
                    last_loss, last_grad_norm = one_step(
                        model, optimizer, method, bx, bm, by,
                        args.perturbations, args.zo_eps, args.grad_clip,
                    )
                step_flops = int(flop_counter.get_total_flops())
            except Exception as error:
                print(json.dumps({"event": "flop_counter_warning", "method": method, "error": repr(error)}), flush=True)
                last_loss, last_grad_norm = one_step(
                    model, optimizer, method, bx, bm, by,
                    args.perturbations, args.zo_eps, args.grad_clip,
                )
        else:
            last_loss, last_grad_norm = one_step(
                model, optimizer, method, bx, bm, by,
                args.perturbations, args.zo_eps, args.grad_clip,
            )
        sync(device)
        train_elapsed += time.perf_counter() - start
        if device.type == "cuda":
            peak_memory_mb = max(peak_memory_mb, torch.cuda.max_memory_allocated(device) / (1024 ** 2))

        if step % args.eval_every == 0 or step == args.steps:
            val_loss, val_accuracy = evaluate(model, eval_x, eval_mask, eval_y)
            row = {
                "method": method,
                "step": step,
                "train_loss": last_loss,
                "val_loss": val_loss,
                "val_accuracy": val_accuracy,
                "elapsed_train_s": train_elapsed,
                "peak_memory_mb": peak_memory_mb,
                "step_flops": step_flops,
                "grad_norm": last_grad_norm,
            }
            rows.append(row)
            print(json.dumps({"event": "eval", **row}), flush=True)
            if first_threshold_step is None and val_accuracy >= args.threshold:
                first_threshold_step = step
                first_threshold_time = train_elapsed

    summary = {
        "method": method,
        "batch_size": method_batch_size,
        "initial_accuracy": initial_accuracy,
        "final_accuracy": rows[-1]["val_accuracy"],
        "best_accuracy": max(row["val_accuracy"] for row in rows),
        "total_train_s": train_elapsed,
        "peak_memory_mb": peak_memory_mb,
        "step_flops": step_flops,
        "threshold": args.threshold,
        "threshold_step": first_threshold_step,
        "threshold_train_s": first_threshold_time,
        "threshold_flops": None if first_threshold_step is None or step_flops is None else first_threshold_step * step_flops,
        "threshold_censored": first_threshold_step is None,
    }
    del model, optimizer, eval_x, eval_mask, eval_y
    gc.collect()
    if device.type == "cuda":
        torch.cuda.empty_cache()
    return rows, summary


def write_outputs(output_dir: Path, rows, summaries, args, hardware):
    output_dir.mkdir(parents=True, exist_ok=True)
    with (output_dir / "metrics.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    by_method = {entry["method"]: entry for entry in summaries}
    bp = by_method["bp_checkpoint"]
    comparisons = {}
    for competitor in ("fmad", "zo"):
        other = by_method[competitor]
        comparison = {
            "final_accuracy_gap_percentage_points": 100 * (bp["final_accuracy"] - other["final_accuracy"]),
            "final_accuracy_relative_gain_percent": 100 * (bp["final_accuracy"] / other["final_accuracy"] - 1),
            "peak_memory_ratio_bp_over_other": (
                bp["peak_memory_mb"] / other["peak_memory_mb"]
                if other["peak_memory_mb"] > 0 else None
            ),
            "total_wall_time_ratio_other_over_bp": other["total_train_s"] / bp["total_train_s"],
        }
        if bp["threshold_train_s"] is not None and other["threshold_train_s"] is not None:
            comparison["bp_faster_to_threshold_percent"] = 100 * (
                1 - bp["threshold_train_s"] / other["threshold_train_s"]
            )
            comparison["threshold_compute_ratio_other_over_bp"] = (
                other["threshold_flops"] / bp["threshold_flops"]
            )
        comparisons[f"bp_vs_{competitor}"] = comparison

    payload = {
        "scope": "scaled AG News proxy; not an exact full-paper replication",
        "dataset": "https://huggingface.co/datasets/fancyzhx/ag_news",
        "official_code": "https://github.com/Astuary/Gradient_Estimation_Methods/tree/bc5798add2339d1317664462648167fc6b987671",
        "hardware": hardware,
        "configuration": vars(args),
        "methods": summaries,
        "comparisons": comparisons,
    }
    with (output_dir / "summary.json").open("w") as handle:
        json.dump(payload, handle, indent=2)
    print("FINAL_SUMMARY=" + json.dumps(payload, sort_keys=True), flush=True)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default="results")
    parser.add_argument("--train-size", type=int, default=8192)
    parser.add_argument("--eval-size", type=int, default=1024)
    parser.add_argument("--vocab-size", type=int, default=12000)
    parser.add_argument("--seq-len", type=int, default=64)
    parser.add_argument("--width", type=int, default=192)
    parser.add_argument("--depth", type=int, default=8)
    parser.add_argument("--rank", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--batch-size-bp", type=int)
    parser.add_argument("--batch-size-fmad", type=int)
    parser.add_argument("--batch-size-zo", type=int)
    parser.add_argument("--steps", type=int, default=300)
    parser.add_argument("--eval-every", type=int, default=25)
    parser.add_argument("--threshold", type=float, default=0.55)
    parser.add_argument("--perturbations", type=int, default=1)
    parser.add_argument("--zo-eps", type=float, default=1e-3)
    parser.add_argument("--grad-clip", type=float, default=5.0)
    parser.add_argument("--lr-bp", type=float, default=3e-3)
    parser.add_argument("--lr-fmad", type=float, default=1e-3)
    parser.add_argument("--lr-zo", type=float, default=1e-3)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--synthetic", action="store_true")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    return parser.parse_args()


def main():
    args = parse_args()
    seed_everything(args.seed)
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but is not available")
    hardware = {
        "device": str(device),
        "gpu": torch.cuda.get_device_name(0) if device.type == "cuda" else None,
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "job_id": os.environ.get("HF_JOB_ID"),
    }
    print(json.dumps({"event": "hardware", **hardware}), flush=True)
    if args.synthetic:
        data = load_synthetic(args.train_size, args.eval_size, args.vocab_size, args.seq_len, args.seed)
    else:
        data = load_ag_news(args.train_size, args.eval_size, args.vocab_size, args.seq_len, args.seed)
    trainable_probe = AdapterTextModel(data[-1], args.width, args.depth, args.rank, args.seed)
    counts = {
        "total_parameters": sum(p.numel() for p in trainable_probe.parameters()),
        "trainable_parameters": sum(p.numel() for p in trainable_probe.parameters() if p.requires_grad),
    }
    print(json.dumps({"event": "model", **counts}), flush=True)
    del trainable_probe

    all_rows = []
    summaries = []
    for method in ("bp_checkpoint", "fmad", "zo"):
        method_rows, summary = run_method(method, args, data, device)
        all_rows.extend(method_rows)
        summaries.append(summary)
    write_outputs(Path(args.output_dir), all_rows, summaries, args, hardware)


if __name__ == "__main__":
    main()