Leonid Pershin
2ccb03f7d2
- Introduced GPU probing functionality to gather and store GPU specifications in `/mnt/swarm_data/.gpu-rent-gpu.json`, aiding in performance tuning. - Updated `install_ollama.sh` and `install_llamacpp.sh` to utilize GPU information for configuring optimal runtime parameters. - Enhanced `provision.py` to include GPU probing and performance tuning logic, ensuring better resource allocation for LLM operations. - Improved documentation in `decisions.md`, `llm.md`, and `swarmui.md` to reflect changes in GPU handling and performance tuning processes. - Added new tests to validate the GPU probing and model resolution logic, ensuring robustness in handling various GPU configurations.
71 lines
1.9 KiB
Python
71 lines
1.9 KiB
Python
from gpu_rent.perf_tiers import (
|
|
GpuInfo,
|
|
compute_cap_at_least,
|
|
ollama_env_lines,
|
|
ollama_tune_for,
|
|
swarm_tune_for,
|
|
tier_for_vram_mib,
|
|
)
|
|
|
|
|
|
def test_tier_boundaries():
|
|
assert tier_for_vram_mib(8 * 1024) == "low"
|
|
assert tier_for_vram_mib(16 * 1024) == "mid"
|
|
assert tier_for_vram_mib(23 * 1024) == "mid"
|
|
assert tier_for_vram_mib(24 * 1024) == "high"
|
|
assert tier_for_vram_mib(40 * 1024) == "high"
|
|
assert tier_for_vram_mib(48 * 1024) == "ultra"
|
|
assert tier_for_vram_mib(80 * 1024) == "ultra"
|
|
|
|
|
|
def test_compute_cap():
|
|
assert compute_cap_at_least("8.0", 8, 0)
|
|
assert compute_cap_at_least("8.9", 8, 0)
|
|
assert not compute_cap_at_least("7.5", 8, 0)
|
|
|
|
|
|
def test_ollama_high_reserves_vram_for_swarm():
|
|
info = GpuInfo(
|
|
name="NVIDIA A100-SXM4-40GB",
|
|
vram_mib=40960,
|
|
compute_cap="8.0",
|
|
uuid="GPU-1",
|
|
tier="high",
|
|
)
|
|
tune = ollama_tune_for(info)
|
|
assert tune.flash_attention
|
|
assert tune.num_parallel == 1
|
|
assert tune.max_loaded_models == 1
|
|
assert tune.kv_cache_type == "q8_0"
|
|
assert tune.gpu_overhead_bytes == 14 * 1024**3
|
|
env = "\n".join(ollama_env_lines(tune))
|
|
assert "OLLAMA_FLASH_ATTENTION=1" in env
|
|
assert "OLLAMA_GPU_OVERHEAD=" in env
|
|
|
|
|
|
def test_swarm_sage_on_ampere_mid():
|
|
info = GpuInfo(
|
|
name="NVIDIA GeForce RTX 4090",
|
|
vram_mib=24576,
|
|
compute_cap="8.9",
|
|
uuid="GPU-2",
|
|
tier="high",
|
|
)
|
|
st = swarm_tune_for(info)
|
|
assert st.use_sage_attention
|
|
assert "--use-sage-attention" in st.comfy_extra_args
|
|
assert st.install_triton_sage
|
|
|
|
|
|
def test_swarm_no_sage_on_low():
|
|
info = GpuInfo(
|
|
name="NVIDIA T4",
|
|
vram_mib=15360,
|
|
compute_cap="7.5",
|
|
uuid="GPU-3",
|
|
tier="low",
|
|
)
|
|
st = swarm_tune_for(info)
|
|
assert not st.use_sage_attention
|
|
assert st.comfy_extra_args == ""
|