trainmeter 0.0.2__py3-none-any.whl

trainmeter/export/files.py

@@ -0,0 +1,24 @@
+"""The passport as files next to the run log. Derived from the log, never the log itself."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from ..passport import to_markdown
+
+
+def write_passport(directory: Path, data: dict[str, Any]) -> list[Path]:
+    """Write `passport.json` and `passport.md` atomically. Returns the paths written."""
+    written = []
+    for name, text in (
+        ("passport.json", json.dumps(data, indent=2, allow_nan=False) + "\n"),
+        ("passport.md", to_markdown(data)),
+    ):
+        tmp = directory / f".{name}.tmp"
+        tmp.write_text(text, encoding="utf-8")
+        os.replace(tmp, directory / name)
+        written.append(directory / name)
+    return written

trainmeter/export/wandb.py

@@ -0,0 +1,87 @@
+"""Export a finished run to Weights & Biases. Explicit, and only when the owner runs it.
+
+`tm export wandb [RUN] --project P` replays the run log and creates one W&B run: the series the
+loop reported under their own keys, the derived metrics (tokens/s, MFU, OFU, the utilization
+ladder, per-GPU counters) under `trainmeter/`, the passport as the run config and summary, and the
+passport files. Nothing leaves the node unless this command is run, and no price is sent.
+
+Verified against wandb 0.30.0 in offline mode and against a fake module in tests. Not run against
+the hosted service.
+"""
+
+from __future__ import annotations
+
+import importlib
+from types import ModuleType
+from typing import Any
+
+from .. import catalog
+from ..replay import Replayed
+
+DERIVED = [
+    m[0] for m in catalog.METRICS if m[0] not in ("train_loss", "val_loss", "lr", "grad_norm")
+]
+
+
+def _flat(prefix: str, obj: Any, out: dict[str, Any]) -> None:
+    if isinstance(obj, dict):
+        for k, v in obj.items():
+            _flat(f"{prefix}.{k}" if prefix else str(k), v, out)
+    elif isinstance(obj, (int, float, str, bool)) or obj is None:
+        out[prefix] = obj
+
+
+def rows(replayed: Replayed) -> dict[int, dict[str, float]]:
+    """Step -> {key: value}. Reported series keep their raw keys, derived ones are prefixed."""
+    tl = replayed.timeline
+    names = [n for n in tl.names() if not n.startswith("gpu.")]
+    reported = [n for n in names if n not in DERIVED]
+    per_gpu = [n for n in tl.names() if n.startswith("gpu.")]
+    table: dict[int, dict[str, float]] = {}
+    for group, prefix in ((reported, ""), (DERIVED, "trainmeter/"), (per_gpu, "trainmeter/")):
+        data = tl.query(group, "step")
+        for name, points in data.items():
+            for step, value in points:
+                table.setdefault(int(step), {})[f"{prefix}{name}"] = value
+    return table
+
+
+def export(
+    replayed: Replayed,
+    project: str,
+    entity: str | None = None,
+    name: str | None = None,
+    wandb: ModuleType | None = None,
+) -> str:
+    wandb = wandb or importlib.import_module("wandb")
+    passport = replayed.passport()
+    config: dict[str, Any] = {}
+    _flat(
+        "",
+        {k: passport[k] for k in ("model", "topology", "versions", "hardware") if k in passport},
+        config,
+    )
+    run = wandb.init(
+        project=project,
+        entity=entity,
+        name=name or replayed.path.name,
+        config=config,
+        job_type="trainmeter-import",
+        tags=["trainmeter"],
+    )
+    for step, row in sorted(rows(replayed).items()):
+        run.log(row, step=step)
+    summary: dict[str, Any] = {}
+    _flat(
+        "",
+        {k: passport[k] for k in ("training", "efficiency", "flops", "results") if k in passport},
+        summary,
+    )
+    run.summary.update({k: v for k, v in summary.items() if v is not None})
+    for fname in ("passport.json", "passport.md"):
+        path = replayed.path / fname
+        if path.is_file():
+            run.save(str(path), base_path=str(replayed.path), policy="now")
+    url = getattr(run, "url", None)
+    run.finish()
+    return url or f"offline run for {replayed.path.name}"

trainmeter/facts.py

@@ -0,0 +1,50 @@
+"""Facts with provenance, and the rule that decides between two values for the same key."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+# Highest precedence first. `measured` values have their own keys and never compete, so it sits
+# with `detected`.
+PRECEDENCE = ("declared", "emitted", "tapped", "detected", "measured", "assumed")
+
+
+def rank(source: str) -> int:
+    return PRECEDENCE.index(source) if source in PRECEDENCE else len(PRECEDENCE)
+
+
+@dataclass(frozen=True)
+class Fact:
+    key: str
+    value: Any
+    source: str
+    note: str | None = None
+
+
+class FactStore:
+    """Keeps, per key, the value from the highest-precedence source; the latest wins on a tie."""
+
+    def __init__(self) -> None:
+        self._facts: dict[str, Fact] = {}
+        self.conflicts: list[tuple[Fact, Fact]] = []  # (kept, ignored) with different values
+
+    def offer(self, key: str, value: Any, source: str, note: str | None = None) -> None:
+        new = Fact(key, value, source, note)
+        old = self._facts.get(key)
+        if old is None or rank(source) <= rank(old.source):
+            if old is not None and old.value != value and rank(source) < rank(old.source):
+                self.conflicts.append((new, old))
+            self._facts[key] = new
+        elif old.value != value:
+            self.conflicts.append((old, new))
+
+    def get(self, key: str) -> Fact | None:
+        return self._facts.get(key)
+
+    def value(self, key: str, default: Any = None) -> Any:
+        fact = self._facts.get(key)
+        return default if fact is None else fact.value
+
+    def items(self) -> dict[str, Fact]:
+        return dict(self._facts)

trainmeter/flops.py

@@ -0,0 +1,52 @@
+"""FLOPs per token for a dense decoder-only transformer, under named conventions.
+
+Frameworks disagree on what to count, and the disagreement is large: on one
+1.04B model the four conventions below span 5.83 to 6.69 GFLOP per token, about
+15%. A number without its convention is not comparable, so every function here
+takes the convention explicitly and there is no silent default in reports.
+
+Symbols: N total parameters, E embedding-table parameters (token and learned
+position tables, not an untied output head), L layers, d model width, T
+sequence length. Training counts forward plus backward, hence the factor 6.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+Convention = Literal["titan", "megatron", "hf", "all_params"]
+
+CONVENTIONS: dict[str, str] = {
+    "titan": "6(N-E) + 12*L*d*T  (TorchTitan, nanochat; full attention square)",
+    "megatron": "6(N-E) + 6*L*d*T  (Megatron-LM; causal mask halves attention)",
+    "hf": "6(N-E)  (HF Trainer total_flos; no attention term)",
+    "all_params": "6N + 12*L*d*T  (nanoGPT-style; embeddings counted as matmuls)",
+}
+
+
+@dataclass(frozen=True)
+class ModelShape:
+    n_params: int
+    n_embedding_params: int
+    n_layers: int
+    d_model: int
+    seq_len: int
+
+
+def flops_per_token(shape: ModelShape, convention: Convention) -> float:
+    dense = 6 * (shape.n_params - shape.n_embedding_params)
+    attn_full = 12 * shape.n_layers * shape.d_model * shape.seq_len
+    if convention == "titan":
+        return dense + attn_full
+    if convention == "megatron":
+        return dense + attn_full / 2
+    if convention == "hf":
+        return dense
+    if convention == "all_params":
+        return 6 * shape.n_params + attn_full
+    raise ValueError(f"unknown convention {convention!r}; expected one of {list(CONVENTIONS)}")
+
+
+def all_conventions(shape: ModelShape) -> dict[str, float]:
+    return {c: flops_per_token(shape, c) for c in CONVENTIONS}  # type: ignore[arg-type]

trainmeter/metrics.py

@@ -0,0 +1,68 @@
+"""Pure metric formulas. No I/O, no GPU: everything here is unit-testable."""
+
+from __future__ import annotations
+
+import math
+
+
+def mfu(flops_per_token: float, tokens_per_s: float, n_gpus: int, peak_flops: float) -> float:
+    """Model FLOPs utilization as a fraction of the dense peak of all GPUs."""
+    return flops_per_token * tokens_per_s / (n_gpus * peak_flops)
+
+
+def ofu(tensor_active: float, sm_clock_mhz: float, max_sm_clock_mhz: float) -> float:
+    """Counter-based MFU proxy: tensor pipe activity scaled by the clock ratio.
+
+    tensor_active is a fraction in [0, 1]. See arXiv 2605.20799. It counts busy
+    tensor-pipe cycles, so it includes recompute and padding that MFU excludes.
+    """
+    return tensor_active * sm_clock_mhz / max_sm_clock_mhz
+
+
+def step_intensity(mfu_value: float, dram_active: float, ridge: float) -> float:
+    """Aggregate useful FLOPs per byte of HBM traffic over a sampling window.
+
+    achieved FLOP/s = mfu * peak_flops; achieved B/s = dram_active * peak_bw,
+    so the ratio is ridge * mfu / dram_active. It is an average over many
+    kernels, not a roofline verdict for any one of them.
+    """
+    if dram_active <= 0:
+        return math.inf
+    return ridge * mfu_value / dram_active
+
+
+def eta_seconds(remaining_tokens: float, tokens_per_s: float) -> float:
+    if tokens_per_s <= 0:
+        return math.inf
+    return max(remaining_tokens, 0.0) / tokens_per_s
+
+
+def cost_usd(price_per_hour: float, billed_seconds: float) -> float:
+    return price_per_hour * billed_seconds / 3600.0
+
+
+def cost_per_million_tokens(price_per_hour: float, billed_seconds: float, tokens: float) -> float:
+    if tokens <= 0:
+        return math.inf
+    return cost_usd(price_per_hour, billed_seconds) / tokens * 1e6
+
+
+class Ema:
+    """Exponential moving average that ignores its first `warmup` samples.
+
+    The first steps include torch.compile and allocator warmup; averaging them
+    in makes the ETA wrong for a long time.
+    """
+
+    def __init__(self, alpha: float = 0.1, warmup: int = 3) -> None:
+        self.alpha = alpha
+        self.warmup = warmup
+        self.seen = 0
+        self.value: float | None = None
+
+    def update(self, x: float) -> float | None:
+        self.seen += 1
+        if self.seen <= self.warmup:
+            return self.value
+        self.value = x if self.value is None else self.alpha * x + (1 - self.alpha) * self.value
+        return self.value

trainmeter/passport.py

@@ -0,0 +1,327 @@
+"""The run passport: what goes into a model card, from a snapshot. Pure.
+
+Physics only: no price, no hostname, no working directory, no environment. The command line is
+the one recorded in the run header, where the values of secret-looking flags are already replaced.
+Every number keeps its convention and provenance, and what could not be computed is listed under
+`not_available` with the reason, never as a zero.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+SCHEMA = 1
+
+
+def _fact(facts: dict[str, Any], key: str) -> dict[str, Any] | None:
+    fact = facts.get(key)
+    if fact is None:
+        return None
+    out = {"value": fact["value"], "source": fact["source"]}
+    if fact.get("note"):
+        out["note"] = fact["note"]
+    return out
+
+
+def build(
+    snap: dict[str, Any],
+    header: dict[str, Any],
+    facts: dict[str, Any],
+    hello: dict[str, Any],
+    exit_info: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """`facts` is `{key: {value, source, note}}` as served by `/api/run`."""
+    passport: dict[str, Any] = {
+        "schema": SCHEMA,
+        "generated_by": f"trainmeter {header.get('tm_version')}",
+    }
+    passport["run"] = {
+        "command": header.get("argv"),
+        "elapsed_s": (snap.get("time") or {}).get("elapsed_s"),
+        "finished": (snap.get("time") or {}).get("finished", False),
+        **({"exit_status": exit_info["exit_status"]} if exit_info else {}),
+    }
+    model = {
+        key: _fact(facts, key)
+        for key in ("n_params", "n_embedding_params", "n_layers", "d_model", "seq_len")
+    }
+    model["moe_experts"] = _fact(facts, "moe_experts")
+    passport["model"] = {k: v for k, v in model.items() if v is not None}
+    fpt = snap.get("flops_per_token")
+    flops: dict[str, Any] = {}
+    if fpt is not None:
+        flops["per_token"] = fpt
+    if "conventions" in snap:
+        flops["conventions"] = snap["conventions"]
+    if "counted" in snap:
+        flops["counted"] = snap["counted"]
+    for key in ("flops_spent", "flops_required", "flops_observed"):
+        if key in snap:
+            flops[key] = snap[key]
+    passport["flops"] = flops
+    passport["training"] = {
+        k: snap[k]
+        for k in ("steps", "tokens_seen", "tokens_per_s", "budget_tokens", "progress")
+        if k in snap
+    }
+    passport["efficiency"] = {
+        k: snap[k] for k in ("mfu", "ofu", "mfu_ofu_gap", "arithmetic_intensity") if k in snap
+    }
+    if "goodput" in (snap.get("time") or {}):
+        passport["efficiency"]["goodput"] = snap["time"]["goodput"]
+    passport["hardware"] = {
+        k: snap[k] for k in ("n_gpus", "devices", "peak", "gpu_versions") if k in snap
+    }
+    passport["topology"] = {
+        "world_size": hello.get("world_size"),
+        **({"dp": _fact(facts, "dp")} if _fact(facts, "dp") else {}),
+    }
+    passport["versions"] = {
+        k: v
+        for k, v in {
+            "trainmeter": header.get("tm_version"),
+            "python": hello.get("python"),
+            "torch": hello.get("torch"),
+            "cuda": hello.get("cuda"),
+        }.items()
+        if v
+    }
+    passport["results"] = {
+        name: entry
+        for name, entry in (snap.get("series") or {}).items()
+        if name in ("train_loss", "val_loss")
+    }
+    if "parts" in (snap.get("time") or {}):
+        passport["time"] = {"parts": snap["time"]["parts"], "goodput": snap["time"].get("goodput")}
+    if "plan" in snap:
+        passport["plan"] = snap["plan"]
+    passport["warnings"] = snap.get("warnings", [])
+    passport["not_available"] = snap.get("absent", {})
+    return passport
+
+
+# ---- Markdown -----------------------------------------------------------------------------
+
+
+def _si(v: float) -> str:
+    for limit, suffix in ((1e15, "P"), (1e12, "T"), (1e9, "G"), (1e6, "M"), (1e3, "k")):
+        if abs(v) >= limit:
+            return f"{v / limit:.3g} {suffix}"
+    return f"{v:.4g}"
+
+
+def _hours(s: float) -> str:
+    """Seconds as the largest sensible unit: 12.3 s, 4.5 min, 1.25 h."""
+    if s < 60:
+        return f"{s:.1f} s"
+    if s < 3600:
+        return f"{s / 60:.1f} min"
+    return f"{s / 3600:.2f} h"
+
+
+def _val(entry: dict[str, Any] | None) -> str:
+    if not entry:
+        return "not available"
+    tag = f" ({entry['source']})" if entry.get("source") else ""
+    value = entry["value"]
+    return f"{_si(value) if isinstance(value, (int, float)) else value}{tag}"
+
+
+def to_markdown(p: dict[str, Any]) -> str:
+    """A model-card section. One sentence per line, tables for numbers."""
+    lines = ["# Run passport", ""]
+    run = p.get("run", {})
+    if run.get("command"):
+        lines += ["Command:", "", "```", " ".join(run["command"]), "```", ""]
+    if run.get("elapsed_s") is not None:
+        state = "finished" if run.get("finished") else "still running"
+        exit_text = f", exit status {run['exit_status']}" if "exit_status" in run else ""
+        lines += [f"Wall time {_hours(run['elapsed_s'])} ({state}{exit_text}).", ""]
+
+    model = p.get("model", {})
+    lines += ["## Model", ""]
+    if model:
+        lines += ["| Item | Value |", "| --- | --- |"]
+        labels = {
+            "n_params": "Parameters",
+            "n_embedding_params": "Embedding parameters",
+            "n_layers": "Layers",
+            "d_model": "Width (d_model)",
+            "seq_len": "Sequence length",
+            "moe_experts": "Experts (mixture of experts)",
+        }
+        lines += [f"| {labels[k]} | {_val(v)} |" for k, v in model.items() if k in labels]
+    else:
+        lines.append("Not available: the agent saw no model, and none was declared.")
+    lines.append("")
+
+    fl = p.get("flops", {})
+    lines += ["## FLOPs", ""]
+    fpt = fl.get("per_token")
+    if fpt:
+        if fpt.get("value") is not None:
+            lines.append(
+                f"FLOPs per token: {_si(fpt['value'])}, convention `{fpt['convention']}` "
+                f"({fpt['source']})."
+            )
+        else:
+            lo, hi = fpt["range"]
+            lines.append(
+                f"FLOPs per token: {_si(lo)} to {_si(hi)} across the four conventions, "
+                "convention not chosen."
+            )
+    else:
+        lines.append("FLOPs per token: not available.")
+    lines.append("")
+    if fl.get("counted"):
+        c = fl["counted"]
+        lines += [
+            f"Counted (`FlopCounterMode`) on one real step: {_si(c['per_token'])} FLOPs per token.",
+            "It counts executed operators, so it includes recompute and is not model FLOPs.",
+            "",
+        ]
+    rows = [
+        (name, label)
+        for name, label in (
+            ("flops_spent", "FLOPs spent (modeled)"),
+            ("flops_required", "FLOPs required (budget)"),
+        )
+        if name in fl
+    ]
+    if rows or "flops_observed" in fl:
+        lines += ["| Quantity | Value |", "| --- | --- |"]
+        lines += [f"| {label} | {_si(fl[name])} |" for name, label in rows]
+        if "flops_observed" in fl:
+            lines.append(
+                f"| FLOPs observed (counters, approximate) | {_si(fl['flops_observed']['value'])} |"
+            )
+        lines.append("")
+
+    tr = p.get("training", {})
+    lines += ["## Training", ""]
+    if tr:
+        lines += ["| Item | Value |", "| --- | --- |"]
+        for key, label in (("steps", "Optimizer steps"), ("tokens_seen", "Tokens seen")):
+            if key in tr:
+                lines.append(f"| {label} | {_si(tr[key])} |")
+        if "tokens_per_s" in tr:
+            lines.append(f"| Tokens per second (moving average) | {_si(tr['tokens_per_s'])} |")
+        if "budget_tokens" in tr:
+            lines.append(f"| Token budget | {_si(tr['budget_tokens'])} |")
+        if "progress" in tr:
+            lines.append(f"| Progress | {tr['progress'] * 100:.1f}% |")
+    else:
+        lines.append("Not available: no steps were seen.")
+    lines.append("")
+
+    res = p.get("results", {})
+    if res:
+        lines += ["## Results", "", "| Series | Last value | Source |", "| --- | --- | --- |"]
+        for name, e in res.items():
+            lines.append(f"| {name} | {e['value']:.4g} | {e.get('source', '')} |")
+        lines.append("")
+
+    eff = p.get("efficiency", {})
+    lines += ["## Efficiency", ""]
+    if eff:
+        lines += ["| Metric | Value | Label |", "| --- | --- | --- |"]
+        if "mfu" in eff:
+            m = eff["mfu"]
+            shown = (
+                f"{m['value'] * 100:.1f}%"
+                if "value" in m
+                else (f"{m['range'][0] * 100:.1f}-{m['range'][1] * 100:.1f}%")
+            )
+            lines.append(
+                f"| MFU | {shown} | convention `{m['convention']}`, peak {m['peak']} dense BF16 |"
+            )
+        if "ofu" in eff:
+            lines.append(f"| OFU | {eff['ofu']['value'] * 100:.1f}% | counters |")
+        if "arithmetic_intensity" in eff:
+            a = eff["arithmetic_intensity"]
+            lines.append(f"| Arithmetic intensity | {_si(a['value'])} FLOP/byte | {a['note']} |")
+        if "goodput" in eff:
+            lines.append(
+                f"| Goodput | {eff['goodput'] * 100:.1f}% | time in counted steps / wall time |"
+            )
+    else:
+        lines.append("Not available.")
+    lines.append("")
+
+    hw = p.get("hardware", {})
+    lines += ["## Hardware", ""]
+    if hw.get("n_gpus"):
+        n = hw["n_gpus"]
+        lines.append(f"{n['value']} GPU(s) ({n['source']}).")
+        names = sorted({d.get("name") for d in hw.get("devices", []) if d.get("name")})
+        if names:
+            lines.append(f"Model: {', '.join(names)}.")
+        if hw.get("peak"):
+            pk = hw["peak"]
+            lines.append(
+                f"Dense BF16 peak {_si(pk['bf16_dense_flops'])}FLOP/s and HBM "
+                f"{_si(pk['hbm_bytes_per_s'])}B/s per GPU, from {pk['source']}."
+            )
+        if hw.get("gpu_versions"):
+            lines.append(
+                "Driver: " + ", ".join(f"{k} {v}" for k, v in hw["gpu_versions"].items()) + "."
+            )
+    else:
+        lines.append("Not available: no GPU was seen.")
+    lines.append("")
+
+    topo, ver = p.get("topology", {}), p.get("versions", {})
+    lines += ["## Software and topology", ""]
+    if topo.get("world_size"):
+        lines.append(f"World size {topo['world_size']}.")
+    lines += [f"- {k}: {v}" for k, v in ver.items()]
+    lines.append("")
+
+    if p.get("time", {}).get("parts"):
+        lines += ["## Where the time went", "", "| Part | Time | Basis |", "| --- | --- | --- |"]
+        for part in p["time"]["parts"]:
+            basis = (
+                "measured"
+                if part["measured"]
+                else ("not measured" if part["s"] is None else "remainder")
+            )
+            secs = "-" if part["s"] is None else _hours(part["s"])
+            lines.append(f"| {part['label']} | {secs} | {basis} |")
+        lines.append("")
+
+    plan = p.get("plan")
+    if plan:
+        lines += [
+            "## Plan against actual",
+            "",
+            "| Item | Planned | Measured |",
+            "| --- | --- | --- |",
+        ]
+        if "mfu" in plan:
+            m = plan["mfu"]["measured"]
+            measured = (
+                "not available" if m is None else f"{m * 100:.1f}% ({plan['mfu']['convention']})"
+            )
+            lines.append(f"| MFU | {plan['mfu']['planned'] * 100:.1f}% | {measured} |")
+        if "tokens_per_s" in plan:
+            t = plan["tokens_per_s"]
+            lines.append(
+                f"| Tokens per second | {_si(t['planned'])} | "
+                f"{'not available' if t['measured'] is None else _si(t['measured'])} |"
+            )
+        if "hours" in plan:
+            h = plan["hours"]
+            proj = (
+                "not available" if h["projected"] is None else f"{h['projected']:.2f} h (projected)"
+            )
+            lines.append(f"| Wall time | {h['planned']:.2f} h | {proj} |")
+        lines.append("")
+
+    if p.get("warnings"):
+        lines += ["## Warnings", ""] + [f"- {w}" for w in p["warnings"]] + [""]
+    if p.get("not_available"):
+        lines += ["## Not available", ""]
+        lines += [f"- {k}: {v}" for k, v in p["not_available"].items()]
+        lines.append("")
+    lines.append(f"Generated by {p.get('generated_by')}.")
+    return "\n".join(lines) + "\n"