"""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)