Leonid Pershin
a882964ce0
- Enhanced the `OllamaTune` class to include a new `context_length` attribute, improving the configuration for different GPU tiers. - Updated performance tuning logic to set appropriate context lengths for low, mid, high, and ultra tiers, ensuring optimal resource allocation. - Modified installation scripts to reflect the new context length settings, enhancing the installation process for Ollama. - Revised documentation to include context length details in the GPU performance table, providing clearer guidance for users. - Added tests to validate the correct context length settings in various scenarios, ensuring robustness in performance tuning.
86 lines
2.4 KiB
Python
86 lines
2.4 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
|
|
assert "OLLAMA_CONTEXT_LENGTH=16384" in env
|
|
assert tune.context_length == 16384
|
|
|
|
|
|
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 == ""
|
|
|
|
|
|
def test_ollama_mid_4090_context_16k():
|
|
info = GpuInfo(
|
|
name="NVIDIA GeForce RTX 4090",
|
|
vram_mib=24564,
|
|
compute_cap="8.9",
|
|
uuid="GPU-4",
|
|
tier="mid",
|
|
)
|
|
tune = ollama_tune_for(info)
|
|
assert tune.context_length == 16384
|
|
assert "OLLAMA_CONTEXT_LENGTH=16384" in "\n".join(ollama_env_lines(tune))
|