wehub-resource-sync
138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import time
|
|
import torch
|
|
from transformers import TrainerControl, TrainerState
|
|
from typing import TYPE_CHECKING
|
|
|
|
from swift.utils import (empty_cache, get_current_device, get_device_count, get_dist_setting, get_env_args, get_logger,
|
|
synchronize)
|
|
from .base import TrainerCallback
|
|
|
|
if TYPE_CHECKING:
|
|
from swift.trainers import Trainer, TrainingArguments
|
|
|
|
logger = get_logger()
|
|
|
|
device_flops_map = {
|
|
'GB200': 2.5e15,
|
|
'B200': 2.25e15,
|
|
'MI300X': 1336e12,
|
|
'H100': 312e12,
|
|
'H800': 312e12,
|
|
'H200': 989e12,
|
|
'A100': 312e12,
|
|
'A800': 312e12,
|
|
'L40S': 362.05e12,
|
|
'L40': 181.05e12,
|
|
'A40': 149.7e12,
|
|
'L20': 119.5e12,
|
|
'H20': 148e12,
|
|
'910B': 354e12,
|
|
'Ascend910': 354e12,
|
|
'RTX 3070 Ti': 21.75e12
|
|
}
|
|
|
|
|
|
class PerfMetricsLogCallback(TrainerCallback):
|
|
"""An callback for perf metrics (MFU etc) log implementation"""
|
|
|
|
def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
|
|
super().__init__(args, trainer)
|
|
self.max_tflops = None
|
|
self.elapsed = 0.0
|
|
self.step_start_time = None
|
|
|
|
def on_init_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
|
|
|
|
# Top priority. Specify by ENV
|
|
tflops = get_env_args('DEVICE_TFLOPS', float, None)
|
|
# `state.total_flos` is summed across all ranks (cluster-global) by HF Trainer,
|
|
# so the theoretical denominator must use the TOTAL number of GPUs in use across the entire cluster.
|
|
_, _, world_size, local_world_size = get_dist_setting()
|
|
local_n_gpu = get_device_count()
|
|
gpus_per_process = max(1, local_n_gpu // max(local_world_size, 1))
|
|
device_count = max(world_size * gpus_per_process, 1)
|
|
if tflops is not None:
|
|
logger.info(f"Specify theoretical max TFLOPS through ENV 'DEVICE_TFLOPS'. [{tflops} TFLOPS]")
|
|
else:
|
|
# Run a estimating test.
|
|
dtype = kwargs.get('model').dtype
|
|
device = torch.device(get_current_device())
|
|
logger.info(f'Estimating device TFLOPS baseline. Device: [{device}] dtype: [{dtype}]')
|
|
tflops = self._estimate_device_tflops_by_dtype(device, dtype)
|
|
logger.info(f'Estimate test finished. [{tflops} TFLOPS] Device count: [{device_count}]')
|
|
|
|
self.max_tflops = tflops * device_count
|
|
|
|
def on_step_begin(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
|
|
self.step_start_time = time.time()
|
|
|
|
def on_step_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
|
|
self.elapsed += time.time() - self.step_start_time
|
|
|
|
def on_log(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, logs=None, **kwargs):
|
|
total_flos = getattr(state, 'total_flos', 0)
|
|
actual_flops = total_flos / self.elapsed
|
|
theoretical_max_flops = self.max_tflops * 1e12
|
|
mfu = actual_flops / theoretical_max_flops
|
|
logger.debug(f'Total_flos[{total_flos}] elapsed_time[{self.elapsed}]sec Average MFU[{mfu}]')
|
|
logs['MFU'] = round(mfu, 6)
|
|
|
|
@staticmethod
|
|
def _estimate_device_tflops_by_dtype(device: torch.device, dtype: torch.dtype, repeats: int = 60, dim: int = 8192):
|
|
# Set matrix dimension
|
|
shape = (dim, dim)
|
|
backend = device.type
|
|
if backend == 'npu':
|
|
import torch_npu
|
|
|
|
# Initialize matrices
|
|
a = torch.randn(*shape, device=device, dtype=dtype)
|
|
b = torch.randn(*shape, device=device, dtype=dtype)
|
|
|
|
# Warm-up
|
|
for _ in range(5):
|
|
c = torch.matmul(a, b)
|
|
synchronize(device)
|
|
|
|
# Run benchmark test
|
|
start = time.time()
|
|
for _ in range(repeats):
|
|
c = torch.matmul(a, b)
|
|
synchronize(device)
|
|
end = time.time()
|
|
total_time = end - start
|
|
avg_time = total_time / repeats
|
|
|
|
# Adjust repeat count and retest if test duration is too short
|
|
if total_time < 3:
|
|
repeats = int(6 / avg_time)
|
|
start = time.time()
|
|
for _ in range(repeats):
|
|
c = torch.matmul(a, b)
|
|
synchronize(device)
|
|
end = time.time()
|
|
total_time = end - start
|
|
avg_time = total_time / repeats
|
|
|
|
del a, b, c
|
|
empty_cache()
|
|
|
|
tflops = (2 * dim**3 / avg_time) / 1e12
|
|
logger.info(f'[Device {device}] Total time: {total_time:.4f}s, dtype: {dtype}, Perf: {tflops:.4f} TFLOPS')
|
|
return tflops
|
|
|
|
@staticmethod
|
|
def _retrieve_flops_from_map(device):
|
|
"""Retrieve theoretical FLOPS from Map. """
|
|
# This function is never used.
|
|
|
|
device_name = device.get_device_name()
|
|
flops = None
|
|
for name, value in device_flops_map.items():
|
|
if name in device_name:
|
|
flops = value
|
|
break
|
|
|
|
return flops
|