mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
192 lines
7.0 KiB
Python
192 lines
7.0 KiB
Python
"""Tests for MFU/MBU efficiency calculator."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from openjarvis.telemetry.efficiency import (
|
|
EfficiencyMetrics,
|
|
compute_efficiency,
|
|
estimate_model_bytes_per_token,
|
|
estimate_model_flops_per_token,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Constants: A100 SXM4 80GB specs
|
|
# ---------------------------------------------------------------------------
|
|
A100_PEAK_TFLOPS = 312.0 # FP16 Tensor Core
|
|
A100_PEAK_BW_GB_S = 2039.0 # HBM2e
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# estimate_model_flops_per_token
|
|
# ---------------------------------------------------------------------------
|
|
class TestEstimateModelFlopsPerToken:
|
|
def test_7b_dense(self) -> None:
|
|
flops = estimate_model_flops_per_token(7.0)
|
|
assert flops == pytest.approx(14e9)
|
|
|
|
def test_70b_dense(self) -> None:
|
|
flops = estimate_model_flops_per_token(70.0)
|
|
assert flops == pytest.approx(140e9)
|
|
|
|
def test_moe_mixtral(self) -> None:
|
|
# Mixtral 8x7B: ~47B total, ~13B active per token
|
|
flops = estimate_model_flops_per_token(47.0, active_params_b=13.0)
|
|
assert flops == pytest.approx(26e9)
|
|
|
|
def test_active_params_none_uses_total(self) -> None:
|
|
result = estimate_model_flops_per_token(7.0, None)
|
|
assert result == estimate_model_flops_per_token(7.0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# estimate_model_bytes_per_token
|
|
# ---------------------------------------------------------------------------
|
|
class TestEstimateModelBytesPerToken:
|
|
def test_7b_fp16(self) -> None:
|
|
bpt = estimate_model_bytes_per_token(7.0)
|
|
assert bpt == pytest.approx(14e9)
|
|
|
|
def test_7b_int8(self) -> None:
|
|
bpt = estimate_model_bytes_per_token(7.0, bytes_per_param=1.0)
|
|
assert bpt == pytest.approx(7e9)
|
|
|
|
def test_70b_fp16(self) -> None:
|
|
bpt = estimate_model_bytes_per_token(70.0)
|
|
assert bpt == pytest.approx(140e9)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# compute_efficiency
|
|
# ---------------------------------------------------------------------------
|
|
class TestComputeEfficiency:
|
|
def test_7b_100tps_single_gpu(self) -> None:
|
|
"""7B model, 100 tok/s, single A100."""
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
)
|
|
# actual_flops = 14e9 * 100 = 1.4e12
|
|
assert m.actual_flops == pytest.approx(1.4e12)
|
|
# peak_flops = 312e12
|
|
assert m.peak_flops == pytest.approx(312e12)
|
|
# MFU = 1.4e12 / 312e12 * 100 ≈ 0.4487%
|
|
expected_mfu = 1.4e12 / 312e12 * 100.0
|
|
assert m.mfu_pct == pytest.approx(expected_mfu)
|
|
|
|
# actual_bandwidth = 14e9 * 100 / 1e9 = 1400 GB/s
|
|
assert m.actual_bandwidth_gb_s == pytest.approx(1400.0)
|
|
# MBU = 1400 / 2039 * 100 ≈ 68.66%
|
|
expected_mbu = 1400.0 / 2039.0 * 100.0
|
|
assert m.mbu_pct == pytest.approx(expected_mbu)
|
|
|
|
def test_multi_gpu_scaling(self) -> None:
|
|
"""2 GPUs should double peak, halving utilization at same throughput."""
|
|
m1 = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
num_gpus=1,
|
|
)
|
|
m2 = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
num_gpus=2,
|
|
)
|
|
assert m2.mfu_pct == pytest.approx(m1.mfu_pct / 2.0)
|
|
assert m2.mbu_pct == pytest.approx(m1.mbu_pct / 2.0)
|
|
assert m2.peak_flops == pytest.approx(m1.peak_flops * 2.0)
|
|
assert m2.peak_bandwidth_gb_s == pytest.approx(m1.peak_bandwidth_gb_s * 2.0)
|
|
|
|
def test_ipj_with_energy(self) -> None:
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
energy_joules=50.0,
|
|
accuracy=0.8,
|
|
)
|
|
assert m.ipj == pytest.approx(0.8 / 50.0)
|
|
|
|
def test_ipj_zero_energy(self) -> None:
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
energy_joules=0.0,
|
|
accuracy=0.8,
|
|
)
|
|
assert m.ipj == 0.0
|
|
|
|
def test_zero_tokens_per_sec(self) -> None:
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=0.0,
|
|
)
|
|
assert m.mfu_pct == 0.0
|
|
assert m.mbu_pct == 0.0
|
|
assert m.actual_flops == 0.0
|
|
assert m.actual_bandwidth_gb_s == 0.0
|
|
|
|
def test_moe_efficiency(self) -> None:
|
|
"""MoE model: FLOPs use active params, bandwidth uses total params."""
|
|
m = compute_efficiency(
|
|
param_count_b=47.0,
|
|
active_params_b=13.0,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=50.0,
|
|
)
|
|
# actual_flops based on active params: 2*13e9*50 = 1.3e12
|
|
assert m.actual_flops == pytest.approx(1.3e12)
|
|
# bandwidth based on total params: 47e9*2*50/1e9 = 4700 GB/s
|
|
assert m.actual_bandwidth_gb_s == pytest.approx(4700.0)
|
|
|
|
def test_returns_dataclass(self) -> None:
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
)
|
|
assert isinstance(m, EfficiencyMetrics)
|
|
|
|
def test_custom_bytes_per_param(self) -> None:
|
|
"""INT4 quantization: 0.5 bytes per param."""
|
|
m = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
bytes_per_param=0.5,
|
|
)
|
|
# bandwidth = 7e9 * 0.5 * 100 / 1e9 = 350 GB/s
|
|
assert m.actual_bandwidth_gb_s == pytest.approx(350.0)
|
|
# MFU should be unchanged (FLOPs don't depend on bytes_per_param)
|
|
m_fp16 = compute_efficiency(
|
|
param_count_b=7.0,
|
|
active_params_b=None,
|
|
gpu_peak_tflops=A100_PEAK_TFLOPS,
|
|
gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S,
|
|
tokens_per_sec=100.0,
|
|
)
|
|
assert m.mfu_pct == pytest.approx(m_fp16.mfu_pct)
|