cosmos-framework/cosmos_framework/callbacks/ofu.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""OFU (Operational FLOPs Utilization) callback for OmniMoT training.

Computes and logs OFU metrics by launching ``nvidia-smi dmon`` as a background
subprocess and parsing the Tensor Core activity (mmaact) and processor clock
(pclk) columns.  OFU is defined as::

    OFU = mmaact * (pclk / max_pclk)

where ``max_pclk`` is the max boost clock for the detected hardware (e.g.
1980 MHz for H100, 2062 MHz for GB200).  The result is in the 0-100 range.
"""

from __future__ import annotations

import subprocess
import threading
from collections import defaultdict
from dataclasses import dataclass

import torch
import wandb

from cosmos_framework.model.attention.utils import is_blackwell_dc
from cosmos_framework.callbacks.every_n import EveryN
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.trainer import ImaginaireTrainer
from cosmos_framework.utils import log
from cosmos_framework.utils.distributed import is_rank0, rank0_only


@dataclass
class HardwareTarget:
    """Hardware-specific constants for OFU normalisation.

    Attributes:
        name: Human-readable name (used as W&B tag, e.g. "H100").
        max_pclk_mhz: Max boost SM clock in MHz used to normalise OFU.
    """

    name: str
    max_pclk_mhz: float


# Pre-defined hardware targets
H100 = HardwareTarget(name="H100", max_pclk_mhz=1980.0)
GB200 = HardwareTarget(name="GB200", max_pclk_mhz=2062.0)


class OFUCallback(EveryN):
    """Callback that computes and logs Operational FLOPs Utilization (OFU) to W&B.

    OFU = mmaact * (pclk / max_pclk), where mmaact is the MMA activity
    percentage and pclk is the current processor clock from ``nvidia-smi dmon``.
    ``max_pclk`` is determined from the detected hardware (H100 or GB200).
    The result is in the 0-100 range.

    The callback launches ``nvidia-smi dmon`` as a background subprocess on
    ``on_train_start`` and a daemon thread continuously reads its output.
    At every logging interval, accumulated samples are consumed, averaged per GPU
    and overall, and logged to W&B under ``ofu/{hardware_name}``.

    Args:
        hit_thres: Number of warm-up training iterations to skip before logging.
    """

    def __init__(
        self,
        *args,
        hit_thres: int = 5,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.hardware_target = GB200 if is_blackwell_dc() else H100
        self.hit_thres = hit_thres

        # Subprocess state
        self._process: subprocess.Popen | None = None
        self._reader_thread: threading.Thread | None = None
        self._stop_event = threading.Event()

        # Buffered samples protected by a lock: list of (gpu_idx, mmaact, pclk)
        self._lock = threading.Lock()
        self._samples: list[tuple[int, float, float]] = []

        # Column indices parsed from the header (set by _reader_loop)
        self._col_gpu: int | None = None
        self._col_mmaact: int | None = None
        self._col_pclk: int | None = None

        # Warm-up counter
        self._hit_counter: int = 0

    # ------------------------------------------------------------------ #
    # Background reader
    # ------------------------------------------------------------------ #

    def _parse_header(self, line: str) -> bool:
        """Parse a dmon header line to locate column indices.

        Called on every ``#`` line because nvidia-smi dmon reprints the header
        every few seconds.  Returns True if ``gpu``, ``mmaact``, and ``pclk``
        columns are all found; warns only when the column-names line (identified
        by the presence of ``gpu``) lacks a required column.  Silently ignores
        the units line (``# Idx  W  C ...``) which does not contain ``gpu``.
        """
        cols = line.lstrip("#").strip().split()
        col_map = {name.lower(): idx for idx, name in enumerate(cols)}
        gpu_idx = col_map.get("gpu")
        mmaact_idx = col_map.get("mmaact")
        pclk_idx = col_map.get("pclk")
        if gpu_idx is not None and mmaact_idx is not None and pclk_idx is not None:
            if self._col_mmaact is None:
                log.info(f"OFUCallback: found mmaact at column {mmaact_idx}, pclk at column {pclk_idx}")
            self._col_gpu = gpu_idx
            self._col_mmaact = mmaact_idx
            self._col_pclk = pclk_idx
            return True
        if gpu_idx is not None:
            missing = [name for name, idx in [("mmaact", mmaact_idx), ("pclk", pclk_idx)] if idx is None]
            log.warning(
                f"OFUCallback: column(s) {missing} not found in nvidia-smi dmon header: {cols}. "
                "OFU metrics will not be available."
            )
        return False

    def _reader_loop(self) -> None:
        """Background thread that reads nvidia-smi dmon output line-by-line."""
        assert self._process is not None and self._process.stdout is not None

        for line in self._process.stdout:
            if self._stop_event.is_set():
                break
            line = line.strip()
            if not line:
                continue

            # Header lines repeat every few seconds — always re-parse so that a
            # missed or failed first parse is recovered on the next occurrence.
            if line.startswith("#"):
                self._parse_header(line)
                continue

            # Skip data lines until we have column indices
            if self._col_gpu is None or self._col_mmaact is None or self._col_pclk is None:
                continue

            parts = line.split()
            try:
                gpu_idx = int(parts[self._col_gpu])
                mmaact = float(parts[self._col_mmaact])
                pclk = float(parts[self._col_pclk])
            except (ValueError, IndexError):
                continue

            with self._lock:
                self._samples.append((gpu_idx, mmaact, pclk))

    # ------------------------------------------------------------------ #
    # Lifecycle hooks
    # ------------------------------------------------------------------ #

    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
        if not is_rank0():
            return

        try:
            # --gpm-metrics 5 means that we access Tensor Activity under mmaact column.
            # -d 5 means that we sample the data every 5 seconds.
            cmd = ["nvidia-smi", "dmon", "--gpm-metrics", "5", "-d", "5"]
            self._process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,  # line-buffered
            )
            self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
            self._reader_thread.start()
            log.info(f"OFUCallback: launched nvidia-smi dmon --gpm-metrics 5")
        except FileNotFoundError:
            log.warning("OFUCallback: nvidia-smi not found, OFU metrics will not be available")
        except Exception as e:
            log.warning(f"OFUCallback: failed to launch nvidia-smi dmon: {e}")

    def on_train_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
        if not is_rank0():
            return
        self._stop_event.set()
        if self._process is not None:
            try:
                self._process.terminate()
                self._process.wait(timeout=5)
            except ProcessLookupError:
                pass  # already exited
            except subprocess.TimeoutExpired:
                self._process.kill()
            self._process = None
        if self._reader_thread is not None:
            self._reader_thread.join(timeout=5)
            self._reader_thread = None

    # ------------------------------------------------------------------ #
    # Per-step gating
    # ------------------------------------------------------------------ #

    def on_training_step_end(
        self,
        model: ImaginaireModel,
        data_batch: dict[str, torch.Tensor],
        output_batch: dict[str, torch.Tensor],
        loss: torch.Tensor,
        iteration: int = 0,
    ) -> None:
        # All ranks must enter super().on_training_step_end() so they reach the
        # distributed barrier inside EveryN.  Only rank 0 has samples to clear.
        if self._hit_counter < self.hit_thres:
            self._hit_counter += 1
            if self._hit_counter == self.hit_thres:
                # Discard samples collected during warm-up (compilation, allocation, etc.)
                with self._lock:
                    self._samples.clear()
            return
        # Delegate to EveryN for the periodic reporting logic
        super().on_training_step_end(model, data_batch, output_batch, loss, iteration)

    # ------------------------------------------------------------------ #
    # Periodic reporting
    # ------------------------------------------------------------------ #

    @rank0_only
    def every_n_impl(
        self,
        trainer: ImaginaireTrainer,
        model: ImaginaireModel,
        data_batch: dict[str, torch.Tensor],
        output_batch: dict[str, torch.Tensor],
        loss: torch.Tensor,
        iteration: int,
    ) -> None:
        if self._process is None:
            return

        # Drain buffered samples
        with self._lock:
            samples = list(self._samples)
            self._samples.clear()

        if not samples:
            log.warning(
                f"OFUCallback: no nvidia-smi samples collected at iteration {iteration}. "
                "Check that the dmon subprocess launched and that the mmaact column is present."
            )
            return

        # Compute per-GPU OFU: mmaact * (pclk / max_pclk)
        max_pclk = self.hardware_target.max_pclk_mhz
        gpu_ofu: dict[int, list[float]] = defaultdict(list)
        gpu_mmaact: dict[int, list[float]] = defaultdict(list)
        gpu_pclk: dict[int, list[float]] = defaultdict(list)
        for gpu_idx, mmaact, pclk in samples:
            gpu_ofu[gpu_idx].append(mmaact * (pclk / max_pclk))
            gpu_mmaact[gpu_idx].append(mmaact)
            gpu_pclk[gpu_idx].append(pclk)

        # Overall averages across all GPUs and samples
        all_ofu = [v for vals in gpu_ofu.values() for v in vals]
        all_mmaact = [v for vals in gpu_mmaact.values() for v in vals]
        all_pclk = [v for vals in gpu_pclk.values() for v in vals]

        log_info: dict[str, float] = {
            f"ofu/{self.hardware_target.name}": sum(all_ofu) / len(all_ofu),
            "ofu/mmaact": sum(all_mmaact) / len(all_mmaact),
            "ofu/avg_pclk_mhz": sum(all_pclk) / len(all_pclk),
            "ofu/num_samples": float(len(samples)),
        }

        if wandb.run is not None:
            wandb.log(log_info, step=iteration)