mfu-tracker 0.1.0__py3-none-any.whl

mfu_tracker/__init__.py

@@ -0,0 +1,18 @@
+"""mfu-tracker: lightweight MFU and MBU tracking for PyTorch models."""
+from .flops import flash_attn_flops, param_bytes, profile_flops
+from .gpu import GPUSpec, get_gpu_spec
+from .optim import MFUOptimizerWrapper
+from .tracker import UtilizationResult, compute_mbu, compute_mfu, track
+
+__all__ = [
+    "track",
+    "compute_mfu",
+    "compute_mbu",
+    "profile_flops",
+    "flash_attn_flops",
+    "param_bytes",
+    "get_gpu_spec",
+    "GPUSpec",
+    "UtilizationResult",
+    "MFUOptimizerWrapper",
+]

mfu_tracker/flops.py

@@ -0,0 +1,157 @@
+"""FLOP counting — FlopCounterMode (PyTorch 2.1+) with thop fallback."""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+
+def flash_attn_flops(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+    *,
+    causal: bool = True,
+    with_backward: bool = False,
+) -> int:
+    """
+    Analytical FLOP count for one flash attention call (escape hatch).
+
+    ``profile_flops()`` already counts ``F.scaled_dot_product_attention``
+    automatically on CUDA via ``FlopCounterMode``. Use this only if your model
+    calls the ``flash_attn`` C extension directly (``flash_attn_func``), which
+    is rare in modern codebases.
+
+    Formula:
+        causal=True  → 2 * batch * seq_len² * num_heads * head_dim FLOPs
+        causal=False → 4 * batch * seq_len² * num_heads * head_dim FLOPs
+    """
+    fwd_flops = 2 * batch * seq_len * seq_len * num_heads * head_dim
+    if not causal:
+        fwd_flops *= 2
+    return fwd_flops * 3 if with_backward else fwd_flops
+
+
+def _profile_with_flop_counter(target: nn.Module, call_args: tuple) -> int:
+    """Profile using torch.utils.flop_counter.FlopCounterMode (PyTorch 2.1+).
+
+    Works at the ATen dispatch level — counts F.scaled_dot_product_attention
+    (SDPA / native flash attention) on CUDA automatically. Returns -1 on failure.
+    """
+    try:
+        from torch.utils.flop_counter import FlopCounterMode
+    except ImportError:
+        return -1
+
+    was_training = target.training
+    target.eval()
+    try:
+        with torch.no_grad(), FlopCounterMode(display=False) as flop_counter:
+            target(*call_args)
+        return int(flop_counter.get_total_flops())
+    except Exception:
+        return -1
+    finally:
+        target.train(was_training)
+
+
+def _profile_with_thop(target: nn.Module, call_args: tuple) -> int:
+    """Profile using thop — fallback for PyTorch < 2.1 or unsupported models."""
+    from thop import profile
+
+    macs, _ = profile(target, inputs=call_args, verbose=False)
+    return int(macs * 2)
+
+
+def profile_flops(
+    model: nn.Module,
+    args: Optional[tuple] = None,
+    kwargs: Optional[dict[str, Any]] = None,
+    *,
+    with_backward: bool = True,
+) -> int:
+    """
+    Count FLOPs for one forward (or forward+backward) pass through *model*.
+
+    Uses ``torch.utils.flop_counter.FlopCounterMode`` (PyTorch 2.1+), which
+    hooks at the ATen operator level and automatically counts
+    ``F.scaled_dot_product_attention`` (SDPA) on CUDA — covering virtually all
+    modern transformer attention implementations without any manual correction.
+    Falls back to ``thop`` on older PyTorch versions.
+
+    **Quantized models (bitsandbytes INT8 / NF4):**
+    Both counters operate on the Python/ATen graph and cannot see inside opaque
+    bitsandbytes CUDA kernels. In practice the FLOPs reported are close to
+    correct because NF4 (used by QLoRA) dequantizes weights to fp16 before the
+    matmul, so the actual computation is a standard fp16 GEMM. INT8 similarly
+    performs an fp16-equivalent matmul after dequantization in most bitsandbytes
+    kernels. Pass the appropriate ``dtype`` to ``track()`` / ``compute_mfu()``
+    to select the correct peak TFLOPS ceiling::
+
+        # QLoRA (NF4 base + fp16 LoRA adapters) — adapters run in fp16
+        flops = profile_flops(model, kwargs=batch, with_backward=False)
+        with track(flops, param_bytes(model), dtype="fp16", spec=spec) as r:
+            ...
+
+        # Pure INT8 inference
+        with track(flops, param_bytes(model), dtype="int8", spec=spec) as r:
+            ...
+
+    **PEFT / LoRA MBU:**
+    ``param_bytes(model)`` counts all parameters including the frozen base,
+    which is correct for the forward pass (all weights are read from memory).
+    For a backward-pass MBU estimate that excludes frozen weights, use
+    ``param_bytes(model, trainable_only=True)``::
+
+        # Backward MBU for a LoRA-finetuned model
+        active_bytes = param_bytes(model, trainable_only=True)
+
+    Args:
+        model:         The nn.Module to profile. Should be on CUDA for accurate
+                       SDPA counts.
+        args:          Positional inputs passed to model(*args).
+        kwargs:        Keyword-only inputs (e.g. HF models). Baked into a thin
+                       wrapper since both profilers call ``model(*inputs)``.
+        with_backward: Include backward-pass FLOPs (default True for training).
+                       Backward ≈ 2× forward → 3× total.
+
+    Returns:
+        Total integer FLOP count for one step.
+    """
+    # kwargs-only models need a wrapper — both profilers call target(*call_args)
+    if kwargs:
+        _kw = kwargs
+
+        class _KwargsAdapter(nn.Module):
+            def forward(self):
+                return model(**(args[0] if args and isinstance(args[0], dict) else _kw))
+
+        target: nn.Module = _KwargsAdapter()
+        call_args: tuple = ()
+    else:
+        target = model
+        call_args = args if args is not None else ()
+
+    forward_flops = _profile_with_flop_counter(target, call_args)
+    if forward_flops < 0:
+        forward_flops = _profile_with_thop(target, call_args)
+
+    return forward_flops * 3 if with_backward else forward_flops
+
+
+def param_bytes(model: nn.Module, *, trainable_only: bool = False) -> int:
+    """
+    Total bytes occupied by model parameters (for MBU calculation).
+
+    Args:
+        trainable_only: If True, count only parameters that require grad.
+                        Useful for PEFT/LoRA backward-pass MBU estimates.
+    """
+    params = (
+        (p for p in model.parameters() if p.requires_grad)
+        if trainable_only
+        else model.parameters()
+    )
+    return sum(p.numel() * p.element_size() for p in params)

mfu_tracker/gpu.py

@@ -0,0 +1,124 @@
+"""GPU capability querying and peak throughput derivation."""
+from __future__ import annotations
+
+import warnings
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+# FP16 dense tensor-core FLOPs per SM per clock cycle.
+# Keyed by (major, minor) compute capability.
+# Values empirically validated against NVIDIA spec-sheet peaks:
+#   A100 SXM4 : 156 TFLOPS = 108 SMs × 1024 × 1.410 GHz ✓
+#   H100 SXM5 : 989 TFLOPS = 132 SMs × 3840 × 1.980 GHz ✓
+#   RTX 4090  : ~660 TFLOPS = 128 SMs × 2048 × 2.520 GHz ✓
+#
+# Data-centre vs. consumer Ampere (8.0 vs. 8.6) differ because GA100's
+# tensor cores are physically larger than GA102's.
+_FP16_FLOPS_PER_SM_PER_CLOCK: dict[tuple[int, int], int] = {
+    (7, 0): 1024,   # Volta         (V100)
+    (7, 5): 1024,   # Turing        (T4, RTX 20xx)
+    (8, 0): 1024,   # Ampere DC     (A100, A30)
+    (8, 6): 512,    # Ampere cons.  (A10, RTX 30xx)
+    (8, 7): 512,    # Ampere Jetson (Orin)
+    (8, 9): 2048,   # Ada Lovelace  (RTX 40xx, L40S)
+    (9, 0): 3840,   # Hopper        (H100, H200)
+    (10, 0): 7680,  # Blackwell     (B100, B200) — estimated ≈ 2× Hopper
+}
+
+# Throughput multiplier relative to FP16 dense for each dtype.
+# Only available from certain compute capabilities onward.
+# (cc_major_min, multiplier)
+_DTYPE_MULTIPLIER: dict[str, tuple[int, float]] = {
+    "fp16":  (7, 1.0),
+    "bf16":  (8, 1.0),   # BF16 tensor cores added in Ampere
+    "int8":  (7, 2.0),   # INT8 tensor cores since Turing
+    "fp8":   (9, 2.0),   # FP8 via Transformer Engine; Ada (8.9) also supports it
+    "int4":  (7, 4.0),   # INT4 since Turing
+    "fp4":   (10, 4.0),  # FP4 native in Blackwell
+}
+# Ada also supports FP8 even though its major is 8
+_ADA_FP8_MINOR = 9
+
+
+@dataclass(frozen=True)
+class GPUSpec:
+    name: str
+    compute_capability: tuple[int, int]
+    num_sms: int
+    clock_rate_hz: float
+    memory_bandwidth_bytes_per_sec: float
+    peak_flops_by_dtype: dict[str, float] = field(default_factory=dict)
+    peak_memory_bandwidth_tbs: float = 0.0
+
+    def peak_tflops(self, dtype: str = "fp16") -> float:
+        """Return peak dense tensor-core TFLOPS for the given dtype."""
+        if dtype not in self.peak_flops_by_dtype:
+            supported = list(self.peak_flops_by_dtype)
+            raise ValueError(
+                f"dtype '{dtype}' not supported on {self.name} "
+                f"(CC {self.compute_capability[0]}.{self.compute_capability[1]}). "
+                f"Supported: {supported}"
+            )
+        return self.peak_flops_by_dtype[dtype] / 1e12
+
+
+def _warn_unknown(cc: tuple[int, int]) -> None:
+    warnings.warn(
+        f"Unknown compute capability {cc[0]}.{cc[1]}; MFU/MBU results may be inaccurate. "
+        "Please open an issue at https://github.com/your-repo/mfu-tracker.",
+        UserWarning,
+        stacklevel=3,
+    )
+
+
+def _fp8_supported(cc: tuple[int, int]) -> bool:
+    major, minor = cc
+    return major >= 9 or (major == 8 and minor >= _ADA_FP8_MINOR)
+
+
+def get_gpu_spec(device: Optional[torch.device] = None) -> GPUSpec:
+    """Query the GPU and derive peak throughput from first principles."""
+    if device is None:
+        device = torch.device("cuda")
+
+    props = torch.cuda.get_device_properties(device)
+    cc = (props.major, props.minor)
+
+    fp16_flops_per_sm = _FP16_FLOPS_PER_SM_PER_CLOCK.get(cc)
+    if fp16_flops_per_sm is None:
+        # Try falling back to the same major with a known minor
+        fallback = next(
+            (v for (maj, _), v in _FP16_FLOPS_PER_SM_PER_CLOCK.items() if maj == props.major),
+            None,
+        )
+        _warn_unknown(cc)
+        fp16_flops_per_sm = fallback if fallback is not None else 1024
+
+    # clock_rate is reported in kHz
+    clock_hz = props.clock_rate * 1_000
+    fp16_peak = props.multi_processor_count * fp16_flops_per_sm * clock_hz
+
+    # Build per-dtype peak FLOPS table for this GPU
+    peak_flops_by_dtype: dict[str, float] = {}
+    for dtype, (min_major, multiplier) in _DTYPE_MULTIPLIER.items():
+        if dtype == "fp8":
+            if not _fp8_supported(cc):
+                continue
+        elif props.major < min_major:
+            continue
+        peak_flops_by_dtype[dtype] = fp16_peak * multiplier
+
+    # Memory bandwidth: memory_clock_rate in kHz, bus width in bits, DDR = ×2
+    mem_bw = props.memory_clock_rate * 1_000 * (props.memory_bus_width / 8) * 2
+
+    return GPUSpec(
+        name=props.name,
+        compute_capability=cc,
+        num_sms=props.multi_processor_count,
+        clock_rate_hz=clock_hz,
+        memory_bandwidth_bytes_per_sec=mem_bw,
+        peak_flops_by_dtype=peak_flops_by_dtype,
+        peak_memory_bandwidth_tbs=mem_bw / 1e12,
+    )

mfu_tracker/integrations/__init__.py

@@ -0,0 +1,3 @@
+from .hf_trainer import MFUCallback
+
+__all__ = ["MFUCallback"]

mfu_tracker/integrations/hf_trainer.py

@@ -0,0 +1,152 @@
+"""HuggingFace Trainer integration via TrainerCallback."""
+from __future__ import annotations
+
+import warnings
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+from transformers import TrainerCallback
+
+from ..flops import param_bytes, profile_flops
+from ..gpu import GPUSpec, get_gpu_spec
+from ..tracker import compute_mbu, compute_mfu
+
+
+class MFUCallback(TrainerCallback):
+    """
+    TrainerCallback that logs MFU and MBU at every Trainer logging step.
+
+    FLOPs per step are ``fwd_flops × (1 + backward_factor)`` where
+    ``backward_factor`` defaults to 2.0 (standard 3× convention). Set it higher
+    when using gradient checkpointing (typical: 3.0–4.0).
+
+    Per-step cost is two non-blocking ``Event.record()`` calls (~10 μs CPU, no
+    GPU stall). The single ``torch.cuda.synchronize()`` is deferred to ``on_log``
+    and amortised across all steps in the logging interval.
+
+    Usage::
+
+        from mfu_tracker.integrations.hf_trainer import MFUCallback
+
+        callback = MFUCallback(sample_batch=next(iter(train_dataloader)), dtype="bf16")
+        # num_gpus is auto-detected from torch.distributed — no manual config needed
+        trainer = Trainer(..., callbacks=[callback])
+
+    **torch.compile**: profile_flops is called at ``on_train_begin``, before the
+    first compiled step. This is correct — ``torch.compile`` does not change the
+    FLOP count (same math, just faster execution). The MFU improvement from
+    compilation is captured automatically in the CUDA event timing of real steps.
+    Do not pass a compiled model to this callback directly; let HF Trainer compile
+    after the callback is registered.
+
+    **DDP / FSDP**: leave ``num_gpus=1`` — per-GPU MFU equals global MFU for
+    data-parallel jobs.
+
+    Args:
+        sample_batch:    A representative batch dict passed as ``**kwargs`` to the
+                         model. Used once at training start to profile forward FLOPs.
+        dtype:           Compute dtype for the peak ceiling — "fp16", "bf16", etc.
+        num_gpus:        GPUs in the peak ceiling (default 1).
+        backward_factor: Multiplier for backward pass cost (default 2.0). Set
+                         higher when using gradient checkpointing (typical: 3.0–4.0).
+        metric_prefix:   Prefix for logged metric names (default ``"throughput"``).
+                         Results in ``throughput/mfu`` and ``throughput/mbu``, which
+                         WandB groups into its own section away from loss/lr. Set to
+                         ``""`` to log bare ``mfu`` / ``mbu`` keys.
+        device:          CUDA device to query. Defaults to current device.
+    """
+
+    def __init__(
+        self,
+        sample_batch: dict[str, Any],
+        dtype: str = "bf16",
+        num_gpus: int = 1,
+        backward_factor: float = 2.0,
+        metric_prefix: str = "throughput",
+        device: Optional[torch.device] = None,
+    ) -> None:
+        self.sample_batch = sample_batch
+        self.dtype = dtype
+        self.num_gpus = num_gpus
+        self.backward_factor = backward_factor
+        self.metric_prefix = metric_prefix
+        self.device = device
+
+        self._model: Optional[nn.Module] = None
+        self._spec: Optional[GPUSpec] = None
+        self._fwd_flops: Optional[int] = None
+        self._param_bytes: Optional[int] = None
+
+        # Each entry: (e_start, e_bwd, e_end, bwd_recorded)
+        # Each entry: (e_start, e_end). Accumulated between on_log calls.
+        self._pending: list[tuple] = []
+
+    def _profile(self, model: nn.Module) -> None:
+        self._spec = get_gpu_spec(self.device)
+        self._param_bytes = param_bytes(model)
+        try:
+            # Move sample batch to the model's device (Trainer may have moved the model).
+            model_device = next(model.parameters()).device
+            batch = {
+                k: v.to(model_device) if isinstance(v, torch.Tensor) else v
+                for k, v in self.sample_batch.items()
+            }
+            self._fwd_flops = profile_flops(
+                model, kwargs=batch, with_backward=False
+            )
+        except Exception as exc:
+            warnings.warn(
+                f"mfu-tracker: FLOP profiling failed ({exc}); MFU will not be logged.",
+                stacklevel=2,
+            )
+
+    # --- TrainerCallback protocol -------------------------------------------
+
+    def on_train_begin(self, args, state, control, model=None, **kwargs):
+        if model is not None and torch.cuda.is_available():
+            self._model = model
+            self._profile(model)
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self._fwd_flops is None or not torch.cuda.is_available():
+            return
+
+        e_start = torch.cuda.Event(enable_timing=True)
+        e_end = torch.cuda.Event(enable_timing=True)
+        e_start.record()
+        self._pending_step = (e_start, e_end)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if not hasattr(self, "_pending_step"):
+            return
+
+        e_start, e_end = self._pending_step
+        e_end.record()
+        self._pending.append((e_start, e_end))
+        del self._pending_step
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is None or not self._pending or self._fwd_flops is None:
+            return
+
+        # Single sync amortised across all accumulated steps.
+        torch.cuda.synchronize(self.device)
+
+        n_steps = len(self._pending)
+        total_ms = sum(e_start.elapsed_time(e_end) for e_start, e_end in self._pending)
+        self._pending.clear()
+
+        elapsed_sec = total_ms / 1000
+        if elapsed_sec <= 0:
+            return
+
+        total_flops = int(self._fwd_flops * n_steps * (1 + self.backward_factor))
+
+        prefix = f"{self.metric_prefix}/" if self.metric_prefix else ""
+        logs[f"{prefix}mfu"] = round(
+            compute_mfu(total_flops, elapsed_sec, dtype=self.dtype, num_gpus=self.num_gpus, spec=self._spec), 4
+        )
+        logs[f"{prefix}mbu"] = round(
+            compute_mbu(self._param_bytes * n_steps, elapsed_sec, num_gpus=self.num_gpus, spec=self._spec), 4
+        )

mfu_tracker/optim.py

@@ -0,0 +1,155 @@
+"""Optimizer wrapper that measures MFU and MBU per training step."""
+from __future__ import annotations
+
+import warnings
+from contextlib import contextmanager
+from typing import Any, Generator, Optional
+
+import torch
+import torch.nn as nn
+
+from .flops import param_bytes, profile_flops
+from .gpu import GPUSpec, get_gpu_spec
+from .tracker import UtilizationResult
+
+
+class MFUOptimizerWrapper:
+    """
+    Wraps any ``torch.optim.Optimizer`` to automatically track MFU and MBU.
+
+    FLOPs are profiled once on the uncompiled model and scaled by
+    ``1 + backward_factor`` (default 2.0, the standard forward + 2× backward
+    convention). Set ``backward_factor`` higher when using gradient checkpointing,
+    which recomputes activations during backward (typical values: 3.0–4.0).
+
+    ``zero_grad()`` is called automatically at the *start* of ``track_step()``.
+    Call ``optimizer.step()`` **after** the block so it is excluded from the
+    timing window::
+
+        optimizer = MFUOptimizerWrapper(
+            torch.optim.AdamW(model.parameters(), lr=1e-4),
+            model=model,
+            sample_batch=sample_batch,
+            dtype="bf16",
+        )
+
+        for batch in dataloader:
+            with optimizer.track_step() as result:
+                output = model(**batch)
+                loss = output.loss
+                loss.backward()
+            optimizer.step()
+            print(f"MFU={result.mfu:.1%}  MBU={result.mbu:.1%}")
+
+    **torch.compile**: profile the uncompiled model first, then compile::
+
+        optimizer = MFUOptimizerWrapper(raw_model, ...)
+        optimizer.profile()          # profile before compile
+        model = torch.compile(model)
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        model: nn.Module,
+        sample_batch: dict[str, Any],
+        dtype: str = "bf16",
+        num_gpus: int = 1,
+        backward_factor: float = 2.0,
+        device: Optional[torch.device] = None,
+    ) -> None:
+        self.optimizer = optimizer
+        self._model = model
+        self._sample_batch = sample_batch
+        self._dtype = dtype
+        self._num_gpus = num_gpus
+        self._backward_factor = backward_factor
+        self._device = device
+
+        self._spec: Optional[GPUSpec] = None
+        self._fwd_flops: Optional[int] = None
+        self._param_bytes: Optional[int] = None
+
+    def profile(self) -> None:
+        """
+        Explicitly profile FLOPs on the current (uncompiled) model.
+
+        Call this before ``torch.compile`` so the FLOP count is measured on the
+        original graph. If not called, profiling happens lazily on the first
+        ``track_step()`` — which may be too late if the model is already compiled::
+
+            optimizer = MFUOptimizerWrapper(raw_optimizer, model, sample_batch, dtype="bf16")
+            optimizer.profile()          # profile uncompiled model
+            model = torch.compile(model) # compile after profiling
+        """
+        if self._spec is None:
+            self._profile_once()
+
+    def _profile_once(self) -> None:
+        self._spec = get_gpu_spec(self._device)
+        try:
+            self._fwd_flops = profile_flops(
+                self._model,
+                kwargs=self._sample_batch,
+                with_backward=False,
+            )
+        except Exception as exc:
+            warnings.warn(
+                f"mfu-tracker: profiling failed ({exc}); MFU will not be populated.",
+                stacklevel=3,
+            )
+        self._param_bytes = param_bytes(self._model)
+
+    @contextmanager
+    def track_step(self) -> Generator[UtilizationResult, None, None]:
+        """
+        Context manager that wraps one training step and populates a
+        :class:`~mfu_tracker.UtilizationResult` with MFU and MBU.
+
+        ``optimizer.zero_grad()`` is called automatically at the *start* of the
+        block. Call ``optimizer.step()`` **after** the block so it is excluded
+        from the timing window::
+
+            with wrapped.track_step() as result:
+                out = model(**batch)
+                out.loss.backward()
+            wrapped.step()
+
+        FLOPs are ``fwd_flops × (1 + backward_factor)`` where ``backward_factor``
+        defaults to 2.0 (standard 3× convention). Set it higher when using
+        gradient checkpointing (typical: 3.0–4.0).
+        """
+        if self._spec is None:
+            self._profile_once()
+
+        self.optimizer.zero_grad()
+
+        e_start = torch.cuda.Event(enable_timing=True)
+        e_end = torch.cuda.Event(enable_timing=True)
+
+        result = UtilizationResult(dtype=self._dtype, gpu_spec=self._spec, num_gpus=self._num_gpus)
+
+        e_start.record()
+        try:
+            yield result
+        finally:
+            e_end.record()
+            result._e_start = e_start
+            result._e_end = e_end
+            result._total_flops = (
+                int(self._fwd_flops * (1 + self._backward_factor))
+                if self._fwd_flops is not None else None
+            )
+            result._param_bytes = self._param_bytes
+            result._device = self._device
+
+    # --- Proxy the underlying optimizer ------------------------------------
+
+    def step(self, *args, **kwargs) -> None:
+        self.optimizer.step(*args, **kwargs)
+
+    def zero_grad(self, *args, **kwargs) -> None:
+        self.optimizer.zero_grad(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.optimizer, name)

mfu_tracker/tracker.py

@@ -0,0 +1,211 @@
+"""MFU and MBU measurement — context manager and standalone functions."""
+from __future__ import annotations
+
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Generator, Optional
+
+import torch
+
+from .gpu import GPUSpec, get_gpu_spec
+
+
+@dataclass
+class UtilizationResult:
+    """
+    Result holder populated after a ``track()`` or ``track_step()`` block exits.
+
+    Fields backed by CUDA events (from ``MFUOptimizerWrapper.track_step()``) are
+    computed lazily on first access — the GPU sync is deferred until the value is
+    actually needed. Fields from the CPU-timed ``track()`` context manager are
+    populated eagerly since ``torch.cuda.synchronize`` is already called there.
+    """
+
+    dtype: str = "fp16"
+    gpu_spec: Optional[GPUSpec] = None
+    num_gpus: int = 1
+
+    # Eagerly-set fields (track()) or lazily-set after _resolve() (track_step()).
+    _mfu: Optional[float] = field(default=None, repr=False)
+    _mbu: Optional[float] = field(default=None, repr=False)
+    _elapsed_sec: Optional[float] = field(default=None, repr=False)
+    _achieved_tflops: Optional[float] = field(default=None, repr=False)
+    _achieved_tbs: Optional[float] = field(default=None, repr=False)
+
+    # Set by MFUOptimizerWrapper for lazy resolution.
+    _e_start: Optional[object] = field(default=None, repr=False)
+    _e_end: Optional[object] = field(default=None, repr=False)
+    _total_flops: Optional[int] = field(default=None, repr=False)
+    _param_bytes: Optional[int] = field(default=None, repr=False)
+    _device: Optional[torch.device] = field(default=None, repr=False)
+
+    def _resolve(self) -> None:
+        """Sync and compute all fields from CUDA events. Called at most once."""
+        if self._e_start is None:
+            return
+        if self._total_flops is None or self._param_bytes is None:
+            return
+        torch.cuda.synchronize(self._device)
+        elapsed = self._e_start.elapsed_time(self._e_end) / 1000
+
+        peak_tflops = self.gpu_spec.peak_tflops(self.dtype) * self.num_gpus
+        peak_tbs = self.gpu_spec.peak_memory_bandwidth_tbs * self.num_gpus
+
+        self._elapsed_sec = elapsed
+        self._achieved_tflops = self._total_flops / elapsed / 1e12
+        self._achieved_tbs = self._param_bytes / elapsed / 1e12
+        self._mfu = self._achieved_tflops / peak_tflops
+        self._mbu = self._achieved_tbs / peak_tbs
+        self._e_start = None  # mark resolved
+
+    @property
+    def mfu(self) -> Optional[float]:
+        self._resolve()
+        return self._mfu
+
+    @mfu.setter
+    def mfu(self, v: Optional[float]) -> None:
+        self._mfu = v
+
+    @property
+    def mbu(self) -> Optional[float]:
+        self._resolve()
+        return self._mbu
+
+    @mbu.setter
+    def mbu(self, v: Optional[float]) -> None:
+        self._mbu = v
+
+    @property
+    def elapsed_sec(self) -> Optional[float]:
+        self._resolve()
+        return self._elapsed_sec
+
+    @elapsed_sec.setter
+    def elapsed_sec(self, v: Optional[float]) -> None:
+        self._elapsed_sec = v
+
+    @property
+    def achieved_tflops(self) -> Optional[float]:
+        self._resolve()
+        return self._achieved_tflops
+
+    @achieved_tflops.setter
+    def achieved_tflops(self, v: Optional[float]) -> None:
+        self._achieved_tflops = v
+
+    @property
+    def achieved_tbs(self) -> Optional[float]:
+        self._resolve()
+        return self._achieved_tbs
+
+    @achieved_tbs.setter
+    def achieved_tbs(self, v: Optional[float]) -> None:
+        self._achieved_tbs = v
+
+
+@contextmanager
+def track(
+    flop_count: int,
+    param_bytes: int,
+    *,
+    dtype: str = "fp16",
+    num_gpus: int = 1,
+    device: Optional[torch.device] = None,
+    spec: Optional[GPUSpec] = None,
+) -> Generator[UtilizationResult, None, None]:
+    """
+    Context manager that measures MFU and MBU for an arbitrary compute block.
+
+    Args:
+        flop_count:  Total FLOPs for the block (use flops.profile_flops or your own).
+        param_bytes: Bytes transferred for weights (num_params * bytes_per_element).
+        dtype:       Compute dtype string — "fp16", "bf16", "int8", "fp8", "int4", "fp4".
+        num_gpus:    Multiplier applied to the peak ceiling (default 1). When using
+                     ``profile_flops`` as the source of *flop_count*, leave this at 1
+                     regardless of parallelism strategy — ``profile_flops`` returns
+                     per-GPU FLOPs, and per-GPU MFU equals global MFU for all standard
+                     parallelism types (the N factors cancel). Only set this when
+                     *flop_count* is the analytically-derived *full-model* FLOP count
+                     (e.g. ``6 × params × tokens``) rather than a profiled per-GPU count.
+        device:      CUDA device to measure against (default: current device).
+        spec:        Pre-queried GPUSpec; fetched once if not provided.
+
+    Yields a :class:`UtilizationResult` whose fields are ``None`` until the block
+    exits, then populated with measured values.
+
+    Example — single GPU::
+
+        flops = profile_flops(model, args=(sample,), with_backward=True)
+        with track(flops, param_bytes(model), dtype="bf16") as result:
+            loss = model(inputs)
+            loss.backward()
+            optimizer.step()
+        print(f"MFU={result.mfu:.1%}  MBU={result.mbu:.1%}")
+
+    Example — any parallelism strategy (DDP, FSDP, tensor, pipeline)::
+
+        # profile_flops on whatever model shard the local rank holds gives
+        # per-GPU FLOPs; per-GPU MFU == global MFU for all standard strategies.
+        with track(flops, param_bytes(model), dtype="bf16") as r:
+            ...
+    """
+    if spec is None:
+        spec = get_gpu_spec(device)
+
+    peak_tflops = spec.peak_tflops(dtype) * num_gpus
+    peak_tbs = spec.peak_memory_bandwidth_tbs * num_gpus
+
+    result = UtilizationResult(dtype=dtype, gpu_spec=spec, num_gpus=num_gpus)
+
+    torch.cuda.synchronize(device)
+    t0 = time.perf_counter()
+    yield result
+    torch.cuda.synchronize(device)
+    elapsed = time.perf_counter() - t0
+
+    result.elapsed_sec = elapsed
+    result.achieved_tflops = flop_count / elapsed / 1e12
+    result.achieved_tbs = param_bytes / elapsed / 1e12
+    result.mfu = result.achieved_tflops / peak_tflops
+    result.mbu = result.achieved_tbs / peak_tbs
+
+
+def compute_mfu(
+    flop_count: int,
+    elapsed_sec: float,
+    *,
+    dtype: str = "fp16",
+    num_gpus: int = 1,
+    device: Optional[torch.device] = None,
+    spec: Optional[GPUSpec] = None,
+) -> float:
+    """
+    Standalone MFU calculation without a context manager.
+
+    Args:
+        num_gpus: GPUs in the peak ceiling. See :func:`track` for guidance.
+    """
+    if spec is None:
+        spec = get_gpu_spec(device)
+    return (flop_count / elapsed_sec / 1e12) / (spec.peak_tflops(dtype) * num_gpus)
+
+
+def compute_mbu(
+    param_bytes: int,
+    elapsed_sec: float,
+    *,
+    num_gpus: int = 1,
+    device: Optional[torch.device] = None,
+    spec: Optional[GPUSpec] = None,
+) -> float:
+    """
+    Standalone MBU calculation without a context manager.
+
+    Args:
+        num_gpus: GPUs in the peak ceiling. See :func:`track` for guidance.
+    """
+    if spec is None:
+        spec = get_gpu_spec(device)
+    return (param_bytes / elapsed_sec / 1e12) / (spec.peak_memory_bandwidth_tbs * num_gpus)

mfu_tracker-0.1.0.dist-info/METADATA

@@ -0,0 +1,235 @@
+Metadata-Version: 2.4
+Name: mfu-tracker
+Version: 0.1.0
+Summary: Lightweight Model FLOPs Utilization and Bandwidth Utilization tracker for PyTorch
+License: MIT License
+        
+        Copyright (c) 2026 Jeremias Lino Ferrao
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.9
+Requires-Dist: numpy>=2.0.2
+Requires-Dist: thop>=0.1.1.post2209072238
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Provides-Extra: hf
+Requires-Dist: transformers>=4.30; extra == 'hf'
+Description-Content-Type: text/markdown
+
+# mfu-tracker
+
+When profiling training runs, I found that most existing tools either lacked MFU/MBU support entirely or dragged in hundreds of megabytes of transitive dependencies. This library is an attempt at a self-contained alternative.
+
+**mfu-tracker** is a PyTorch library for measuring Model FLOPs Utilization (MFU) and Model Bandwidth Utilization (MBU). It supports bare PyTorch training loops, an optimizer wrapper, and a HuggingFace Trainer callback.
+
+- **Minimal dependencies** — PyTorch and `thop` only
+- **Profiled FLOPs, not formula estimates** — uses `FlopCounterMode` to count the FLOPs your model actually executes rather than a formula like `6 × params × tokens`. For Mixture-of-Experts models this means only active experts are counted, giving a more accurate numerator than parameter-based estimates.
+- **Three integration styles** — context manager, optimizer wrapper, HF Trainer callback
+- **WandB / TensorBoard / MLflow** — metrics are logged through HF Trainer's existing pipeline when using `MFUCallback`
+
+MFU as a training efficiency metric was introduced in the [PaLM paper](https://arxiv.org/abs/2204.02311) (Chowdhery et al., 2022).
+
+---
+
+## What MFU and MBU measure
+
+**MFU (Model FLOPs Utilization)** is the ratio of observed FLOP throughput to the GPU's theoretical peak for the given dtype. A value of 0.50 means the model is executing at half the GPU's rated peak. Well-optimized large models on modern hardware typically fall in the 0.40–0.60 range; small models often land much lower due to kernel dispatch overhead relative to compute time.
+
+**MBU (Model Bandwidth Utilization)** as computed here is a proxy, not a direct DRAM measurement. It is defined as:
+
+```
+MBU = (param_bytes / elapsed_sec) / peak_memory_bandwidth
+```
+
+where `param_bytes` is the total size of model parameters and `elapsed_sec` is wall time. This assumes one full pass through model weights per step and does not account for activation memory, gradients, optimizer state, or data layout effects. It is most useful as a relative indicator across runs rather than an absolute efficiency measure.
+
+If both MFU and MBU are low simultaneously, the GPU is underutilized. Two common causes: kernel dispatch overhead (the CPU cannot issue kernels fast enough to keep the GPU busy — `torch.compile` reduces this by fusing operations), or CPU-side pipeline stalls (slow DataLoader, heavy host preprocessing, or host-to-device transfers in the hot path).
+
+---
+
+## Installation
+
+```bash
+pip install mfu-tracker
+```
+
+HuggingFace Trainer integration requires no extra install — if you are already running HF Trainer, `transformers` is already available. Import `MFUCallback` directly.
+
+---
+
+## Usage
+
+### Context manager (bare PyTorch)
+
+```python
+from mfu_tracker import track, profile_flops, param_bytes
+
+# Profile once on the uncompiled model before training begins
+sample = {"input_ids": batch["input_ids"][:1]}
+flops = profile_flops(model, kwargs=sample, with_backward=True)
+p_bytes = param_bytes(model)
+
+for batch in dataloader:
+    optimizer.zero_grad()
+    with track(flops, p_bytes, dtype="bf16") as result:
+        loss = model(**batch).loss
+        loss.backward()
+    optimizer.step()
+
+    print(f"MFU: {result.mfu:.3f}  MBU: {result.mbu:.3f}  {result.elapsed_sec*1000:.0f} ms/step")
+```
+
+### Optimizer wrapper
+
+```python
+from mfu_tracker import MFUOptimizerWrapper
+
+base_optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+optimizer = MFUOptimizerWrapper(
+    base_optimizer, model,
+    sample_batch={"input_ids": sample_ids},
+    dtype="bf16",
+)
+
+# Profile before compiling — FlopCounterMode may not trace compiled graphs
+optimizer.profile()
+model = torch.compile(model)
+
+for batch in dataloader:
+    with optimizer.track_step() as result:  # calls zero_grad() at block entry
+        loss = model(**batch).loss
+        loss.backward()
+    optimizer.step()                        # outside the timing window
+
+    if step % 10 == 0:
+        print(f"MFU {result.mfu:.3f}  MBU {result.mbu:.3f}")
+```
+
+### HuggingFace Trainer
+
+```python
+from mfu_tracker.integrations.hf_trainer import MFUCallback
+
+sample_batch = {k: v[:batch_size] for k, v in next(iter(train_dataloader)).items()}
+
+callback = MFUCallback(
+    sample_batch=sample_batch,
+    dtype="bf16",
+    metric_prefix="throughput",  # logs throughput/mfu and throughput/mbu
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    callbacks=[callback],
+)
+trainer.train()
+```
+
+`throughput/mfu` and `throughput/mbu` are added to the Trainer log dict at each logging step and forwarded automatically to any configured integrations (WandB, TensorBoard, MLflow). WandB groups metrics by the `/` separator, so these appear in a distinct "throughput" section rather than alongside loss and learning rate.
+
+---
+
+## FLOP counting
+
+```python
+from mfu_tracker import profile_flops, flash_attn_flops, param_bytes
+
+# Standard models — FlopCounterMode counts SDPA automatically on CUDA
+flops = profile_flops(model, kwargs=batch, with_backward=True)
+
+# Models calling flash_attn_func directly (rare; older HF with use_flash_attention_2=True)
+# need a manual correction since the C extension is opaque to FlopCounterMode:
+flops += flash_attn_flops(batch_size=B, seq_len=S, num_heads=H, head_dim=D)
+
+# PEFT / LoRA — restrict param_bytes to trainable parameters only
+p_bytes = param_bytes(model, trainable_only=True)
+```
+
+`with_backward=True` applies the standard 3× convention (1× forward + 2× backward). For gradient checkpointing, pass `backward_factor=3.0` or `4.0` to `MFUOptimizerWrapper` or `MFUCallback`.
+
+---
+
+## GPU spec
+
+```python
+from mfu_tracker import get_gpu_spec
+
+spec = get_gpu_spec()
+print(spec.name)                        # e.g. "NVIDIA GeForce RTX 4080"
+print(spec.peak_tflops("fp16"))         # e.g. 97.6
+print(spec.peak_tflops("fp8"))          # Ada Lovelace (CC 8.9) and Hopper (CC 9.0)+
+print(spec.peak_memory_bandwidth_tbs)   # e.g. 0.717
+```
+
+Supported dtypes: `fp32`, `fp16`, `bf16`, `int8`, `fp8`, `int4`, `fp4`. Unrecognized compute capabilities fall back to the nearest known major version with a `UserWarning`.
+
+---
+
+## Benchmark (RTX 4080, GPT-2 124M, fp16)
+
+| Configuration | MFU | ms/step |
+|---|---|---|
+| batch=1 · eager | ~0.027 | ~40 ms |
+| batch=8 · eager | ~0.09 | ~93 ms |
+| batch=8 · sdpa | ~0.12 | ~74 ms |
+| batch=8 · sdpa + compile | ~0.17 | ~50 ms |
+| batch=16 · sdpa + compile | ~0.16 | ~104 ms |
+
+GPT-2 (124M) is a small model relative to the compute capacity of a modern GPU, so low MFU is expected — the model spends a large fraction of step time waiting for kernel dispatch rather than doing arithmetic. Larger models (e.g. LLaMA-70B) typically reach 0.40–0.60 MFU. The improvement from `torch.compile` reflects kernel fusion reducing dispatch overhead. I'll add some testing on this later.
+
+```bash
+python examples/benchmark_mfu.py --help
+python examples/hf_trainer_mfu.py --dtype bf16 --batch-size 16
+```
+
+---
+
+## Multi-GPU
+
+Leave `num_gpus=1` (the default) when using `profile_flops` as the FLOP source. For data-parallel strategies (DDP, FSDP), per-GPU FLOPs equal total FLOPs divided by N and wall time is the same on all ranks, so per-GPU MFU equals global MFU and the N factors cancel. Set `num_gpus > 1` only when pairing an analytically-derived full-model FLOP count (e.g. `6 × params × tokens`) with a total-job peak ceiling.
+
+---
+
+## Limitations
+
+- **SDPA on CPU is not counted** — `FlopCounterMode` does not intercept flash attention dispatch on CPU. Profile with a CUDA model.
+- **bitsandbytes quantized layers** — INT8/NF4 kernels are opaque to `FlopCounterMode`. NF4 dequantizes to fp16 before the matmul, so FLOP counts are approximately correct. Pass the appropriate dtype to use the right peak ceiling.
+- **`flash_attn_func` direct calls** — models bypassing `F.scaled_dot_product_attention` need a manual `flash_attn_flops()` correction (see above).
+- **Peak ceilings from spec sheets** — these are not independently measured. MFU > 1.0 indicates the ceiling is underestimated.
+- **MBU is a proxy** — the formula uses parameter bytes as a stand-in for memory traffic; actual DRAM traffic (activations, gradients, optimizer state) is higher and not measured.
+- I have not tested the library extensively yet; please open an issue if you encounter any bugs or unexpected behavior.
+
+---
+
+## Requirements
+
+- Python 3.9+
+- PyTorch 2.0+ (2.1+ recommended for `FlopCounterMode`)
+- A CUDA GPU is required for meaningful results; CPU timing works but MFU will be near zero for any realistic model
+
+---
+
+## License
+
+MIT

mfu_tracker-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+mfu_tracker/__init__.py,sha256=gvoCPl6mEHBRV7Y_q_FLJjK16JQMsQ11nfgrpQiF8-0,503
+mfu_tracker/flops.py,sha256=NwCrPXRNXPJMHSWEcb5SSWxEZEqkGocwdNY0NJ-ZeiU,5762
+mfu_tracker/gpu.py,sha256=60ZSbSe1cAyfjPm-uqZR7ptRBYDFno6SJJO32VGePwA,4689
+mfu_tracker/optim.py,sha256=kOfQfDMC7oPpWGGRHROs4Cwv9vxFAoe3GpYlV_2sGOw,5466
+mfu_tracker/tracker.py,sha256=fDGpFm7XKEkwtSb6CHgWsrrS01ic2l8RxVOrvuDtaUY,7362
+mfu_tracker/integrations/__init__.py,sha256=KmKZNbpniR3qqNG5Zloe6Z2tGwi7fgH1zCb6p_8a0po,63
+mfu_tracker/integrations/hf_trainer.py,sha256=aP7ld5HGgOgDu_PQt0vcoOU2hx7ocqp__sHDPnmblgY,6164
+mfu_tracker-0.1.0.dist-info/METADATA,sha256=wQ9f1PHsrcgYa-kB53wlTYht6luRnbS20PFkX-SR3Us,10365
+mfu_tracker-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+mfu_tracker-0.1.0.dist-info/licenses/LICENSE,sha256=aXaBx4UYHF1tx67IIN-nh5BIggjU8q9WWYyN2HupHXA,1077
+mfu_tracker-0.1.0.dist-info/RECORD,,

mfu_tracker-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mfu_tracker-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jeremias Lino Ferrao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.