trainmeter 0.0.2__py3-none-any.whl

trainmeter/__init__.py

@@ -0,0 +1,17 @@
+"""trainmeter: where your training FLOPs go, and how well the GPUs are used."""
+
+from .emit import emit
+from .flops import CONVENTIONS, ModelShape, all_conventions, flops_per_token
+from .peaks import Peak, lookup
+
+__version__ = "0.0.2"
+
+__all__ = [
+    "CONVENTIONS",
+    "ModelShape",
+    "Peak",
+    "all_conventions",
+    "emit",
+    "flops_per_token",
+    "lookup",
+]

trainmeter/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import main
+
+sys.exit(main())

trainmeter/agent/__init__.py

@@ -0,0 +1 @@
+"""The passive agent that `tm` injects into the training process."""

trainmeter/agent/bootstrap/sitecustomize.py

@@ -0,0 +1,28 @@
+"""Loaded by Python at startup because `tm` puts this directory first on PYTHONPATH.
+
+It starts the agent, then runs any other `sitecustomize` that the environment already had.
+Every failure here is swallowed: the training process must start exactly as it would without tm.
+"""
+
+import os
+import sys
+
+try:
+    from trainmeter.agent.core import install
+
+    install()
+except Exception:  # noqa: BLE001
+    pass
+
+try:
+    import importlib.machinery
+    import importlib.util
+
+    _here = os.path.dirname(os.path.abspath(__file__))
+    _rest = [p for p in sys.path if os.path.abspath(p or ".") != _here]
+    _spec = importlib.machinery.PathFinder.find_spec("sitecustomize", _rest)
+    if _spec is not None and _spec.loader is not None:
+        _module = importlib.util.module_from_spec(_spec)
+        _spec.loader.exec_module(_module)
+except Exception:  # noqa: BLE001
+    pass

trainmeter/agent/core.py

@@ -0,0 +1,578 @@
+"""The agent: passive probes that run inside the training process.
+
+Standard library only, plus the records codec through the sender. It never imports torch itself:
+it waits until the user's code does, then attaches. It never synchronizes a device, never reads a
+tensor value, never allocates on a GPU, and never raises: a probe that fails disables itself,
+warns once on stderr, and nothing else changes.
+
+Verified against real torch 2.14.0 on CPU. The device probe is verified only against a fake
+`torch.cuda`, not on real hardware.
+"""
+
+from __future__ import annotations
+
+import atexit
+import importlib
+import importlib.machinery
+import os
+import sys
+import time
+from typing import Any
+
+from . import taps
+from .sender import Sender, get_sender
+from .wrap import patch
+
+PROBES_ENV = "TRAINMETER_PROBES"  # comma-separated allow-list; unset means every probe
+TAP_ENV = "TRAINMETER_TAP"  # "0" turns the logger taps off
+COUNT_STEP_ENV = "TRAINMETER_COUNT_STEP"  # optimizer steps to skip before counting FLOPs
+COUNT_AFTER_STEP = 5  # warmup, compile and allocator effects are over by then
+
+# Attributes that name the layer count and the width in the configs of well-known model families.
+LAYER_ATTRS = ("num_hidden_layers", "n_layer", "num_layers", "n_layers")
+WIDTH_ATTRS = ("hidden_size", "n_embd", "d_model", "dim")
+EXPERT_ATTRS = ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")
+
+FLUSH_INTERVAL = 0.05  # at most about 20 step records per second per rank
+
+
+class Agent:
+    def __init__(self, sender: Sender) -> None:
+        self.sender = sender
+        self.rank = int(os.environ["RANK"]) if "RANK" in os.environ else None
+        self.world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        self.torch: Any = None
+        self.disabled: set[str] = set()
+        self._seq_len_sent = False
+        # step state
+        self._saw_forward = False
+        self._forward_tracking = False
+        self._step_index = 0
+        self._micro = 0
+        self._tokens: int | None = 0  # None once any micro-batch of the step had unknown tokens
+        # batch state
+        self._root: Any = None
+        self._global_handle: Any = None
+        self._post_handle: Any = None
+        self._barrier = True
+        # aggregation
+        self._pending_k = 0
+        self._pending_micro = 0
+        self._pending_tokens: int | None = 0
+        self._last_flush = 0.0
+        # phases
+        self.attached: set[str] = set()
+        self._data_s = 0.0
+        self._data_n = 0
+        self._eval_start: float | None = None
+        self._tapped = False
+        # counted FLOPs: idle -> active (one whole optimizer step) -> done
+        self._count_state = "idle"
+        self._count_mode: Any = None
+        self._count_after = int(os.environ.get(COUNT_STEP_ENV, COUNT_AFTER_STEP))
+
+    # ---- plumbing ----------------------------------------------------------------------
+
+    def _fail(self, probe: str, exc: BaseException) -> None:
+        if probe not in self.disabled:
+            self.disabled.add(probe)
+            print(f"trainmeter agent: probe {probe!r} disabled: {exc!r}", file=sys.stderr)
+
+    def _send(self, kind: str, d: dict[str, Any]) -> None:
+        self.sender.send(kind, "agent", d, self.rank if self.rank is not None else 0)
+
+    def _fact(self, key: str, value: Any) -> None:
+        self._send("fact", {"key": key, "value": value, "source": "detected"})
+
+    def _enabled(self, probe: str) -> bool:
+        allowed = os.environ.get(PROBES_ENV)
+        if allowed is not None and probe not in {p.strip() for p in allowed.split(",")}:
+            return False
+        return not (probe == "tap" and os.environ.get(TAP_ENV) == "0")
+
+    # ---- attaching ---------------------------------------------------------------------
+
+    def attach(self, torch: Any) -> None:
+        self.torch = torch
+        try:
+            optimizer_module = importlib.import_module("torch.optim.optimizer")
+            optimizer_module.register_optimizer_step_post_hook(self._on_optimizer_step)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("steps", exc)
+            return
+        try:
+            from torch.nn.modules.module import register_module_forward_pre_hook
+
+            self._global_handle = register_module_forward_pre_hook(self._on_any_forward)
+            self._forward_tracking = True
+            from torch.nn.modules.module import register_module_forward_hook
+
+            # The global pre-hook cannot see keyword arguments; this one can, after forward.
+            self._post_handle = register_module_forward_hook(self._on_first_post, with_kwargs=True)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("batch", exc)
+        atexit.register(self.flush)
+        self._attach_phases(torch)
+        if self._enabled("flops"):
+            self.attached.add("flops")
+
+    # ---- phases: data wait, checkpoints and eval spans ----------------------------------
+
+    def _attach_phases(self, torch: Any) -> None:
+        for probe, attach in (
+            ("dataloader", self._attach_dataloader),
+            ("checkpoint", self._attach_checkpoint),
+            ("eval", self._attach_eval),
+        ):
+            if not self._enabled(probe):
+                continue
+            try:
+                attach(torch)
+                self.attached.add(probe)
+            except Exception as exc:  # noqa: BLE001
+                self._fail(probe, exc)
+
+    def _attach_dataloader(self, torch: Any) -> None:
+        iterator = importlib.import_module("torch.utils.data.dataloader")._BaseDataLoaderIter
+
+        def factory(orig: Any) -> Any:
+            def __next__(it: Any) -> Any:  # noqa: N807
+                start = time.monotonic()
+                try:
+                    return orig(it)
+                finally:
+                    self._add_data_wait(time.monotonic() - start)
+
+            return __next__
+
+        if patch(iterator, "__next__", factory) is None:
+            raise RuntimeError("DataLoader iterator not found or already wrapped")
+
+    def _add_data_wait(self, seconds: float) -> None:
+        try:
+            if "dataloader" not in self.disabled and self._eval_start is None:
+                self._data_s += seconds  # waits inside an eval span belong to the eval span
+                self._data_n += 1
+        except Exception as exc:  # noqa: BLE001
+            self._fail("dataloader", exc)
+
+    def _attach_checkpoint(self, torch: Any) -> None:
+        def factory(orig: Any) -> Any:
+            def save(*args: Any, **kwargs: Any) -> Any:
+                start = time.monotonic()
+                try:
+                    return orig(*args, **kwargs)
+                finally:
+                    self._span("checkpoint", time.monotonic() - start, "checkpoint")
+
+            return save
+
+        original = patch(torch, "save", factory)
+        if original is None:
+            raise RuntimeError("torch.save not found or already wrapped")
+        serialization = getattr(torch, "serialization", None)
+        if getattr(serialization, "save", None) is original:
+            serialization.save = torch.save  # one wrapper, so a save is timed once
+
+    def _attach_eval(self, torch: Any) -> None:
+        def factory(orig: Any) -> Any:
+            def train(module: Any, mode: bool = True) -> Any:
+                self._on_train_toggle(module, mode)
+                return orig(module, mode)
+
+            return train
+
+        if patch(torch.nn.Module, "train", factory) is None:
+            raise RuntimeError("Module.train not found or already wrapped")
+
+    def _on_train_toggle(self, module: Any, mode: Any) -> None:
+        """An eval span runs from `eval()` on the root module to the next `train()` on it."""
+        try:
+            if "eval" in self.disabled or module is not self._root:
+                return
+            if not mode:
+                if self._eval_start is None:
+                    self._eval_start = time.monotonic()
+                    if self._count_state == "active":
+                        self._cancel_count()  # eval FLOPs are not training FLOPs
+            elif self._eval_start is not None:
+                start, self._eval_start = self._eval_start, None
+                self._span("eval", time.monotonic() - start, "eval")
+        except Exception as exc:  # noqa: BLE001
+            self._fail("eval", exc)
+
+    def _span(self, name: str, seconds: float, probe: str) -> None:
+        try:
+            if probe not in self.disabled:
+                self._send("phase", {"name": name, "s": seconds, "n": 1})
+        except Exception as exc:  # noqa: BLE001
+            self._fail(probe, exc)
+
+    # ---- logger taps ------------------------------------------------------------------
+
+    def _report(self, key: str, value: float, step: int | None, via: str) -> None:
+        if self.rank not in (None, 0):
+            return  # one copy of a series, whichever rank the loop logs from
+        d: dict[str, Any] = {"key": key, "value": value, "via": via}
+        if step is not None:
+            d["step"] = step
+        self.sender.send("metric", "tap", d, self.rank if self.rank is not None else 0)
+
+    def _tap_fail(self, via: str, exc: BaseException) -> None:
+        self._fail(f"tap:{via}", exc)
+
+    def attach_tracker(self, module_name: str) -> None:
+        """Called once the tracker's module has finished importing."""
+        if not self._enabled("tap"):
+            return
+        try:
+            taps.ADAPTERS[module_name](self._report, self._tap_fail, self._tap_disabled)
+            self.attached.add("tap")
+        except Exception as exc:  # noqa: BLE001
+            self._fail(f"tap:{module_name}", exc)
+
+    def on_import(self, name: str, module: Any) -> None:
+        if name == "torch":
+            self.attach(module)
+        else:
+            self.attach_tracker(name)
+
+    def _tap_disabled(self, via: str) -> bool:
+        return f"tap:{via}" in self.disabled
+
+    # ---- batch probe -------------------------------------------------------------------
+
+    def _training(self, module: Any) -> bool:
+        return bool(getattr(module, "training", False)) and self.torch.is_grad_enabled()
+
+    def _on_any_forward(self, module: Any, args: Any) -> None:
+        """Global pre-hook, used only until the root module is found."""
+        try:
+            if "batch" in self.disabled or self._root is not None:
+                return
+            training = self._training(module)
+            if training:
+                self._saw_forward = True
+            if self._barrier:
+                self._barrier = False
+                if training:
+                    self._found_root(module, args)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("batch", exc)
+
+    def _found_root(self, module: Any, args: Any) -> None:
+        """The first training-mode module called after a barrier is the outermost one."""
+        self._root = module
+        if self._global_handle is not None:
+            self._global_handle.remove()
+            self._global_handle = None
+        module.register_forward_pre_hook(self._on_root_forward, with_kwargs=True)
+        self._micro += 1  # its tokens are counted in the post-hook, which sees keyword arguments
+        try:
+            self._params(module)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("params", exc)
+
+    def _on_first_post(
+        self, module: Any, args: Any, kwargs: Any = None, output: Any = None
+    ) -> None:
+        """Counts the tokens of the very first call of the root module, then removes itself."""
+        try:
+            if module is not self._root or self._post_handle is None:
+                return
+            self._post_handle.remove()
+            self._post_handle = None
+            self._count_tokens(args, kwargs or {})
+        except Exception as exc:  # noqa: BLE001
+            self._fail("batch", exc)
+
+    def _on_root_forward(self, module: Any, args: Any, kwargs: Any) -> None:
+        try:
+            if "batch" in self.disabled or not self._training(module):
+                return
+            self._saw_forward = True
+            self._micro += 1
+            if self._micro == 1:
+                self._maybe_start_count()
+            self._count_tokens(args, kwargs)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("batch", exc)
+
+    # ---- counted FLOPs ----------------------------------------------------------------
+
+    def _maybe_start_count(self) -> None:
+        """Count one whole optimizer step, from its first forward to the step. Never raises.
+
+        The mode counts operators that execute, from their shapes: no device synchronization, no
+        allocation, no change to any value. It uses `FlopCounterMode` and swallows errors in the
+        FLOPs formulas, so a formula that cannot handle an operator cannot reach the loop.
+        """
+        if (
+            self._count_state != "idle"
+            or "flops" in self.disabled
+            or not self._enabled("flops")
+            or self._step_index < self._count_after
+            or self._eval_start is not None
+        ):
+            return
+        try:
+            from torch.utils.flop_counter import FlopCounterMode
+
+            class Guarded(FlopCounterMode):
+                def _count_flops(self, func_packet: Any, out: Any, args: Any, kwargs: Any) -> Any:
+                    try:
+                        return super()._count_flops(func_packet, out, args, kwargs)
+                    except Exception:  # noqa: BLE001 - the operator already ran
+                        return out
+
+            mode = Guarded(display=False)
+            mode.__enter__()
+            self._count_mode, self._count_state = mode, "active"
+        except Exception as exc:  # noqa: BLE001
+            self._count_state = "done"
+            self._fail("flops", exc)
+
+    def _cancel_count(self) -> None:
+        mode, self._count_mode = self._count_mode, None
+        self._count_state = "idle"
+        if mode is not None:
+            try:
+                mode.__exit__(None, None, None)
+            except Exception as exc:  # noqa: BLE001
+                self._fail("flops", exc)
+
+    def _finish_count(self) -> None:
+        """Called at the optimizer step that ends the counted step."""
+        mode, self._count_mode = self._count_mode, None
+        self._count_state = "done"
+        try:
+            mode.__exit__(None, None, None)
+            total = int(mode.get_total_flops())
+            tokens = self._tokens
+            if total <= 0 or not tokens:
+                print(
+                    "trainmeter agent: no FLOPs counted for the sampled step (torch.compile, "
+                    "custom kernels, or unknown token count); the counted convention stays absent",
+                    file=sys.stderr,
+                )
+                return
+            self._send(
+                "fact",
+                {
+                    "key": "counted_flops_per_token",
+                    "value": total / tokens,
+                    "source": "measured",
+                    "note": f"FlopCounterMode, step {self._step_index + 1}, {self._micro} "
+                    f"micro-batches, {tokens} tokens on this rank",
+                },
+            )
+        except Exception as exc:  # noqa: BLE001
+            self._fail("flops", exc)
+
+    def _count_tokens(self, args: Any, kwargs: Any) -> None:
+        torch = self.torch
+        for value in (*args, *kwargs.values()):
+            if (
+                torch.is_tensor(value)
+                and value.ndim == 2
+                and not value.dtype.is_floating_point
+                and not value.dtype.is_complex
+                and value.dtype != torch.bool
+            ):
+                if self._tokens is not None:
+                    self._tokens += int(value.numel())
+                if not self._seq_len_sent:
+                    self._seq_len_sent = True
+                    self._fact("seq_len", int(value.shape[1]))
+                return
+        self._tokens = None
+
+    def _params(self, root: Any) -> None:
+        torch = self.torch
+        total = sum(p.numel() for p in root.parameters())  # shared weights are counted once
+        seen: set[int] = set()
+        embedding = 0
+        for m in root.modules():
+            if isinstance(m, torch.nn.Embedding) and id(m.weight) not in seen:
+                seen.add(id(m.weight))
+                embedding += m.weight.numel()
+        self._fact("n_params", int(total))
+        self._fact("n_embedding_params", int(embedding))
+        if self._enabled("config"):
+            try:
+                self._config(root)
+            except Exception as exc:  # noqa: BLE001
+                self._fail("config", exc)
+
+    def _config(self, root: Any) -> None:
+        """Layers and width from a model config, when the family has one. Low confidence."""
+        model = root
+        while True:  # look through DDP, FSDP and torch.compile wrappers
+            inner = getattr(model, "module", None) or getattr(model, "_orig_mod", None)
+            if inner is None or inner is model:
+                break
+            model = inner
+        config = getattr(model, "config", None)
+        if config is None:
+            return
+        for key, names in (("n_layers", LAYER_ATTRS), ("d_model", WIDTH_ATTRS)):
+            for name in names:
+                value = getattr(config, name, None)
+                if isinstance(value, int) and not isinstance(value, bool) and value > 0:
+                    self._send(
+                        "fact",
+                        {
+                            "key": key,
+                            "value": value,
+                            "source": "detected",
+                            "note": f"model.config.{name}",
+                        },
+                    )
+                    break
+        for name in EXPERT_ATTRS:
+            value = getattr(config, name, None)
+            if isinstance(value, int) and not isinstance(value, bool) and value > 1:
+                self._fact("moe_experts", value)
+                break
+
+    # ---- steps probe -------------------------------------------------------------------
+
+    def _on_optimizer_step(self, optimizer: Any, args: Any, kwargs: Any) -> None:
+        try:
+            if "steps" in self.disabled:
+                return
+            tracking = self._forward_tracking and "batch" not in self.disabled
+            if tracking and not self._saw_forward:
+                return  # optimizer steps back to back count as one step
+            self._saw_forward = False
+            if self._root is None:
+                self._barrier = True
+            if self._count_state == "active":
+                self._finish_count()
+            self._step_index += 1
+            if self._step_index == 1:
+                self._first_step()
+            self._record_step()
+        except Exception as exc:  # noqa: BLE001
+            self._fail("steps", exc)
+
+    def _first_step(self) -> None:
+        torch = self.torch
+        hello: dict[str, Any] = {
+            "python": sys.version.split()[0],
+            "torch": torch.__version__,
+            "cuda": getattr(torch.version, "cuda", None),
+            "world_size": self.world_size,
+            "probes": sorted(self.attached),
+        }
+        if self.rank is not None:
+            hello["local_rank"] = int(os.environ.get("LOCAL_RANK", "0"))
+        self._send("hello", hello)
+        try:
+            if torch.cuda.is_initialized():  # never initializes CUDA itself
+                index = torch.cuda.current_device()
+                props = torch.cuda.get_device_properties(index)
+                device = {"name": props.name, "index": index}
+                uuid = getattr(props, "uuid", None)
+                if uuid:
+                    device["uuid"] = str(uuid)
+                self._fact("device", device)
+        except Exception as exc:  # noqa: BLE001
+            self._fail("device", exc)
+
+    def _record_step(self) -> None:
+        self._pending_k += 1
+        self._pending_micro += self._micro
+        if self._tokens is None or self._pending_tokens is None:
+            self._pending_tokens = None
+        else:
+            self._pending_tokens += self._tokens
+        self._micro, self._tokens = 0, 0
+        if time.monotonic() - self._last_flush >= FLUSH_INTERVAL:
+            self.flush()
+
+    def flush(self) -> None:
+        if self._data_n:
+            self._send("phase", {"name": "data", "s": self._data_s, "n": self._data_n})
+            self._data_s, self._data_n = 0.0, 0
+        if not self._pending_k:
+            return
+        d: dict[str, Any] = {
+            "n": self._step_index,
+            "k": self._pending_k,
+            "micro": self._pending_micro,
+        }
+        if self._pending_tokens is not None and self._pending_tokens > 0:
+            d["tokens"] = self._pending_tokens
+        self._send("step", d)
+        self._pending_k = self._pending_micro = 0
+        self._pending_tokens = 0
+        self._last_flush = time.monotonic()
+
+
+class _Loader:
+    """Wraps a module's own loader so the agent can attach right after the import finishes."""
+
+    def __init__(self, loader: Any, finder: _Finder, name: str) -> None:
+        self._loader = loader
+        self._finder = finder
+        self._name = name
+
+    def create_module(self, spec: Any) -> Any:
+        return self._loader.create_module(spec)
+
+    def exec_module(self, module: Any) -> None:
+        self._loader.exec_module(module)
+        self._finder.imported(self._name, module)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._loader, name)
+
+
+class _Finder:
+    """Meta path finder that runs the agent's attach step once for each module it waits for."""
+
+    def __init__(self, agent: Agent, names: set[str]) -> None:
+        self._agent = agent
+        self._waiting = set(names)
+
+    def find_spec(self, fullname: str, path: Any = None, target: Any = None) -> Any:
+        if fullname not in self._waiting:
+            return None
+        spec = importlib.machinery.PathFinder.find_spec(fullname, path)
+        if spec is None or spec.loader is None:
+            return None
+        spec.loader = _Loader(spec.loader, self, fullname)
+        return spec
+
+    def imported(self, name: str, module: Any) -> None:
+        self._waiting.discard(name)
+        if not self._waiting:
+            sys.meta_path[:] = [f for f in sys.meta_path if f is not self]
+        try:
+            self._agent.on_import(name, module)
+        except Exception as exc:  # noqa: BLE001
+            self._agent._fail(f"attach:{name}", exc)
+
+
+def install() -> None:
+    """Called from the bootstrap `sitecustomize`. Does nothing outside `tm` or when switched off."""
+    if os.environ.get("TRAINMETER_AGENT") == "0":
+        return
+    sender = get_sender()
+    if sender is None:
+        return
+    agent = Agent(sender)
+    waiting: set[str] = set()
+    torch = sys.modules.get("torch")
+    if torch is not None and hasattr(torch, "optim"):
+        agent.attach(torch)
+    else:
+        waiting.add("torch")
+    for name in taps.ADAPTERS:
+        if name in sys.modules:
+            agent.attach_tracker(name)
+        else:
+            waiting.add(name)
+    if waiting:
+        sys.meta_path.insert(0, _Finder(agent, waiting))