[power] Add initial power chart support (meta-pytorch#420)

Author: xuzhao9

.github/workflows/_linux-test-h100.yml

@@ -30,8 +30,10 @@ jobs:
       - name: Install Tritonbench
         run: |
           # speedup install and skip compile by reusing the docker .so files
+          . "${SETUP_SCRIPT}"
           mkdir -p /workspace/tritonbench/.data
           ln -s /workspace/tritonbench/.data .
+          pip install -r requirements.txt
       - name: Test Tritonbench operators on H100 GPU
         run: |
           bash ./.ci/tritonbench/test-gpu.sh

requirements.txt

@@ -1,5 +1,6 @@
 packaging
-pynvml
+nvidia-ml-py
 psutil
 tabulate
+matplotlib
 transformers==4.46.1

test/test_gpu/skip_tests_h100_pytorch.yaml

@@ -26,6 +26,9 @@ fp8_attention:
 # fp8_fused_quant_gemm_rowwise requires fb-only kernels
 fp8_fused_quant_gemm_rowwise:
 gemm:
+  # torch._inductor.exc.InductorError: LoweringException
+  # NoValidChoicesError: No choices to select. 
+  - pt2_cutlass_matmul
   # internal only kernels
   - hstu_triton_matmul
 # jagged tests are slow, so disable them in OSS

test/test_gpu/skip_tests_h100_triton_main.yaml

@@ -30,6 +30,7 @@ gemm:
   - hstu_triton_matmul
   # No need to test cutlass on triton main
   - pt2_cutlass_matmul
+int4_gemm:
 # jagged tests are slow, so disable them in OSS
 jagged_layer_norm:
 jagged_mean:

tritonbench/components/power/__init__.py

@@ -0,0 +1 @@
+from .chart import power_chart_begin, power_chart_end

tritonbench/components/power/chart.py

@@ -0,0 +1,134 @@
+import csv
+import logging
+import os
+import signal
+import subprocess
+import time
+
+import matplotlib.pyplot as plt
+import torch
+
+# query every 10 ms
+QUERY_FREQUENCY = 10
+QUERY_STDOUT_FILE = "power.csv"
+QUERY_STDERR_FILE = "power.log"
+QUERY_COMMAND = """nvidia-smi -lms {QUERY_FREQUENCY} -i {QUERY_DEVICE} --query-gpu=power.draw.average,power.draw.instant,power.max_limit,temperature.gpu,temperature.memory,clocks.current.sm,clocks.current.memory,clocks_throttle_reasons.hw_thermal_slowdown,clocks_throttle_reasons.sw_thermal_slowdown --format=csv,nounits"""
+global QUERY_PROC
+global POWER_OUTPUT_DIR
+
+QUERY_PROC = None
+POWER_OUTPUT_DIR = None
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def _get_cuda_device_id():
+    return torch.cuda.current_device()
+
+
+def _gen_power_charts(benchmark_name: str, device_name: str, power_csv_file: str):
+    # Read CSV
+    with open(power_csv_file) as f:
+        reader = csv.reader(f)
+        header = next(reader)  # first row as header
+        header = [col.strip() for col in header]
+        data = {col: [] for col in header}
+
+        for row in reader:
+            for col, value in zip(header, row):
+                if value == "[N/A]":
+                    logger.warning(
+                        f"[tritonbench][power] {col} is not available, skipping"
+                    )
+                    value = 0.0
+                else:
+                    value = (
+                        float(value)
+                        if col
+                        not in [
+                            "clocks_event_reasons.hw_thermal_slowdown",
+                            "clocks_event_reasons.sw_thermal_slowdown",
+                        ]
+                        else value
+                    )
+                data[col].append(value)
+
+    # Generate synthetic time axis (100 ms per sample)
+    n_samples = len(next(iter(data.values())))
+    time = [i * 0.1 for i in range(n_samples)]  # seconds (0.1s = 100 ms)
+
+    # Plot power chart
+    plt.figure(figsize=(10, 6))
+    for power_col in header[:3]:
+        plt.plot(time, data[power_col], label=power_col)
+    plt.xlabel("Time (s)")
+    plt.ylabel("Power (W)")
+    plt.legend()
+    plt.title(
+        f"[tritonbench] {benchmark_name} power consumption over time on {device_name}"
+    )
+    plt.savefig(
+        os.path.join(POWER_OUTPUT_DIR, "power.png"), dpi=300, bbox_inches="tight"
+    )
+    # Plot temp chart
+    plt.figure(figsize=(10, 6))
+    for temp_col in header[3:5]:
+        plt.plot(time, data[temp_col], label=temp_col)
+        plt.xlabel("Time (s)")
+        plt.ylabel("Temperature (C)")
+    plt.legend()
+    plt.title(f"[tritonbench] {benchmark_name} temperature over time on {device_name}")
+    plt.savefig(
+        os.path.join(POWER_OUTPUT_DIR, "temp.png"), dpi=300, bbox_inches="tight"
+    )
+    # Plot frequency chart
+    plt.figure(figsize=(10, 6))
+    for temp_col in header[5:7]:
+        plt.plot(time, data[temp_col], label=temp_col)
+        plt.xlabel("Time (s)")
+        plt.ylabel("Frequency (MHz)")
+    plt.legend()
+    plt.title(f"[tritonbench] {benchmark_name} frequency over time on {device_name}")
+    plt.savefig(
+        os.path.join(POWER_OUTPUT_DIR, "freq.png"), dpi=300, bbox_inches="tight"
+    )
+
+
+def power_chart_begin(benchmark_name, output_dir):
+    # check no other proc is running
+    global QUERY_PROC, POWER_OUTPUT_DIR
+    assert QUERY_PROC is None, "Power query process must be None to start a new one"
+    # clean up the directory
+    POWER_OUTPUT_DIR = os.path.join(output_dir, benchmark_name)
+    if not os.path.exists(POWER_OUTPUT_DIR):
+        os.mkdir(POWER_OUTPUT_DIR)
+    stdout_file_path = os.path.join(POWER_OUTPUT_DIR, QUERY_STDOUT_FILE)
+    stderr_file_path = os.path.join(POWER_OUTPUT_DIR, QUERY_STDERR_FILE)
+    # Run the command
+    query_cmd = QUERY_COMMAND.format(
+        QUERY_FREQUENCY=QUERY_FREQUENCY, QUERY_DEVICE=_get_cuda_device_id()
+    ).split(" ")
+    with open(stdout_file_path, "w") as stdout_file, open(
+        stderr_file_path, "w"
+    ) as stderr_file:
+        QUERY_PROC = subprocess.Popen(
+            query_cmd, stdout=stdout_file, stderr=stderr_file, start_new_session=True
+        )
+
+
+def power_chart_end():
+    global QUERY_PROC, POWER_OUTPUT_DIR
+    assert QUERY_PROC is not None, "Power query process cannot be None"
+    # Kill the process
+    QUERY_PROC.send_signal(signal.SIGINT)
+    time.sleep(0.2)
+    assert (
+        QUERY_PROC.poll() is not None
+    ), "Power query process must be killed to proceed"
+    # generate the chart based on csv
+    stdout_file_path = os.path.join(POWER_OUTPUT_DIR, QUERY_STDOUT_FILE)
+    benchmark_name = os.path.basename(POWER_OUTPUT_DIR)
+    device_name = torch.cuda.get_device_name(_get_cuda_device_id())
+    _gen_power_charts(benchmark_name, device_name, stdout_file_path)
+    logger.warning(f"[tritonbench][power] Power chart saved to {POWER_OUTPUT_DIR}.")

tritonbench/operators/softmax/operator.py

@@ -42,6 +42,7 @@ def parse_op_args(args: List[str]):
 
 
 class Operator(BenchmarkOperator):
+    DEFAULT_PRECISION = "fp16"
     is_compute_bound = False
 
     def __init__(

tritonbench/utils/parser.py

@@ -180,6 +180,12 @@ def get_parser(args=None):
         default=None,
         help="Dump Triton IR to specific directory.",
     )
+    parser.add_argument(
+        "--power-chart",
+        type=str,
+        default=None,
+        help="Dump GPU power chart to specific directory.",
+    )
     parser.add_argument(
         "--gpu-lockdown",
         action="store_true",

tritonbench/utils/triton_op.py

@@ -28,11 +28,11 @@
 import tabulate
 import torch
 import triton
-from torch.utils._pytree import tree_flatten, tree_map
 
 from tritonbench.components.do_bench import do_bench_wrapper, Latency
 from tritonbench.components.export import export_data
 
+from tritonbench.components.power.chart import power_chart_begin, power_chart_end
 from tritonbench.utils.constants import (
     DEFAULT_QUANTILES,
     DEFAULT_REP,
@@ -873,6 +873,8 @@ def run(
     ) -> None:
         """Benchmarking the operator and returning its metrics."""
         metrics = []
+        if self.tb_args.power_chart:
+            power_chart_begin(self.benchmark_name, self.tb_args.power_chart)
         try:
             if "proton" in self.required_metrics:
                 import triton.profiler as proton
@@ -998,6 +1000,8 @@ def _reduce_benchmarks(acc, bm_name: str):
                 os._exit(1)
             raise
         finally:
+            if self.tb_args.power_chart:
+                power_chart_end()
             self.output = BenchmarkOperatorResult(
                 benchmark_name=self.tb_args.benchmark_name,
                 op_name=self.name,