From b4e817f65e0f27ce22f07f8ccc74f2399a81641e Mon Sep 17 00:00:00 2001
From: aivrar <42790594+aivrar@users.noreply.github.com>
Date: Thu, 20 Aug 2026 19:54:11 -0700
Subject: [PATCH 1/3] Harden TriAttention model calibration preflight
---
README.md | 29 +-
docs/context-extension.md | 8 +
docs/manual.md | 20 +-
docs/ui-workspace.md | 25 +-
multi_turboquant/calibration/__init__.py | 8 +
.../calibration/godzilla_triattention.py | 263 +++++++++++++++---
.../calibration/triattention_runner.py | 144 ++++++++++
.../integration/godzilla_workspace.py | 114 +++++++-
multi_turboquant/ui/runtime.py | 4 +
run_ui.py | 39 ++-
tests/test_godzilla_triattention.py | 104 ++++++-
tests/test_godzilla_workspace.py | 33 +++
tests/test_run_ui.py | 25 ++
tests/test_triattention_runner.py | 66 +++++
14 files changed, 810 insertions(+), 72 deletions(-)
create mode 100644 multi_turboquant/calibration/triattention_runner.py
create mode 100644 tests/test_triattention_runner.py
diff --git a/README.md b/README.md
index 41cc9b3..9e6be19 100644
--- a/README.md
+++ b/README.md
@@ -220,7 +220,7 @@ mtq-godzilla-triattention calibrate \
--input calibration.txt \
--output model.triattention \
--max-length 2048 \
- --device cuda \
+ --device cuda:1 \
--attn-implementation sdpa
```
@@ -285,6 +285,15 @@ not offered because the current Godzilla binary does not expose a real calibrati
command. Calibration still needs the exact Hugging Face model or a compatible
source; a GGUF alone does not contain the pre-RoPE query statistics. The official
script uses `trust_remote_code=True`, so only calibrate model sources you trust.
+`IQ4_XS`, `Q4_K_M`, and similar labels describe the selected GGUF's weight
+encoding; they do not make the Transformers calibration load quantized. Before
+weights are downloaded, the workflow now reads the matching model config,
+prefers authoritative nested `rope_parameters.rope_theta` over conflicting
+legacy defaults, and rejects a requested sequence longer than the model's
+declared context. `--device` accepts `cuda:N`, and the UI lists each GPU and
+defaults to the one with the most free VRAM. A model-config Transformers version
+different from the managed runtime is reported for qualification rather than
+silently treated as compatible.
The older Godzilla checkout-owned PowerShell workflow remains available as an
explicit fallback for checkouts that provide it. `mtq-triattention-stats` writes
a different `.pt` schema for Multi-TurboQuant's Python/vLLM path and cannot be
@@ -315,14 +324,16 @@ mtq-godzilla-triattention domvox \
The conversion is deliberately opt-in and lossy: Godzilla v1 has no fields
for domvox layer-budget scales or attention scale, so those fields are reported
-as dropped. Calibration lengths from 128 through 200,000 tokens are accepted;
-anything above 32,768 requires `--allow-long-calibration` and is processed as
-one upstream sequence, with substantially higher memory and runtime risk. The
-default remains conservative. The local UI accepts only one calibration job at
-a time and reports current CUDA free/total VRAM, but neither measure reduces or
-predicts the memory required by the one long sequence. System RAM and GPU VRAM
-are shown separately; their optional combined figure is capacity inventory,
-not interchangeable calibration memory. The UI can also create deterministic
+as dropped. The upstream guide targets enough coherent text to approach its
+32,768-token default; it does not establish 200,000 tokens as an optimal
+calibration. Multi-TurboQuant retains 200,000 only as a global input ceiling.
+The effective limit is the smaller of that ceiling and the matching model's
+declared context, and anything above 32,768 still requires
+`--allow-long-calibration`. The official path also estimates its retained Q
+tensors, BF16 weights, and transient state and fails before model download when
+that conservative floor exceeds the selected GPU's free VRAM. The estimate is
+a lower bound, not a promise that a run will fit. System RAM and GPU VRAM remain
+separate capacity domains. The UI can also create deterministic
offline starter text inside the saved model root without overwriting unrelated
files. Corpus files carry a schema and completion marker, and simultaneous
requests cannot clobber one another or reuse a partial file. Use representative
diff --git a/docs/context-extension.md b/docs/context-extension.md
index 94a15a6..6a3085e 100644
--- a/docs/context-extension.md
+++ b/docs/context-extension.md
@@ -126,6 +126,14 @@ older checkout-owned PowerShell workflow remains an explicit fallback for
compatible checkouts. Multi-TurboQuant never synthesizes unverified model
statistics and reuses existing artifacts only after strict validation.
+Calibration performs a Transformers BF16 model load; a small `IQ4_XS` GGUF does
+not reduce that memory. The preflight resolves the matching Hugging Face config,
+prefers nested `rope_parameters.rope_theta` over conflicting legacy defaults,
+bounds the sequence by `max_position_embeddings`, estimates the official
+one-shot memory floor, and supports explicit `cuda:N` selection. The upstream
+32,768-token default remains the evidence-backed target; 200,000 is only the
+application's global ceiling for models and hardware that pass those checks.
+
The UI selects calibration Python only after a bounded isolated import probe,
does not combine packages from different environments, and rechecks the exact
interpreter immediately before launch. domvox calibration and conversion stay
diff --git a/docs/manual.md b/docs/manual.md
index e108ead..08d2405 100644
--- a/docs/manual.md
+++ b/docs/manual.md
@@ -826,14 +826,18 @@ script's sibling `triattention_common.py`, header, and exact file length, and
requires an explicit acknowledgement because
Godzilla v1 cannot store domvox layer-budget scales or attention scale. Those
fields are reported as dropped; this is not a lossless format conversion.
-Calibration lengths from 128 through 200,000 are accepted. Above 32,768 the
-operator must enable long calibration explicitly; the upstream script processes
-one long sequence, so the UI warns about memory and runtime instead of assuming
-chunked aggregation. The UI runs at most one calibration job at a time and
-reports the selected CUDA device's current free/total VRAM after dependency
-preflight. That snapshot cannot predict the long sequence's peak usage, and
-system RAM is not a substitute for discrete VRAM. A GGUF alone remains
-insufficient for calibration.
+The 200,000-token value is a global input ceiling, not a recommendation or a
+guarantee. Upstream documents a 32,768-token default. The actual limit is also
+bounded by the matching Hugging Face model's declared context, and the planner
+rejects an oversized request before downloading weights. Above 32,768 the
+operator must still enable long calibration explicitly; the upstream script
+processes one sequence rather than treating the value as total tokens across a
+chunked corpus. The official path reports a conservative memory floor for its
+retained Q tensors, BF16 weights, and transient state and blocks a request that
+cannot fit the selected GPU's currently free VRAM. Select GPUs explicitly as
+`cuda:N`; the UI initially chooses the device with the most free VRAM. System
+RAM is not a substitute for discrete VRAM, and a GGUF alone remains insufficient
+for calibration. GGUF labels such as `IQ4_XS` describe inference weights only.
Selecting Gigatoken runs domvox through the same fail-closed parity wrapper as
the official script and forwards only domvox-supported arguments afterward.
The domvox forward pass and Godzilla conversion both run inside the exact
diff --git a/docs/ui-workspace.md b/docs/ui-workspace.md
index af5ccef..010b45e 100644
--- a/docs/ui-workspace.md
+++ b/docs/ui-workspace.md
@@ -97,13 +97,24 @@ When Gigatoken is selected, domvox is launched through the same fail-closed
parity wrapper as the official calibrator; its script receives only supported
arguments after parity succeeds.
-Calibration lengths from 128 through 200,000 tokens are supported. Values above
-32,768 require the explicit **Allow long calibration** checkbox and produce a
-one-shot memory/runtime warning; the UI does not silently chunk or aggregate
-the upstream calibrator's sequence. It permits one calibration job at a time so
-two model loads cannot overlap. A successful dependency preflight reports the
-selected CUDA device's free and total VRAM, but that snapshot is not a memory
-estimate or guarantee for a 200,000-token run.
+`IQ4_XS` and similar names are GGUF weight encodings, not calibration model IDs.
+Both Python calibrators load the matching Hugging Face checkpoint rather than
+the selected quantized GGUF. The plan now validates that checkpoint's config
+before weight download, uses nested `rope_parameters.rope_theta` when it
+conflicts with a legacy fallback, and reports the model context and RoPE source.
+It also reports when the checkpoint's recorded Transformers version differs from
+the managed calibration runtime.
+
+The global input ceiling remains 200,000 tokens, but the effective one-shot
+limit is also bounded by the matching model's `max_position_embeddings`. Values
+above 32,768 require **Allow long calibration**; the upstream guide recommends
+approaching its 32,768-token default and does not establish 200,000 as best.
+The UI does not silently reinterpret 200k total corpus tokens as one or more
+safe sequences. For the official path it estimates retained Q tensors, BF16
+weights, and transient state and blocks a run when that conservative floor
+already exceeds free VRAM. Individual `cuda:N` devices are selectable, with the
+largest currently free GPU selected initially. The estimate remains a lower
+bound and one-job concurrency does not reduce a single sequence's peak use.
If the automatically selected managed interpreter is missing `accelerate` or
another declared dependency, the plan offers **Repair TriAttention
diff --git a/multi_turboquant/calibration/__init__.py b/multi_turboquant/calibration/__init__.py
index a154cd1..62b1c96 100644
--- a/multi_turboquant/calibration/__init__.py
+++ b/multi_turboquant/calibration/__init__.py
@@ -19,7 +19,11 @@
inspect_domvox_triattention_file,
inspect_official_triattention_calibrator,
inspect_official_triattention_checkout,
+ load_huggingface_model_metadata,
+ normalize_calibration_device,
select_compatible_calibration_python,
+ estimate_official_calibration_bytes,
+ validate_model_calibration_length,
)
__all__ = [
@@ -41,5 +45,9 @@
"inspect_domvox_triattention_file",
"inspect_official_triattention_calibrator",
"inspect_official_triattention_checkout",
+ "load_huggingface_model_metadata",
+ "normalize_calibration_device",
"select_compatible_calibration_python",
+ "estimate_official_calibration_bytes",
+ "validate_model_calibration_length",
]
diff --git a/multi_turboquant/calibration/godzilla_triattention.py b/multi_turboquant/calibration/godzilla_triattention.py
index ad46982..3c1b6f1 100644
--- a/multi_turboquant/calibration/godzilla_triattention.py
+++ b/multi_turboquant/calibration/godzilla_triattention.py
@@ -69,6 +69,122 @@
"TRIA",
"triattention_common",
)
+_CUDA_DEVICE = re.compile(r"cuda(?::(\d+))?\Z")
+
+
+def normalize_calibration_device(value: str) -> str:
+ """Return a supported CPU or explicitly indexed CUDA device."""
+ normalized = value.strip().lower()
+ if normalized == "cpu":
+ return normalized
+ match = _CUDA_DEVICE.fullmatch(normalized)
+ if match is None:
+ raise ValueError("Calibration device must be 'cpu', 'cuda', or 'cuda:N'")
+ return "cuda" if match.group(1) is None else f"cuda:{int(match.group(1))}"
+
+
+def calibration_device_index(value: str) -> int | None:
+ normalized = normalize_calibration_device(value)
+ if normalized == "cpu":
+ return None
+ return int(normalized.split(":", 1)[1]) if ":" in normalized else 0
+
+
+def _config_value(config: object, name: str) -> object | None:
+ value = getattr(config, name, None)
+ if value is not None:
+ return value
+ if isinstance(config, Mapping):
+ return config.get(name)
+ return None
+
+
+def _nested_rope_theta(config: object) -> float | None:
+ parameters = _config_value(config, "rope_parameters")
+ if parameters is None:
+ parameters = _config_value(config, "rope_scaling")
+ value = _config_value(parameters, "rope_theta") if parameters is not None else None
+ return float(value) if value is not None else None
+
+
+def estimate_official_calibration_bytes(
+ metadata: Mapping[str, object], max_length: int
+) -> dict[str, int | None]:
+ """Estimate a conservative memory floor for the upstream one-shot calibrator."""
+ layers = _require_positive_int(metadata.get("num_layers"), "metadata.num_layers")
+ heads = _require_positive_int(
+ metadata.get("num_attention_heads"), "metadata.num_attention_heads"
+ )
+ head_dim = _require_positive_int(metadata.get("head_dim"), "metadata.head_dim")
+ hidden_size = metadata.get("hidden_size")
+ captured_q = layers * heads * max_length * head_dim * 2
+ transient = (
+ max_length * _require_positive_int(hidden_size, "metadata.hidden_size") * 2 * 4
+ if hidden_size is not None
+ else None
+ )
+ weights = metadata.get("estimated_bf16_weight_bytes")
+ floor = captured_q
+ if isinstance(weights, int):
+ floor += weights
+ if isinstance(transient, int):
+ floor += transient
+ return {
+ "captured_q_bytes": captured_q,
+ "estimated_bf16_weight_bytes": weights if isinstance(weights, int) else None,
+ "estimated_transient_bytes": transient,
+ "estimated_floor_bytes": floor,
+ }
+
+
+def validate_official_calibration_memory(
+ metadata: Mapping[str, object], max_length: int, device: str
+) -> dict[str, int | str | None]:
+ """Reject a CUDA run when its known memory floor cannot fit."""
+ normalized = normalize_calibration_device(device)
+ estimate = estimate_official_calibration_bytes(metadata, max_length)
+ report: dict[str, int | str | None] = {**estimate, "device": normalized}
+ if normalized == "cpu":
+ return report
+ try:
+ import torch
+ except ImportError as exc:
+ raise RuntimeError("torch is required to inspect CUDA calibration capacity") from exc
+ index = calibration_device_index(normalized)
+ assert index is not None
+ count = torch.cuda.device_count()
+ if index >= count:
+ raise ValueError(f"CUDA device {index} was requested, but only {count} device(s) are visible")
+ free_bytes, total_bytes = torch.cuda.mem_get_info(index)
+ report.update(
+ {
+ "device_index": index,
+ "device_name": torch.cuda.get_device_name(index),
+ "free_bytes": int(free_bytes),
+ "total_bytes": int(total_bytes),
+ }
+ )
+ floor = int(estimate["estimated_floor_bytes"] or 0)
+ if floor > int(free_bytes):
+ gib = 1024**3
+ raise ValueError(
+ "Official one-shot TriAttention calibration cannot fit on "
+ f"cuda:{index} ({torch.cuda.get_device_name(index)}): the conservative memory "
+ f"floor is {floor / gib:.1f} GiB but only {int(free_bytes) / gib:.1f} GiB is free. "
+ "Choose another CUDA device or reduce the one-shot calibration length."
+ )
+ return report
+
+
+def validate_model_calibration_length(
+ max_length: int, metadata: Mapping[str, object]
+) -> None:
+ context = metadata.get("max_position_embeddings")
+ if isinstance(context, int) and max_length > context:
+ raise ValueError(
+ f"Requested calibration length {max_length} exceeds the model's declared "
+ f"context limit of {context} tokens"
+ )
def inspect_calibration_python(
@@ -82,10 +198,14 @@ def inspect_calibration_python(
) -> dict[str, object]:
"""Verify one calibration interpreter without importing it into this process."""
interpreter = lexical_absolute_path(python)
- normalized_device = device.strip().lower()
+ issues: list[str] = []
+ try:
+ normalized_device = normalize_calibration_device(device)
+ except ValueError as exc:
+ normalized_device = device.strip().lower()
+ issues.append(str(exc))
normalized_tokenizer = tokenizer_backend.strip().lower()
normalized_attention = attention_implementation.strip().lower()
- issues: list[str] = []
required_modules = list(_CALIBRATION_REQUIRED_MODULES)
if normalized_attention == "flash_attention_2":
required_modules.append("flash_attn")
@@ -103,8 +223,7 @@ def unavailable() -> dict[str, object]:
"issues": issues,
}
- if normalized_device not in {"cuda", "cpu"}:
- issues.append("Calibration device must be 'cuda' or 'cpu'")
+ if issues:
return unavailable()
if normalized_tokenizer not in {"transformers", "gigatoken"}:
issues.append("Tokenizer backend must be 'transformers' or 'gigatoken'")
@@ -112,7 +231,7 @@ def unavailable() -> dict[str, object]:
if normalized_attention not in {"eager", "sdpa", "flash_attention_2"}:
issues.append("Attention implementation must be 'eager', 'sdpa', or 'flash_attention_2'")
return unavailable()
- if normalized_attention == "flash_attention_2" and normalized_device != "cuda":
+ if normalized_attention == "flash_attention_2" and not normalized_device.startswith("cuda"):
issues.append("flash_attention_2 calibration requires the CUDA device")
return unavailable()
if not interpreter.is_file():
@@ -147,6 +266,9 @@ def unavailable() -> dict[str, object]:
" report['cuda_available'] = torch.cuda.is_available()\n"
" if report['cuda_available']:\n"
" try:\n"
+ f" requested_index = {calibration_device_index(normalized_device)!r}\n"
+ " if requested_index is not None:\n"
+ " torch.cuda.set_device(requested_index)\n"
" device_index = torch.cuda.current_device()\n"
" free_bytes, total_bytes = torch.cuda.mem_get_info(device_index)\n"
" report.update({'cuda_device': torch.cuda.get_device_name(device_index), "
@@ -207,13 +329,25 @@ def unavailable() -> dict[str, object]:
)
if (
isinstance(report, dict)
- and normalized_device == "cuda"
+ and normalized_device.startswith("cuda")
and not any(issue.startswith("torch import failed:") for issue in issues)
):
if not report.get("torch_cuda"):
issues.append("Calibration Python has a CPU-only PyTorch build")
elif report.get("cuda_available") is not True:
issues.append("Calibration Python cannot access CUDA")
+ else:
+ expected_index = calibration_device_index(normalized_device)
+ if report.get("cuda_memory_error"):
+ issues.append(
+ "Calibration Python could not select the requested CUDA device: "
+ f"{report['cuda_memory_error']}"
+ )
+ elif report.get("cuda_device_index") != expected_index:
+ issues.append(
+ f"Calibration Python selected CUDA device {report.get('cuda_device_index')}, "
+ f"expected {expected_index}"
+ )
return {
"python": str(interpreter),
"python_resolved": str(interpreter.resolve()),
@@ -955,13 +1089,16 @@ def inspect_godzilla_triattention_file(path: str | Path) -> dict[str, object]:
}
-def load_huggingface_model_metadata(model: str) -> dict[str, object]:
+def load_huggingface_model_metadata(
+ model: str, *, trust_remote_code: bool = True
+) -> dict[str, object]:
"""Load only the matching Hugging Face config needed by the converter."""
try:
+ import transformers
from transformers import AutoConfig
except ImportError as exc:
raise RuntimeError("transformers is required to load model calibration metadata") from exc
- config = AutoConfig.from_pretrained(model, trust_remote_code=True)
+ config = AutoConfig.from_pretrained(model, trust_remote_code=trust_remote_code)
text_config = getattr(config, "text_config", config)
num_layers = _require_positive_int(
getattr(text_config, "num_hidden_layers", None), "config.num_hidden_layers"
@@ -981,15 +1118,60 @@ def load_huggingface_model_metadata(model: str) -> dict[str, object]:
if hidden_size % num_attention_heads:
raise ValueError("config.hidden_size is not divisible by num_attention_heads")
head_dim = hidden_size // num_attention_heads
- rope_theta = getattr(text_config, "rope_theta", None)
+ head_dim = _require_positive_int(head_dim, "config.head_dim")
+ legacy_rope_theta = getattr(text_config, "rope_theta", None)
+ nested_rope_theta = _nested_rope_theta(text_config)
+ rope_theta = nested_rope_theta if nested_rope_theta is not None else legacy_rope_theta
if rope_theta is None:
raise ValueError("The Hugging Face config does not declare rope_theta")
+ max_positions = getattr(text_config, "max_position_embeddings", None)
+ if max_positions is not None:
+ max_positions = _require_positive_int(
+ max_positions, "config.max_position_embeddings"
+ )
+ hidden_size = _require_positive_int(
+ getattr(text_config, "hidden_size", None), "config.hidden_size"
+ )
+ intermediate_size = getattr(text_config, "intermediate_size", None)
+ vocab_size = getattr(text_config, "vocab_size", None)
+ estimated_weights = None
+ if intermediate_size is not None and vocab_size is not None:
+ intermediate_size = _require_positive_int(
+ intermediate_size, "config.intermediate_size"
+ )
+ vocab_size = _require_positive_int(vocab_size, "config.vocab_size")
+ attention_parameters = hidden_size * (
+ num_attention_heads * head_dim
+ + 2 * num_key_value_heads * head_dim
+ + hidden_size
+ )
+ mlp_parameters = 3 * hidden_size * intermediate_size
+ embedding_parameters = vocab_size * hidden_size
+ if not bool(getattr(text_config, "tie_word_embeddings", True)):
+ embedding_parameters *= 2
+ estimated_weights = 2 * (
+ num_layers * (attention_parameters + mlp_parameters) + embedding_parameters
+ )
return {
- "head_dim": _require_positive_int(head_dim, "config.head_dim"),
+ "head_dim": head_dim,
"num_layers": num_layers,
"num_attention_heads": num_attention_heads,
"num_key_value_heads": num_key_value_heads,
"rope_theta": float(rope_theta),
+ "rope_theta_source": "rope_parameters" if nested_rope_theta is not None else "legacy",
+ "legacy_rope_theta": (
+ float(legacy_rope_theta) if legacy_rope_theta is not None else None
+ ),
+ "rope_theta_conflict": (
+ nested_rope_theta is not None
+ and legacy_rope_theta is not None
+ and not math.isclose(float(nested_rope_theta), float(legacy_rope_theta))
+ ),
+ "max_position_embeddings": max_positions,
+ "hidden_size": hidden_size,
+ "estimated_bf16_weight_bytes": estimated_weights,
+ "declared_transformers_version": getattr(config, "transformers_version", None),
+ "runtime_transformers_version": transformers.__version__,
}
@@ -1009,6 +1191,7 @@ def calibrate_official_triattention_for_godzilla(
) -> dict[str, object]:
"""Run the official calibrator, then convert and verify its output."""
max_length = _validate_calibration_length(max_length, allow_long=allow_long_calibration)
+ normalized_device = normalize_calibration_device(device)
script = Path(calibrator).expanduser().resolve()
script_report = inspect_official_triattention_calibrator(script)
if not script_report["valid"]:
@@ -1030,10 +1213,22 @@ def calibrate_official_triattention_for_godzilla(
normalized_tokenizer = tokenizer_backend.strip().lower()
if normalized_tokenizer not in {"transformers", "gigatoken"}:
raise ValueError("Tokenizer backend must be 'transformers' or 'gigatoken'")
- executable = [sys.executable, str(script)]
- if normalized_tokenizer == "gigatoken":
- wrapper = Path(__file__).with_name("gigatoken_runner.py")
- executable = [sys.executable, str(wrapper), "--calibrator", str(script)]
+ model_metadata = load_huggingface_model_metadata(model, trust_remote_code=True)
+ validate_model_calibration_length(max_length, model_metadata)
+ memory_report = validate_official_calibration_memory(
+ model_metadata, max_length, normalized_device
+ )
+ wrapper = Path(__file__).with_name("triattention_runner.py")
+ executable = [
+ sys.executable,
+ str(wrapper),
+ "--kind",
+ "official",
+ "--tokenizer-backend",
+ normalized_tokenizer,
+ "--calibrator",
+ str(script),
+ ]
command = [
*executable,
"--model",
@@ -1045,7 +1240,7 @@ def calibrate_official_triattention_for_godzilla(
"--max-length",
str(max_length),
"--device",
- device,
+ normalized_device,
"--attn-implementation",
attention_implementation,
]
@@ -1056,7 +1251,6 @@ def calibrate_official_triattention_for_godzilla(
)
if not stats_output.is_file():
raise RuntimeError(f"Official calibrator did not create {stats_output}")
- model_metadata = load_huggingface_model_metadata(model)
payload = _load_official_payload(stats_output)
metadata = payload.get("metadata")
if not isinstance(metadata, Mapping):
@@ -1081,6 +1275,7 @@ def calibrate_official_triattention_for_godzilla(
"command": command,
"official_stats": str(stats_output),
"tokenizer_backend": normalized_tokenizer,
+ "memory_estimate": memory_report,
**report,
}
@@ -1103,6 +1298,7 @@ def calibrate_domvox_triattention_for_godzilla(
) -> dict[str, object]:
"""Run domvox calibration and explicitly adapt TRIA v2 to Godzilla v1."""
max_length = _validate_calibration_length(max_length, allow_long=allow_long_calibration)
+ normalized_device = normalize_calibration_device(device)
script = Path(calibrator).expanduser().resolve()
script_report = inspect_domvox_triattention_calibrator(script)
if not script_report["valid"]:
@@ -1127,17 +1323,19 @@ def calibrate_domvox_triattention_for_godzilla(
normalized_tokenizer = tokenizer_backend.strip().lower()
if normalized_tokenizer not in {"transformers", "gigatoken"}:
raise ValueError("Tokenizer backend must be 'transformers' or 'gigatoken'")
- executable = [str(python_path), str(script)]
- if normalized_tokenizer == "gigatoken":
- wrapper = Path(__file__).with_name("gigatoken_runner.py")
- executable = [
- str(python_path),
- str(wrapper),
- "--kind",
- "domvox",
- "--calibrator",
- str(script),
- ]
+ model_metadata = load_huggingface_model_metadata(model, trust_remote_code=True)
+ validate_model_calibration_length(max_length, model_metadata)
+ wrapper = Path(__file__).with_name("triattention_runner.py")
+ executable = [
+ str(python_path),
+ str(wrapper),
+ "--kind",
+ "domvox",
+ "--tokenizer-backend",
+ normalized_tokenizer,
+ "--calibrator",
+ str(script),
+ ]
command = [
*executable,
"--model",
@@ -1149,7 +1347,7 @@ def calibrate_domvox_triattention_for_godzilla(
"--max-length",
str(max_length),
"--device",
- device,
+ normalized_device,
]
result = runner(command, check=False)
if result.returncode != 0:
@@ -1158,7 +1356,6 @@ def calibrate_domvox_triattention_for_godzilla(
)
if not stats_output.is_file():
raise RuntimeError(f"domvox calibrator did not create {stats_output}")
- model_metadata = load_huggingface_model_metadata(model)
display_name = model.rstrip("/\\").replace("\\", "/").rsplit("/", 1)[-1]
report = convert_domvox_triattention_stats(
stats_output,
@@ -1231,7 +1428,9 @@ def build_parser() -> argparse.ArgumentParser:
action="store_true",
help=f"Allow one-shot calibration above {LONG_CALIBRATION_THRESHOLD} tokens (maximum {MAX_CALIBRATION_TOKENS})",
)
- calibrate.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
+ calibrate.add_argument(
+ "--device", default="cuda", help="Calibration device: cpu, cuda, or cuda:N"
+ )
calibrate.add_argument(
"--attn-implementation",
choices=("eager", "sdpa", "flash_attention_2"),
@@ -1255,7 +1454,9 @@ def build_parser() -> argparse.ArgumentParser:
domvox.add_argument("--stats-output")
domvox.add_argument("--max-length", type=int, default=2048)
domvox.add_argument("--allow-long-calibration", action="store_true")
- domvox.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
+ domvox.add_argument(
+ "--device", default="cuda", help="Calibration device: cpu, cuda, or cuda:N"
+ )
domvox.add_argument("--accept-lossy", action="store_true")
domvox.add_argument(
"--tokenizer-backend",
diff --git a/multi_turboquant/calibration/triattention_runner.py b/multi_turboquant/calibration/triattention_runner.py
new file mode 100644
index 0000000..9aeaa1f
--- /dev/null
+++ b/multi_turboquant/calibration/triattention_runner.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: MIT
+"""Run reviewed TriAttention calibrators with model/device compatibility patches."""
+
+from __future__ import annotations
+
+import argparse
+import runpy
+import sys
+from pathlib import Path
+from typing import Callable, Mapping
+
+
+def _value(source: object, name: str) -> object | None:
+ value = getattr(source, name, None)
+ if value is not None:
+ return value
+ return source.get(name) if isinstance(source, Mapping) else None
+
+
+def _authoritative_rope_theta(config: object) -> float | None:
+ text_config = getattr(config, "text_config", config)
+ parameters = _value(text_config, "rope_parameters")
+ if parameters is None:
+ parameters = _value(text_config, "rope_scaling")
+ nested = _value(parameters, "rope_theta") if parameters is not None else None
+ return float(nested) if nested is not None else None
+
+
+def patch_auto_config(*, stderr: object = sys.stderr) -> Callable[..., object]:
+ """Promote nested RoPE metadata for calibrators written against older Transformers."""
+ from transformers import AutoConfig
+
+ original = AutoConfig.from_pretrained
+
+ def compatible_from_pretrained(*args: object, **kwargs: object) -> object:
+ config = original(*args, **kwargs)
+ text_config = getattr(config, "text_config", config)
+ authoritative = _authoritative_rope_theta(config)
+ legacy = getattr(text_config, "rope_theta", None)
+ if authoritative is not None and (
+ legacy is None or float(legacy) != authoritative
+ ):
+ setattr(text_config, "rope_theta", authoritative)
+ print(
+ "TriAttention compatibility: using nested rope_parameters.rope_theta="
+ f"{authoritative:g} instead of legacy value {legacy!r}.",
+ file=stderr,
+ )
+ return config
+
+ AutoConfig.from_pretrained = staticmethod(compatible_from_pretrained)
+ return original
+
+
+def patch_auto_model_device(device: str) -> Callable[..., object]:
+ """Force reviewed calibrators' ``device_map='auto'`` onto the selected CUDA device."""
+ from transformers import AutoModelForCausalLM
+
+ original = AutoModelForCausalLM.from_pretrained
+
+ def selected_from_pretrained(*args: object, **kwargs: object) -> object:
+ if device.startswith("cuda") and kwargs.get("device_map") == "auto":
+ kwargs["device_map"] = device
+ return original(*args, **kwargs)
+
+ AutoModelForCausalLM.from_pretrained = staticmethod(selected_from_pretrained)
+ return original
+
+
+def _select_cuda_device(device: str) -> None:
+ if not device.startswith("cuda"):
+ return
+ import torch
+
+ index = int(device.split(":", 1)[1]) if ":" in device else 0
+ if index >= torch.cuda.device_count():
+ raise RuntimeError(
+ f"CUDA device {index} was requested, but only {torch.cuda.device_count()} device(s) "
+ "are visible"
+ )
+ torch.cuda.set_device(index)
+
+
+def build_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(
+ description="Run a reviewed TriAttention calibrator with compatibility guards"
+ )
+ parser.add_argument("--kind", choices=("official", "domvox"), required=True)
+ parser.add_argument(
+ "--tokenizer-backend", choices=("transformers", "gigatoken"), required=True
+ )
+ parser.add_argument("--calibrator", required=True)
+ parser.add_argument("--model", required=True)
+ parser.add_argument("--input", required=True)
+ parser.add_argument("--output", required=True)
+ parser.add_argument("--max-length", type=int, required=True)
+ parser.add_argument("--device", required=True)
+ parser.add_argument("--attn-implementation")
+ return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+ args = build_parser().parse_args(argv)
+ calibrator = Path(args.calibrator).expanduser().resolve()
+ calibration_input = Path(args.input).expanduser().resolve()
+ if not calibrator.is_file():
+ raise ValueError(f"{args.kind} TriAttention calibrator not found: {calibrator}")
+ if not calibration_input.is_file() or calibration_input.stat().st_size == 0:
+ raise ValueError(f"Calibration input must be non-empty: {calibration_input}")
+ if args.kind == "official" and not args.attn_implementation:
+ raise ValueError("Official calibration requires --attn-implementation")
+ if args.kind == "domvox" and args.attn_implementation:
+ raise ValueError("domvox calibration does not accept --attn-implementation")
+
+ _select_cuda_device(args.device)
+ patch_auto_config()
+ patch_auto_model_device(args.device)
+ if args.tokenizer_backend == "gigatoken":
+ from gigatoken_runner import _patch_auto_tokenizer
+
+ text = calibration_input.read_text(encoding="utf-8")
+ _patch_auto_tokenizer(validation_text=text, max_length=args.max_length)
+
+ sys.argv = [
+ str(calibrator),
+ "--model",
+ args.model,
+ "--input",
+ str(calibration_input),
+ "--output",
+ args.output,
+ "--max-length",
+ str(args.max_length),
+ "--device",
+ args.device,
+ ]
+ if args.kind == "official":
+ sys.argv.extend(("--attn-implementation", args.attn_implementation))
+ runpy.run_path(str(calibrator), run_name="__main__")
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/multi_turboquant/integration/godzilla_workspace.py b/multi_turboquant/integration/godzilla_workspace.py
index cf09225..e0560f7 100644
--- a/multi_turboquant/integration/godzilla_workspace.py
+++ b/multi_turboquant/integration/godzilla_workspace.py
@@ -23,6 +23,9 @@
inspect_domvox_triattention_calibrator,
inspect_godzilla_triattention_file,
inspect_official_triattention_calibrator,
+ estimate_official_calibration_bytes,
+ normalize_calibration_device,
+ validate_model_calibration_length,
)
@@ -179,6 +182,8 @@ class GodzillaCalibrationPlan:
issues: tuple[GodzillaIssue, ...]
tokenizer_backend: str = "transformers"
python_discovery: Mapping[str, object] | None = None
+ model_metadata: Mapping[str, object] | None = None
+ memory_estimate: Mapping[str, object] | None = None
@property
def ready(self) -> bool:
@@ -212,6 +217,8 @@ def to_dict(self) -> dict[str, object]:
"dependency_validation": self.dependency_validation,
"dependency_override": self.dependency_override,
"python_discovery": self.python_discovery,
+ "model_metadata": self.model_metadata,
+ "memory_estimate": self.memory_estimate,
"command": list(self.command),
"environment": dict(self.environment),
"issues": [issue.to_dict() for issue in self.issues],
@@ -243,6 +250,7 @@ def plan_godzilla_triattention(
python_discovery: Mapping[str, object] | None = None,
dependency_runner=subprocess.run,
shell_executable: str | None = None,
+ model_metadata_loader: Callable[[str], Mapping[str, object]] | None = None,
) -> GodzillaCalibrationPlan:
"""Plan a validated official-Python or checkout-owned calibration workflow."""
checkout_path = Path(checkout).expanduser().resolve()
@@ -253,7 +261,12 @@ def plan_godzilla_triattention(
else checkout_path / "calibrations" / f"{gguf_path.stem}.triattention"
)
normalized_hf = hf_model.strip() if hf_model and hf_model.strip() else None
- normalized_device = device.strip().lower()
+ device_error = None
+ try:
+ normalized_device = normalize_calibration_device(device)
+ except ValueError as exc:
+ normalized_device = device.strip().lower()
+ device_error = str(exc)
normalized_mode = mode.strip().lower()
normalized_attention = attention_implementation.strip().lower()
dependency_attention = normalized_attention if normalized_mode == "official_python" else "sdpa"
@@ -301,6 +314,8 @@ def plan_godzilla_triattention(
output_exists = output_path.is_file()
issues: list[GodzillaIssue] = []
dependency_validation: Mapping[str, object] | None = None
+ model_metadata: Mapping[str, object] | None = None
+ memory_estimate: Mapping[str, object] | None = None
if not inspection["valid"]:
issues.append(
@@ -485,8 +500,8 @@ def plan_godzilla_triattention(
"The upstream calibrator processes one long sequence; no chunked aggregation is being assumed. Expect substantially higher memory and runtime.",
)
)
- if normalized_device not in {"cuda", "cpu"}:
- issues.append(GodzillaIssue("error", "invalid_device", "Device must be 'cuda' or 'cpu'."))
+ if device_error is not None:
+ issues.append(GodzillaIssue("error", "invalid_device", device_error))
if normalized_mode == "official_python" and normalized_attention not in {
"eager",
"sdpa",
@@ -502,7 +517,7 @@ def plan_godzilla_triattention(
elif (
normalized_mode == "official_python"
and normalized_attention == "flash_attention_2"
- and normalized_device != "cuda"
+ and not normalized_device.startswith("cuda")
):
issues.append(
GodzillaIssue(
@@ -542,6 +557,64 @@ def plan_godzilla_triattention(
"Provide the matching Hugging Face model because this checkout has no resolver.",
)
)
+ if (
+ not output_exists
+ and normalized_hf is not None
+ and normalized_mode in {"official_python", "official_convert", "domvox"}
+ and model_metadata_loader is not None
+ ):
+ try:
+ loaded_metadata = model_metadata_loader(normalized_hf)
+ model_metadata = dict(loaded_metadata)
+ if normalized_mode in {"official_python", "domvox"}:
+ validate_model_calibration_length(n_tokens, model_metadata)
+ if normalized_mode == "official_python":
+ memory_estimate = estimate_official_calibration_bytes(model_metadata, n_tokens)
+ except Exception as exc:
+ issues.append(
+ GodzillaIssue(
+ "error",
+ "model_metadata_incompatible",
+ f"Matching Hugging Face model metadata could not be validated: {exc}",
+ )
+ )
+ else:
+ context = model_metadata.get("max_position_embeddings")
+ theta = model_metadata.get("rope_theta")
+ issues.append(
+ GodzillaIssue(
+ "info",
+ "model_metadata_validated",
+ f"Validated model metadata: context={context or 'unknown'}, "
+ f"RoPE theta={theta or 'unknown'}.",
+ )
+ )
+ if model_metadata.get("rope_theta_conflict"):
+ issues.append(
+ GodzillaIssue(
+ "warning",
+ "nested_rope_theta_selected",
+ "The model declares conflicting legacy and nested RoPE values; the "
+ "calibration wrapper will use rope_parameters.rope_theta.",
+ )
+ )
+ declared_transformers = model_metadata.get("declared_transformers_version")
+ runtime_transformers = model_metadata.get("runtime_transformers_version")
+ if (
+ declared_transformers
+ and runtime_transformers
+ and declared_transformers != runtime_transformers
+ ):
+ issues.append(
+ GodzillaIssue(
+ "warning",
+ "model_transformers_version_differs",
+ "The model config was saved by Transformers "
+ f"{declared_transformers}, while calibration uses {runtime_transformers}. "
+ "The compatibility wrapper normalizes reviewed RoPE metadata, but model "
+ "code compatibility must still be validated.",
+ )
+ )
issues.append(
GodzillaIssue(
"info",
@@ -651,10 +724,35 @@ def plan_godzilla_triattention(
"info",
"calibration_device_memory",
f"{device_name}: {free_gib:.1f} GiB free of {total_gib:.1f} GiB VRAM "
- "at preflight time. This is capacity information, not a 200k-token "
- "memory guarantee.",
+ "at preflight time. This capacity snapshot is not a peak-memory "
+ "guarantee.",
)
)
+ if memory_estimate is not None:
+ floor = memory_estimate.get("estimated_floor_bytes")
+ free_bytes = dependency_validation.get("cuda_free_memory_bytes")
+ if isinstance(floor, int) and isinstance(free_bytes, int):
+ if floor > free_bytes:
+ issues.append(
+ GodzillaIssue(
+ "error",
+ "calibration_memory_floor_exceeded",
+ "The official one-shot calibration memory floor is "
+ f"{floor / gib:.1f} GiB, exceeding the selected device's "
+ f"{free_bytes / gib:.1f} GiB currently free. Select another "
+ "CUDA device or reduce calibration tokens.",
+ )
+ )
+ else:
+ issues.append(
+ GodzillaIssue(
+ "info",
+ "calibration_memory_floor",
+ "The official one-shot calibration memory floor is "
+ f"approximately {floor / gib:.1f} GiB. This is a lower-bound "
+ "estimate, not a guarantee.",
+ )
+ )
elif dependency_validation and dependency_validation.get("cuda_memory_error"):
issues.append(
GodzillaIssue(
@@ -861,6 +959,8 @@ def plan_godzilla_triattention(
issues=tuple(issues),
tokenizer_backend=normalized_tokenizer,
python_discovery=python_discovery,
+ model_metadata=model_metadata,
+ memory_estimate=memory_estimate,
)
@@ -1048,6 +1148,8 @@ def collect_godzilla_calibration_diagnostics(
"tokenizer_backend": plan.tokenizer_backend,
"attention_implementation": plan.attention_implementation,
"n_tokens": plan.n_tokens,
+ "model_metadata": plan.model_metadata,
+ "memory_estimate": plan.memory_estimate,
"dependency_override_requested": plan.dependency_override,
"ready_at_submission": plan.ready,
"issues": [issue.to_dict() for issue in plan.issues],
diff --git a/multi_turboquant/ui/runtime.py b/multi_turboquant/ui/runtime.py
index 89bf627..b05c2a9 100644
--- a/multi_turboquant/ui/runtime.py
+++ b/multi_turboquant/ui/runtime.py
@@ -359,6 +359,7 @@ def start(
dependency_override: bool = False,
python_discovery: Mapping[str, object] | None = None,
) -> dict[str, object]:
+ from ..calibration.godzilla_triattention import load_huggingface_model_metadata
from ..integration.godzilla_workspace import plan_godzilla_triattention
plan = plan_godzilla_triattention(
@@ -381,6 +382,9 @@ def start(
verify_dependencies=True,
dependency_override=dependency_override,
python_discovery=python_discovery,
+ model_metadata_loader=lambda model_id: load_huggingface_model_metadata(
+ model_id, trust_remote_code=False
+ ),
)
if not plan.ready:
errors = "; ".join(issue.message for issue in plan.issues if issue.severity == "error")
diff --git a/run_ui.py b/run_ui.py
index 9b82d14..d3c837d 100644
--- a/run_ui.py
+++ b/run_ui.py
@@ -37,6 +37,7 @@
from multi_turboquant.calibration import (
CALIBRATION_CORPUS_SCHEMA_VERSION,
generate_calibration_text,
+ load_huggingface_model_metadata,
select_compatible_calibration_python,
)
from multi_turboquant._paths import lexical_absolute_path
@@ -101,6 +102,7 @@ def api_status():
plat = detect_platform()
gpus = [
{
+ "index": g.index,
"name": g.name,
"vram_mb": g.vram_total_mb,
"vram_used_mb": g.vram_used_mb,
@@ -870,6 +872,9 @@ def _godzilla_plan_from_params(params):
verify_dependencies=True,
dependency_override=_truthy(params.get("dependency_override")),
python_discovery=python_discovery,
+ model_metadata_loader=lambda model_id: load_huggingface_model_metadata(
+ model_id, trust_remote_code=False
+ ),
)
@@ -914,7 +919,8 @@ def api_plan_godzilla(params):
"max_concurrent_calibrations": GODZILLA_JOBS.max_concurrent_jobs,
"message": (
"The local UI permits one calibration job at a time to prevent overlapping "
- "model loads. This does not reduce the memory used by one long sequence."
+ "model loads. Model context and a conservative one-shot memory floor are checked "
+ "before weights are loaded."
),
}
return result
@@ -1589,14 +1595,14 @@ def api_runtime_status():
-
+
No preparation plan checked.
Python: ${escapeHtml(plan.python || 'not found')}
` +
discoveryHtml +
`Output: ${escapeHtml(plan.output)}
` +
+ (plan.model_metadata ? `Model preflight: context ${escapeHtml(plan.model_metadata.max_position_embeddings || 'unknown')}, ` +
+ `RoPE theta ${escapeHtml(plan.model_metadata.rope_theta || 'unknown')} (${escapeHtml(plan.model_metadata.rope_theta_source || 'unknown source')}).
` : '') +
+ (plan.memory_estimate?.estimated_floor_bytes ? `Official one-shot memory floor: ` +
+ `${escapeHtml((plan.memory_estimate.estimated_floor_bytes / (1024 ** 3)).toFixed(1))} GiB.
` : '') +
(plan.resource_policy ? `${escapeHtml(plan.resource_policy.message)}
` : '') +
(plan.official_stats ? `Official .pt stats: ${escapeHtml(plan.official_stats)}
` : '') +
(plan.issues || []).map(issue =>
diff --git a/tests/test_godzilla_triattention.py b/tests/test_godzilla_triattention.py
index db020e9..fa1083d 100644
--- a/tests/test_godzilla_triattention.py
+++ b/tests/test_godzilla_triattention.py
@@ -4,6 +4,7 @@
import struct
import subprocess
from pathlib import Path
+from types import SimpleNamespace
import pytest
import torch
@@ -78,6 +79,7 @@ def _python_probe_report(
cuda_available: bool = True,
gigatoken: str | None = None,
module_errors: dict[str, str] | None = None,
+ cuda_device_index: int = 0,
) -> str:
versions = {
"torch": "2.7.1",
@@ -113,6 +115,10 @@ def _python_probe_report(
"modules": modules,
"torch_cuda": torch_cuda,
"cuda_available": cuda_available,
+ "cuda_device": "Test GPU",
+ "cuda_device_index": cuda_device_index,
+ "cuda_free_memory_bytes": 24 * 1024**3,
+ "cuda_total_memory_bytes": 24 * 1024**3,
}
)
@@ -216,7 +222,7 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
monkeypatch.setattr(
calibration,
"load_huggingface_model_metadata",
- lambda model: {
+ lambda model, **_kwargs: {
"head_dim": 4,
"num_layers": 1,
"num_attention_heads": 2,
@@ -266,7 +272,7 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
monkeypatch.setattr(
calibration,
"load_huggingface_model_metadata",
- lambda model: {
+ lambda model, **_kwargs: {
"head_dim": 4,
"num_layers": 1,
"num_attention_heads": 2,
@@ -288,8 +294,10 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
runner=runner,
)
- assert Path(calls[0][1]).name == "gigatoken_runner.py"
- assert calls[0][2:6] == ["--kind", "domvox", "--calibrator", str(script.resolve())]
+ assert Path(calls[0][1]).name == "triattention_runner.py"
+ assert calls[0][2:8] == [
+ "--kind", "domvox", "--tokenizer-backend", "gigatoken", "--calibrator", str(script.resolve())
+ ]
assert "--attn-implementation" not in calls[0]
assert report["tokenizer_backend"] == "gigatoken"
@@ -380,7 +388,7 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
monkeypatch.setattr(
calibration,
"load_huggingface_model_metadata",
- lambda model: {
+ lambda model, **_kwargs: {
"head_dim": 4,
"num_layers": 2,
"num_attention_heads": 2,
@@ -455,6 +463,79 @@ def runner(command, **kwargs):
assert "torch.cuda.mem_get_info" in commands[0][-1]
+def test_calibration_python_preflight_selects_explicit_cuda_device(tmp_path: Path):
+ python = tmp_path / "python"
+ python.write_bytes(b"python")
+ commands = []
+
+ def runner(command, **kwargs):
+ commands.append(command)
+ return subprocess.CompletedProcess(
+ command,
+ 0,
+ stdout=_python_probe_report(cuda_device_index=1) + "\n",
+ stderr="",
+ )
+
+ report = calibration.inspect_calibration_python(
+ python, device="cuda:1", runner=runner
+ )
+
+ assert report["valid"] is True
+ assert report["report"]["cuda_device_index"] == 1
+ assert "requested_index = 1" in commands[0][-1]
+
+
+def test_model_metadata_prefers_nested_rope_and_exposes_context(monkeypatch):
+ import transformers
+
+ config = SimpleNamespace(
+ num_hidden_layers=36,
+ num_attention_heads=16,
+ num_key_value_heads=2,
+ head_dim=128,
+ hidden_size=2048,
+ intermediate_size=11008,
+ vocab_size=151936,
+ tie_word_embeddings=True,
+ max_position_embeddings=131072,
+ rope_theta=10_000.0,
+ rope_parameters={"rope_theta": 1_000_000.0, "rope_type": "default"},
+ )
+ monkeypatch.setattr(
+ transformers.AutoConfig,
+ "from_pretrained",
+ staticmethod(lambda *_args, **_kwargs: config),
+ )
+
+ metadata = calibration.load_huggingface_model_metadata("org/model")
+
+ assert metadata["rope_theta"] == 1_000_000.0
+ assert metadata["legacy_rope_theta"] == 10_000.0
+ assert metadata["rope_theta_conflict"] is True
+ assert metadata["rope_theta_source"] == "rope_parameters"
+ assert metadata["max_position_embeddings"] == 131_072
+
+
+def test_mythos_shape_rejects_200k_and_reports_memory_floor():
+ metadata = {
+ "head_dim": 128,
+ "num_layers": 36,
+ "num_attention_heads": 16,
+ "num_key_value_heads": 2,
+ "hidden_size": 2048,
+ "estimated_bf16_weight_bytes": 6_171_394_048,
+ "max_position_embeddings": 131_072,
+ }
+
+ with pytest.raises(ValueError, match="131072"):
+ calibration.validate_model_calibration_length(200_000, metadata)
+
+ estimate = calibration.estimate_official_calibration_bytes(metadata, 200_000)
+ assert estimate["captured_q_bytes"] == 29_491_200_000
+ assert estimate["estimated_floor_bytes"] == 38_939_394_048
+
+
@pytest.mark.parametrize(
("arguments", "message"),
[
@@ -678,7 +759,7 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
monkeypatch.setattr(
calibration,
"load_huggingface_model_metadata",
- lambda model: {
+ lambda model, **_kwargs: {
"head_dim": 4,
"num_layers": 2,
"num_attention_heads": 2,
@@ -698,6 +779,13 @@ def runner(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]:
runner=runner,
)
- assert Path(calls[0][1]).name == "gigatoken_runner.py"
- assert calls[0][2:4] == ["--calibrator", str(script.resolve())]
+ assert Path(calls[0][1]).name == "triattention_runner.py"
+ assert calls[0][2:8] == [
+ "--kind",
+ "official",
+ "--tokenizer-backend",
+ "gigatoken",
+ "--calibrator",
+ str(script.resolve()),
+ ]
assert report["tokenizer_backend"] == "gigatoken"
diff --git a/tests/test_godzilla_workspace.py b/tests/test_godzilla_workspace.py
index af65048..200065f 100644
--- a/tests/test_godzilla_workspace.py
+++ b/tests/test_godzilla_workspace.py
@@ -472,6 +472,35 @@ def test_long_calibration_requires_acknowledgement_and_has_hard_cap(tmp_path: Pa
assert any(issue.code == "invalid_token_count" for issue in capped.issues)
+def test_plan_rejects_tokens_beyond_matching_model_context(tmp_path: Path):
+ checkout = _godzilla_checkout(tmp_path)
+ model = tmp_path / "Mythos-nano-heretic.IQ4_XS.gguf"
+ model.write_bytes(b"gguf")
+
+ plan = plan_godzilla_triattention(
+ checkout,
+ model,
+ hf_model="richardyoung/Mythos-nano-heretic",
+ n_tokens=200_000,
+ allow_long_calibration=True,
+ model_metadata_loader=lambda _model: {
+ "head_dim": 128,
+ "num_layers": 36,
+ "num_attention_heads": 16,
+ "num_key_value_heads": 2,
+ "rope_theta": 1_000_000.0,
+ "rope_theta_source": "rope_parameters",
+ "rope_theta_conflict": True,
+ "max_position_embeddings": 131_072,
+ "hidden_size": 2048,
+ "estimated_bf16_weight_bytes": 6_171_394_048,
+ },
+ )
+
+ assert any(issue.code == "model_metadata_incompatible" for issue in plan.issues)
+ assert not plan.ready
+
+
def test_gigatoken_plan_forwards_backend_and_wraps_domvox(tmp_path: Path):
checkout = _godzilla_checkout(tmp_path)
model = tmp_path / "model.gguf"
@@ -766,6 +795,10 @@ def runner(argv, **kwargs):
"sentencepiece": "0.2.2",
"torch_cuda": "12.6",
"cuda_available": True,
+ "cuda_device": "Test GPU",
+ "cuda_device_index": 0,
+ "cuda_free_memory_bytes": 24 * 1024**3,
+ "cuda_total_memory_bytes": 24 * 1024**3,
"modules": {
"torch": {"status": "ok", "version": "2.7.1"},
"transformers": {"status": "ok", "version": "4.57.6"},
diff --git a/tests/test_run_ui.py b/tests/test_run_ui.py
index fe8434a..3f4405a 100644
--- a/tests/test_run_ui.py
+++ b/tests/test_run_ui.py
@@ -23,6 +23,26 @@
)
+@pytest.fixture(autouse=True)
+def _model_metadata(monkeypatch):
+ monkeypatch.setattr(
+ run_ui,
+ "load_huggingface_model_metadata",
+ lambda _model, **_kwargs: {
+ "head_dim": 128,
+ "num_layers": 2,
+ "num_attention_heads": 2,
+ "num_key_value_heads": 1,
+ "rope_theta": 10_000.0,
+ "rope_theta_source": "legacy",
+ "rope_theta_conflict": False,
+ "max_position_embeddings": 32_768,
+ "hidden_size": 256,
+ "estimated_bf16_weight_bytes": 1024,
+ },
+ )
+
+
def _valid_calibration_python(path) -> dict[str, object]:
return {
"python": str(Path(path).absolute()),
@@ -48,6 +68,10 @@ def _valid_calibration_python(path) -> dict[str, object]:
"accelerate": "1.14.0",
"torch_cuda": "12.6",
"cuda_available": True,
+ "cuda_device": "Test GPU",
+ "cuda_device_index": 0,
+ "cuda_free_memory_bytes": 24 * 1024**3,
+ "cuda_total_memory_bytes": 24 * 1024**3,
"modules": {
name: {"status": "ok", "version": version}
for name, version in {
@@ -377,6 +401,7 @@ def test_api_status_reports_ram_and_vram_separately(monkeypatch):
assert result["total_vram_gb"] == 24
assert result["available_vram_gb"] == 16
assert result["combined_memory_gb"] == 88
+ assert result["gpus"][0]["index"] == 0
def test_api_inspects_informational_addon_source(tmp_path):
diff --git a/tests/test_triattention_runner.py b/tests/test_triattention_runner.py
new file mode 100644
index 0000000..76a97a8
--- /dev/null
+++ b/tests/test_triattention_runner.py
@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from multi_turboquant.calibration import triattention_runner
+
+
+def test_runner_promotes_nested_rope_theta(monkeypatch):
+ import transformers
+
+ config = SimpleNamespace(
+ rope_theta=10_000.0,
+ rope_parameters={"rope_theta": 1_000_000.0, "rope_type": "default"},
+ )
+ original = transformers.AutoConfig.from_pretrained
+ monkeypatch.setattr(
+ transformers.AutoConfig,
+ "from_pretrained",
+ staticmethod(lambda *_args, **_kwargs: config),
+ )
+ messages: list[str] = []
+
+ try:
+ triattention_runner.patch_auto_config(stderr=SimpleNamespace(
+ write=lambda value: messages.append(value)
+ ))
+ loaded = transformers.AutoConfig.from_pretrained("org/model")
+ finally:
+ transformers.AutoConfig.from_pretrained = original
+
+ assert loaded.rope_theta == 1_000_000.0
+ assert any("rope_parameters.rope_theta=1e+06" in message for message in messages)
+
+
+def test_runner_selects_requested_cuda_index(monkeypatch):
+ import torch
+
+ selected: list[int] = []
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: 2)
+ monkeypatch.setattr(torch.cuda, "set_device", lambda index: selected.append(index))
+
+ triattention_runner._select_cuda_device("cuda:1")
+
+ assert selected == [1]
+
+
+def test_runner_forces_auto_device_map_to_selected_cuda(monkeypatch):
+ import transformers
+
+ observed: list[dict[str, object]] = []
+ original = transformers.AutoModelForCausalLM.from_pretrained
+ monkeypatch.setattr(
+ transformers.AutoModelForCausalLM,
+ "from_pretrained",
+ staticmethod(lambda *_args, **kwargs: observed.append(kwargs) or object()),
+ )
+
+ try:
+ triattention_runner.patch_auto_model_device("cuda:1")
+ transformers.AutoModelForCausalLM.from_pretrained(
+ "org/model", device_map="auto"
+ )
+ finally:
+ transformers.AutoModelForCausalLM.from_pretrained = original
+
+ assert observed[0]["device_map"] == "cuda:1"
From fe62bd2a53a6de1c986dfbe88d437d4736eedc2a Mon Sep 17 00:00:00 2001
From: aivrar <42790594+aivrar@users.noreply.github.com>
Date: Thu, 20 Aug 2026 19:57:28 -0700
Subject: [PATCH 2/3] Keep calibration tests independent of optional
Transformers
---
tests/test_godzilla_triattention.py | 13 ++++----
tests/test_triattention_runner.py | 49 +++++++++++++----------------
2 files changed, 29 insertions(+), 33 deletions(-)
diff --git a/tests/test_godzilla_triattention.py b/tests/test_godzilla_triattention.py
index fa1083d..b9b7d71 100644
--- a/tests/test_godzilla_triattention.py
+++ b/tests/test_godzilla_triattention.py
@@ -3,6 +3,7 @@
import json
import struct
import subprocess
+import sys
from pathlib import Path
from types import SimpleNamespace
@@ -487,8 +488,6 @@ def runner(command, **kwargs):
def test_model_metadata_prefers_nested_rope_and_exposes_context(monkeypatch):
- import transformers
-
config = SimpleNamespace(
num_hidden_layers=36,
num_attention_heads=16,
@@ -502,11 +501,13 @@ def test_model_metadata_prefers_nested_rope_and_exposes_context(monkeypatch):
rope_theta=10_000.0,
rope_parameters={"rope_theta": 1_000_000.0, "rope_type": "default"},
)
- monkeypatch.setattr(
- transformers.AutoConfig,
- "from_pretrained",
- staticmethod(lambda *_args, **_kwargs: config),
+ fake_transformers = SimpleNamespace(
+ __version__="test",
+ AutoConfig=SimpleNamespace(
+ from_pretrained=staticmethod(lambda *_args, **_kwargs: config)
+ ),
)
+ monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
metadata = calibration.load_huggingface_model_metadata("org/model")
diff --git a/tests/test_triattention_runner.py b/tests/test_triattention_runner.py
index 76a97a8..7f4ff20 100644
--- a/tests/test_triattention_runner.py
+++ b/tests/test_triattention_runner.py
@@ -1,32 +1,30 @@
from __future__ import annotations
+import sys
from types import SimpleNamespace
from multi_turboquant.calibration import triattention_runner
def test_runner_promotes_nested_rope_theta(monkeypatch):
- import transformers
-
config = SimpleNamespace(
rope_theta=10_000.0,
rope_parameters={"rope_theta": 1_000_000.0, "rope_type": "default"},
)
- original = transformers.AutoConfig.from_pretrained
- monkeypatch.setattr(
- transformers.AutoConfig,
- "from_pretrained",
- staticmethod(lambda *_args, **_kwargs: config),
+ auto_config = SimpleNamespace(
+ from_pretrained=staticmethod(lambda *_args, **_kwargs: config)
+ )
+ monkeypatch.setitem(
+ sys.modules, "transformers", SimpleNamespace(AutoConfig=auto_config)
)
messages: list[str] = []
- try:
- triattention_runner.patch_auto_config(stderr=SimpleNamespace(
+ triattention_runner.patch_auto_config(
+ stderr=SimpleNamespace(
write=lambda value: messages.append(value)
- ))
- loaded = transformers.AutoConfig.from_pretrained("org/model")
- finally:
- transformers.AutoConfig.from_pretrained = original
+ )
+ )
+ loaded = auto_config.from_pretrained("org/model")
assert loaded.rope_theta == 1_000_000.0
assert any("rope_parameters.rope_theta=1e+06" in message for message in messages)
@@ -45,22 +43,19 @@ def test_runner_selects_requested_cuda_index(monkeypatch):
def test_runner_forces_auto_device_map_to_selected_cuda(monkeypatch):
- import transformers
-
observed: list[dict[str, object]] = []
- original = transformers.AutoModelForCausalLM.from_pretrained
- monkeypatch.setattr(
- transformers.AutoModelForCausalLM,
- "from_pretrained",
- staticmethod(lambda *_args, **kwargs: observed.append(kwargs) or object()),
+ auto_model = SimpleNamespace(
+ from_pretrained=staticmethod(
+ lambda *_args, **kwargs: observed.append(kwargs) or object()
+ )
+ )
+ monkeypatch.setitem(
+ sys.modules,
+ "transformers",
+ SimpleNamespace(AutoModelForCausalLM=auto_model),
)
- try:
- triattention_runner.patch_auto_model_device("cuda:1")
- transformers.AutoModelForCausalLM.from_pretrained(
- "org/model", device_map="auto"
- )
- finally:
- transformers.AutoModelForCausalLM.from_pretrained = original
+ triattention_runner.patch_auto_model_device("cuda:1")
+ auto_model.from_pretrained("org/model", device_map="auto")
assert observed[0]["device_map"] == "cuda:1"
From c83eb6e00f2de3a9927ee859ce983028b95d1daa Mon Sep 17 00:00:00 2001
From: aivrar <42790594+aivrar@users.noreply.github.com>
Date: Thu, 20 Aug 2026 20:03:28 -0700
Subject: [PATCH 3/3] Credit Mythos calibration report
---
README.md | 3 +++
1 file changed, 3 insertions(+)
diff --git a/README.md b/README.md
index 9e6be19..2ad1834 100644
--- a/README.md
+++ b/README.md
@@ -935,6 +935,7 @@ contracts and do not imply runtime compatibility.
| Proposed the JetSpec, Lucebox, Proxima, Jet-Long, ChunkLlama, RaBitQCache, ScoPE, DuoAttention, IceCache, PFlash/KVFlash, and Resonance-JetLong review, prompting pinned source profiles, read-only discovery contracts, runtime capability scanning, and fail-closed composition metadata | [@jawadala](https://github.com/jawadala) | Issue [#43](https://github.com/aivrar/multi-turboquant/issues/43) |
| Requested a safe PFlash/KVFlash composition path for the exact Godzilla `09214b160` baseline, prompting the pinned overlay, request and runtime guardrails, and build verification | [@jawadala](https://github.com/jawadala) | Issue [#44](https://github.com/aivrar/multi-turboquant/issues/44) |
| Requested full guarded treatment of the reviewed add-on catalog, workload routing, simulation, LuceBox composition research, UI coverage, and SM86/SM89 qualification for the canonical Godzilla baseline | [@jawadala](https://github.com/jawadala) | Issue [#46](https://github.com/aivrar/multi-turboquant/issues/46) |
+| Reported model-specific TriAttention calibration failures with Mythos-nano-heretic `IQ4_XS`, prompting source-model context validation, nested RoPE correction, memory preflight, and explicit multi-GPU device selection | [@jawadala](https://github.com/jawadala) | Community testing report (August 2026) |
| ForgeAttention — fused MLX kernels for Apple Silicon (`multi_turboquant/kernels/metal/`): packed-3-bit fused QK, tiled SV, flash decode, sparse SV with phase-1/2 early exit, per-head attention budget calibration | [@user-23xyz](https://github.com/user-23xyz) | PR [#1](https://github.com/aivrar/multi-turboquant/pull/1) · sibling project [user-23xyz/forgeattention](https://github.com/user-23xyz/forgeattention) |
Thank you to [@jawadala](https://github.com/jawadala) for the sustained issue
@@ -946,6 +947,8 @@ correction, parity-checked Gigatoken option, and the broader issue #43 research
catalog with explicit safety boundaries, including the exact-commit composition
workflow prompted by issue #44 and the full guarded composition, routing,
simulation, LuceBox, UI, and SM86/SM89 qualification follow-up in issue #46.
+Their subsequent Mythos-nano-heretic calibration report also prompted the
+model-aware context, RoPE, memory, and CUDA-device safeguards.
The Metal path is community-maintained — the maintainer does not have Apple Silicon hardware, so issues specific to MLX/Metal should tag the contributor for context.