opal-optimizer 0.3.5__py3-none-any.whl

opal/defaults.py

@@ -0,0 +1,54 @@
+"""OPAL's defaults are fixed: the default prompt, the starting programs, and the code that writes them.
+
+The blog's and the README's results were made with these files as released. A coding agent that edits them inside an
+(editable) install changes every later run without a word, so their sha256 hashes ship in defaults.json, and a run
+(rank 0) compares the files it is about to use with them when it starts: an edited default prints a warning, and the
+run's records (the start event of events.jsonl, report.html) say which files differ. Customise through wrap()'s
+arguments instead: prompt=<a template file>, program=<a program folder>, task_facts=, and the rest.
+
+tests/test_defaults.py fails when defaults.json is stale; after a deliberate change to one of these files, write it
+again with `python -m opal.defaults --write`.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).parent
+MANIFEST = HERE / "defaults.json"
+PROMPT_FILES = ("prompts/neutral.md", "misc/aggressive.py")  # the named prompts (not used with prompt=<a path>)
+PROGRAM_FILES = ("prompts/example_updater.py", "misc/neutral_analysis.py")  # starting analysis, example (not with program=)
+# what fills the prompt and writes the messages and the starting updater, whatever prompt or program is used
+SHARED_FILES = ("misc/setup.py", "misc/taskfacts.py", "programs.py", "agent.py", "context.py")
+FILES = PROMPT_FILES + PROGRAM_FILES + SHARED_FILES
+
+
+def hashes(names=FILES) -> dict:
+    return {n: hashlib.sha256((HERE / n).read_bytes()).hexdigest() for n in names}
+
+
+def modified(names=FILES) -> list:
+    """The files among `names` whose content differs from the released one (defaults.json), or that are missing."""
+    try:
+        shipped = json.loads(MANIFEST.read_text())
+    except (OSError, ValueError):
+        return ["defaults.json"]
+    out = []
+    for n in names:
+        try:
+            now = hashlib.sha256((HERE / n).read_bytes()).hexdigest()
+        except OSError:
+            now = None
+        if now != shipped.get(n):
+            out.append(n)
+    return out
+
+
+if __name__ == "__main__":
+    if sys.argv[1:] != ["--write"]:
+        sys.exit("usage: python -m opal.defaults --write   (after a deliberate change to OPAL's defaults)")
+    MANIFEST.write_text(json.dumps(hashes(), indent=1, sort_keys=True) + "\n")
+    print(f"wrote {MANIFEST}")

opal/estimate.py

@@ -0,0 +1,120 @@
+"""Before a run: how long training will wait for the agent, and what the agent will cost.
+
+`opal estimate --total-steps N --sec-per-step S` prints it before launch; a run prints the same at its first check,
+measured from its own time per update. Both use the per-decision times and costs below, which come from our runs.
+
+When `decisions` is not set, a blocking run schedules at most as many decisions as keep the expected waiting within
+MAX_WAIT_FACTOR (4) times the run's own training time, at least one; a `decisions` the user sets is never changed, and a line says when its waiting is above that.
+
+Per decision of claude-opus-5-5 at effort high, the agent's `seconds` and `cost_usd`, 567 decisions in all:
+  the first decision, which reads the training code and writes the first program: 101 to 432 s over 20 runs (median
+  251 s), $0.33 to $1.62 (median $1.08);
+  each later decision, averaged over a run: 10 to 51 s over the 15 runs with at least 10 of them (median 27 s), $0.058
+  to $0.377 (median $0.23).
+Through a mailbox, add the broker's poll and the ssh transfer, a few seconds per decision. Other models and efforts were not measured.
+"""
+
+from __future__ import annotations
+
+MODEL, EFFORT = "claude-opus-5-5", "high"  # what the numbers below were measured with
+FIRST_SECONDS = (100.0, 250.0, 430.0)  # (low, typical, high) seconds of the first decision
+LATER_SECONDS = (10.0, 27.0, 50.0)  # seconds of each later decision (a run's average)
+FIRST_COST = (0.30, 1.08, 1.60)  # dollars
+LATER_COST = (0.05, 0.23, 0.38)
+MAX_WAIT_FACTOR = 4.0  # with decisions not set: expected waiting at most this many times the run's own training time
+
+
+def waiting(n, blocking=True, sec_per_step=None, max_lag=50, which=1, first=None, later=None):
+    """Seconds training waits for n decisions: which=0 low, 1 typical, 2 high. blocking=False: none (training waits
+    only when the loss turns non-finite). "auto": the first decision, and the later ones whose answer would land more
+    than max_lag updates late."""
+    if n <= 0 or blocking is False:
+        return 0.0
+    f = FIRST_SECONDS[which] if first is None else first
+    g = LATER_SECONDS[which] if later is None else later
+    if blocking == "auto" and sec_per_step and g / sec_per_step <= max_lag:
+        g = 0.0
+    return f + (n - 1) * g
+
+
+def cost(n, which=1):
+    return 0.0 if n <= 0 else FIRST_COST[which] + (n - 1) * LATER_COST[which]
+
+
+def capped_decisions(n, run_seconds, blocking=True, sec_per_step=None, max_lag=50, factor=MAX_WAIT_FACTOR,
+                     first=None, later=None):  # fmt: skip
+    """The most decisions, at most n and at least 1, whose typical waiting stays within factor x run_seconds."""
+    if n <= 1:
+        return n
+    budget = factor * run_seconds
+    k = n
+    while k > 1 and waiting(k, blocking, sec_per_step, max_lag, 1, first, later) > budget:
+        k -= 1
+    return k
+
+
+def _t(s):
+    if s < 90:
+        return f"{s:.0f} s"
+    if s < 90 * 60:
+        return f"{s / 60:.0f} min"
+    return f"{s / 3600:.1f} h"
+
+
+def _x(a, b):
+    if b <= 0:
+        return "?"
+    r = a / b
+    return f"{r:.2f}x" if r < 1 else f"{r:.1f}x" if r < 10 else f"{r:,.0f}x"
+
+
+def _usd(x):
+    return f"${x:.2f}" if x < 10 else f"${x:.0f}"
+
+
+def lines(total_steps, sec_per_step, decisions, blocking=True, max_lag=50, model=MODEL, effort=EFFORT, mailbox=False,
+          capped_from=None, explicit=True):  # fmt: skip
+    """The estimate, a few lines. decisions: the scheduled count. capped_from: the count before the cap of
+    MAX_WAIT_FACTOR (decisions not set). explicit: the user set decisions (then a waiting above the factor is said)."""
+    run = total_steps * sec_per_step
+    lo, typ, hi = (waiting(decisions, blocking, sec_per_step, max_lag, w) for w in (0, 1, 2))
+    mode = (
+        "training waits for each (blocking=True)"
+        if blocking is True
+        else "training goes on while the agent thinks (blocking=False)"
+        if blocking is False
+        else 'training waits at the first and whenever an answer would land late (blocking="auto")'
+    )
+    out = [f"training without the agent: {total_steps:,} updates x {sec_per_step:.3g} s = {_t(run)}"]
+    if blocking is False:
+        wait = "no waiting for the agent, except when the training loss turns non-finite"
+    else:
+        wait = (
+            f"about {_t(typ)} of waiting for the agent ({_t(lo)} to {_t(hi)}), {_x(typ, run)} the run's own time "
+            f"({_x(lo, run)} to {_x(hi, run)})"
+            + (", plus a few seconds per decision through the mailbox" if mailbox else "")
+        )
+    out.append(f"{decisions} decision{'s' * (decisions != 1)}, {mode}: {wait}")
+    out.append(f"agent cost: about {_usd(cost(decisions))} ({_usd(cost(decisions, 0))} to {_usd(cost(decisions, 2))})")
+    over = blocking is not False and typ > MAX_WAIT_FACTOR * run
+    if capped_from is not None and capped_from != decisions:
+        out.append(
+            f"decisions not set: {decisions} scheduled instead of {capped_from}, so that the expected waiting stays "
+            f"within {MAX_WAIT_FACTOR:g}x the run's own time (set decisions= to choose)"
+        )
+    if over and (decisions <= 1 or not explicit):
+        out.append(
+            f"even {'one decision' if decisions <= 1 else f'{decisions} decisions'} wait{'s' * (decisions <= 1)} more "
+            f"than {MAX_WAIT_FACTOR:g}x the run's own time; blocking=False lets training go on while the agent thinks"
+        )
+    elif over:
+        out.append(
+            f"the expected waiting is more than {MAX_WAIT_FACTOR:g}x the run's own time: ask fewer times "
+            f"(decisions={capped_decisions(decisions, run, blocking, sec_per_step, max_lag)} keeps it within "
+            f"{MAX_WAIT_FACTOR:g}x) or let training go on while the agent thinks (blocking=False)"
+        )
+    if (model, effort) != (MODEL, EFFORT):
+        out.append(
+            f"these numbers are from {MODEL} at effort {EFFORT}; {model} at effort {effort} was not measured"
+        )
+    return out

opal/hf.py

@@ -0,0 +1,109 @@
+"""OPAL in a Hugging Face Trainer: OpalTrainer is the Trainer with its optimizer wrapped by opal.wrap.
+
+    from opal.hf import OpalTrainer
+
+    trainer = OpalTrainer(model=model, args=args, train_dataset=train, eval_dataset=dev,
+                          opal=dict(goal="...", max_cost_usd=20))   # opal.wrap's keyword arguments
+    trainer.train()
+
+The Trainer builds its optimizer and learning-rate scheduler as usual; OpalTrainer builds the scheduler on the Trainer's
+own torch optimizer and then wraps that optimizer with it, so the scheduler stays on the Trainer's own optimizer (a scheduler built on the wrapper would overwrite the
+rates the agent's program sets) and total_steps is the Trainer's number of updates. The Trainer calls optimizer.step()
+without the loss, so OpalTrainer adds up the loss training_step returns for each micro-batch of an update and hands it
+to OPAL at that update's step(): the update's mean loss (also at the end of an epoch, where an update can have fewer or
+more micro-batches than gradient_accumulation_steps). The Trainer's evaluation values (eval_loss, ...) go into OPAL's
+report and are not shown to the agent.
+
+Clipping: the Trainer clips the gradients (max_grad_norm, 1.0 by default) before the optimizer's step; that clipping
+stays the Trainer's, outside the optimizer, and the agent's program gets clipped gradients and cannot change it. This
+keeps the run bit for bit the plain Trainer's. To hand the clipping to the agent, turn the Trainer's off and give it to
+OPAL: TrainingArguments(max_grad_norm=0) and opal=dict(clip_grad_norm=1.0). The run then starts the same up to rounding
+(accelerate computes the Trainer's gradient norm its own way, so not bit for bit: about 3e-4 apart after 6 updates of
+the tiny test model). clip_grad_norm in opal= while max_grad_norm is set raises: the gradients would be clipped twice.
+
+Needs transformers (pip install "opal-optimizer[hf]"). Tested on CPU with transformers 5.20 and accelerate 1.15.
+"""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from transformers import Trainer
+except ImportError as exc:  # pragma: no cover - the optional extra
+    raise ImportError('opal.hf needs transformers: pip install "opal-optimizer[hf]"') from exc
+
+from .wrapper import AgentOptimizer, wrap
+
+
+class OpalTrainer(Trainer):
+    def __init__(self, *args, opal=None, **kwargs):
+        opal = dict(opal or {})
+        targs = kwargs.get("args") if "args" in kwargs else (args[1] if len(args) > 1 else None)
+        if opal.get("clip_grad_norm") is not None and (targs is None or getattr(targs, "max_grad_norm", None)):
+            raise ValueError(
+                "OpalTrainer: opal= sets clip_grad_norm and args.max_grad_norm is set too, so the gradients would be "
+                "clipped twice; to hand the clipping to the agent, also pass TrainingArguments(max_grad_norm=0)"
+            )
+        super().__init__(*args, **kwargs)
+        self.opal_kwargs = opal
+        self._opal_losses = []  # the losses training_step returned for this update's micro-batches
+        self._opal_loss_step = None
+
+    def create_scheduler(self, num_training_steps, optimizer=None):
+        """Build the Trainer's schedule on its own torch optimizer, then wrap that optimizer with it. Transformers 4
+        calls this from create_optimizer_and_scheduler; 5 calls it after accelerate has wrapped the optimizer, so the
+        AgentOptimizer takes the torch optimizer's place inside accelerate's wrapper."""
+        holder = optimizer if optimizer is not None else self.optimizer
+        raw = holder
+        while isinstance(getattr(raw, "optimizer", None), torch.optim.Optimizer):  # accelerate's wrapper
+            raw = raw.optimizer
+        if isinstance(raw, AgentOptimizer):  # wrapped already
+            return super().create_scheduler(num_training_steps, optimizer=optimizer)
+        scheduler = super().create_scheduler(num_training_steps, optimizer=raw)
+        agent = wrap(raw, self.model, scheduler=scheduler, total_steps=num_training_steps, **self.opal_kwargs)
+        agent._loss_hook = self._update_loss  # step() takes the update's loss from here
+        if self.optimizer is raw:
+            self.optimizer = agent
+        else:
+            outer = self.optimizer
+            while getattr(outer, "optimizer", None) is not raw:
+                outer = outer.optimizer
+            outer.optimizer = agent
+        return scheduler
+
+    def training_step(self, *args, **kwargs):
+        if self._opal_loss_step != self.state.global_step:
+            self._opal_losses = []  # a GradScaler-skipped update never called the optimizer's loss hook
+            self._opal_loss_step = self.state.global_step
+        loss = super().training_step(*args, **kwargs)
+        self._opal_losses.append(loss.detach())
+        return loss
+
+    def _update_loss(self):
+        """The update's mean loss from its micro-batches' losses as training_step returns them: each is already scaled
+        for the update in recent Transformers (by current_gradient_accumulation_steps, or by the number of items), so
+        they add up to it; older ones divide by gradient_accumulation_steps whatever the update's micro-batches are."""
+        losses, self._opal_losses = self._opal_losses, []
+        if not losses:
+            return None
+        total = torch.stack([x.float().reshape(()) for x in losses]).sum()
+        if not hasattr(self, "current_gradient_accumulation_steps"):
+            total = total * self.args.gradient_accumulation_steps / len(losses)
+        return total
+
+    def log(self, logs, *args, **kwargs):
+        super().log(logs, *args, **kwargs)
+        opt = agent_optimizer(self)
+        if opt is not None:
+            shown = {k: v for k, v in logs.items() if k.startswith("eval_") and isinstance(v, (int, float))}
+            if shown:
+                opt.log(**shown)  # in OPAL's report; not shown to the agent
+
+
+def agent_optimizer(trainer):
+    """The AgentOptimizer inside a Trainer's optimizer (accelerate wraps it), or None before it is built."""
+    opt = trainer.optimizer
+    while opt is not None and not isinstance(opt, AgentOptimizer):
+        opt = getattr(opt, "optimizer", None)
+    return opt

opal/misc/__init__.py

@@ -0,0 +1,2 @@
+"""Side machinery the main loop in wrapper.py delegates to: checkpoints and resume, recovery from a failing
+program, DDP collectives, the user's schedulers, prompt and hook setup, tensor helpers."""

opal/misc/aggressive.py

@@ -0,0 +1,157 @@
+"""OPAL 0.1.2-aggressive (prompt="aggressive"): the neutral prompt, except that each
+decision starts with a review of every parameter group's hyperparameters, the rationale has five parts (observed,
+review, changed, prediction, rules), the decision message no longer says that keeping the programs is a valid decision,
+and reference/example_updater.py has a feedback controller of the learning rates (off by default). Everything else is
+neutral.md word for word; this module holds only the differences, applied to neutral.md and example_updater.py."""
+
+from __future__ import annotations
+
+# ----- the system prompt: neutral.md's DECISIONS section, old -> new ------------------------------------------------
+
+PROMPT_EDITS = [
+    (
+        "At each decision you may edit your programs or keep them; keeping them is a valid decision. Before changing "
+        "anything,\ncheck what you predicted at your previous decision against what happened. End every decision with a "
+        "short rationale:\n",
+        "At each decision, first check what you predicted at your previous decision against what happened. Then review, "
+        "for\neach parameter group, whether its current hyperparameters (the update rule, the learning rate and its "
+        "schedule,\nmomentum or betas, epsilon, weight decay, clipping) suit the next stretch of training, and which "
+        "adjustments, in which\ndirection and by how much, would improve performance. Rules in your programs can also "
+        "make such adjustments between\ndecisions, at every update, from signals they measure. End every decision with a "
+        "short rationale:\n",
+    ),
+    (
+        "2. Changed or kept: what you edited and why; or why you kept the programs and what evidence you are waiting "
+        "for.\n3. Prediction: the effect you expect, by roughly when, and what result would show you were wrong.\n4. "
+        "Rules: for each rule in your programs, whether it is a fixed schedule or reacts to a signal.\n",
+        "2. Review: for each parameter group, whether its hyperparameters suit the next stretch, and which adjustments\n"
+        "   (direction and size) you expect to help.\n3. Changed: what you edited and why.\n4. Prediction: the effect you "
+        "expect, by roughly when, and what result would show you were wrong.\n5. Rules: for each rule in your programs, "
+        "whether it is a fixed schedule or reacts to a signal.\n",
+    ),
+]
+
+# ----- the decision message (agent.format_message) ------------------------------------------------------------------
+
+EDIT_SENTENCE = "Edit the programs in place. "
+RATIONALE_SENTENCE = (
+    "End with your rationale in the five parts of DECISIONS: observed, review, changed, prediction, rules."
+)
+
+
+def _apply(text, edits, what):
+    for old, new in edits:
+        if text.count(old) != 1:
+            raise RuntimeError(f"opal: the aggressive variant no longer matches {what}: {old[:80]!r}")
+        text = text.replace(old, new)
+    return text
+
+
+def template(neutral: str) -> str:
+    return _apply(neutral, PROMPT_EDITS, "prompts/neutral.md")
+
+
+# ----- reference/example_updater.py: a feedback controller of the learning rates, off by default ---------------------
+
+CONSTANTS = """
+
+# A feedback controller of the learning rates, off by default: at every update it measures, on the device, a running
+# mean of each AdamW group's update-to-weight ratio, and lr_rule() below turns that signal into a multiplier of the
+# group's learning rate. To use it, set LR_CONTROLLER = True and SIGNAL_DECAY (the weight of the newest measurement in
+# the running mean), and write lr_rule() from your own measurements.
+LR_CONTROLLER = False
+SIGNAL_DECAY = None
+"""
+
+CONTROLLER = '''
+
+# ----- a feedback controller of the learning rates (off by default) -------------------------------------------------
+
+
+def lr_rule(signal, multiplier, update, total):
+    """Your rule: the new learning-rate multiplier of one group, from its signal (a 0-d tensor on the device: the
+    running mean of the group's update-to-weight ratio) and its current multiplier (a 0-d tensor). Keep it on the
+    device (torch.where, torch.clamp, arithmetic; no .item(), no if on a tensor). As written it changes nothing."""
+    return multiplier
+
+
+class LRController:
+    """Per AdamW group (keyed like the group): the running mean of its update-to-weight ratio and its learning-rate
+    multiplier, 0-d float32 tensors on the device. The multiplier scales the group's whole update after the rule's
+    step, which for AdamW is the same as scaling its learning rate, so nothing is read on the host."""
+
+    def __init__(self, rule, carried):
+        if SIGNAL_DECAY is None:
+            raise ValueError("set SIGNAL_DECAY to use the learning-rate controller")
+        self.rule = rule
+        self.keys = [str(g["key"]) for g in rule.groups]
+        self.signal, self.mult = {}, {}
+        for key, g in zip(self.keys, rule.groups):
+            device = g["params"][0].device
+            self.signal[key] = carried.get("signal", {}).get(key, torch.zeros((), device=device)).clone()
+            self.mult[key] = carried.get("mult", {}).get(key, torch.ones((), device=device)).clone()
+
+    def before(self):
+        return [torch._foreach_mul([p.detach() for p in g["params"]], 1) for g in self.rule.groups]  # copies
+
+    def after(self, before, update, total):
+        for key, g, old in zip(self.keys, self.rule.groups, before):
+            params = [p.detach() for p in g["params"]]
+            delta = torch._foreach_sub(params, old)
+            moved = torch.linalg.vector_norm(torch.stack(torch._foreach_norm(delta)).float())
+            size = torch.linalg.vector_norm(torch.stack(torch._foreach_norm(old)).float())
+            ratio = moved / size.clamp_min(torch.finfo(size.dtype).tiny)
+            self.signal[key].mul_(1 - SIGNAL_DECAY).add_(ratio * SIGNAL_DECAY)
+            self.mult[key] = lr_rule(self.signal[key], self.mult[key], update, total)
+            # Interpolate from the saved weights: subtracting the update back can lose
+            # their low bits even at multiplier zero. lerp preserves both endpoints.
+            torch._foreach_lerp_(old, params, [self.mult[key].to(params[0].dtype)] * len(old))
+            torch._foreach_copy_(params, old)
+
+    def state_dict(self):
+        return {"signal": dict(self.signal), "mult": dict(self.mult)}
+
+    def load_state_dict(self, state):
+        for name in ("signal", "mult"):
+            mine = getattr(self, name)
+            mine.update({k: v for k, v in state.get(name, {}).items() if k in mine})
+'''
+
+EXAMPLE_EDITS = [
+    (
+        "Nothing\nhere is a recommendation: the multipliers, knots and rule table are left at values that change "
+        "nothing.",
+        "Nothing\nhere is a recommendation: the multipliers, knots and rule table are left at values that change "
+        "nothing, and the\nfeedback controller of the learning rates is off.",
+    ),
+    ("\n\n# ----- grouping by role and block", CONSTANTS + "\n\n# ----- grouping by role and block"),
+    ("\n\n# ----- the updater -----", CONTROLLER + "\n\n# ----- the updater -----"),
+    (
+        "        self.rules = {r: RULES[r](entries.get(r, []), carried) for r in RULES}\n",
+        "        self.rules = {r: RULES[r](entries.get(r, []), carried) for r in RULES}\n"
+        '        old = getattr(previous, "controller", None)  # its signals and multipliers carry over\n'
+        "        self.controller = (\n"
+        '            LRController(self.rules["adamw"], old.state_dict() if old else {}) if LR_CONTROLLER else None\n'
+        "        )\n",
+    ),
+    (
+        "        for rule in self.rules.values():\n            rule.set_lrs(update, self.total)\n            rule.step()\n",
+        "        before = self.controller.before() if self.controller else None\n"
+        "        for rule in self.rules.values():\n            rule.set_lrs(update, self.total)\n            rule.step()\n"
+        "        if self.controller:\n            self.controller.after(before, update, self.total)\n",
+    ),
+    (
+        "    def state_dict(self):\n        return {r: rule.state_by_name() for r, rule in self.rules.items()}  # keyed by "
+        "parameter name\n\n    def load_state_dict(self, state):\n        for r, rule in self.rules.items():\n"
+        "            rule.load_state_by_name(state.get(r, {}))\n",
+        "    def state_dict(self):\n        out = {r: rule.state_by_name() for r, rule in self.rules.items()}  # keyed by "
+        'parameter name\n        if self.controller:\n            out["lr_controller"] = self.controller.state_dict()\n'
+        "        return out\n\n    def load_state_dict(self, state):\n        for r, rule in self.rules.items():\n"
+        '            rule.load_state_by_name(state.get(r, {}))\n        if self.controller and "lr_controller" in state:\n'
+        '            self.controller.load_state_dict(state["lr_controller"])\n',
+    ),
+]
+
+
+def example(base: str) -> str:
+    return _apply(base, EXAMPLE_EDITS, "prompts/example_updater.py")

opal/misc/audit.py

@@ -0,0 +1,84 @@
+"""One audited update per new updater (rank 0, second update of each version): operator calls per update, and host
+synchronisations per update on a GPU. These show when an updater is slow because of how it is written (values read
+on the host, Python loops over tensors) rather than because of the algorithm it implements.
+
+On a GPU the CUDA profiler is not used: synchronisations are counted with torch's sync debug mode, a flag checked at each sync, and
+operators with the CPU-only profiler."""
+
+from __future__ import annotations
+
+import warnings
+
+COUNT_OPS_ON_CUDA = True  # the CPU-only profiler on a GPU
+
+
+def audited(fn, device):
+    """Call fn() once and return its counts, or None if they cannot be taken (fn still runs; its exceptions
+    propagate)."""
+    if device.type != "cuda":
+        return _profiled(fn)
+    import torch
+
+    try:
+        prev = torch.cuda.get_sync_debug_mode()
+    except Exception:
+        fn()
+        return None
+    ops = {}
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        torch.cuda.set_sync_debug_mode(1)
+        try:
+            if COUNT_OPS_ON_CUDA:
+                ops = _profiled(fn) or {}
+            else:
+                fn()
+        finally:
+            torch.cuda.set_sync_debug_mode(prev)
+    syncs = sum(1 for w in caught if "synchroniz" in str(w.message))
+    return dict(ops or {}, host_syncs=int(syncs))
+
+
+def _profiled(fn):
+    """fn() under the CPU-only profiler: outermost operator calls (an operator called from Python, not one called by
+    another operator)."""
+    try:
+        from torch.profiler import ProfilerActivity, profile
+
+        try:
+            prof = profile(activities=[ProfilerActivity.CPU], acc_events=True)  # no end-of-cycle warning
+        except TypeError:  # torch < 2.4
+            prof = profile(activities=[ProfilerActivity.CPU])
+        prof.__enter__()
+    except Exception:
+        fn()
+        return None
+    try:
+        fn()
+    finally:
+        try:
+            prof.__exit__(None, None, None)
+        except Exception:
+            prof = None
+    if prof is None:
+        return None
+    try:
+        ops = sum(
+            1
+            for e in prof.events()
+            if e.name.startswith("aten::") and (e.cpu_parent is None or not e.cpu_parent.name.startswith("aten::"))
+        )
+    except Exception:
+        return None
+    return dict(operator_calls=int(ops))
+
+
+def describe(audit):
+    if not audit:
+        return "not measured"
+    parts = []
+    if "operator_calls" in audit:
+        parts.append(f"{audit['operator_calls']} operator calls")
+    if "host_syncs" in audit:
+        parts.append(f"{audit['host_syncs']} host synchronisations")
+    return " and ".join(parts) + " per update"