mfu-tracker 0.1.0__tar.gz

mfu_tracker-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,22 @@
+name: pypi
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write  # required for trusted publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+
+      - run: uv build
+
+      - uses: pypa/gh-action-pypi-publish@release/v1

mfu_tracker-0.1.0/.gitignore

@@ -0,0 +1,210 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Vibe Coding Stuff
+.serena/

mfu_tracker-0.1.0/CLAUDE.md

@@ -0,0 +1,75 @@
+# mfu-tracker
+
+PyPI library for tracking Model FLOPs Utilization (MFU) and Model Bandwidth Utilization (MBU).
+
+## Architecture
+
+- [src/mfu_tracker/gpu.py](src/mfu_tracker/gpu.py) — queries `torch.cuda.get_device_properties()` to derive peak TFLOPS and memory bandwidth from first principles. Uses `_FP16_FLOPS_PER_SM_PER_CLOCK` keyed by `(major, minor)` compute capability tuple (empirically validated against spec sheets). Supports per-dtype peak ceilings (fp16, bf16, int8, fp8, int4, fp4).
+- [src/mfu_tracker/flops.py](src/mfu_tracker/flops.py) — FLOP counting via `torch.utils.flop_counter.FlopCounterMode` (PyTorch 2.1+), with `thop` as fallback. `FlopCounterMode` hooks at the ATen dispatch level, so it counts `F.scaled_dot_product_attention` (SDPA / native flash attention) automatically on CUDA — no manual correction needed for modern transformer models. SDPA is NOT counted on CPU (the kernel dispatches differently); profile with a CUDA model for accurate counts. `thop` fallback handles PyTorch < 2.1. For the rare `flash_attn` C extension (direct `flash_attn_func` calls), use `flash_attn_flops(B, S, H, D)` to compute the missing FLOPs manually. For kwargs-only models (e.g. HF), both paths wrap in `_KwargsAdapter`. `param_bytes()` accepts `trainable_only=True` for PEFT/LoRA backward MBU estimates.
+- [src/mfu_tracker/tracker.py](src/mfu_tracker/tracker.py) — `track()` context manager and `compute_mfu`/`compute_mbu` standalone functions. All accept a `dtype` parameter. `UtilizationResult` is a mutable dataclass yielded by `track()`; fields start as `None` and are populated after the block exits. CUDA-event-backed fields (from `track_step()`) are resolved lazily on first attribute access; `_resolve()` is idempotent and calls `synchronize` exactly once.
+- [src/mfu_tracker/optim.py](src/mfu_tracker/optim.py) — `MFUOptimizerWrapper`. Wraps any `torch.optim.Optimizer` and exposes a `track_step()` context manager. Uses a fixed `backward_factor` (default 2.0, the standard 3× convention) — no gradient hook. `zero_grad()` is called automatically at the **start** of `track_step()`; call `optimizer.step()` **after** the block to keep it outside the timing window. Profile the uncompiled model via `wrapper.profile()` before calling `torch.compile`.
+- [src/mfu_tracker/integrations/hf_trainer.py](src/mfu_tracker/integrations/hf_trainer.py) — `MFUCallback(TrainerCallback)`. Profiles the model once at `on_train_begin` (moves sample batch to model device automatically). Records two non-blocking CUDA events per step (`on_step_begin` / `on_step_end`), defers `torch.cuda.synchronize()` to `on_log` amortised across the logging interval. Logs `throughput/mfu` and `throughput/mbu` (configurable prefix). Does NOT read `state.total_flos` — HF Trainer uses the dense 6ND formula for all models including MoE, overcounting MoE by up to 4×. Skips silently when CUDA is unavailable.
+
+## Key design decisions
+
+- `(major, minor)` tuple keys for GPU lookup — CC 8.0 (A100) and CC 8.6 (RTX 3090) have genuinely different per-SM throughput (1024 vs 512 FP16 FLOPs/SM/clock) despite both being Ampere. Major-version-only keys would be wrong.
+- Ada Lovelace is CC 8.9 — gets FP8 support via a special case in `_fp8_supported()` even though its major version is 8 (below the FP8 min_major of 9 for Hopper).
+- `thop` over `calflops` — calflops unconditionally imports `transformers` in `__init__.py`, making it a 600MB transitive dep that defeats the lightweight goal.
+- **Fixed backward factor, not dynamic measurement.** The original design used a gradient hook on `trainable[-1]` to measure the forward/backward time split dynamically. This was abandoned because: (1) gradient hooks fire on the CPU autograd thread based on scheduling, not GPU completion — making the "start of backward" event unreliable; (2) with `optimizer.step()` inside the timing window, the measured factor absorbed optimizer time and gave ~4× instead of ~2×; (3) `torch.compile` restructures the backward graph so the hook often never fires, producing inconsistent FLOP multipliers between compiled and uncompiled runs. Fixed `backward_factor=2.0` (3× convention) with a user-overridable parameter is simpler and consistent. Users with gradient checkpointing set `backward_factor=3.0–4.0` explicitly.
+- `optimizer.step()` belongs **outside** `track_step()`. `zero_grad()` is called at the start of the block so gradients are valid until the block exits. This way timing captures forward + backward only.
+- `UtilizationResult` is mutable (no `frozen=True`) so the context manager pattern works correctly — yield first, populate after block exits.
+- `MFUOptimizerWrapper.track_step()` is lazy — no `synchronize` in `finally`, only on first attribute access of the result. Skipping `result.mfu` on some steps incurs zero sync cost for those steps.
+- HF integration uses `TrainerCallback`, not monkey-patch — cleaner, composable, and avoids patching internal Trainer methods. `MFUCallback` inherits from `TrainerCallback` so HF Trainer can call all callback events on it.
+- `metric_prefix="throughput"` on `MFUCallback` — logs `throughput/mfu` and `throughput/mbu`. WandB groups metrics by `/` separator, placing these in their own "throughput" section away from `loss`/`lr`. Set `metric_prefix=""` for bare keys.
+- HF Trainer forwards the `logs` dict from `on_log` to all configured integrations (WandB, TensorBoard, MLflow) automatically — no extra configuration needed.
+- Graceful degradation: unknown compute capability emits a `UserWarning` and falls back to the closest known major version.
+- MBU is always reported alongside MFU.
+- `num_gpus` parameter on `track()`, `compute_mfu()`, `compute_mbu()`, `MFUCallback`, and `MFUOptimizerWrapper` scales the peak ceiling. Default 1, correct for all parallelism strategies when using `profile_flops` — per-GPU MFU equals global MFU for DDP, FSDP, tensor, and pipeline parallelism because the N factors cancel (per-GPU FLOPs = total/N, wall time is the same across all GPUs). Only set `num_gpus > 1` when pairing analytically-derived full-model FLOPs (e.g. `6 × params × tokens`) with a total-job peak.
+- `torch.compile` does not change FLOP count (same math, faster execution). Profile the *uncompiled* model — `FlopCounterMode` may not trace compiled graphs correctly. The MFU improvement from compilation is captured automatically via CUDA event timing of real steps.
+- `src/` layout for correct PyPI packaging (hatchling build backend).
+
+## Benchmark findings (RTX 4080, GPT-2 124M, fp16)
+
+From `examples/benchmark_mfu.py` — uses `track()` with `profile_flops(with_backward=True)` (fixed 3× convention) for consistent results across all configurations:
+
+| Configuration         | MFU   | ms/step |
+|-----------------------|-------|---------|
+| batch=1  \| eager     | ~2.7% | ~40ms   |
+| batch=8  \| eager     | ~9%   | ~93ms   |
+| batch=8  \| sdpa      | ~12%  | ~74ms   |
+| batch=8  \| sdpa+compile | ~17% | ~50ms |
+| batch=16 \| sdpa+compile | ~16% | ~104ms |
+
+Key observations:
+- Low MFU (~2–17%) is expected for GPT-2 on modern hardware — the model is too small to saturate tensor cores. Large models (LLaMA-70B) reach 40–60% MFU.
+- Low MBU (~0.2–0.7%) at batch≥4 means memory bandwidth is **not** the bottleneck — the model is compute-bound (or kernel-launch-bound). MBU is more meaningful for inference at batch=1.
+- Both low MFU and low MBU simultaneously indicates **kernel launch overhead**: the GPU idles between small operations waiting for the CPU to dispatch the next kernel. `torch.compile` addresses this by fusing kernels, giving +5–8pp MFU improvement.
+- `sdpa` over `eager` attention: +2–3pp MFU from avoiding materialising the full B×H×S×S attention matrix (flash attention tiling).
+
+## Testing
+
+Two test tiers:
+
+- **Mock-based** (`test_flops.py`, `test_gpu.py`, `test_tracker.py`, `test_hf_callback.py`, `test_optim.py`) — no GPU required, run anywhere.
+- **GPU integration** (`test_integration_gpu.py`) — skipped automatically without CUDA. Validated on RTX 4080 (CC 8.9). Covers: spec detection without warnings, thop FLOP counts matching theory within 1% for `Linear` and `Conv2d`, MFU/MBU in `(0, 1]` on real hardware, `compute_mfu` agreeing with `track()`, larger batch → higher MFU.
+
+Known faithfulness limitations:
+- Peak ceiling is from NVIDIA spec sheets, not our own measurements.
+- `F.scaled_dot_product_attention` (SDPA) is counted automatically on CUDA via `FlopCounterMode`. Models using `flash_attn_func` directly (rare — older HF with `use_flash_attention_2=True`) still need `flash_attn_flops()` correction.
+- SDPA is not counted when profiling on CPU — profile the CUDA model for accurate counts.
+- bitsandbytes INT8/NF4 quantized layers (QLoRA) are opaque CUDA kernels not visible to either counter. NF4 dequantizes to fp16 before matmul so FLOPs are approximately correct. Pass `dtype="int8"` to get the right peak ceiling.
+- CUDA event timing is accurate; CPU-timer `track()` requires a `synchronize` at block boundaries.
+- The dynamic backward factor measurement (gradient hook on `trainable[-1]`) was removed — it gave unreliable results (~4× instead of ~2×) due to CPU/GPU async timing and broke comparisons between compiled/uncompiled models.
+
+`transformers` and `accelerate` are dev dependencies (needed to test `MFUCallback` and run the HF Trainer example).
+
+```bash
+uv sync --group dev
+.venv/bin/pytest tests/ -v                          # mock tests only (no GPU needed)
+.venv/bin/pytest tests/test_integration_gpu.py -v  # GPU tests
+```
+
+## Examples
+
+- `examples/benchmark_mfu.py` — benchmarks MFU/MBU across batch size, attention implementation (`eager` vs `sdpa`), and `torch.compile` using GPT-2 (124M). Uses `track()` with pre-profiled FLOPs for consistent results.
+- `examples/hf_trainer_mfu.py` — demonstrates `MFUCallback` with HF Trainer on synthetic data. Metrics appear as `throughput/mfu` / `throughput/mbu` in training logs and WandB. Run with `--wandb` to enable WandB logging.

mfu_tracker-0.1.0/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.

mfu_tracker-0.1.0/PKG-INFO

@@ -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