import time
from collections import deque
from contextlib import nullcontext
from typing import Any, Callable, Deque, Dict, Optional
import torch
from lightning import Callback, Fabric, LightningModule, Trainer
from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1
from lightning.fabric.plugins import (
BitsandbytesPrecision,
DoublePrecision,
FSDPPrecision,
HalfPrecision,
MixedPrecision,
Precision,
TransformerEnginePrecision,
XLAPrecision,
)
from lightning.fabric.utilities.rank_zero import rank_zero_only as fabric_rank_zero_only
from lightning.pytorch.plugins import (
DoublePrecisionPlugin,
FSDPPrecisionPlugin,
HalfPrecisionPlugin,
MixedPrecisionPlugin,
XLAPrecisionPlugin,
)
from lightning.pytorch.utilities.rank_zero import rank_zero_only as trainer_rank_zero_only
from torch.utils.flop_counter import FlopCounterMode
from tsai_gpt import GPT
from tsai_gpt.utils import num_parameters
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": {
torch.float64: 67e12,
torch.float32: 67e12,
torch.bfloat16: 1.979e15 / 2,
torch.float16: 1.979e15 / 2,
torch.int8: 3.958e15 / 2,
},
"h100-pcie": {
torch.float64: 51e12,
torch.float32: 51e12,
torch.bfloat16: 1.513e15 / 2,
torch.float16: 1.513e15 / 2,
torch.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": {torch.float64: 19.5e12, torch.float32: 19.5e12, torch.bfloat16: 312e12, torch.float16: 312e12},
# source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a10/pdf/a10-datasheet.pdf
"a10g": {torch.float32: 31.2e12, torch.bfloat16: 125e12, torch.float16: 125e12},
# source: https://images.nvidia.com/content/technologies/volta/pdf/volta-v100-datasheet-update-us-1165301-r5.pdf
"v100-sxm": {torch.float64: 7.8e12, torch.float32: 15.7e12, torch.float16: 125e12},
"v100-pcie": {torch.float64: 7e12, torch.float32: 14e12, torch.float16: 112e12},
"v100s-pcie": {torch.float64: 8.2e12, torch.float32: 16.4e12, torch.float16: 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": {torch.float32: 8.1e12, torch.float16: 65e12, torch.int8: 130e12},
# https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/quadro-product-literature/quadro-rtx-5000-data-sheet-us-nvidia-704120-r4-web.pdf
"quadro rtx 5000": {torch.float32: 11.2e12, torch.float16: 89.2e12},
}
TPU_AVAILABLE_FLOPS = {
# flop count for each TPU generation is the same for all precisions
# since bfloat16 precision is always used for performing matrix operations
# for more info: https://cloud.google.com/tpu/docs/bfloat16#choosing_bfloat16
# source: https://arxiv.org/pdf/1907.10701.pdf
"v2": 45e12,
# source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v3
"v3": 123e12,
# source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v4
"v4": 275e12,
# source: https://cloud.google.com/tpu/docs/v5e-training
"v5litepod": 197e12,
}
def get_flops_available(device: torch.device, dtype: torch.dtype) -> Optional[float]:
if device.type == "cuda":
device_name = torch.cuda.get_device_name(device).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 "a10g" in device_name:
device_name = "a10g"
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 "quadro rtx 5000" in device_name:
device_name = "quadro rtx 5000"
else:
device_name = None
if device_name is not None:
try:
return int(GPU_AVAILABLE_FLOPS[device_name][dtype])
except KeyError:
raise KeyError(
f"flop count not found for {device_name} with dtype: {dtype}; "
"MFU cannot be calculated and reported."
)
elif device.type == "xla":
if _XLA_GREATER_EQUAL_2_1:
from torch_xla._internal import tpu
else:
from torch_xla.experimental import tpu
device_name = tpu.get_tpu_env()["TYPE"].lower()
try:
return int(TPU_AVAILABLE_FLOPS[device_name])
except KeyError:
raise KeyError(
f"flop count not found for {device_name} with dtype: {dtype}; MFU cannot be calculated and reported."
)
return None
# Adapted from https://github.com/mosaicml/composer/blob/f2a2dc820cb75023b9eb7c46fdfd25273712abd0/composer/callbacks/speed_monitor.py
class SpeedMonitorBase:
"""Logs the training throughput and utilization.
+-------------------------------------+-----------------------------------------------------------+
| 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. |
| | This may include padding depending on dataset |
+-------------------------------------+-----------------------------------------------------------+
| | Estimates flops by `flops_per_batch * batches_per_sec` |
| `throughput/flops_per_sec` | |
| | |
+-------------------------------------+-----------------------------------------------------------+
| `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. This |
| `throughput/device/tokens_per_sec` | may include pad tokens depending on dataset |
| | |
+-------------------------------------+-----------------------------------------------------------+
| | `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` | |
| | |
+-------------------------------------+-----------------------------------------------------------+
| `time/train` | Total elapsed training time |
+-------------------------------------+-----------------------------------------------------------+
| `time/val` | Total elapsed validation time |
+-------------------------------------+-----------------------------------------------------------+
| `time/total` | Total elapsed time (time/train + time/val) |
+-------------------------------------+-----------------------------------------------------------+
Notes:
- The implementation assumes that devices are homogeneous as it normalizes by the world size.
- Tokens/sec, flops/sec and MFU do not account for padding tokens if present. We suggest using samples/sec or
batches/sec to measure throughput under this circumstance.
- Be careful when comparing MFU numbers across projects, as this will highly depend on the ``flops_per_batch``.
There is no widespread, realistic, and reliable implementation to compute them.
We suggest using our ``measure_flops`` function, but many other works will use ``estimated_flops`` which
will almost always be an overestimate when compared to the true value.
Args:
window_size (int, optional): Number of batches to use for a rolling average of throughput.
Defaults to 100.
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,
flops_available: float,
log_dict: Callable[[Dict, int], None],
window_size: int = 100,
time_unit: str = "hours",
):
self.flops_available = flops_available
self.log_dict = log_dict
# 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_wct: Deque[float] = deque(maxlen=window_size + 1)
self.history_lengths: Deque[int] = deque(maxlen=window_size + 1)
self.history_flops: Deque[int] = deque(maxlen=window_size + 1)
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
self.step = -1
def on_train_batch_end(
self,
samples: int, # total samples seen (per device)
train_elapsed: float, # total training time (seconds)
world_size: int,
flops_per_batch: Optional[int] = None, # (per device)
lengths: Optional[int] = None, # total length of the samples seen (per device)
) -> None:
self.step += 1
step = self.step
metrics = {}
self.history_samples.append(samples)
if lengths is not None:
self.history_lengths.append(lengths)
# if lengths are passed, there should be as many values as samples
assert len(self.history_samples) == len(self.history_lengths)
self.history_wct.append(train_elapsed)
if len(self.history_wct) == self.history_wct.maxlen:
elapsed_batches = len(self.history_samples) - 1
elapsed_samples = self.history_samples[-1] - self.history_samples[0]
elapsed_wct = self.history_wct[-1] - self.history_wct[0]
samples_per_sec = elapsed_samples * world_size / elapsed_wct
dev_samples_per_sec = elapsed_samples / elapsed_wct
metrics.update(
{
"throughput/batches_per_sec": elapsed_batches * world_size / elapsed_wct,
"throughput/samples_per_sec": samples_per_sec,
"throughput/device/batches_per_sec": elapsed_batches / elapsed_wct,
"throughput/device/samples_per_sec": dev_samples_per_sec,
}
)
if lengths is not None:
elapsed_lengths = int(self.history_lengths[-1]) - int(self.history_lengths[0])
avg_length = elapsed_lengths / elapsed_batches
metrics.update(
{
"throughput/tokens_per_sec": samples_per_sec * avg_length,
"throughput/device/tokens_per_sec": dev_samples_per_sec * avg_length,
}
)
if flops_per_batch is not None:
# sum of flops per batch across ranks
self.history_flops.append(flops_per_batch * world_size)
if len(self.history_flops) == self.history_flops.maxlen:
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
metrics.update(
{"throughput/flops_per_sec": flops_per_sec, "throughput/device/flops_per_sec": device_flops_per_sec}
)
if self.flops_available:
metrics["throughput/device/mfu"] = device_flops_per_sec / self.flops_available
metrics.update(
{
"time/train": train_elapsed / self.divider,
"time/val": self.total_eval_wct / self.divider,
"time/total": (train_elapsed + self.total_eval_wct) / self.divider,
"samples": samples,
}
)
self.log_dict(metrics, step)
def eval_end(self, eval_elapsed: float) -> None:
self.total_eval_wct += eval_elapsed # seconds
def plugin_to_compute_dtype(plugin: Precision) -> torch.dtype:
if isinstance(plugin, BitsandbytesPrecision):
return plugin.dtype
if isinstance(plugin, (HalfPrecision, MixedPrecision, HalfPrecisionPlugin)):
return plugin._desired_input_dtype
if isinstance(plugin, MixedPrecisionPlugin):
return torch.bfloat16 if plugin.precision == "bf16-mixed" else torch.half
if isinstance(plugin, (DoublePrecision, DoublePrecisionPlugin)):
return torch.double
if isinstance(plugin, (XLAPrecision, XLAPrecisionPlugin)):
return plugin._desired_dtype
if isinstance(plugin, TransformerEnginePrecision):
return torch.int8
if isinstance(plugin, (FSDPPrecision, FSDPPrecisionPlugin)):
return plugin.mixed_precision_config.reduce_dtype
if isinstance(plugin, Precision):
return torch.float32
raise NotImplementedError(plugin)
class SpeedMonitorFabric(SpeedMonitorBase):
def __init__(self, fabric: Fabric, *args: Any, **kwargs: Any) -> None:
dtype = plugin_to_compute_dtype(fabric.strategy.precision)
flops_available = get_flops_available(fabric.device, dtype)
super().__init__(flops_available, fabric.log_dict, *args, **kwargs)
@fabric_rank_zero_only
def on_train_batch_end(self, *args: Any, **kwargs: Any) -> None:
super().on_train_batch_end(*args, **kwargs)
class SpeedMonitorCallback(Callback):
def __init__(self, length_fn: Callable[[Any], int], batch_size: int, **kwargs: Any) -> None:
super().__init__()
self.speed_monitor: Optional[SpeedMonitorBase] = None
self.speed_monitor_kwargs = kwargs
self.length_fn = length_fn
self.batch_size = batch_size
self.eval_t0: int = 0
self.train_t0: int = 0
self.total_lengths: int = 0
def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
if self.speed_monitor is not None:
return # already setup
dtype = plugin_to_compute_dtype(trainer.precision_plugin)
flops_available = get_flops_available(trainer.strategy.root_device, dtype)
self.speed_monitor = SpeedMonitorBase(flops_available, trainer.logger.log_metrics, **self.speed_monitor_kwargs)
@trainer_rank_zero_only
def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
if trainer.fit_loop._should_accumulate():
return
self.train_t0 = time.perf_counter()
@trainer_rank_zero_only
def on_train_batch_end(
self, trainer: Trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int
) -> None:
self.total_lengths += self.length_fn(batch)
if trainer.fit_loop._should_accumulate():
return
train_elapsed = time.perf_counter() - self.train_t0
assert self.speed_monitor is not None
iter_num = trainer.fit_loop.total_batch_idx
assert (measured_flops := pl_module.measured_flops) is not None
self.speed_monitor.on_train_batch_end(
(iter_num + 1) * self.batch_size,
train_elapsed,
# this assumes that device FLOPs are the same and that all devices have the same batch size
trainer.world_size,
flops_per_batch=measured_flops,
lengths=self.total_lengths,
)
@trainer_rank_zero_only
def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
self.eval_t0 = time.perf_counter()
@trainer_rank_zero_only
def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
eval_elapsed = time.perf_counter() - self.eval_t0
assert self.speed_monitor is not None
self.speed_monitor.eval_end(eval_elapsed)
def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
flops_per_token = 2 * n_params # each parameter is used for a MAC (2 FLOPS) per network operation
# this assumes that all samples have a fixed length equal to the block size
# which is most likely false during finetuning
flops_per_seq = flops_per_token * max_seq_length
attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2))
return flops_per_seq + attn_flops_per_seq
def estimate_flops(model: GPT) -> int:
"""Measures estimated FLOPs for MFU.
Refs:
* https://ar5iv.labs.arxiv.org/html/2205.05198#A1
* https://ar5iv.labs.arxiv.org/html/2204.02311#A2
"""
# using all parameters for this is a naive over estimation because not all model parameters actually contribute to
# this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage
# (~10%) compared to the measured FLOPs, making those lower but more realistic.
# For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper.
n_trainable_params = num_parameters(model, requires_grad=True)
trainable_flops = flops_per_param(
model.max_seq_length, model.config.n_layer, model.config.n_embd, n_trainable_params
)
# forward + backward + gradients (assumes no gradient accumulation)
ops_per_step = 3 if model.training else 1
n_frozen_params = num_parameters(model, requires_grad=False)
frozen_flops = flops_per_param(model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params)
# forward + backward
frozen_ops_per_step = 2 if model.training else 1
return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops
def measure_flops(model: GPT, x: torch.Tensor) -> int:
"""Measures real FLOPs for HFU"""
flop_counter = FlopCounterMode(model, display=False)
ctx = nullcontext() if model.training else torch.no_grad()
with ctx, flop_counter:
y = model(x)
if model.training:
y.sum().backward()
return flop_counter.get_total_flops()