"""Train the model matrix identically on one device, for comparison on another. python benchmarks/pretrain.py --out "MPS bench" # on the Mac Studio python benchmarks/pretrain.py --out "CUDA bench" # on the 5090 Every model is built on **CPU** under a fixed seed and only then moved to the device, and every batch is drawn on CPU from a seeded generator. That is the whole design: identical initial weights and identical data on both machines, so the only thing that differs between two runs is the arithmetic. Compare the two directories with ``benchmarks/compare.py``. **What a CUDA-vs-MPS comparison can and cannot show.** It cannot show bitwise agreement — different devices reduce in different orders, and cuDNN, Metal and a pure-torch loop are three different implementations of the same formula. It *can* show whether they agree to the precision the arithmetic allows, whether they train to the same place, and how much faster one is. The interesting rows are the vendored models. On CUDA they take the authors' fused kernels; on MPS they take the reference path. So those rows are not really "CUDA vs MPS" — they are **fused vs reference**, which is the comparison PLAN.md fixes as the rule for any fast path: the portable path is the reference, and the fused path must agree with it. """ from __future__ import annotations import argparse import json import platform import time from pathlib import Path import init_weights import precision import torch import torch.nn as nn import torch_dimensions as td SEED = 20260803 def device_name(device: str) -> str: if device == "cuda": return torch.cuda.get_device_name(0) if device == "mps": return f"Apple {platform.machine()} (MPS)" return platform.processor() or platform.machine() def pick_device(requested: str | None) -> str: if requested: return requested if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available(): return "mps" return "cpu" def sync(device: str) -> None: if device == "cuda": torch.cuda.synchronize() elif device == "mps": torch.mps.synchronize() # One recipe for every model, which is only possible because `td.param_groups` # reads the tags the upstream authors put on their own parameters: s4's # `_optim` and Mamba's `_no_weight_decay`. Training an SSM's `A` and `dt` at # the projection rate with weight decay is what made the earlier runs diverge # at 1e-2 and forced hand-picked per-family rates — which were themselves a # confound, since a device comparison should hold everything it is not # measuring fixed. See BENCHMARK-DESIGN.md. LR = 3e-3 BETAS = (0.9, 0.95) WEIGHT_DECAY = 0.1 GRAD_CLIP = 1.0 WARMUP_FRACTION = 0.1 # --- the matrix ------------------------------------------------------------ # Small on purpose: the point is coverage of every family and composition, not # a leaderboard. Each model has to fit the same task so the losses are # comparable, and stay small enough that fifteen of them are minutes rather # than hours — and that the checkpoints are a sane upload. def sparse_2d() -> td.Lattice: gen = torch.Generator().manual_seed(11) valid = torch.rand(6, 8, generator=gen) > 0.25 valid[0, 0] = True return td.Lattice(shape=(6, 8), names=("h", "w"), valid=valid, time=True) def dense_3d() -> td.Lattice: return td.Lattice(shape=(4, 5, 6), names=("d", "h", "w"), time=True) ZOO: dict[str, dict] = { # --- recurrent "lstm_2d_sparse": {"build": lambda lat: td.LSTM(32, 6, lat, d_input=1), "lat": sparse_2d}, "gru_2d_sparse": {"build": lambda lat: td.GRU(32, 6, lat, d_input=1), "lat": sparse_2d}, # --- state space, portable vs the authors' own code "s4d_portable_2d": { "build": lambda lat: td.S4D(32, 4, lat, d_input=1, portable=True, d_state=16), "lat": sparse_2d, }, "s4d_upstream_2d": { "build": lambda lat: td.S4D(32, 4, lat, d_input=1, d_state=16), "lat": sparse_2d, }, "s4_upstream_2d": { "build": lambda lat: td.S4(32, 4, lat, d_input=1, d_state=16), "lat": sparse_2d, }, "mamba_portable_2d": { "build": lambda lat: td.Mamba(32, 4, lat, d_input=1, portable=True, d_state=8), "lat": sparse_2d, }, "mamba_upstream_2d": { "build": lambda lat: td.Mamba(32, 4, lat, d_input=1, d_state=8), "lat": sparse_2d, }, "mamba2_2d": { "build": lambda lat: td.Mamba2( 64, 4, lat, d_input=1, mixer_kwargs={"d_state": 16, "headdim": 32} ), "lat": sparse_2d, }, "mamba3_2d": { "build": lambda lat: td.Mamba3( 64, 4, lat, d_input=1, mixer_kwargs={"d_state": 64, "headdim": 32} ), "lat": sparse_2d, }, # --- attention, across all three compositions "transformer_scan_2d": { "build": lambda lat: td.Transformer(32, 4, lat, d_input=1), "lat": sparse_2d, }, "transformer_cafa_2d": { "build": lambda lat: td.Transformer(32, 4, lat, d_input=1, method=td.cafa), "lat": sparse_2d, }, "transformer_flatten_2d": { "build": lambda lat: td.Transformer(32, 4, lat, d_input=1, method=td.flatten), "lat": sparse_2d, }, # --- convolutional "cnn_2d_sparse": {"build": lambda lat: td.CNN(32, 4, lat, d_input=1), "lat": sparse_2d}, "tcn_2d_sparse": {"build": lambda lat: td.TCN(32, 4, lat, d_input=1), "lat": sparse_2d}, # --- rank 3, where the composition has more to do "mamba_upstream_3d": { "build": lambda lat: td.Mamba(32, 6, lat, d_input=1, d_state=8), "lat": dense_3d, }, "lstm_3d": {"build": lambda lat: td.LSTM(32, 6, lat, d_input=1), "lat": dense_3d}, } def train_one( name: str, cfg: dict, device: str, steps: int, batch: int, t_len: int, init: Path | None = None, ) -> dict: """Build on CPU under a fixed seed, move, train, and record. Building on CPU first is most of what makes two machines comparable. It is not all of it: S4 and S4D diagonalise HiPPO with `torch.linalg.eigh`, whose eigenvectors are only defined up to a phase, so their `B` and `P` differ between macOS and Linux under the same seed. Pass `--init` to share one set of starting weights and remove that difference from the measurement — see benchmarks/init_weights.py. """ lat = cfg["lat"]() torch.manual_seed(SEED) model = cfg["build"](lat) head = nn.Linear(model.config["d_model"], 1) weights_from = init_weights.sync(model, init, name) init_weights.sync(head, init, f"{name}.head") # On CPU in float64: MPS has no float64 at all, and a norm computed in the # device's own precision would differ between machines for that reason # rather than for the reason being measured. def weight_norm(m: nn.Module) -> float: total = sum((p.detach().cpu().double() ** 2).sum() for p in m.parameters()) return float(torch.sqrt(total)) init_norm = weight_norm(model) n_params = sum(p.numel() for p in model.parameters()) model, head = model.to(device), head.to(device) lr = cfg.get("lr", LR) opt = torch.optim.AdamW( td.param_groups(model, lr=lr, weight_decay=WEIGHT_DECAY) + [{"params": list(head.parameters()), "lr": lr, "weight_decay": WEIGHT_DECAY}], lr=lr, betas=BETAS, ) sched = td.warmup_cosine(opt, warmup=max(1, int(steps * WARMUP_FRACTION)), total=steps) mask = lat.mask(torch.float32).to(device) w_dim = lat.tensor_dim(lat.axis_names[-1]) gen = torch.Generator().manual_seed(SEED + 1) def draw(): # Drawn on CPU so both machines see the same numbers, then moved. x = torch.randn(batch, t_len, *lat.shape, 1, generator=gen) return x.to(device) * mask, x.to(device).cumsum(dim=w_dim) * mask losses: list[float] = [] sync(device) started = time.perf_counter() for _ in range(steps): x, y = draw() loss = (head(model(x)) - y).pow(2).mean() opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_([*model.parameters(), *head.parameters()], GRAD_CLIP) opt.step() sched.step() losses.append(float(loss.detach())) sync(device) seconds = time.perf_counter() - started # A fixed probe batch, so the two machines can be compared on outputs and # not only on losses. Same input on both, by construction. probe_gen = torch.Generator().manual_seed(SEED + 2) probe = torch.randn(1, t_len, *lat.shape, 1, generator=probe_gen).to(device) * mask model.eval() with torch.no_grad(): out = head(model(probe)).float().cpu() model.train() return ( { "name": name, "n_params": n_params, "weights_from": weights_from, "lr": lr, "recipe": "adamw+param_groups+warmup_cosine+clip", "steps": steps, "init_weight_norm": init_norm, "final_weight_norm": weight_norm(model), "loss_first": losses[0], "loss_final": losses[-1], "loss_min": min(losses), "losses": [round(v, 8) for v in losses], "seconds": round(seconds, 3), "steps_per_second": round(steps / seconds, 3), "probe_mean": float(out.mean()), "probe_absmax": float(out.abs().max()), "probe_sum": float(out.double().sum()), "mixer": type(model.nd.mixers[0]).__name__ if getattr(model.nd, "mixers", None) else None, }, model, head, ) def main() -> int: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--out", required=True, help='output directory, e.g. "MPS bench"') ap.add_argument("--device", default=None, help="cuda | mps | cpu (default: best available)") ap.add_argument("--steps", type=int, default=300) ap.add_argument("--batch", type=int, default=4) ap.add_argument("--t-len", type=int, default=6) ap.add_argument("--only", default=None, help="comma-separated subset of model names") init_weights.add_argument(ap) precision.add_arguments(ap, tf32_default="torch") args = ap.parse_args() device = pick_device(args.device) prec = precision.apply(args) out = Path(args.out) out.mkdir(parents=True, exist_ok=True) init = Path(args.init) if args.init else None names = args.only.split(",") if args.only else list(ZOO) print(f"device: {device} ({device_name(device)})") print(f"models: {len(names)} steps: {args.steps}") print(f"weights: {init if init else 'from the seed (see --init)'}") print(f"{precision.describe(prec)}\n") def manifest_of(records: list) -> dict: return { "device": device, "device_name": device_name(device), "platform": platform.platform(), "python": platform.python_version(), "torch": torch.__version__, "torch_cuda": torch.version.cuda, "torch_dimensions": getattr(td, "__version__", "unknown"), "seed": SEED, "precision": prec, "steps": args.steps, "batch": args.batch, "t_len": args.t_len, "complete": len(records) == len(names), "requested": names, "models": records, } results: list[dict] = [] for i, name in enumerate(names, 1): cfg = ZOO[name] print(f"[{i:2d}/{len(names)}] {name:26s} ", end="", flush=True) try: record, model, head = train_one( name, cfg, device, args.steps, args.batch, args.t_len, init ) except Exception as exc: # noqa: BLE001 - one model failing is a result print(f"FAILED — {type(exc).__name__}: {exc}") results.append({"name": name, "error": f"{type(exc).__name__}: {exc}"}) continue model_dir = out / name model_dir.mkdir(exist_ok=True) # Weights on CPU so the file is identical in layout across devices and # loadable anywhere. torch.save( { "model": {k: v.cpu() for k, v in model.state_dict().items()}, "head": {k: v.cpu() for k, v in head.state_dict().items()}, }, model_dir / "weights.pt", ) (model_dir / "metrics.json").write_text(json.dumps(record, indent=2) + "\n") results.append(record) # Rewritten after every model, not once at the end. A CPU pass over # this zoo is hours — the authors' reference Mamba scan is a Python # loop over sequence length — and a run that keeps everything in memory # until the last model loses all of it to one interrupt. `complete` # says whether the file is the whole run or a prefix of it. (out / "manifest.json").write_text(json.dumps(manifest_of(results), indent=2) + "\n") print( f"loss {record['loss_first']:.4f} -> {record['loss_final']:.5f} " f"{record['steps_per_second']:6.1f} steps/s {record['n_params']:,} params" ) (out / "manifest.json").write_text(json.dumps(manifest_of(results), indent=2) + "\n") ok = sum(1 for r in results if "error" not in r) print(f"\nwrote {out}/manifest.json — {ok}/{len(results)} models trained") return 0 if ok == len(results) else 1 if __name__ == "__main__": raise SystemExit(main())