| """Private N1-only performance probes; no captured input receives optimizer updates.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import defaultdict |
| from contextlib import contextmanager |
| import functools |
| import json |
| from pathlib import Path |
| import time |
| from typing import Any, Callable, Iterator, TypedDict, cast |
|
|
| import torch |
| from torch import Tensor |
|
|
| from .loss import distillation_loss |
| from .model import Drafter, Stage, load |
| from .ops import Dense, PackedLinear |
| from .safety import assert_gpu_lease |
|
|
|
|
| def n1_inputs(root: Path, batch: int, device: str) -> dict[str, Tensor]: |
| paths = sorted(root.glob("n1-*.pt")) |
| if len(paths) != 4 or batch < 1: |
| raise ValueError("require all four N1 synthetic fixtures and positive batch") |
| windows = [torch.load(p, map_location="cpu", weights_only=False)["inputs"] for p in paths] |
| return {k: torch.cat([windows[i % 4][k] for i in range(batch)], 0).to(device) for k in windows[0]} |
|
|
|
|
| class Targets(TypedDict): |
| previous: Tensor |
| top_ids: Tensor |
| top_logp: Tensor |
| tail_mass: Tensor |
| labels: Tensor |
| present: Tensor |
| position_mask: Tensor |
|
|
|
|
| def synthetic_targets(model: Drafter, inp: dict[str, Tensor]) -> Targets: |
| b = inp["anchor"].numel() |
| device = inp["anchor"].device |
| previous = (inp["anchor"][:, None] + torch.arange(5, device=device)) % model.c.vocab |
| present = torch.ones(b, 5, device=device, dtype=torch.bool) |
| return {"previous": previous, "top_ids": torch.arange(20, device=device)[None, None].expand(b, 5, -1), |
| "top_logp": torch.full((b, 5, 20), .04, device=device).log(), |
| "tail_mass": torch.full((b, 5), .2, device=device), "labels": previous, |
| "present": present, "position_mask": present} |
|
|
|
|
| class PhaseTimes: |
| """Nested synchronized exclusive wall times. Instrumentation overhead is explicit.""" |
|
|
| def __init__(self, synchronize: Callable[[], None]): |
| self.synchronize = synchronize |
| self.seconds: dict[str, float] = defaultdict(float) |
| self.stack: list[list[Any]] = [] |
|
|
| @contextmanager |
| def phase(self, name: str) -> Iterator[None]: |
| self.synchronize() |
| start = time.perf_counter() |
| frame: list[Any] = [name, 0.] |
| self.stack.append(frame) |
| try: |
| yield |
| finally: |
| self.synchronize() |
| elapsed = time.perf_counter() - start |
| self.stack.pop() |
| self.seconds[name] += elapsed - frame[1] |
| if self.stack: |
| self.stack[-1][1] += elapsed |
|
|
| def wrap(self, fn: Callable[..., Any], name: str) -> Callable[..., Any]: |
| @functools.wraps(fn) |
| def timed(*args: Any, **kwargs: Any) -> Any: |
| with self.phase(name): |
| return fn(*args, **kwargs) |
| return timed |
|
|
|
|
| @contextmanager |
| def instrument(times: PhaseTimes) -> Iterator[None]: |
| changes = [(Drafter, "conditioning", "conditioning"), (Stage, "attention", "attention"), |
| (Stage, "moe", "experts"), (Dense, "forward", "dense"), |
| (PackedLinear, "forward", "experts"), (Drafter, "logit_chunk", "markov_head")] |
| originals = [] |
| for cls, attr, name in changes: |
| original = getattr(cls, attr) |
| originals.append((cls, attr, original)) |
| setattr(cls, attr, times.wrap(original, name)) |
| try: |
| yield |
| finally: |
| for cls, attr, original in originals: |
| setattr(cls, attr, original) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(__doc__) |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--fixtures", type=Path, required=True) |
| ap.add_argument("--batch", type=int, required=True) |
| ap.add_argument("--steps", type=int, default=2) |
| ap.add_argument("--out", type=Path, required=True) |
| ap.add_argument("--reference-math", action="store_true") |
| ap.add_argument("--fast-training", action="store_true") |
| ap.add_argument("--operator-profile", action="store_true", help="copy operator evidence (adds profiler overhead)") |
| a = ap.parse_args() |
| assert_gpu_lease() |
| torch.set_num_threads(8) |
| model = load(a.ckpt, "cuda", reference_math=a.reference_math) |
| if a.fast_training: |
| model.enable_fast_training() |
| model.train() |
| inp = n1_inputs(a.fixtures, a.batch, model.activation_device) |
| target = synthetic_targets(model, inp) |
| parameters = [p for n, p in model.named_parameters() if "confidence_head" not in n] |
| optimizer = torch.optim.AdamW(parameters, lr=2e-4, foreach=False, fused=a.fast_training) |
| rows = [] |
| for step in range(a.steps): |
| times = PhaseTimes(torch.cuda.synchronize) |
| optimizer.zero_grad(set_to_none=True) |
| begin = time.perf_counter() |
| with instrument(times), times.phase("other_forward"): |
| h, _ = model.hidden_from(inp) |
| with times.phase("loss"): |
| loss, _ = distillation_loss(model, h, **target) |
| with times.phase("backward_other"): |
| cast(Callable[[], None], loss.backward)() |
| with times.phase("optimizer"): |
| torch.nn.utils.clip_grad_norm_(parameters, 1., error_if_nonfinite=True) |
| optimizer.step() |
| seconds = time.perf_counter() - begin |
| row = {"step": step, "seconds": seconds, "blocks_per_second": a.batch / seconds, |
| "exclusive_seconds": dict(times.seconds), "loss": float(loss.detach())} |
| rows.append(row) |
| print(json.dumps(row), flush=True) |
| copies = [] |
| if a.operator_profile: |
| optimizer.zero_grad(set_to_none=True) |
| with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, |
| torch.profiler.ProfilerActivity.CUDA]) as prof: |
| h, _ = model.hidden_from(inp) |
| loss, _ = distillation_loss(model, h, **target) |
| cast(Callable[[], None], loss.backward)() |
| optimizer.step() |
| torch.cuda.synchronize() |
| copies = [{"name": e.key, "calls": e.count, "self_cpu_seconds": e.self_cpu_time_total / 1e6, |
| "self_device_seconds": e.self_device_time_total / 1e6} for e in prof.key_averages() |
| if "copy" in e.key.lower() or "memcpy" in e.key.lower() or e.key == "aten::to"] |
| report = {"batch": a.batch, "reference_math": a.reference_math, "fast_training": a.fast_training, "steps": rows, |
| "copy_operators_nonadditive": copies, "peak_reserved_bytes": torch.cuda.max_memory_reserved(), |
| "peak_allocated_bytes": torch.cuda.max_memory_allocated(), |
| "note": "Synchronized phase wall times include overhead; backward recomputation appears in its phase. " |
| "Copy operator times overlap phase times and must not be added to them."} |
| a.out.parent.mkdir(parents=True, exist_ok=True, mode=0o700) |
| a.out.write_text(json.dumps(report, indent=2) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
