Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/quantpreflight/planner.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

from .models import CheckpointMeta, Finding, ScanReport, SetupPlan, TargetSpec
from .profiles import resolve_gpu_profile
from .rules import _effective_text_config, _pos_int, analyze_checkpoint, estimate_per_gpu_vram_gb
from .rules import _effective_text_config, analyze_checkpoint, context_limit, estimate_per_gpu_vram_gb

TP_CHOICES = (1, 2, 4, 8)
# They say a model is risky to load, not that it won't start
Expand Down Expand Up @@ -147,7 +147,7 @@ def plan_setup(
) -> SetupPlan:
"""Find the closest setup to `target` that clears every finding a setting can clear."""
tcfg = _effective_text_config(meta.config)
limit = _pos_int(tcfg.get("max_position_embeddings"))
limit = context_limit(meta.config, meta.tokenizer_config)
gpu = resolve_gpu_profile(target.gpu) if target.gpu else None
vram = target.vram_gb_override if target.vram_gb_override is not None else (gpu.vram_gb if gpu else None)
usable = vram * target.gpu_memory_utilization if vram else None
Expand Down
42 changes: 39 additions & 3 deletions src/quantpreflight/rules.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,42 @@ def _effective_text_config(config: dict[str, Any]) -> dict[str, Any]:
return config


_CONTEXT_LENGTH_KEYS = ("max_position_embeddings", "n_positions", "max_seq_len", "seq_length")
_ROPE_TYPES_WITHOUT_FACTOR = frozenset(
{"su", "longrope", "llama3", "yarn", "deepseek_yarn", "deepseek_llama_scaling"}
)


def context_limit(config: dict[str, Any], tokenizer_config: dict[str, Any] | None = None) -> int | None:
"""Derive the serving context limit using vLLM-compatible config precedence."""
text_config = _effective_text_config(config)
limits: list[int] = []
for key in _CONTEXT_LENGTH_KEYS:
value = _pos_int(text_config.get(key))
if value is not None:
limits.append(value)
model_max = _pos_int(text_config.get("model_max_length"))
if model_max is not None:
limits.append(model_max)
if not limits:
return None
limit = min(limits)
rope = text_config.get("rope_scaling")
model_type = str(text_config.get("model_type", "")).lower()
architectures = text_config.get("architectures")
architecture_names = architectures if isinstance(architectures, list) else ()
is_gemma3 = model_type == "gemma3" or any("gemma3" in str(name).lower() for name in architecture_names)
if isinstance(rope, dict) and not is_gemma3:
rope_type = str(rope.get("rope_type", rope.get("type", "default"))).lower()
factor = rope.get("factor")
if rope_type not in _ROPE_TYPES_WITHOUT_FACTOR and isinstance(factor, (int, float)) and factor > 0:
limit = int(limit * factor)
tokenizer_max = _pos_int((tokenizer_config or {}).get("model_max_length"))
if tokenizer_max is not None:
limit = min(limit, tokenizer_max)
return limit


def _to_int_set(val: Any) -> set[int]:
if isinstance(val, int) and not isinstance(val, bool):
return {val}
Expand Down Expand Up @@ -1347,8 +1383,8 @@ def _version_in_range(version: tuple[int, ...], item: dict[str, Any]) -> bool:
return True


def _check_qpf016_context_limit(tcfg: dict[str, Any], target: TargetSpec) -> list[Finding]:
limit = _pos_int(tcfg.get("max_position_embeddings"))
def _check_qpf016_context_limit(meta: CheckpointMeta, target: TargetSpec) -> list[Finding]:
limit = context_limit(meta.config, meta.tokenizer_config)
if target.engine != "vllm" or not target.context_len or not limit or target.context_len <= limit:
return []
return [
Expand Down Expand Up @@ -1460,7 +1496,7 @@ def analyze_checkpoint(
findings.extend(_check_qpf013_safetensors_spans(meta))
findings.extend(_check_qpf014_legacy_pickle_weights(meta))
findings.extend(_check_qpf015_engine_advisories(tgt, engine_prof))
findings.extend(_check_qpf016_context_limit(tcfg, tgt))
findings.extend(_check_qpf016_context_limit(meta, tgt))

est_vram, vram_findings = _check_qpf012_vram_budget(meta, tcfg, tgt, gpu_prof)
findings.extend(vram_findings)
Expand Down
7 changes: 7 additions & 0 deletions tests/test_planner.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,13 @@ def test_a_model_that_fits_is_ready_as_asked() -> None:
assert plan.command == "vllm serve org/model"


def test_setup_uses_tokenizer_context_cap_and_rope_scaling() -> None:
config = {**LLAMA, "max_position_embeddings": 4096, "rope_scaling": {"type": "linear", "factor": 4}}
meta = _meta(config=config, tokenizer_config={"model_max_length": 12000})
plan = plan_setup(meta, TargetSpec(gpu="H100", context_len=20000))
assert plan.context_len == 12000


def test_heads_that_do_not_split_move_to_the_closest_gpu_count_with_fewer_gpus_on_a_tie() -> None:
plan = plan_setup(_meta(10), TargetSpec(gpu="H100", tp=3))
assert plan.status == "changed" and plan.tp == 2
Expand Down
16 changes: 15 additions & 1 deletion tests/test_rule_boundaries.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from quantpreflight.formatters import _fmt_bytes, format_json_report, format_sarif_report, format_text_report
from quantpreflight.models import CheckpointMeta, QuantSpec, Severity, TargetSpec, TensorSpec, TokenizerMeta
from quantpreflight.profiles import load_advisories
from quantpreflight.rules import analyze_checkpoint, estimate_per_gpu_vram_gb
from quantpreflight.rules import analyze_checkpoint, context_limit, estimate_per_gpu_vram_gb

LLAMA = {"architectures": ["LlamaForCausalLM"], "model_type": "llama"}
AWQ4 = QuantSpec(method="awq", bits=4.0)
Expand Down Expand Up @@ -752,6 +752,20 @@ def test_qpf016_names_a_setting_that_works() -> None:
assert report.findings[0].remediation == "Set --context-len to 2048 or lower."


@pytest.mark.parametrize(
("config", "tokenizer", "expected"),
[
({"max_position_embeddings": 8192, "n_positions": 4096}, {}, 4096),
({"max_position_embeddings": 8192, "model_max_length": 2048}, {}, 2048),
({"max_position_embeddings": 8192}, {"model_max_length": 4096}, 4096),
({"max_position_embeddings": 4096, "rope_scaling": {"type": "linear", "factor": 4}}, {}, 16384),
({"max_position_embeddings": 4096, "rope_scaling": {"type": "yarn", "factor": 4}}, {}, 4096),
],
)
def test_context_limit_matches_vllm_precedence(config: dict[str, Any], tokenizer: dict[str, Any], expected: int) -> None:
assert context_limit(config, tokenizer) == expected


# Corrections from scanning 500 public models


Expand Down
Loading