mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 05:12:26 +00:00
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
"""Tests for FLOPs estimation and MFU computation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from openjarvis.telemetry.flops import (
|
|
GPU_PEAK_TFLOPS_BF16,
|
|
MODEL_PARAMS_B,
|
|
compute_mfu,
|
|
estimate_flops,
|
|
estimate_flops_no_kv_cache,
|
|
)
|
|
|
|
|
|
class TestEstimateFlops:
|
|
def test_known_model(self):
|
|
total, per_tok = estimate_flops("qwen3:8b", 100, 50)
|
|
# 2 * 8e9 * 150 = 2.4e12
|
|
assert total == pytest.approx(2.4e12)
|
|
# 2 * 8e9 = 16e9
|
|
assert per_tok == pytest.approx(16e9)
|
|
|
|
def test_unknown_model_zero(self):
|
|
total, per_tok = estimate_flops("totally-unknown-model", 100, 50)
|
|
assert total == 0.0
|
|
assert per_tok == 0.0
|
|
|
|
def test_prefix_matching(self):
|
|
# "llama-3.1-8b-instruct" should match prefix "llama-3.1-8b"
|
|
total, per_tok = estimate_flops("llama-3.1-8b-instruct", 10, 10)
|
|
assert total > 0
|
|
# 2 * 8e9 * 20 = 3.2e11
|
|
assert total == pytest.approx(3.2e11)
|
|
|
|
def test_zero_tokens(self):
|
|
total, per_tok = estimate_flops("qwen3:8b", 0, 0)
|
|
assert total == 0.0
|
|
assert per_tok == 0.0
|
|
|
|
def test_flops_proportional_to_tokens(self):
|
|
total_100, _ = estimate_flops("qwen3:8b", 50, 50)
|
|
total_200, _ = estimate_flops("qwen3:8b", 100, 100)
|
|
assert total_200 == pytest.approx(total_100 * 2.0)
|
|
|
|
|
|
class TestEstimateFlopsNoKvCache:
|
|
def test_known_model(self):
|
|
total, per_tok = estimate_flops_no_kv_cache("qwen3:8b", 100, 50)
|
|
# P * N * (N+1) = 8e9 * 150 * 151 = 1.812e14
|
|
assert total == pytest.approx(8e9 * 150 * 151)
|
|
assert per_tok == pytest.approx(total / 150)
|
|
|
|
def test_zero_tokens(self):
|
|
total, per_tok = estimate_flops_no_kv_cache("qwen3:8b", 0, 0)
|
|
assert total == 0.0
|
|
assert per_tok == 0.0
|
|
|
|
def test_unknown_model_zero(self):
|
|
total, per_tok = estimate_flops_no_kv_cache("unknown-model", 100, 50)
|
|
assert total == 0.0
|
|
assert per_tok == 0.0
|
|
|
|
|
|
class TestComputeMfu:
|
|
def test_known_gpu(self):
|
|
# 100 TFLOPS actual for 1s on H100 → 100 / 989 * 100 ≈ 10.1%
|
|
flops = 100e12
|
|
mfu = compute_mfu(flops, 1.0, "H100")
|
|
assert mfu == pytest.approx(100.0 / 989.0 * 100.0, rel=1e-3)
|
|
|
|
def test_unknown_gpu_zero(self):
|
|
mfu = compute_mfu(100e12, 1.0, "QuantumGPU")
|
|
assert mfu == 0.0
|
|
|
|
def test_zero_duration(self):
|
|
mfu = compute_mfu(100e12, 0.0, "H100")
|
|
assert mfu == 0.0
|
|
|
|
def test_multi_gpu(self):
|
|
flops = 100e12
|
|
mfu_single = compute_mfu(flops, 1.0, "H100", num_gpus=1)
|
|
mfu_dual = compute_mfu(flops, 1.0, "H100", num_gpus=2)
|
|
assert mfu_dual == pytest.approx(mfu_single / 2.0)
|
|
|
|
def test_substring_matching(self):
|
|
# "NVIDIA H100 80GB" should match "H100"
|
|
mfu = compute_mfu(100e12, 1.0, "NVIDIA H100 80GB")
|
|
assert mfu > 0.0
|
|
|
|
def test_tables_nonempty(self):
|
|
assert len(GPU_PEAK_TFLOPS_BF16) > 0
|
|
assert len(MODEL_PARAMS_B) > 0
|