add optional clock lock

Author: drisspg

test/test_utils.py

@@ -1,13 +1,84 @@
+import ctypes
+import os
 import tempfile
 import unittest.mock as mock
 from pathlib import Path
 
 import pytest
 import torch
+from transformer_nuggets.utils.benchmark import (
+    _get_max_sm_clock,
+    _get_nvml_handle,
+    _nvml,
+    locked_clocks,
+    benchmark_cuda_function_in_microseconds_triton,
+)
 from transformer_nuggets.utils.shape_trace import open_logs, ShapeLog
 from transformer_nuggets.utils.tracing import NanInfDetect
 
 
+requires_root = pytest.mark.skipif(os.geteuid() != 0, reason="Locking GPU clocks requires root")
+
+
+def _read_sm_clock() -> int:
+    nvml = _nvml()
+    nvml.nvmlInit()
+    handle = _get_nvml_handle(nvml)
+    clk = ctypes.c_uint()
+    nvml.nvmlDeviceGetClockInfo(handle, ctypes.c_uint(1), ctypes.byref(clk))
+    val = clk.value
+    nvml.nvmlShutdown()
+    return val
+
+
+@requires_root
+def test_locked_clocks_context_manager():
+    with locked_clocks() as target_mhz:
+        assert isinstance(target_mhz, int)
+        assert target_mhz > 0
+
+        x = torch.randn(1024, 1024, device="cuda")
+        for _ in range(10):
+            x @ x
+        torch.cuda.synchronize()
+
+        assert _read_sm_clock() == target_mhz
+
+
+@requires_root
+def test_locked_clocks_custom_frequency():
+    nvml = _nvml()
+    nvml.nvmlInit()
+    handle = _get_nvml_handle(nvml)
+    max_clk = _get_max_sm_clock(nvml, handle)
+    nvml.nvmlShutdown()
+
+    with locked_clocks(clock_mhz=max_clk) as target_mhz:
+        assert target_mhz == max_clk
+
+
+@requires_root
+def test_locked_clocks_resets_on_exception():
+    with pytest.raises(ValueError, match="intentional"):
+        with locked_clocks():
+            raise ValueError("intentional")
+
+
+@requires_root
+def test_benchmark_lock_clocks_kwarg():
+    x = torch.randn(1024, 1024, device="cuda")
+    t = benchmark_cuda_function_in_microseconds_triton(lambda: x @ x, LOCK_CLOCKS=True)
+    assert t > 0
+
+
+def test_locked_clocks_no_permission():
+    if os.geteuid() == 0:
+        pytest.skip("Running as root, cannot test permission failure")
+    with pytest.raises(RuntimeError, match="Requires root"):
+        with locked_clocks():
+            pass
+
+
 def test_nan():
     a = torch.tensor(
         [

transformer_nuggets/utils/__init__.py

@@ -2,6 +2,7 @@
     benchmark_torch_function_in_microseconds,
     benchmark_cuda_function_in_microseconds,
     benchmark_cuda_function_in_microseconds_triton,
+    locked_clocks,
     max_memory_usage,
     cuda_memory_usage,
     profile_function,

transformer_nuggets/utils/benchmark.py

@@ -18,6 +18,44 @@
 logger.addHandler(logging.NullHandler())
 
 
+@contextmanager
+def locked_clocks(device: int = 0, clock_mhz: int | None = None):
+    """Lock GPU SM clocks for stable benchmarking.
+
+    Uses ``sudo nvidia-smi -lgc`` to lock and ``sudo nvidia-smi -rgc`` to
+    reset. Will prompt for a password in interactive terminals.
+
+    Args:
+        device: CUDA device index.
+        clock_mhz: SM clock frequency in MHz. If None, locks to the GPU's max SM clock.
+    """
+    import subprocess
+
+    if clock_mhz is None:
+        clock_mhz = int(
+            subprocess.check_output(
+                [
+                    "nvidia-smi",
+                    "-i",
+                    str(device),
+                    "--query-gpu=clocks.max.sm",
+                    "--format=csv,noheader,nounits",
+                ],
+                text=True,
+            ).strip()
+        )
+
+    subprocess.check_call(
+        ["sudo", "nvidia-smi", "-i", str(device), "-lgc", f"{clock_mhz},{clock_mhz}"]
+    )
+    logger.info(f"Locked GPU {device} SM clocks to {clock_mhz} MHz")
+    try:
+        yield clock_mhz
+    finally:
+        subprocess.call(["sudo", "nvidia-smi", "-i", str(device), "-rgc"])
+        logger.info(f"Reset GPU {device} SM clocks")
+
+
 def lazy_import_error(error_msg: str):
     """Decorator that allows functions with imports to be defined without the dependency"""
 
@@ -60,40 +98,48 @@ class ProfileConfig:
 
 
 def benchmark_torch_function_in_microseconds(func: Callable, *args, **kwargs) -> float:
-    # warmup
-    for _ in range(5):
-        func(*args, **kwargs)
-    t0 = benchmark.Timer(
-        stmt="func(*args, **kwargs)",
-        globals={"args": args, "kwargs": kwargs, "func": func},
-    )
-    return t0.adaptive_autorange(min_run_time=0.1).median * 1e6
+    lock = kwargs.pop("LOCK_CLOCKS", False)
+    ctx = locked_clocks() if lock else nullcontext()
+    with ctx:
+        for _ in range(5):
+            func(*args, **kwargs)
+        t0 = benchmark.Timer(
+            stmt="func(*args, **kwargs)",
+            globals={"args": args, "kwargs": kwargs, "func": func},
+        )
+        return t0.adaptive_autorange(min_run_time=0.1).median * 1e6
 
 
 def benchmark_cuda_function_in_microseconds(func: Callable, *args, **kwargs) -> float:
     """Thin wrapper around do_bench_using_profiling.
 
-    Accepts NUM_ITERS as a kwarg but removes it before calling func so it
-    never leaks into the benchmarked callable.
+    Accepts NUM_ITERS, IS_VETTED_BENCHMARKING, and lock_clocks as kwargs but
+    removes them before calling func so they never leak into the benchmarked callable.
     """
     num_iters = kwargs.pop("NUM_ITERS", 100)
     is_vetted_benchmarking = kwargs.pop("IS_VETTED_BENCHMARKING", False)
-    no_args = lambda: func(*args, **kwargs)
-    time = do_bench_using_profiling(
-        no_args, rep=num_iters, is_vetted_benchmarking=is_vetted_benchmarking
-    )
-    return time * 1e3
+    lock = kwargs.pop("LOCK_CLOCKS", False)
+    ctx = locked_clocks() if lock else nullcontext()
+    with ctx:
+        no_args = lambda: func(*args, **kwargs)
+        return (
+            do_bench_using_profiling(
+                no_args, rep=num_iters, is_vetted_benchmarking=is_vetted_benchmarking
+            )
+            * 1e3
+        )
 
 
 @lazy_import_error("This function requires Triton. Please install it with: pip install triton")
 def benchmark_cuda_function_in_microseconds_triton(func: Callable, *args, **kwargs) -> float:
     """Thin wrapper around do_bench"""
-    from triton.testing import do_bench  # Python caches this automatically
-
-    no_args = lambda: func(*args, **kwargs)
-    time = do_bench(no_args)
+    from triton.testing import do_bench
 
-    return time * 1e3
+    lock = kwargs.pop("LOCK_CLOCKS", False)
+    ctx = locked_clocks() if lock else nullcontext()
+    with ctx:
+        no_args = lambda: func(*args, **kwargs)
+        return do_bench(no_args) * 1e3
 
 
 def profile_function(