from __future__ import annotations import os import shutil import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer ROOT = Path(__file__).resolve().parents[1] EVAL_DIR = ROOT / "eval" DUMP_PATH = EVAL_DIR / "probe_dump.npz" CANDIDATE_MODELS = [ "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", "Qwen/Qwen2.5-3B-Instruct", "openbmb/MiniCPM3-4B", ] TRANSCRIPT_CHUNKS = [ "so basically", "our startup uses", "ai to help", "small businesses", "manage inventory", "and we think", "the market is huge", "and honestly", "we already have", "like a thousand", "users and", "growing fast", ] MIN_MODEL_DOWNLOAD_FREE_BYTES = 6 * 1024**3 def configure_local_caches() -> None: os.environ.setdefault("HF_HOME", str(ROOT / ".hf-cache")) os.environ.setdefault("TRANSFORMERS_CACHE", str(ROOT / ".hf-cache" / "transformers")) os.environ.setdefault("TORCH_HOME", str(ROOT / ".torch-cache")) def cuda_summary() -> torch.device: print(f"torch.__version__ = {torch.__version__}") print(f"torch.version.cuda = {torch.version.cuda}") print(f"torch.cuda.is_available() = {torch.cuda.is_available()}") if torch.cuda.is_available(): print(f"torch.cuda.get_device_name(0) = {torch.cuda.get_device_name(0)}") return torch.device("cuda:0") print("LOUD CUDA FALLBACK: CUDA/Blackwell is not available in this torch environment; using CPU.") return torch.device("cpu") def common_prefix_len(previous: list[int], current: list[int]) -> int: length = 0 for left, right in zip(previous, current): if left != right: break length += 1 return length def save_failure(failure: str) -> None: EVAL_DIR.mkdir(parents=True, exist_ok=True) np.savez( DUMP_PATH, nll_series=np.asarray([], dtype=np.float32), hidden_states=np.empty((0, 0), dtype=np.float32), update_ms=np.asarray([], dtype=np.float32), added_text=np.asarray([], dtype=object), model=np.asarray("", dtype=object), device=np.asarray("cpu", dtype=object), dtype=np.asarray("", dtype=object), failure=np.asarray(failure, dtype=object), ) def load_first_model(device: torch.device) -> tuple[object, object, str, float, float] | None: free_bytes = shutil.disk_usage(ROOT).free local_files_only = free_bytes < MIN_MODEL_DOWNLOAD_FREE_BYTES if local_files_only: free_gib = free_bytes / 1024**3 needed_gib = MIN_MODEL_DOWNLOAD_FREE_BYTES / 1024**3 print( "LOUD MODEL DOWNLOAD SKIP: only " f"{free_gib:.2f} GiB free; need at least {needed_gib:.1f} GiB to attempt these 3B/4B model downloads. " "Trying repo-local cache only." ) if device.type == "cuda": dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 else: dtype = torch.float32 failures: list[str] = [] for model_id in CANDIDATE_MODELS: print(f"Attempting model: {model_id}") if device.type == "cuda": torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats(0) start = time.perf_counter() try: tokenizer = AutoTokenizer.from_pretrained( model_id, trust_remote_code=True, local_files_only=local_files_only, ) load_kwargs = { "trust_remote_code": True, "torch_dtype": dtype, "low_cpu_mem_usage": True, "local_files_only": local_files_only, } if device.type == "cuda": load_kwargs["device_map"] = {"": 0} model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) if device.type == "cpu": model.to(device) model.eval() load_seconds = time.perf_counter() - start actual_device = next(model.parameters()).device actual_dtype = next(model.parameters()).dtype vram_gib = 0.0 if device.type == "cuda": torch.cuda.synchronize() vram_gib = torch.cuda.memory_allocated(0) / 1024**3 print( "LOADED " f"model={model_id} device={actual_device} dtype={actual_dtype} " f"load_seconds={load_seconds:.2f} vram_used_gib={vram_gib:.2f}" ) return tokenizer, model, model_id, load_seconds, vram_gib except Exception as exc: # noqa: BLE001 - spike should continue through model fallbacks. elapsed = time.perf_counter() - start message = f"{model_id} failed after {elapsed:.2f}s: {type(exc).__name__}: {exc}" print(message) failures.append(message) failure = "No candidate model loaded. " + " | ".join(failures) print(f"LOUD PROBE FAILURE: {failure}") save_failure(failure) return None def run_updates(tokenizer: object, model: object, model_id: str) -> None: device = next(model.parameters()).device previous_ids: list[int] = [] prefixes: list[str] = [] running = "" for chunk in TRANSCRIPT_CHUNKS: running = f"{running} {chunk}".strip() prefixes.append(running) nll_series: list[float] = [] hidden_rows: list[np.ndarray] = [] update_ms: list[float] = [] added_text: list[str] = [] print("step | added_text | mean_NLL | hidden_dim | update_ms") print("-----|------------|----------|------------|----------") for step, (chunk, prefix) in enumerate(zip(TRANSCRIPT_CHUNKS, prefixes), start=1): if device.type == "cuda": torch.cuda.synchronize() start = time.perf_counter() encoded = tokenizer(prefix, return_tensors="pt", add_special_tokens=False) current_ids = encoded["input_ids"][0].tolist() new_start = common_prefix_len(previous_ids, current_ids) inputs = {name: tensor.to(device) for name, tensor in encoded.items()} with torch.inference_mode(): outputs = model(**inputs, output_hidden_states=True) input_ids = inputs["input_ids"] logits = outputs.logits[:, :-1, :].float() targets = input_ids[:, 1:] token_nll = F.cross_entropy( logits.reshape(-1, logits.shape[-1]), targets.reshape(-1), reduction="none", ).reshape(targets.shape) nll_start = max(new_start, 1) - 1 new_nll = token_nll[0, nll_start:] mean_nll = float(new_nll.mean().detach().cpu()) if new_nll.numel() else float("nan") last_hidden = outputs.hidden_states[-1][0] new_hidden = last_hidden[new_start:, :] mean_hidden = new_hidden.float().mean(dim=0).detach().cpu() if device.type == "cuda": torch.cuda.synchronize() elapsed_ms = (time.perf_counter() - start) * 1000.0 hidden_vec = mean_hidden.numpy().astype(np.float32) nll_series.append(mean_nll) hidden_rows.append(hidden_vec) update_ms.append(elapsed_ms) added_text.append(chunk) previous_ids = current_ids print(f"{step:>4} | {chunk} | {mean_nll:.4f} | {hidden_vec.shape[0]} | {elapsed_ms:.2f}") EVAL_DIR.mkdir(parents=True, exist_ok=True) hidden_matrix = np.vstack(hidden_rows).astype(np.float32) actual_dtype = str(next(model.parameters()).dtype) np.savez( DUMP_PATH, nll_series=np.asarray(nll_series, dtype=np.float32), hidden_states=hidden_matrix, update_ms=np.asarray(update_ms, dtype=np.float32), added_text=np.asarray(added_text, dtype=object), model=np.asarray(model_id, dtype=object), device=np.asarray(str(device), dtype=object), dtype=np.asarray(actual_dtype, dtype=object), failure=np.asarray("", dtype=object), ) print(f"Saved {DUMP_PATH}") def main() -> None: configure_local_caches() EVAL_DIR.mkdir(parents=True, exist_ok=True) device = cuda_summary() loaded = load_first_model(device) if loaded is None: return tokenizer, model, model_id, _load_seconds, _vram_gib = loaded run_updates(tokenizer, model, model_id) if __name__ == "__main__": main()