Source code for composer.callbacks.speed_monitor

# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""Monitor throughput during training."""
from __future__ import annotations

import warnings
from collections import deque
from typing import Any, Callable, Deque, Dict, Optional, Union

import torch

from composer.core import Callback, State
from composer.loggers import Logger
from composer.models.base import ComposerModel
from composer.utils import dist, is_xla_installed

if is_xla_installed():
    import torch_xla.core.xla_model as xm

__all__ = ['SpeedMonitor']

GPU_AVAILABLE_FLOPS = {
    # source: https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor-core-gpu-datasheet
    # nvidia publishes spec sheet with a 2x sparsity factor
    'h100-sxm': {
        'fp64': 67e12,
        'fp32': 67e12,
        'tf32': 989e12 / 2,
        'fp16': 1.979e15 / 2,
        'amp_fp16': 1.979e15 / 2,
        'bf16': 1.979e15 / 2,
        'amp_bf16': 1.979e15 / 2,
        'fp8': 3.958e15 / 2,
        'amp_fp8': 3.958e15 / 2,
        'int8': 3.958e15 / 2,
    },
    'h100-pcie': {
        'fp64': 51e12,
        'fp32': 51e12,
        'tf32': 756e12 / 2,
        'fp16': 1.513e15 / 2,
        'amp_fp16': 1.513e15 / 2,
        'bf16': 1.513e15 / 2,
        'amp_bf16': 1.513e15 / 2,
        'fp8': 3.026e15 / 2,
        'amp_fp8': 3.026e15 / 2,
        'int8': 3.026e15 / 2,
    },
    # source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf
    # sxm and pcie have same flop counts
    'a100': {
        'fp64': 19.5e12,
        'fp32': 19.5e12,
        'tf32': 156e12,
        'fp16': 312e12,
        'amp_fp16': 312e12,
        'bf16': 312e12,
        'amp_bf16': 312e12,
    },
    # source: https://images.nvidia.com/content/technologies/volta/pdf/volta-v100-datasheet-update-us-1165301-r5.pdf
    'v100-sxm': {
        'fp64': 7.8e12,
        'fp32': 15.7e12,
        'fp16': 125e12,
        'amp_fp16': 125e12,
    },
    'v100-pcie': {
        'fp64': 7e12,
        'fp32': 14e12,
        'fp16': 112e12,
        'amp_fp16': 112e12,
    },
    'v100s-pcie': {
        'fp64': 8.2e12,
        'fp32': 16.4e12,
        'fp16': 130e12,
        'amp_fp16': 130e12,
    },
    # source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/tesla-t4/t4-tensor-core-datasheet-951643.pdf
    # sxm and pcie have same flop counts
    't4': {
        'fp32': 8.1e12,
        'fp16': 65e12,
        'amp_fp16': 65e12,
        'int8': 130e12,
        'int4': 260e12,
    },
    # source: https://aws.amazon.com/blogs/machine-learning/aws-inferentia2-builds-on-aws-inferentia1-by-delivering-4x-higher-throughput-and-10x-lower-latency/
    # Numbers are halved as the above flops is per chip and each chip appears as 2 devices.
    'trn1': {
        'fp32': 47.5e12 / 2,
        'tf32': 47.5e12 / 2,
        'fp16': 190e12 / 2,
        'amp_fp16': 190e12 / 2,
        'bf16': 190e12 / 2,
        'amp_bf16': 190e12 / 2,
        'int8': 380e12 / 2,
    }
}


def get_gpu_flops_available(state: State):
    gpu_flops_available = None

    # Return 0 if no CUDA device (e.g., when running with CPU only)
    if torch.cuda.is_available():
        # torch.cuda.get_device_name() ex output: 'NVIDIA A100-SXM4-40GB'
        device_name = torch.cuda.get_device_name().lower()
        if 'h100' in device_name and 'hbm3' in device_name:
            device_name = 'h100-sxm'
        elif 'h100' in device_name and ('pcie' in device_name or 'hbm2e' in device_name):
            device_name = 'h100-pcie'
        elif 'a100' in device_name:
            device_name = 'a100'
        elif 'v100-sxm' in device_name:
            device_name = 'v100-sxm'
        elif 'v100-pcie' in device_name:
            device_name = 'v100-pcie'
        elif 't4' in device_name:
            device_name = 't4'
    elif is_xla_installed():
        if xm.xla_device_hw(xm.xla_device()) == 'NEURON':
            device_name = 'trn1'
        else:
            # For TPU return 0
            return 0
    else:
        # When running on CPU, return 0 without warning
        return 0

    if device_name in GPU_AVAILABLE_FLOPS and state.precision.value in GPU_AVAILABLE_FLOPS[device_name]:
        gpu_flops_available = int(GPU_AVAILABLE_FLOPS[device_name][state.precision.value])
    else:
        gpu_flops_available = None

    if gpu_flops_available is None:
        warnings.warn(
            f'gpu_flop count not found for {device_name} with precision={state.precision.value} ' +\
            f'so MFU cannot be calculated and reported. gpu_flops_available can be manually ' +\
            f'overridden by setting gpu_flops_available in SpeedMonitor or {device_name} can ' +\
            f'be added to GPU_AVAILABLE_FLOPS in composer/callbacks/speed_monitor.py',
            stacklevel=2,
        )
        # Setting to 0 will disable MFU computation and prevent
        # the speed monitor from running this helper every batch
        gpu_flops_available = 0

    return gpu_flops_available


[docs]class SpeedMonitor(Callback): """Logs the training throughput and utilization. The training throughput is logged on the :attr:`.Event.BATCH_END` event once we have reached the `window_size` threshold. If a model has `flops_per_batch` attribute, then flops per second is also logged. If running on a known GPU type or if `gpu_flops_available` is set, then MFU is also logged. All metrics are also logged as per device by dividing by world size. To compute `flops_per_sec`, the model attribute `flops_per_batch` should be set to a callable which accepts a batch and returns the number of flops for that batch. Typically, this should be flops per sample times the batch size unless pad tokens are used. The wall clock time is logged on every :attr:`.Event.BATCH_END` event. Example: .. doctest:: >>> from composer import Trainer >>> from composer.callbacks import SpeedMonitor >>> # constructing trainer object with this callback >>> trainer = Trainer( ... model=model, ... train_dataloader=train_dataloader, ... eval_dataloader=eval_dataloader, ... optimizers=optimizer, ... max_duration='1ep', ... callbacks=[SpeedMonitor(window_size=100)], ... ) The training throughput is logged by the :class:`.Logger` to the following keys as described below. +-------------------------------------+-----------------------------------------------------------+ | Key | Logged data | +=====================================+===========================================================+ | | Rolling average (over `window_size` most recent | | `throughput/batches_per_sec` | batches) of the number of batches processed per second | | | | +-------------------------------------+-----------------------------------------------------------+ | | Rolling average (over `window_size` most recent | | `throughput/samples_per_sec` | batches) of the number of samples processed per second | | | | +-------------------------------------+-----------------------------------------------------------+ | | Rolling average (over `window_size` most recent | | `throughput/tokens_per_sec` | batches) of the number of tokens processed per second. | | | Only logged if dataspec returns tokens per batch | +-------------------------------------+-----------------------------------------------------------+ | | Estimates flops by `flops_per_batch * batches_per_sec` | | `throughput/flops_per_sec` | if model has attribute `flops_per_batch` | | | | +-------------------------------------+-----------------------------------------------------------+ | `throughput/device/batches_per_sec` | `throughput/batches_per_sec` divided by world size | +-------------------------------------+-----------------------------------------------------------+ | `throughput/device/samples_per_sec` | `throughput/samples_per_sec` divided by world size | +-------------------------------------+-----------------------------------------------------------+ | | `throughput/tokens_per_sec` divided by world size. Only | | `throughput/device/tokens_per_sec` | logged if dataspec returns tokens per batch | | | | +-------------------------------------+-----------------------------------------------------------+ | | `throughput/flops_per_sec` divided by world size. Only | | `throughput/device/flops_per_sec` | logged when model has attribute `flops_per_batch` | | | | +-------------------------------------+-----------------------------------------------------------+ | | `throughput/device/flops_per_sec` divided by world size. | | `throughput/device/mfu` | Only logged when model has attribute `flops_per_batch` | | | and `gpu_flops_available`, which can be passed as an | | | argument if not automatically determined by SpeedMonitor | +-------------------------------------+-----------------------------------------------------------+ | `time/train` | Total elapsed training time | +-------------------------------------+-----------------------------------------------------------+ | `time/val` | Total elapsed validation time | +-------------------------------------+-----------------------------------------------------------+ | `time/total` | Total elapsed time (time/train + time/val) | +-------------------------------------+-----------------------------------------------------------+ Args: window_size (int, optional): Number of batches to use for a rolling average of throughput. Defaults to 100. gpu_flops_available (float, optional): Number of flops available on the GPU. If not set, SpeedMonitor will attempt to determine this automatically. Defaults to None. time_unit (str, optional): Time unit to use for `time` logging. Can be one of 'seconds', 'minutes', 'hours', or 'days'. Defaults to 'hours'. """ def __init__( self, window_size: int = 100, gpu_flops_available: Optional[Union[float, int]] = None, time_unit: str = 'hours', ): # Track the batch num samples and wct to compute throughput over a window of batches self.history_samples: Deque[int] = deque(maxlen=window_size + 1) self.history_tokens: Deque[int] = deque(maxlen=window_size + 1) self.history_wct: Deque[float] = deque(maxlen=window_size + 1) self.history_flops: Deque[float] = deque(maxlen=window_size + 1) self.gpu_flops_available = gpu_flops_available self.divider = 1 if time_unit == 'seconds': self.divider = 1 elif time_unit == 'minutes': self.divider = 60 elif time_unit == 'hours': self.divider = 60 * 60 elif time_unit == 'days': self.divider = 60 * 60 * 24 else: raise ValueError( f'Invalid time_unit: {time_unit}. Must be one of "seconds", "minutes", "hours", or "days".') # Keep track of time spent evaluating self.total_eval_wct = 0.0 def state_dict(self) -> Dict[str, Any]: return { 'total_eval_wct': self.total_eval_wct, } def load_state_dict(self, state: Dict[str, Any]) -> None: self.total_eval_wct = state['total_eval_wct'] def init(self, state: State, logger: Logger) -> None: del logger # unused if self.gpu_flops_available is None: self.gpu_flops_available = get_gpu_flops_available(state) def batch_end(self, state: State, logger: Logger): # Add the new element self.history_samples.append(state.timestamp.sample.value) self.history_tokens.append(state.timestamp.token.value) self.history_wct.append(state.timestamp.total_wct.total_seconds()) # Log the throughput if len(self.history_wct) == self.history_wct.maxlen: world_size = dist.get_world_size() elapsed_batches = len(self.history_samples) - 1 elapsed_samples = int(self.history_samples[-1]) - int(self.history_samples[0]) elapsed_tokens = int(self.history_tokens[-1]) - int(self.history_tokens[0]) elapsed_wct = self.history_wct[-1] - self.history_wct[0] batches_per_sec = elapsed_batches / elapsed_wct samples_per_sec = elapsed_samples / elapsed_wct dev_batches_per_sec = batches_per_sec / world_size dev_samples_per_sec = samples_per_sec / world_size logger.log_metrics({ 'throughput/batches_per_sec': batches_per_sec, 'throughput/samples_per_sec': samples_per_sec, 'throughput/device/batches_per_sec': dev_batches_per_sec, 'throughput/device/samples_per_sec': dev_samples_per_sec, }) if elapsed_tokens > 0: tokens_per_sec = elapsed_tokens / elapsed_wct dev_tokens_per_sec = tokens_per_sec / world_size logger.log_metrics({ 'throughput/tokens_per_sec': tokens_per_sec, 'throughput/device/tokens_per_sec': dev_tokens_per_sec, }) # Compute flops stats if model has flops_per_batch composer_model = state.model if not isinstance(composer_model, ComposerModel): composer_model = composer_model.module # Pass through DDP wrapping if hasattr(composer_model, 'flops_per_batch'): model_flops_per_batch = composer_model.flops_per_batch # type: ignore if not isinstance(model_flops_per_batch, Callable): raise TypeError('flops_per_batch must a callable accepting a batch and ' f'returning an int or float. Instead, got {type(model_flops_per_batch)}.') device_flops_per_batch = model_flops_per_batch(state.batch) # Sum flops across all ranks since each rank computes the flops for its own batch flops_per_batch_tensor = state.device.tensor_to_device( torch.tensor(device_flops_per_batch, dtype=torch.float)) dist.all_reduce(flops_per_batch_tensor, reduce_operation='SUM') flops_per_batch = flops_per_batch_tensor.item() self.history_flops.append(flops_per_batch) # Log the flops throughput if len(self.history_flops) == self.history_flops.maxlen: world_size = dist.get_world_size() elapsed_flops = sum(self.history_flops) - self.history_flops[0] elapsed_wct = self.history_wct[-1] - self.history_wct[0] flops_per_sec = elapsed_flops / elapsed_wct device_flops_per_sec = flops_per_sec / world_size logger.log_metrics({ 'throughput/flops_per_sec': flops_per_sec, 'throughput/device/flops_per_sec': device_flops_per_sec, }) if self.gpu_flops_available: mfu = device_flops_per_sec / self.gpu_flops_available logger.log_metrics({'throughput/device/mfu': mfu}) # Log the time # `state.timestamp` excludes any time spent in evaluation train_wct = state.timestamp.total_wct.total_seconds() logger.log_metrics({ 'time/train': train_wct / self.divider, 'time/val': self.total_eval_wct / self.divider, 'time/total': (train_wct + self.total_eval_wct) / self.divider, }) def eval_end(self, state: State, logger: Logger): del logger # unused self.total_eval_wct += state.eval_timestamp.total_wct.total_seconds()