Coverage for transformer_lens/tools/model_registry/verify_models.py: 33%
737 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Batch model verification tool for the TransformerLens model registry.
3Iterates through supported models, estimates memory requirements, runs benchmarks
4phase-by-phase, and updates the registry with status, phase scores, and notes.
6Usage:
7 python -m transformer_lens.tools.model_registry.verify_models [options]
9Examples:
10 # Dry run to see what would be tested
11 python -m transformer_lens.tools.model_registry.verify_models --dry-run
13 # Verify top 10 models per architecture on CPU
14 python -m transformer_lens.tools.model_registry.verify_models --device cpu
16 # Verify only GPT2 models, limit to 3
17 python -m transformer_lens.tools.model_registry.verify_models --architectures GPT2LMHeadModel --limit 3
19 # Resume from a previous interrupted run
20 python -m transformer_lens.tools.model_registry.verify_models --resume
22 # Re-verify already-tested models for a specific architecture
23 python -m transformer_lens.tools.model_registry.verify_models --reverify --architectures Olmo2ForCausalLM
24"""
26import argparse
27import gc
28import json
29import logging
30import re
31import signal
32import time
33from dataclasses import dataclass, field
34from datetime import datetime
35from pathlib import Path
36from typing import Optional
38from transformer_lens.benchmarks.text_quality_profiles import (
39 P4_SCORING_VERSION,
40 p4_pass_threshold,
41)
42from transformer_lens.utilities.heterogeneous_config import het_safe_view
44# Exit code used for graceful interrupts (Ctrl+C). The wrapper script
45# recognises this and stops without marking the in-flight model as failed.
46_EXIT_GRACEFUL_INTERRUPT = 42
48# Module-level flag set by the SIGINT handler so the main loop can stop
49# between models without corrupting state.
50_interrupt_requested = False
52from .registry_io import (
53 QUANTIZED_NOTE,
54 STATUS_FAILED,
55 STATUS_PROVISIONAL,
56 STATUS_SKIPPED,
57 STATUS_UNVERIFIED,
58 STATUS_VERIFIED,
59 add_verification_record,
60 is_incompatible_quantized,
61 load_supported_models_raw,
62 required_quant_library_for_model,
63 update_model_status,
64)
66logger = logging.getLogger(__name__)
68# Architectures added via the TransformerBridge system that need trust_remote_code=True.
69# These are not in the legacy NEED_REMOTE_CODE_MODELS tuple (loading_from_pretrained.py).
70_BRIDGE_REMOTE_CODE_PREFIXES: tuple[str, ...] = (
71 "baichuan-inc/", # BaichuanForCausalLM — ships own modeling_baichuan.py
72 "ByteDance/Ouro-", # OuroForCausalLM — ships own modeling_ouro.py
73 "internlm/", # InternLM2ForCausalLM — ships own modeling_internlm2.py
74 "GSAI-ML/LLaDA", # LLaDAModelLM — ships configuration_llada.py/modeling_llada.py
75 "kuleshov-group/", # BD3LM — ships own custom modeling_d_dit.py
76 "Dream-org/", # DreamModel — ships own modeling_dream.py
77 "dvruette/", # GiddForDiffusionLM — ships own modeling_gidd.py
78 "LongSafari/", # HyenaDNAForCausalLM — ships own modeling_hyena.py
79 "inclusionAI/", # LLaDA2MoeModelLM — ships own modeling_llada2_moe.py
80 "poolside/", # LagunaForCausalLM — ships own modeling_laguna.py
81 "apple/DiffuCoder", # DreamModel (DiffuCoder) — same remote code family
82 "LGAI-EXAONE/", # ExaoneForCausalLM (EXAONE-3.x) — ships own modeling_exaone.py
83)
85# Data directory for registry files
86_DATA_DIR = Path(__file__).parent / "data"
87_CHECKPOINT_PATH = _DATA_DIR / "verification_checkpoint.json"
90def _handle_sigint(signum, frame): # noqa: ARG001
91 """Handle Ctrl+C by setting a flag instead of raising immediately.
93 The main verification loop checks this flag between models so it can
94 save the checkpoint cleanly and exit without marking the current model
95 as failed.
96 """
97 global _interrupt_requested # noqa: PLW0603
98 if _interrupt_requested:
99 # Second Ctrl+C — force exit immediately
100 print("\nForce quit.")
101 raise SystemExit(1)
102 _interrupt_requested = True
103 print("\n\nInterrupt received — finishing current model before stopping.")
104 print("(Press Ctrl+C again to force quit immediately.)\n")
107# Pattern matching HuggingFace API tokens (hf_ followed by 20+ alphanumeric chars)
108_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}")
111def _sanitize_note(note: Optional[str]) -> Optional[str]:
112 """Sanitize a note string to remove sensitive information.
114 Strips HuggingFace tokens and replaces verbose gated-repo error messages
115 with a concise summary.
116 """
117 if not note:
118 return note
119 # Replace any HF tokens that leaked into the message
120 note = _HF_TOKEN_RE.sub("HF_TOKEN", note)
121 # Replace verbose gated-repo 401 errors with a clean summary
122 if "gated repo" in note:
123 url_match = re.search(r"https://huggingface\.co/([^\s.]+)", note)
124 model_ref = url_match.group(1) if url_match else "unknown"
125 return f"Config unavailable: Gated repo ({model_ref})"
126 return note
129def _phases_to_run(arch: str, phases: list[int]) -> list[int]:
130 """Restrict requested phases to those the adapter supports.
132 An adapter's ``applicable_phases`` declares which text phases (1-4) it covers. Phases
133 7/8/9 are gated separately by ``is_multimodal``/``is_audio_model``/``is_visual_model``
134 in the benchmark, so they are never filtered out here. An empty result means none of
135 the requested phases apply to this architecture (SSM / recurrent families run all four).
136 """
137 from transformer_lens.factories.architecture_adapter_factory import (
138 SUPPORTED_ARCHITECTURES,
139 )
141 applicable = getattr(SUPPORTED_ARCHITECTURES.get(arch), "applicable_phases", [1, 2, 3, 4])
142 return [p for p in phases if p in applicable or p in (7, 8, 9)]
145def _full_and_core_phases(arch: str) -> tuple[set[int], set[int]]:
146 """``(full verification set, core subset)`` for this architecture -- the single
147 source for both default phase selection and the status-writing decision."""
148 from transformer_lens.utilities.architectures import AUDIO_TEXT_ARCHITECTURES
150 kind = classify_architecture(arch)
151 if kind == "audio":
152 return {1, 8}, {1, 8}
153 if kind == "vision":
154 # Vision encoders have no tokenizer and no text tower: Phases 2/3 need
155 # HookedTransformer, Phase 4 needs text generation, and Phase 7 covers
156 # vision+text multimodal models, not these. Phase 1 (HF parity) plus
157 # Phase 9 (pixel forward/cache/stability) are the whole story.
158 return {1, 9}, {1, 9}
159 if kind == "multimodal":
160 return {1, 2, 3, 4, 7}, {1, 4, 7}
161 if arch in AUDIO_TEXT_ARCHITECTURES: 161 ↛ 164line 161 didn't jump to line 164 because the condition on line 161 was never true
162 # Phase 8 (audio-conditioned forward) out of the core set so a partial {1,4}
163 # run still verifies; a full run records and gates it via _check_phase_scores.
164 return {1, 2, 3, 4, 8}, {1, 4}
165 return {1, 2, 3, 4}, {1, 4}
168def _default_phases_for_architecture(arch: str) -> list[int]:
169 """Phases to run when the caller names none — a full verification."""
170 return sorted(_full_and_core_phases(arch)[0])
173def _pass_status(use_hf_reference: bool) -> int:
174 """Status for a passing run: VERIFIED with an HF reference, else PROVISIONAL
175 (a --no-hf-reference structural-only pass is recorded but not counted verified)."""
176 return STATUS_VERIFIED if use_hf_reference else STATUS_PROVISIONAL
179def _get_current_model_status(model_id: str, arch_id: str) -> int:
180 """Look up a model's current status in the registry.
182 Returns STATUS_UNVERIFIED (0) if the model is not found.
183 """
184 data = load_supported_models_raw()
185 for entry in data.get("models", []):
186 if not isinstance(entry, dict):
187 continue
188 if entry.get("model_id") == model_id and entry.get("architecture_id") == arch_id:
189 return entry.get("status", STATUS_UNVERIFIED)
190 return STATUS_UNVERIFIED
193@dataclass
194class ModelCandidate:
195 """A model selected for verification."""
197 model_id: str
198 architecture_id: str
199 estimated_params: Optional[int] = None
200 estimated_memory_gb: Optional[float] = None
203@dataclass
204class VerificationProgress:
205 """Tracks progress across a verification run."""
207 tested: list[str] = field(default_factory=list)
208 skipped: list[str] = field(default_factory=list)
209 failed: list[str] = field(default_factory=list)
210 verified: list[str] = field(default_factory=list)
211 # Structural-only (--no-hf-reference) passes; kept out of the verified tally.
212 provisional: list[str] = field(default_factory=list)
213 start_time: Optional[str] = None
215 def to_dict(self) -> dict:
216 return {
217 "tested": self.tested,
218 "skipped": self.skipped,
219 "failed": self.failed,
220 "verified": self.verified,
221 "provisional": self.provisional,
222 "start_time": self.start_time,
223 }
225 @classmethod
226 def from_dict(cls, data: dict) -> "VerificationProgress":
227 return cls(
228 tested=data.get("tested", []),
229 skipped=data.get("skipped", []),
230 failed=data.get("failed", []),
231 verified=data.get("verified", []),
232 provisional=data.get("provisional", []),
233 start_time=data.get("start_time"),
234 )
237def published_param_count(model_id: str) -> Optional[int]:
238 """Exact parameter count from the hub's safetensors metadata, or None.
240 Metadata only -- no weights are fetched. Preferred over the config formula
241 below, which assumes every layer carries full attention and an MLP and so
242 over-counts a hybrid Mamba/attention stack roughly fourfold
243 (NVIDIA-Nemotron-Nano-9B-v2: 36.6B estimated against 8.89B published, enough
244 to skip the model as too large for memory it does not need).
245 """
246 from transformer_lens.utilities.hf_utils import get_hf_token
248 try:
249 from huggingface_hub import HfApi
251 info = HfApi().model_info(model_id, expand=["safetensors"], token=get_hf_token())
252 except Exception:
253 # Unpublished metadata, a gated repo or a network blip: fall back rather
254 # than fail, since the config formula needs no hub metadata.
255 return None
256 total = getattr(info.safetensors, "total", None) if info.safetensors else None
257 return int(total) if total else None
260def estimate_model_params(model_id: str) -> int:
261 """Parameter count for this model: published metadata first, else a config estimate.
263 Fetches only metadata and the config JSON (~KB), never weights.
265 Args:
266 model_id: HuggingFace model ID
268 Returns:
269 Number of parameters — exact when the hub publishes it, else estimated
270 from config dimensions using the same formula as
271 HookedTransformerConfig.__post_init__.
273 Raises:
274 Exception: If config cannot be fetched or parsed
275 """
276 published = published_param_count(model_id)
277 if published:
278 return published
280 from transformer_lens.loading_from_pretrained import NEED_REMOTE_CODE_MODELS
282 _all_remote_prefixes = NEED_REMOTE_CODE_MODELS + _BRIDGE_REMOTE_CODE_PREFIXES
283 trust_remote_code = any(model_id.startswith(prefix) for prefix in _all_remote_prefixes)
284 from transformer_lens.utilities.hf_utils import (
285 autoconfig_with_remote_post_init_compat,
286 get_hf_token,
287 )
289 config = autoconfig_with_remote_post_init_compat(
290 model_id, trust_remote_code=trust_remote_code, token=get_hf_token()
291 )
293 # For multimodal models (LLaVA, Gemma3 multimodal), the language model config
294 # is nested under text_config. Fall through to the top-level config otherwise.
295 lang_config = getattr(config, "text_config", config)
297 # Encoder-decoder models (e.g. T5Gemma) nest dimensions under decoder/encoder
298 # subconfigs rather than the top level; prefer the decoder for the estimate.
299 if not (hasattr(lang_config, "hidden_size") or hasattr(lang_config, "d_model")): 299 ↛ 300line 299 didn't jump to line 300 because the condition on line 299 was never true
300 for _sub in ("decoder", "encoder"):
301 _subcfg = getattr(config, _sub, None)
302 if _subcfg is not None and (
303 hasattr(_subcfg, "hidden_size") or hasattr(_subcfg, "d_model")
304 ):
305 lang_config = _subcfg
306 break
308 # Heterogeneous configs (transformers>=5.15 Gemma 4) raise on global reads of
309 # per-layer fields like head_dim; the view resolves them to majority values.
310 lang_config = het_safe_view(lang_config)
312 # Extract dimensions from config (different models use different attribute names)
313 d_model = (
314 getattr(lang_config, "hidden_size", None)
315 or getattr(lang_config, "d_model", None)
316 or getattr(lang_config, "model_dim", None) # OpenELM
317 or 0
318 )
319 n_heads_raw = (
320 getattr(lang_config, "num_attention_heads", None)
321 or getattr(lang_config, "n_head", None)
322 or getattr(lang_config, "num_query_heads", None) # OpenELM (may be per-layer list)
323 or getattr(lang_config, "num_heads", None) # Mamba-2 SSM heads
324 or 0
325 )
326 # OpenELM uses per-layer lists for heads; take the max for estimation
327 n_heads = max(n_heads_raw) if isinstance(n_heads_raw, (list, tuple)) else n_heads_raw
328 n_layers = (
329 getattr(lang_config, "num_hidden_layers", None)
330 or getattr(lang_config, "n_layer", None)
331 or getattr(lang_config, "num_transformer_layers", None) # OpenELM
332 or 0
333 )
334 d_mlp = (
335 getattr(lang_config, "intermediate_size", None)
336 or getattr(lang_config, "d_inner", None)
337 or getattr(lang_config, "n_inner", None)
338 or getattr(lang_config, "ffn_dim", None) # OPT
339 or getattr(lang_config, "d_ff", None) # T5
340 )
341 # Gemma 3n exposes a per-layer intermediate_size list (uniform in all released
342 # checkpoints); collapse to max for the scalar param estimate.
343 if isinstance(d_mlp, (list, tuple)): 343 ↛ 344line 343 didn't jump to line 344 because the condition on line 343 was never true
344 d_mlp = max(d_mlp) if d_mlp else None
345 # OpenELM uses per-layer ffn_multipliers instead of a fixed intermediate_size
346 if not d_mlp and d_model: 346 ↛ 347line 346 didn't jump to line 347 because the condition on line 346 was never true
347 ffn_multipliers = getattr(lang_config, "ffn_multipliers", None)
348 if isinstance(ffn_multipliers, (list, tuple)):
349 d_mlp = int(max(ffn_multipliers) * d_model)
350 else:
351 # Many architectures (GPT-2, Bloom, GPT-Neo, GPT-J) leave d_mlp/n_inner
352 # as None and default to 4 * hidden_size internally.
353 d_mlp = 4 * d_model
354 d_vocab = getattr(lang_config, "vocab_size", None) or 0
356 if d_model == 0 or n_layers == 0: 356 ↛ 357line 356 didn't jump to line 357 because the condition on line 356 was never true
357 raise ValueError(f"Could not extract model dimensions from config for {model_id}")
359 # Attention-less architectures (Mamba SSMs) have no heads. Use nominal
360 # values so the estimate doesn't attribute phantom attention params.
361 is_attention_less = n_heads == 0
362 if is_attention_less: 362 ↛ 363line 362 didn't jump to line 363 because the condition on line 362 was never true
363 n_heads = 1
364 d_head = d_model
365 else:
366 d_head = getattr(lang_config, "head_dim", None) or (d_model // n_heads)
368 # Attention parameters: W_Q, W_K, W_V, W_O per layer (skipped for SSMs)
369 if is_attention_less: 369 ↛ 370line 369 didn't jump to line 370 because the condition on line 369 was never true
370 n_params = 0
371 else:
372 n_params = n_layers * (d_model * d_head * n_heads * 4)
374 # MLP parameters (if present)
375 if d_mlp is not None and d_mlp > 0: 375 ↛ 418line 375 didn't jump to line 418 because the condition on line 375 was always true
376 # Check for gated MLP (LLaMA, Gemma, Mistral, Qwen, T5 gated-gelu, etc.)
377 has_gate = getattr(lang_config, "is_gated_act", False) or (
378 hasattr(lang_config, "intermediate_size")
379 and (
380 getattr(lang_config, "hidden_act", None) in ("silu", "gelu", "swiglu")
381 or getattr(lang_config, "model_type", None)
382 in (
383 "llama",
384 "gemma",
385 "gemma2",
386 "gemma3",
387 "mistral",
388 "mixtral",
389 "qwen2",
390 "qwen3",
391 "qwen3_moe",
392 "phi3",
393 "stablelm",
394 )
395 )
396 )
397 mlp_multiplier = 3 if has_gate else 2
398 n_params += n_layers * (d_model * d_mlp * mlp_multiplier)
400 # MoE expert scaling
401 num_experts = (
402 getattr(lang_config, "num_local_experts", None)
403 or getattr(lang_config, "num_experts", None)
404 or getattr(lang_config, "n_routed_experts", None) # DeepSeek-V2/V3
405 )
406 if num_experts and num_experts > 1: 406 ↛ 409line 406 didn't jump to line 409 because the condition on line 406 was never true
407 # Qwen3MoE and similar store per-expert hidden size in moe_intermediate_size;
408 # intermediate_size refers to a dense fallback MLP that we don't use here.
409 moe_d_mlp = getattr(lang_config, "moe_intermediate_size", None) or d_mlp
410 # MLP params scale with num_experts; add gate params per expert
411 mlp_per_layer = d_model * moe_d_mlp * mlp_multiplier
412 moe_per_layer = (mlp_per_layer + d_model) * num_experts
413 # Replace the non-MoE MLP contribution
414 n_params -= n_layers * (d_model * d_mlp * mlp_multiplier)
415 n_params += n_layers * moe_per_layer
417 # Embedding parameters (not in HookedTransformerConfig formula but relevant for memory)
418 n_params += d_vocab * d_model
420 return n_params
423def estimate_benchmark_memory_gb(
424 n_params: int,
425 dtype: str = "float32",
426 phases: Optional[list[int]] = None,
427 use_hf_reference: bool = True,
428 device: str = "cpu",
429) -> float:
430 """Estimate peak memory needed for benchmark suite.
432 Phases run sequentially, so peak memory is the maximum of any single phase,
433 not the sum. The multiplier represents how many model copies exist at peak:
435 Phase 1 (HF ref on): HF ref + Bridge → 2.0x peak
436 Phase 1 (HF ref off): Bridge only → 1.0x peak
437 Phase 2: Bridge + HookedTransformer (separate copy) → 2.0x model + overhead
438 Phase 3: Same as Phase 2 (processed versions) → 2.0x model + overhead
439 Phase 4: Bridge + GPT-2 scorer (~500MB) → ~1.0x model + 0.5 GB
441 Args:
442 n_params: Number of model parameters
443 dtype: Data type for memory calculation
444 phases: Which phases will be run (None = all phases)
445 use_hf_reference: Whether Phase 1 loads an HF reference alongside the
446 Bridge. Mirrors the ``--no-hf-reference`` CLI flag.
448 Returns:
449 Estimated peak memory in GB
450 """
451 bytes_per_param = {"float32": 4, "float16": 2, "bfloat16": 2}
452 bpp = bytes_per_param.get(dtype, 4)
453 model_size_gb = n_params * bpp / (1024**3)
455 # Phase-4 judge overhead: measured 2.33 GB RSS loading Qwen2.5-0.5B fp32
456 # on CPU (494M params). Kept slightly above the measurement; over-counting
457 # is the safe direction.
458 # The CPU-pinned judge never occupies accelerator memory; charging it to
459 # a cuda budget produces spurious VRAM skips.
460 judge_overhead_gb = 2.5 if device == "cpu" else 0.0
462 # Activation/framework overhead as a fraction of model size
463 overhead_fraction = 0.2
465 # Determine peak memory across all requested phases
466 phase_peaks = []
468 if phases is None: 468 ↛ 469line 468 didn't jump to line 469 because the condition on line 468 was never true
469 phases = [1, 2, 3, 4]
471 for p in phases:
472 if p == 1:
473 # HF ref + Bridge (2 copies) or Bridge alone
474 multiplier = 2.0 if use_hf_reference else 1.0
475 phase_peaks.append(model_size_gb * multiplier * (1 + overhead_fraction))
476 elif p in (2, 3): 476 ↛ 478line 476 didn't jump to line 478 because the condition on line 476 was never true
477 # Bridge + HookedTransformer = 2 copies
478 phase_peaks.append(model_size_gb * 2.0 * (1 + overhead_fraction))
479 elif p == 4: 479 ↛ 471line 479 didn't jump to line 471 because the condition on line 479 was always true
480 # Bridge + judge
481 phase_peaks.append(model_size_gb * (1 + overhead_fraction) + judge_overhead_gb)
483 return max(phase_peaks) if phase_peaks else model_size_gb
486def get_available_memory_gb(device: str) -> float:
487 """Detect available memory on the target device.
489 Args:
490 device: "cpu" or "cuda"
492 Returns:
493 Available memory in GB
494 """
495 if device.startswith("cuda"):
496 try:
497 import torch
499 if torch.cuda.is_available():
500 device_idx = 0
501 if ":" in device:
502 device_idx = int(device.split(":")[1])
503 props = torch.cuda.get_device_properties(device_idx)
504 return props.total_memory / (1024**3)
505 except Exception:
506 pass
507 return 8.0 # Conservative default for GPU
509 # CPU: use psutil if available, else conservative default
510 try:
511 import psutil
513 return psutil.virtual_memory().available / (1024**3)
514 except ImportError:
515 return 16.0 # Conservative default for CPU
518def select_models_for_verification(
519 per_arch: int = 10,
520 architectures: Optional[list[str]] = None,
521 limit: Optional[int] = None,
522 resume_progress: Optional[VerificationProgress] = None,
523 retry_failed: bool = False,
524 reverify: bool = False,
525) -> list[ModelCandidate]:
526 """Select models for verification from the registry.
528 Loads supported_models.json (already sorted by downloads).
529 Takes the top N unverified models per architecture.
531 Args:
532 per_arch: Maximum models to verify per architecture
533 architectures: Filter to specific architectures (None = all)
534 limit: Total model cap (None = no cap)
535 resume_progress: If resuming, skip already-tested models
536 retry_failed: If True, include previously failed models for re-testing
537 reverify: If True, ignore previous status and re-test all matching models
539 Returns:
540 List of ModelCandidate objects to verify
541 """
542 already_tested: set[str] = set()
543 if resume_progress and not reverify:
544 already_tested = set(resume_progress.tested)
545 if retry_failed:
546 # Remove failed models from already_tested so they get re-selected
547 failed_set = set(resume_progress.failed)
548 already_tested -= failed_set
550 data = load_supported_models_raw()
551 models = data.get("models", [])
553 # Group by architecture
554 by_arch: dict[str, list[dict]] = {}
555 for model in models:
556 arch = model["architecture_id"]
557 by_arch.setdefault(arch, []).append(model)
559 # Determine which architectures to scan
560 if architectures:
561 arch_ids = architectures
562 else:
563 arch_ids = sorted(by_arch.keys())
565 candidates: list[ModelCandidate] = []
567 for arch in arch_ids:
568 arch_models = by_arch.get(arch, [])
569 count = 0
571 for model in arch_models:
572 model_id = model["model_id"]
574 # Skip already-verified or already-tested models
575 if not reverify:
576 model_status = model.get("status", 0)
577 if model_status == STATUS_VERIFIED or model_status == STATUS_SKIPPED:
578 continue
579 if model_status == STATUS_FAILED and not retry_failed:
580 continue
581 if model_id in already_tested:
582 continue
584 # Check per-arch limit
585 if count >= per_arch:
586 break
588 count += 1
589 candidates.append(ModelCandidate(model_id=model_id, architecture_id=arch))
591 # Check total limit
592 if limit and len(candidates) >= limit:
593 return candidates
595 return candidates
598def _extract_phase_scores(results: list) -> dict[int, Optional[float]]:
599 """Extract phase scores from benchmark results.
601 Mirrors the logic in update_model_registry() from main_benchmark.py.
603 Args:
604 results: List of BenchmarkResult objects
606 Returns:
607 Dict mapping phase number to score (0-100) or None
608 """
609 from transformer_lens.benchmarks.utils import BenchmarkSeverity
611 phase_results: dict[int, list[bool]] = {1: [], 2: [], 3: [], 4: [], 7: [], 8: [], 9: []}
612 for result in results:
613 if result.phase in phase_results and result.severity != BenchmarkSeverity.SKIPPED:
614 phase_results[result.phase].append(result.passed)
616 scores: dict[int, Optional[float]] = {}
617 for phase, passed_list in phase_results.items():
618 if passed_list:
619 scores[phase] = round(sum(passed_list) / len(passed_list) * 100, 1)
620 # Omit phases with no results — they weren't run, so their
621 # existing registry scores should be preserved.
623 # Phase 4 (text quality): store the actual 0-100 quality score from the
624 # benchmark details instead of a binary pass/fail percentage.
625 if 4 in scores:
626 for result in results: 626 ↛ 631line 626 didn't jump to line 631 because the loop on line 626 didn't complete
627 if result.phase == 4 and result.details and "score" in result.details:
628 scores[4] = round(result.details["score"], 1)
629 break
631 return scores
634def _extract_prompt_profile(results: list) -> Optional[str]:
635 """Effective Phase-4 prompt profile from the benchmark details, or None
636 when no Phase-4 result exists. The default "continuation" is reported so
637 the registry write can clear a stale non-default key."""
638 for result in results:
639 if result.phase == 4 and result.details:
640 profile = result.details.get("prompt_profile")
641 if isinstance(profile, str): 641 ↛ 638line 641 didn't jump to line 638 because the condition on line 641 was always true
642 return profile
643 return None
646# Per-phase minimum score thresholds (0-100).
647# Phase 1: Core correctness (bridge vs HF) — must pass everything.
648# Phase 2: Hook/cache/gradient tests — most should pass.
649# Phase 3: Weight processing tests — most should pass.
650# Phase 4: Text quality — inherently fuzzy, keep lenient.
651_MIN_PHASE_SCORES: dict[int, float] = {
652 1: 100.0,
653 2: 75.0,
654 3: 75.0,
655 # Phase 4 floor == the benchmark pass line; a gap between them lets a
656 # failing score carry a clean "completed" note.
657 4: p4_pass_threshold(),
658 7: 75.0,
659 8: 75.0,
660 9: 75.0,
661}
662_DEFAULT_MIN_PHASE_SCORE = 50.0
664# Architectures that include a vision encoder and require Phase 7 (multimodal
665# benchmarks) as part of core verification.
666from transformer_lens.utilities.architectures import classify_architecture
668_AUDIO_ARCHITECTURES = {
669 "HubertForCTC",
670 "HubertModel",
671 "HubertForSequenceClassification",
672}
674# Tests that MUST pass for a phase to be considered passing, regardless of
675# the overall percentage score. If any required test fails, the phase fails
676# even if the score is above the minimum threshold.
677_REQUIRED_PHASE_TESTS: dict[int, list[str]] = {
678 2: ["logits_equivalence", "loss_equivalence"],
679 3: ["logits_equivalence", "loss_equivalence"],
680 7: ["multimodal_forward"],
681 8: ["audio_forward", "audio_text_forward"],
682 9: ["vision_forward", "vision_cache"],
683}
685# Failure text for a modality phase that produced no score — either absent from
686# phase_scores (all tests skipped) or explicitly NULL.
687_MODALITY_NULL_MESSAGES: dict[int, str] = {
688 7: "P7=NULL (multimodal tests skipped — processor unavailable)",
689 8: "P8=NULL (audio tests skipped — no results)",
690 9: "P9=NULL (vision tests skipped — no results)",
691}
694def _measured_nothing(phase_scores: dict) -> bool:
695 """True when no phase produced a score, so the run verified nothing.
697 An adapter's ``applicable_phases`` can prune the requested phases to empty.
698 That leaves ``required_phases`` empty as well, so no score check fires and a
699 run that measured nothing would otherwise be recorded as VERIFIED.
701 ``is not None`` rather than truthiness: 0.0 is a real score.
702 """
703 return not any(score is not None for score in phase_scores.values())
706def _check_phase_scores(
707 phase_scores: dict[int, Optional[float]],
708 all_results: list,
709 required_phases: Optional[set[int]] = None,
710) -> Optional[str]:
711 """Check phase scores against per-phase minimum thresholds and required tests.
713 A phase fails if:
714 1. Its overall score is below the minimum threshold, OR
715 2. Any of its required tests (per _REQUIRED_PHASE_TESTS) failed.
717 Phase 4 (text quality) is excluded — it is a quality metric, not a
718 correctness check. Low text quality is surfaced in the verification
719 note via _build_verified_note() but never causes a model to fail.
721 Args:
722 phase_scores: Per-phase scores; a phase whose tests all skipped is
723 absent entirely (``extract_phase_scores`` omits empty phases).
724 all_results: Benchmark results, used to name the failing tests.
725 required_phases: Phases this architecture must produce a score for —
726 the core set from ``_full_and_core_phases``. A required modality
727 phase that is absent counts as NULL, not as a pass.
729 Returns an error message if any phase fails, or None if all phases pass.
730 The message includes the names of failed tests.
731 """
732 from transformer_lens.benchmarks.utils import BenchmarkSeverity
734 failing_phases: list[str] = []
736 # A required modality phase whose tests all skipped never reaches phase_scores.
737 for phase in sorted((required_phases or set()) - set(phase_scores)):
738 if phase in _MODALITY_NULL_MESSAGES: 738 ↛ 737line 738 didn't jump to line 737 because the condition on line 738 was always true
739 failing_phases.append(_MODALITY_NULL_MESSAGES[phase])
741 for phase, score in sorted(phase_scores.items()):
742 if score is None: 742 ↛ 746line 742 didn't jump to line 746 because the condition on line 742 was never true
743 # Phase 7 (multimodal), Phase 8 (audio), or 9 (vision) with a NULL score means
744 # the modality tests never ran. This is a verification failure,
745 # not something to silently skip.
746 if phase in _MODALITY_NULL_MESSAGES:
747 failing_phases.append(_MODALITY_NULL_MESSAGES[phase])
748 continue
750 # Phase 4 is a quality metric, not a pass/fail check — skip it here.
751 # Low text quality is reported in the note by _build_verified_note().
752 if phase == 4:
753 continue
755 # Check 1: overall score threshold
756 threshold = _MIN_PHASE_SCORES.get(phase, _DEFAULT_MIN_PHASE_SCORE)
757 if score < threshold:
758 failed_tests = [
759 r.name
760 for r in all_results
761 if r.phase == phase and not r.passed and r.severity != BenchmarkSeverity.SKIPPED
762 ]
763 tests_str = ", ".join(failed_tests) if failed_tests else "unknown"
764 failing_phases.append(f"P{phase}={score}% < {threshold}% (failed: {tests_str})")
765 continue # Already failing; no need to also check required tests
767 # Check 2: required tests must pass
768 required_tests = _REQUIRED_PHASE_TESTS.get(phase, [])
769 if required_tests:
770 failed_required = [
771 r.name
772 for r in all_results
773 if r.phase == phase
774 and r.name in required_tests
775 and not r.passed
776 and r.severity != BenchmarkSeverity.SKIPPED
777 ]
778 if failed_required: 778 ↛ 779line 778 didn't jump to line 779 because the condition on line 778 was never true
779 tests_str = ", ".join(failed_required)
780 failing_phases.append(f"P{phase}={score}% but required tests failed: {tests_str}")
782 if failing_phases:
783 return f"Below threshold: {'; '.join(failing_phases)}"
784 return None
787def _build_verified_note(
788 phase_scores: dict[int, Optional[float]],
789 all_results: list,
790) -> str:
791 """Build a verification note summarizing phase scores.
793 Phase 4 (text quality) is excluded from the score summary since it's a
794 quality metric, not a pass/fail comparison. It only contributes a "low
795 text quality" flag when below threshold.
796 """
797 from transformer_lens.benchmarks.utils import BenchmarkSeverity
799 issue_parts: list[str] = []
800 low_text_quality = False
802 for phase in sorted(phase_scores):
803 score = phase_scores[phase]
804 if score is None: 804 ↛ 805line 804 didn't jump to line 805 because the condition on line 804 was never true
805 continue
806 # Phase 4 is a quality score, not a pass/fail comparison — don't
807 # include it in the normal score summary.
808 if phase == 4:
809 threshold = _MIN_PHASE_SCORES.get(4, _DEFAULT_MIN_PHASE_SCORE)
810 if score < threshold: 810 ↛ 811line 810 didn't jump to line 811 because the condition on line 810 was never true
811 low_text_quality = True
812 continue
814 if score < 100.0: 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true
815 failed_tests = [
816 r.name
817 for r in all_results
818 if r.phase == phase and not r.passed and r.severity != BenchmarkSeverity.SKIPPED
819 ]
820 if failed_tests:
821 issue_parts.append(f"P{phase}={score}% (failed: {', '.join(failed_tests)})")
822 else:
823 issue_parts.append(f"P{phase}={score}%")
825 p4_uncovered = next(
826 (
827 r.message
828 for r in all_results
829 if r.phase == 4
830 and r.severity == BenchmarkSeverity.SKIPPED
831 and r.message.startswith("P4 skipped:")
832 ),
833 None,
834 )
835 suffix = ""
836 if p4_uncovered: 836 ↛ 838line 836 didn't jump to line 838 because the condition on line 836 was never true
837 # Keep the gap visible in the registry until prompt coverage is added.
838 reason = p4_uncovered.split("—")[0].replace("P4 skipped:", "").strip()
839 suffix = f"; P4 skipped (uncovered: {reason} — file a coverage issue)"
841 if issue_parts and low_text_quality: 841 ↛ 842line 841 didn't jump to line 842 because the condition on line 841 was never true
842 return (
843 f"Full verification completed with issues, low text quality: {'; '.join(issue_parts)}"
844 + suffix
845 )
846 if issue_parts: 846 ↛ 847line 846 didn't jump to line 847 because the condition on line 846 was never true
847 return f"Full verification completed with issues: {'; '.join(issue_parts)}" + suffix
848 if low_text_quality: 848 ↛ 849line 848 didn't jump to line 849 because the condition on line 848 was never true
849 return "Full verification completed with issues, low text quality" + suffix
850 return "Full verification completed" + suffix
853def _preserved_issue_suffix(model_id: str, eff_phases) -> str:
854 """Sub-100 scores from phases not re-run this pass stay visible in the
855 note; a partial pass must not overwrite tracked residue."""
856 from transformer_lens.tools.model_registry.registry_io import (
857 load_supported_models_raw,
858 )
860 try:
861 entry = next(
862 (
863 m
864 for m in load_supported_models_raw().get("models", [])
865 if m.get("model_id") == model_id
866 ),
867 None,
868 )
869 except OSError:
870 return ""
871 if entry is None: 871 ↛ 872line 871 didn't jump to line 872 because the condition on line 871 was never true
872 return ""
873 residue = []
874 for phase in (2, 3, 7, 8, 9):
875 if phase in (eff_phases or []):
876 continue
877 score = entry.get(f"phase{phase}_score")
878 if score is not None and score < 100.0:
879 residue.append(f"P{phase}={score}%")
880 if not residue:
881 return ""
882 return f" (prior issues retained: {', '.join(residue)})"
885def _p1_only_core_note(p4_score, all_results: list) -> str:
886 """Note for a core run where P1 passed but P4 did not contribute a pass.
888 A skipped P4 is a coverage gap, not a quality failure — the stale
889 (possibly old-scale) score must not be relabeled "poor"."""
890 from transformer_lens.benchmarks.utils import BenchmarkSeverity
892 p4_skip_msg = next(
893 (
894 r.message
895 for r in all_results
896 if r.phase == 4
897 and r.severity == BenchmarkSeverity.SKIPPED
898 and r.message.startswith("P4 skipped:")
899 ),
900 None,
901 )
902 if p4_skip_msg is not None:
903 reason = p4_skip_msg.split("—")[0].replace("P4 skipped:", "").strip()
904 return f"Core verification passed; P4 skipped ({reason})"
905 if p4_score is None:
906 return "Core verification passed, but text quality benchmark errored. Needs review"
907 return f"Core verification passed, but text quality poor (P4={p4_score}). Needs review"
910def _clear_hf_cache(quiet: bool = False) -> None:
911 """Remove downloaded model weights from the HuggingFace cache to free disk."""
912 from pathlib import Path
914 cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
915 if not cache_dir.exists(): 915 ↛ 916line 915 didn't jump to line 916 because the condition on line 915 was never true
916 return
918 from transformer_lens.benchmarks.text_quality import JUDGE_MODEL_ID
920 # The pinned Phase-4 judge is needed by every run; deleting it here would
921 # force a re-download per family.
922 judge_dir = "models--" + JUDGE_MODEL_ID.replace("/", "--")
924 freed = 0
925 for blobs_dir in cache_dir.glob("models--*/blobs"):
926 if blobs_dir.parent.name == judge_dir:
927 continue
928 for blob in blobs_dir.iterdir():
929 try:
930 size = blob.stat().st_size
931 blob.unlink()
932 freed += size
933 except OSError:
934 pass
936 if not quiet and freed > 0: 936 ↛ 937line 936 didn't jump to line 937 because the condition on line 936 was never true
937 print(f" Cleared {freed / (1024**3):.1f} GB from HuggingFace cache")
940def _save_checkpoint(progress: VerificationProgress) -> None:
941 """Save verification progress to checkpoint file."""
942 with open(_CHECKPOINT_PATH, "w") as f:
943 json.dump(progress.to_dict(), f, indent=2)
944 f.write("\n")
947def _skip_model(
948 model_id: str, arch: str, note: str, progress: VerificationProgress, quiet: bool
949) -> None:
950 """Record a model as skipped with ``note``, preserving an existing verified or provisional
951 status, and checkpoint. Callers ``continue`` the loop afterwards.
952 """
953 if not quiet:
954 print(f" SKIP: {note}")
955 if _get_current_model_status(model_id, arch) not in (STATUS_VERIFIED, STATUS_PROVISIONAL):
956 update_model_status(model_id, arch, STATUS_SKIPPED, note=note, sanitize_fn=_sanitize_note)
957 elif not quiet:
958 print(" (preserving existing verified/provisional status)")
959 progress.skipped.append(model_id)
960 _save_checkpoint(progress)
963def _load_checkpoint() -> Optional[VerificationProgress]:
964 """Load verification progress from checkpoint file."""
965 if not _CHECKPOINT_PATH.exists():
966 return None
967 try:
968 with open(_CHECKPOINT_PATH) as f:
969 data = json.load(f)
970 return VerificationProgress.from_dict(data)
971 except (json.JSONDecodeError, KeyError):
972 return None
975def verify_models(
976 candidates: list[ModelCandidate],
977 device: str = "cpu",
978 max_memory_gb: Optional[float] = None,
979 dtype: str = "float32",
980 use_hf_reference: bool = True,
981 use_ht_reference: bool = True,
982 phases: Optional[list[int]] = None,
983 quiet: bool = False,
984 progress: Optional[VerificationProgress] = None,
985) -> VerificationProgress:
986 """Run verification benchmarks on a list of model candidates.
988 Args:
989 candidates: Models to verify
990 device: Device for benchmarks
991 max_memory_gb: Memory limit (auto-detected if None)
992 dtype: Dtype for memory estimation
993 use_hf_reference: Whether to compare against HuggingFace model
994 use_ht_reference: Whether to compare against HookedTransformer
995 phases: Which benchmark phases to run
996 quiet: Suppress verbose output
997 progress: Existing progress for resume
999 Returns:
1000 VerificationProgress with results
1001 """
1002 from transformer_lens.benchmarks.main_benchmark import run_benchmark_suite
1004 if progress is None:
1005 progress = VerificationProgress(start_time=datetime.now().isoformat())
1007 if max_memory_gb is None:
1008 max_memory_gb = get_available_memory_gb(device)
1009 if not quiet:
1010 print(f"Auto-detected available memory: {max_memory_gb:.1f} GB")
1012 # phases stays None = full verification for the model.
1014 # Pre-load the Phase-4 judge so it persists across all models in the batch
1015 # instead of being loaded and destroyed for each one.
1016 _judge_model = None
1017 _judge_tokenizer = None
1018 if phases is None or 4 in phases:
1019 try:
1020 from transformer_lens.benchmarks.text_quality import (
1021 JUDGE_MODEL_ID,
1022 JUDGE_REVISION,
1023 load_judge,
1024 )
1026 _judge_model, _judge_tokenizer = load_judge()
1027 if not quiet:
1028 print(f"Pre-loaded Phase 4 judge {JUDGE_MODEL_ID}@{JUDGE_REVISION[:8]}")
1029 except Exception as e:
1030 if not quiet:
1031 print(f"Warning: Could not pre-load Phase 4 judge: {e}")
1032 print(" Phase 4 will load its own judge per model.")
1034 total = len(candidates)
1035 for i, candidate in enumerate(candidates, 1):
1036 # Check for graceful interrupt between models
1037 if _interrupt_requested:
1038 if not quiet:
1039 print(
1040 f"\nStopping gracefully. Progress saved "
1041 f"({len(progress.verified)} verified, "
1042 f"{len(progress.provisional)} provisional)."
1043 )
1044 _save_checkpoint(progress)
1045 raise SystemExit(_EXIT_GRACEFUL_INTERRUPT)
1047 model_id = candidate.model_id
1048 arch = candidate.architecture_id
1050 if not quiet:
1051 print(f"\n{'='*70}")
1052 print(f"[{i}/{total}] {model_id} ({arch})")
1053 print(f"{'='*70}")
1055 progress.tested.append(model_id)
1057 # Step 0: Skip formats with no HF loader path (GGUF / MLX / FP4 / FP8).
1058 if is_incompatible_quantized(model_id):
1059 _skip_model(model_id, arch, QUANTIZED_NOTE, progress, quiet)
1060 continue
1062 # Step 0a: skip HF-loadable quantized models when their loader lib is missing.
1063 required_lib = required_quant_library_for_model(model_id)
1064 if required_lib is not None:
1065 import importlib.util
1067 if importlib.util.find_spec(required_lib) is None:
1068 note = f"Skipped: {required_lib} not installed (required to load this quantized format)"
1069 _skip_model(model_id, arch, note, progress, quiet)
1070 continue
1072 # Step 0b: Check adapter-level phase applicability. applicable_phases=[]
1073 # means no phase applies (a genuinely unsupported architecture — rare);
1074 # SSM / recurrent families run the full [1, 2, 3, 4].
1075 from transformer_lens.factories.architecture_adapter_factory import (
1076 SUPPORTED_ARCHITECTURES,
1077 )
1079 adapter_cls = SUPPORTED_ARCHITECTURES.get(arch)
1080 eff_phases = phases if phases is not None else _default_phases_for_architecture(arch)
1081 phases_to_run = _phases_to_run(arch, eff_phases)
1082 if adapter_cls is not None and not phases_to_run:
1083 applicable = getattr(adapter_cls, "applicable_phases", [1, 2, 3, 4])
1084 note = (
1085 f"Architecture {arch} has applicable_phases={applicable}; "
1086 f"verify_models coverage is deferred. Verification lives "
1087 f"in integration tests."
1088 )
1089 _skip_model(model_id, arch, note, progress, quiet)
1090 continue
1092 # Step 1: Estimate parameters
1093 try:
1094 n_params = estimate_model_params(model_id)
1095 candidate.estimated_params = n_params
1096 if not quiet:
1097 print(f" Estimated parameters: {n_params:,}")
1098 except Exception as e:
1099 _skip_model(model_id, arch, f"Config unavailable: {str(e)[:200]}", progress, quiet)
1100 continue
1102 # Step 2: Check memory
1103 estimated_mem = estimate_benchmark_memory_gb(
1104 n_params, dtype, phases=phases_to_run, use_hf_reference=use_hf_reference, device=device
1105 )
1106 candidate.estimated_memory_gb = estimated_mem
1107 if not quiet:
1108 print(
1109 f" Estimated benchmark memory: {estimated_mem:.1f} GB (limit: {max_memory_gb:.1f} GB)"
1110 )
1112 if estimated_mem > max_memory_gb:
1113 note = f"Estimated {estimated_mem:.1f} GB exceeds {max_memory_gb:.1f} GB limit"
1114 _skip_model(model_id, arch, note, progress, quiet)
1115 continue
1117 # Step 3: Run benchmarks (all phases in a single call to share models)
1118 all_results: list = []
1119 error_msg: Optional[str] = None
1121 from transformer_lens.loading_from_pretrained import NEED_REMOTE_CODE_MODELS
1123 _all_remote_prefixes = NEED_REMOTE_CODE_MODELS + _BRIDGE_REMOTE_CODE_PREFIXES
1124 needs_remote_code = any(model_id.startswith(prefix) for prefix in _all_remote_prefixes)
1126 # Convert string dtype to torch.dtype for benchmark suite
1127 import torch
1129 _dtype_map = {
1130 "float32": torch.float32,
1131 "float16": torch.float16,
1132 "bfloat16": torch.bfloat16,
1133 }
1134 torch_dtype = _dtype_map[dtype]
1136 from transformer_lens.benchmarks.text_quality_profiles import resolve_profile
1137 from transformer_lens.tools.model_registry.registry_io import (
1138 registry_prompt_profile,
1139 )
1141 resolved_profile = str(resolve_profile(model_id, arch, registry_prompt_profile(model_id)))
1142 if not quiet:
1143 print(f" Prompt profile: {resolved_profile}")
1144 print(f" Running phases {phases} in a single benchmark call...")
1145 try:
1146 all_results = run_benchmark_suite(
1147 model_id,
1148 device=device,
1149 dtype=torch_dtype,
1150 use_hf_reference=use_hf_reference,
1151 use_ht_reference=use_ht_reference,
1152 verbose=not quiet,
1153 phases=phases_to_run,
1154 trust_remote_code=needs_remote_code,
1155 judge_model=_judge_model,
1156 judge_tokenizer=_judge_tokenizer,
1157 prompt_profile=resolved_profile,
1158 )
1159 except Exception as e:
1160 error_msg = str(e)
1161 if not quiet:
1162 print(f" Benchmark failed: {error_msg[:200]}")
1164 phase_scores = _extract_phase_scores(all_results)
1166 if not error_msg:
1167 # Only require the core phases this run actually requested, so a
1168 # partial run (e.g. --phases 1 2) isn't failed for a missing P7.
1169 _, core_for_arch = _full_and_core_phases(arch)
1170 score_error = _check_phase_scores(
1171 phase_scores, all_results, required_phases=core_for_arch & set(eff_phases)
1172 )
1173 if score_error:
1174 error_msg = score_error
1176 if error_msg:
1177 is_oom = "out of memory" in error_msg.lower() or "oom" in error_msg.lower()
1178 if is_oom:
1179 note = "OOM during benchmark"
1180 else:
1181 # Include the specific error from failed results (e.g., tokenizer
1182 # errors, load failures) so the note explains WHY it failed.
1183 root_errors = [r.message for r in all_results if not r.passed and r.message]
1184 if root_errors:
1185 # Deduplicate and use first unique error as the detail
1186 unique_errors = list(dict.fromkeys(root_errors))
1187 detail = unique_errors[0][:150]
1188 note = f"{error_msg[:100]} — {detail}"
1189 else:
1190 note = error_msg[:200]
1191 final_status = STATUS_FAILED
1192 else:
1193 note = _build_verified_note(phase_scores, all_results)
1194 final_status = STATUS_VERIFIED
1196 # When running a partial phase set (e.g., --phases 4 for backfill),
1197 # only update the phase scores that were run. Don't change the
1198 # model's overall status or note — those reflect the full
1199 # verification and should only be set by a complete run.
1200 is_multimodal = classify_architecture(arch) == "multimodal"
1201 is_audio = classify_architecture(arch) == "audio"
1202 full_phases, core_required = _full_and_core_phases(arch)
1203 is_partial_run = set(eff_phases) != full_phases
1205 if is_partial_run and phase_scores:
1206 # Only write scores for phases that were actually requested.
1207 # Bridge load failures can produce Phase 1-tagged error results
1208 # even during Phase 4-only runs — don't let those corrupt
1209 # existing scores for unrequested phases.
1210 filtered_scores = {p: s for p, s in phase_scores.items() if p in eff_phases}
1211 if filtered_scores:
1212 if not quiet:
1213 score_parts = [f"P{p}={s}%" for p, s in sorted(filtered_scores.items())]
1214 print(f" Partial phase update: {', '.join(score_parts)}")
1216 # Core verification: P1+P4 for text-only, P1+P4+P7 for multimodal.
1217 is_core_verification = set(eff_phases) >= core_required
1218 partial_status = None
1219 partial_note = None
1221 if is_core_verification:
1222 p1 = filtered_scores.get(1)
1223 p4 = filtered_scores.get(4)
1224 p1_pass = p1 is not None and p1 >= _MIN_PHASE_SCORES.get(
1225 1, _DEFAULT_MIN_PHASE_SCORE
1226 )
1227 p4_pass = p4 is not None and p4 >= _MIN_PHASE_SCORES.get(
1228 4, _DEFAULT_MIN_PHASE_SCORE
1229 )
1231 # For multimodal, Phase 7 is required. A score below 75%
1232 # or a missing score (NULL — processor unavailable) both
1233 # count as failures.
1234 p7_pass = True
1235 if is_multimodal:
1236 p7 = filtered_scores.get(7)
1237 if p7 is not None:
1238 p7_pass = p7 >= _MIN_PHASE_SCORES.get(7, _DEFAULT_MIN_PHASE_SCORE)
1239 else:
1240 p7_pass = False
1242 # For audio models, Phase 8 is required; Phase 4 is not applicable
1243 p8_pass = True
1244 if is_audio:
1245 p4_pass = True # Audio models skip text quality
1246 p8 = filtered_scores.get(8)
1247 if p8 is not None:
1248 p8_pass = p8 >= _MIN_PHASE_SCORES.get(8, _DEFAULT_MIN_PHASE_SCORE)
1249 else:
1250 p8_pass = False
1252 if p1_pass and p4_pass and p7_pass and p8_pass:
1253 partial_status = STATUS_VERIFIED
1254 partial_note = "Core verification completed" + _preserved_issue_suffix(
1255 model_id, eff_phases
1256 )
1257 elif p1_pass and p4_pass and not p7_pass:
1258 p7_score = filtered_scores.get(7)
1259 if p7_score is None:
1260 partial_status = STATUS_FAILED
1261 partial_note = (
1262 "Core verification failed: multimodal tests skipped "
1263 "(processor unavailable)"
1264 )
1265 else:
1266 partial_status = STATUS_FAILED
1267 partial_note = (
1268 f"Core verification failed: multimodal tests "
1269 f"scored {p7_score}% (requires >= 75%)"
1270 )
1271 elif p1_pass:
1272 partial_status = STATUS_VERIFIED
1273 partial_note = _p1_only_core_note(p4, all_results)
1274 else:
1275 # P1 failed — build a descriptive failure note
1276 partial_status = STATUS_FAILED
1277 if error_msg:
1278 partial_note = f"CORE FAILED: {error_msg[:200]}"
1279 else:
1280 # Score-based failure — include details
1281 from transformer_lens.benchmarks.utils import (
1282 BenchmarkSeverity,
1283 )
1285 failed_tests = [
1286 r.name
1287 for r in all_results
1288 if r.phase == 1
1289 and not r.passed
1290 and r.severity != BenchmarkSeverity.SKIPPED
1291 ]
1292 tests_str = ", ".join(failed_tests) if failed_tests else "unknown"
1293 partial_note = f"CORE FAILED: P1={p1}% (failed: {tests_str})"
1295 # A structural-only core pass (--no-hf-reference) is provisional,
1296 # never numerically compared to HF, so it must not count as verified.
1297 if partial_status == STATUS_VERIFIED and not use_hf_reference:
1298 partial_status = STATUS_PROVISIONAL
1299 partial_note = f"Structural only (no HF reference): {partial_note}"
1301 if not quiet:
1302 print(f" {partial_note}")
1304 update_model_status(
1305 model_id,
1306 arch,
1307 status=partial_status,
1308 phase_scores=filtered_scores,
1309 note=partial_note,
1310 prompt_profile=_extract_prompt_profile(all_results),
1311 )
1312 # A provisional run was not numerically verified; do not write a
1313 # verification-history record (VerificationHistory.is_verified()
1314 # treats any record as verified — a second "counts as verified" path).
1315 if partial_status != STATUS_PROVISIONAL:
1316 add_verification_record(
1317 model_id,
1318 arch,
1319 notes=partial_note,
1320 sanitize_fn=_sanitize_note,
1321 prompt_profile=_extract_prompt_profile(all_results),
1322 p4_scoring_version=(P4_SCORING_VERSION if 4 in filtered_scores else None),
1323 )
1324 if partial_status == STATUS_FAILED:
1325 progress.failed.append(model_id)
1326 elif partial_status == STATUS_PROVISIONAL:
1327 progress.provisional.append(model_id)
1328 elif partial_status is None:
1329 # Scores were written but the status deliberately was not:
1330 # this phase set cannot establish core verification for
1331 # this architecture class. Reporting it as verified is how
1332 # a model ends up with all-pass scores and status 0.
1333 if not quiet:
1334 print(
1335 f" Scores updated, status unchanged — core verification for "
1336 f"{arch} requires phases {sorted(core_required)}, "
1337 f"this run had {sorted(eff_phases)}."
1338 )
1339 progress.skipped.append(model_id)
1340 else:
1341 progress.verified.append(model_id)
1342 else:
1343 if not quiet:
1344 print(f" No results for requested phases {eff_phases} — skipping update")
1345 progress.skipped.append(model_id)
1346 elif final_status == STATUS_VERIFIED and _measured_nothing(phase_scores):
1347 if not quiet:
1348 print(
1349 f" No phase produced a score (requested {eff_phases}) — "
1350 f"status left unchanged"
1351 )
1352 progress.skipped.append(model_id)
1353 elif final_status == STATUS_VERIFIED:
1354 # A passing run is VERIFIED only if it was numerically compared to an
1355 # HF reference; a --no-hf-reference (structural-only) pass is PROVISIONAL.
1356 written_status = _pass_status(use_hf_reference)
1357 is_provisional = written_status == STATUS_PROVISIONAL
1358 if is_provisional:
1359 note = f"Structural only (no HF reference): {note}"
1360 if not quiet:
1361 label = "PROVISIONAL" if is_provisional else "VERIFIED"
1362 print(
1363 f" {label}: P1={phase_scores.get(1)}%, "
1364 f"P2={phase_scores.get(2)}%, P3={phase_scores.get(3)}%, "
1365 f"P4={phase_scores.get(4)}%, P7={phase_scores.get(7)}%, "
1366 f"P8={phase_scores.get(8)}%, P9={phase_scores.get(9)}%"
1367 )
1368 update_model_status(
1369 model_id,
1370 arch,
1371 written_status,
1372 phase_scores=phase_scores,
1373 note=note,
1374 prompt_profile=_extract_prompt_profile(all_results),
1375 )
1376 # Provisional runs are not numerically verified — no history record
1377 # (is_verified() would otherwise report them as verified).
1378 if not is_provisional:
1379 add_verification_record(
1380 model_id,
1381 arch,
1382 notes=note,
1383 prompt_profile=_extract_prompt_profile(all_results),
1384 p4_scoring_version=(P4_SCORING_VERSION if 4 in phase_scores else None),
1385 )
1386 if is_provisional:
1387 progress.provisional.append(model_id)
1388 else:
1389 progress.verified.append(model_id)
1390 else:
1391 if not quiet:
1392 print(f" FAILED: {note}")
1393 if any(v is not None for v in phase_scores.values()):
1394 print(
1395 f" Partial scores saved: P1={phase_scores.get(1)}%, "
1396 f"P2={phase_scores.get(2)}%, P3={phase_scores.get(3)}%, "
1397 f"P4={phase_scores.get(4)}%, P7={phase_scores.get(7)}%, "
1398 f"P8={phase_scores.get(8)}%, P9={phase_scores.get(9)}%"
1399 )
1400 update_model_status(
1401 model_id,
1402 arch,
1403 STATUS_FAILED,
1404 note=note,
1405 phase_scores=phase_scores,
1406 sanitize_fn=_sanitize_note,
1407 prompt_profile=_extract_prompt_profile(all_results),
1408 )
1409 add_verification_record(
1410 model_id,
1411 arch,
1412 notes=note,
1413 sanitize_fn=_sanitize_note,
1414 prompt_profile=_extract_prompt_profile(all_results),
1415 p4_scoring_version=(P4_SCORING_VERSION if 4 in phase_scores else None),
1416 )
1417 progress.failed.append(model_id)
1419 # Post-model cleanup
1420 gc.collect()
1421 try:
1422 import torch
1424 if torch.cuda.is_available():
1425 torch.cuda.empty_cache()
1426 torch.cuda.synchronize()
1427 if device == "mps" and hasattr(torch, "mps") and torch.backends.mps.is_available():
1428 torch.mps.synchronize()
1429 torch.mps.empty_cache()
1431 # Log MPS memory state for debugging long runs
1432 if device == "mps" and not quiet and hasattr(torch.mps, "current_allocated_memory"):
1433 alloc_mb = torch.mps.current_allocated_memory() / (1024 * 1024)
1434 driver_mb = torch.mps.driver_allocated_memory() / (1024 * 1024)
1435 print(f" MPS memory: {alloc_mb:.0f} MB allocated, " f"{driver_mb:.0f} MB driver")
1436 except ImportError:
1437 pass
1439 # Brief pause to let the OS and MPS reclaim memory between models
1440 if device in ("mps", "cuda"):
1441 time.sleep(3)
1443 # Periodically clear the HuggingFace cache to prevent disk exhaustion
1444 if i % 50 == 0:
1445 _clear_hf_cache(quiet)
1447 _save_checkpoint(progress)
1449 # Clean up pre-loaded scoring model
1450 if _judge_model is not None:
1451 del _judge_model
1452 del _judge_tokenizer
1453 gc.collect()
1455 return progress
1458def _print_dry_run(
1459 candidates: list[ModelCandidate],
1460 dtype: str,
1461 max_memory_gb: float,
1462 phases: Optional[list[int]] = None,
1463 use_hf_reference: bool = True,
1464 device: str = "cpu",
1465) -> None:
1466 """Print what would be tested in a dry run."""
1467 print(f"\nDry run: {len(candidates)} models would be tested")
1468 print(f"Memory limit: {max_memory_gb:.1f} GB | Dtype: {dtype}")
1469 print()
1471 # Group by architecture
1472 by_arch: dict[str, list[ModelCandidate]] = {}
1473 for c in candidates:
1474 by_arch.setdefault(c.architecture_id, []).append(c)
1476 skippable = 0
1477 testable = 0
1479 for arch in sorted(by_arch.keys()):
1480 models = by_arch[arch]
1481 # Same per-architecture default as the real run, so the dry run's
1482 # memory estimates cover the phases that will actually execute.
1483 eff_phases = phases if phases is not None else _default_phases_for_architecture(arch)
1484 phases_to_run = _phases_to_run(arch, eff_phases)
1485 print(f" {arch} ({len(models)} models):")
1486 for c in models:
1487 try:
1488 n_params = estimate_model_params(c.model_id)
1489 mem = estimate_benchmark_memory_gb(
1490 n_params,
1491 dtype,
1492 phases=phases_to_run,
1493 use_hf_reference=use_hf_reference,
1494 device=device,
1495 )
1496 status = "OK" if mem <= max_memory_gb else "SKIP (too large)"
1497 if mem > max_memory_gb:
1498 skippable += 1
1499 else:
1500 testable += 1
1501 print(f" {c.model_id}: ~{n_params/1e6:.0f}M params, ~{mem:.1f} GB [{status}]")
1502 except Exception as e:
1503 skippable += 1
1504 print(f" {c.model_id}: config error ({e})")
1505 print()
1507 print(f"Summary: {testable} testable, {skippable} would be skipped")
1510def _print_summary(progress: VerificationProgress) -> None:
1511 """Print a summary of the verification run."""
1512 total = len(progress.tested)
1513 print(f"\n{'='*70}")
1514 print("Verification Summary")
1515 print(f"{'='*70}")
1516 print(f" Total tested: {total}")
1517 print(f" Verified: {len(progress.verified)}")
1518 print(f" Provisional: {len(progress.provisional)}")
1519 print(f" Skipped: {len(progress.skipped)}")
1520 print(f" Failed: {len(progress.failed)}")
1522 if progress.verified:
1523 print(f"\n Verified models:")
1524 for m in progress.verified:
1525 print(f" - {m}")
1527 if progress.provisional:
1528 print(f"\n Provisional models (structural only, no HF reference):")
1529 for m in progress.provisional:
1530 print(f" - {m}")
1532 if progress.failed:
1533 print(f"\n Failed models:")
1534 for m in progress.failed:
1535 print(f" - {m}")
1537 if progress.skipped:
1538 print(f"\n Skipped models:")
1539 for m in progress.skipped[:20]:
1540 print(f" - {m}")
1541 if len(progress.skipped) > 20:
1542 print(f" ... and {len(progress.skipped) - 20} more")
1545def main() -> None:
1546 """CLI entry point for batch model verification."""
1547 parser = argparse.ArgumentParser(
1548 description="Batch verify models in the TransformerLens registry",
1549 formatter_class=argparse.RawDescriptionHelpFormatter,
1550 epilog="""
1551Examples:
1552 %(prog)s --dry-run Show what would be tested
1553 %(prog)s --limit 3 Test 3 models total
1554 %(prog)s --architectures GPT2LMHeadModel --per-arch 5
1555 %(prog)s --device cuda --max-memory 24
1556 %(prog)s --resume Resume from checkpoint
1557 %(prog)s --reverify --architectures Olmo2ForCausalLM Re-verify already-tested models
1558 %(prog)s --model google/gemma-2b Verify a single model by ID
1559 """,
1560 )
1561 parser.add_argument(
1562 "--per-arch",
1563 type=int,
1564 default=10,
1565 help="Max models to verify per architecture (default: 10)",
1566 )
1567 parser.add_argument(
1568 "--device",
1569 type=str,
1570 default="cpu",
1571 help="Device for benchmarks (default: cpu)",
1572 )
1573 parser.add_argument(
1574 "--max-memory",
1575 type=float,
1576 default=None,
1577 help="Memory limit in GB (default: auto-detect)",
1578 )
1579 parser.add_argument(
1580 "--architectures",
1581 nargs="+",
1582 default=None,
1583 help="Filter to specific architectures",
1584 )
1585 parser.add_argument(
1586 "--limit",
1587 type=int,
1588 default=None,
1589 help="Total model cap",
1590 )
1591 parser.add_argument(
1592 "--resume",
1593 action="store_true",
1594 help="Resume from checkpoint",
1595 )
1596 parser.add_argument(
1597 "--dry-run",
1598 action="store_true",
1599 help="Show what would be tested without running benchmarks",
1600 )
1601 parser.add_argument(
1602 "--no-hf-reference",
1603 action="store_true",
1604 help=(
1605 "Skip HuggingFace reference comparison (Phase 1 is structural-only). "
1606 "A passing run is recorded as PROVISIONAL, not verified — re-run without "
1607 "this flag for a real HF-compared verification."
1608 ),
1609 )
1610 parser.add_argument(
1611 "--no-ht-reference",
1612 action="store_true",
1613 help="Skip HookedTransformer reference comparison",
1614 )
1615 parser.add_argument(
1616 "--phases",
1617 nargs="+",
1618 type=int,
1619 default=None,
1620 help=(
1621 "Which benchmark phases to run (default: a full verification for each "
1622 "model's architecture — 1 2 3 4 for text, 1 2 3 4 7 for multimodal, "
1623 "1 8 for audio, 1 9 for vision)"
1624 ),
1625 )
1626 parser.add_argument(
1627 "--dtype",
1628 type=str,
1629 default="float32",
1630 choices=["float32", "float16", "bfloat16"],
1631 help="Dtype for memory estimation (default: float32)",
1632 )
1633 parser.add_argument(
1634 "--quiet",
1635 action="store_true",
1636 help="Suppress verbose output",
1637 )
1638 parser.add_argument(
1639 "--retry-failed",
1640 action="store_true",
1641 help="Re-run previously failed models instead of skipping them",
1642 )
1643 parser.add_argument(
1644 "--reverify",
1645 action="store_true",
1646 help="Re-run verification for already-verified/skipped/failed models. "
1647 "Ignores previous status and re-tests matching models from scratch.",
1648 )
1649 parser.add_argument(
1650 "--model",
1651 type=str,
1652 nargs="+",
1653 default=None,
1654 help="Verify one or more models by HuggingFace model ID. "
1655 "Looks up architecture from the registry automatically.",
1656 )
1658 args = parser.parse_args()
1660 # Setup logging
1661 logging.basicConfig(
1662 level=logging.WARNING if args.quiet else logging.INFO,
1663 format="%(asctime)s [%(levelname)s] %(message)s",
1664 )
1666 # Auto-detect memory
1667 max_memory_gb = args.max_memory
1668 if max_memory_gb is None:
1669 max_memory_gb = get_available_memory_gb(args.device)
1671 # Load checkpoint if resuming
1672 progress = None
1673 if args.resume:
1674 progress = _load_checkpoint()
1675 if progress:
1676 print(f"Resuming from checkpoint: {len(progress.tested)} models already tested")
1677 else:
1678 print("No checkpoint found, starting fresh")
1680 # If retrying failed, clean them from checkpoint and reset status in registry
1681 if args.retry_failed and progress and not args.dry_run:
1682 failed_set = set(progress.failed)
1683 if failed_set:
1684 # Reset status in supported_models.json
1685 registry_data = load_supported_models_raw()
1686 for entry in registry_data.get("models", []):
1687 if entry["model_id"] in failed_set and entry.get("status") == STATUS_FAILED:
1688 update_model_status(
1689 entry["model_id"],
1690 entry["architecture_id"],
1691 STATUS_UNVERIFIED,
1692 )
1693 # Clean checkpoint
1694 progress.tested = [m for m in progress.tested if m not in failed_set]
1695 progress.failed = []
1696 _save_checkpoint(progress)
1697 print(f" Cleared {len(failed_set)} failed models for retry")
1699 # Select models — either --model list or the normal batch selection
1700 if args.model:
1701 # Look up architecture for each model from the registry
1702 registry_data = load_supported_models_raw()
1703 candidates = []
1704 for model_id in args.model:
1705 arch_id = None
1706 for entry in registry_data.get("models", []):
1707 if entry["model_id"] == model_id:
1708 arch_id = entry["architecture_id"]
1709 break
1710 if arch_id is None:
1711 print(f"Model '{model_id}' not found in supported_models.json, skipping")
1712 continue
1713 candidates.append(ModelCandidate(model_id=model_id, architecture_id=arch_id))
1714 if not candidates:
1715 print("No valid models found in registry")
1716 return
1717 print(f"Model list mode: {len(candidates)} model(s)")
1718 else:
1719 candidates = select_models_for_verification(
1720 per_arch=args.per_arch,
1721 architectures=args.architectures,
1722 limit=args.limit,
1723 resume_progress=progress,
1724 retry_failed=args.retry_failed,
1725 reverify=args.reverify,
1726 )
1728 if not candidates:
1729 print("No models to verify (all matching models already tested)")
1730 return
1732 print(f"Selected {len(candidates)} models for verification")
1734 # Dry run
1735 if args.dry_run:
1736 _print_dry_run(
1737 candidates,
1738 args.dtype,
1739 max_memory_gb,
1740 phases=args.phases,
1741 use_hf_reference=not args.no_hf_reference,
1742 device=args.device,
1743 )
1744 return
1746 # Install graceful interrupt handler (Ctrl+C stops between models)
1747 signal.signal(signal.SIGINT, _handle_sigint)
1749 # Run verification
1750 start = time.time()
1751 progress = verify_models(
1752 candidates,
1753 device=args.device,
1754 max_memory_gb=max_memory_gb,
1755 dtype=args.dtype,
1756 use_hf_reference=not args.no_hf_reference,
1757 use_ht_reference=not args.no_ht_reference,
1758 phases=args.phases,
1759 quiet=args.quiet,
1760 progress=progress,
1761 )
1762 elapsed = time.time() - start
1764 _print_summary(progress)
1765 print(f"\nTotal time: {elapsed:.1f}s")
1767 # Clean up checkpoint on successful completion
1768 if _CHECKPOINT_PATH.exists():
1769 _CHECKPOINT_PATH.unlink()
1770 print("Checkpoint cleared (run complete)")
1773if __name__ == "__main__": 1773 ↛ 1774line 1773 didn't jump to line 1774 because the condition on line 1773 was never true
1774 main()