#!/usr/bin/env python3 """Trace all 27 Chat decoder layers and all 26 MoE gates. The probe reuses the fixed 16-source, eight-condition completion cohort but runs prompt-only full-depth forwards. It stores hashes and summaries rather than raw hidden tensors. Target-content comparisons use only byte-identical interior content tokens; boundary-crossing tokens are explicitly excluded. """ from __future__ import annotations import argparse import hashlib import json import math import os import platform import resource import time from datetime import datetime, timezone from pathlib import Path from statistics import mean from typing import Any import accelerate import safetensors import torch import torch.nn.functional as F import transformers from transformers import AutoTokenizer import v2_lite_chat_special_token_behavior_probe as behavior import v2_lite_routing_special_token_family_control as special EDGE_PAIRS = { "system_eos": ("s0_eos", "s1_eos"), "system_bos": ("s0_bos", "s1_bos"), "system_x": ("s0_x", "s1_x"), "system_period": ("s0_period", "s1_period"), "bos_at_s0": ("s0_eos", "s0_bos"), "bos_at_s1": ("s1_eos", "s1_bos"), "x_at_s0": ("s0_eos", "s0_x"), "x_at_s1": ("s1_eos", "s1_x"), "period_at_s0": ("s0_eos", "s0_period"), "period_at_s1": ("s1_eos", "s1_period"), } SCOPES = ("target_content", "full_input") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--artifact-dir", type=Path, required=True) parser.add_argument("--reference-routing-json", type=Path, required=True) parser.add_argument("--human-eval", type=Path, required=True) parser.add_argument("--gsm8k", type=Path, required=True) parser.add_argument("--tnews", type=Path, required=True) parser.add_argument("--tnews-archive", type=Path, required=True) parser.add_argument("--wikitext", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument( "--domains", nargs="+", choices=behavior.DEFAULT_DOMAINS, default=list(behavior.DEFAULT_DOMAINS), ) parser.add_argument("--per-domain", type=int, default=1) parser.add_argument("--gpu-memory", default="28GiB") parser.add_argument("--cpu-memory", default="80GiB") parser.add_argument("--captured-at", default=None) return parser.parse_args() def canonical_hash(value: Any) -> str: payload = json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ).encode() return hashlib.sha256(payload).hexdigest() def tensor_hash(tensor: torch.Tensor) -> str: value = tensor.detach().cpu().contiguous() header = json.dumps( { "shape": list(value.shape), "dtype": str(value.dtype), }, separators=(",", ":"), ).encode() raw = value.view(torch.uint8).numpy().tobytes() return hashlib.sha256(header + b"\0" + raw).hexdigest() def validate_args(args: argparse.Namespace) -> None: if args.per_domain <= 0: raise ValueError("--per-domain must be positive") for path in ( args.artifact_dir, args.reference_routing_json, args.human_eval, args.gsm8k, args.tnews, args.tnews_archive, args.wikitext, ): if not path.exists(): raise FileNotFoundError(path) missing = [ name for name in behavior.MODEL_FILES if not (args.artifact_dir / name).is_file() ] if missing: raise FileNotFoundError( f"artifact directory is incomplete: {missing}" ) def condition_positions( variants: dict[str, dict[str, Any]], batch_tokens: int, ) -> dict[str, dict[str, list[int]]]: result = {} for condition in special.CONDITIONS: variant = variants[condition] padding = batch_tokens - variant["tokens"] result[condition] = { "target_content": [ padding + int(position) for position in variant["content_positions"] ], "full_input": list(range(padding, batch_tokens)), } return result def hidden_stats(value: torch.Tensor) -> dict[str, Any]: fp32 = value.float() return { "shape": list(value.shape), "dtype": str(value.dtype), "tensor_sha256": tensor_hash(value), "mean": fp32.mean().item(), "std": fp32.std(unbiased=False).item(), "rms": fp32.square().mean().sqrt().item(), "max_abs": fp32.abs().max().item(), } def compare_hidden( left: torch.Tensor, right: torch.Tensor, ) -> dict[str, Any]: if left.shape != right.shape: raise RuntimeError( f"hidden comparison shape mismatch: {left.shape} != " f"{right.shape}" ) left_fp32 = left.float() right_fp32 = right.float() delta = right_fp32 - left_fp32 denominator = left_fp32.norm(dim=-1).clamp_min(1e-12) relative = delta.norm(dim=-1) / denominator cosine = F.cosine_similarity( left_fp32, right_fp32, dim=-1, eps=1e-12, ).clamp(-1.0, 1.0) exact = (left == right).all(dim=-1) return { "tokens": left.shape[0], "exact_hidden_rows": exact.sum().item(), "mean_cosine_similarity": cosine.mean().item(), "min_cosine_similarity": cosine.min().item(), "mean_relative_l2": relative.mean().item(), "max_abs_delta": delta.abs().max().item(), } def route_distribution( route_ids: torch.Tensor, route_weights: torch.Tensor, experts: int, ) -> dict[str, Any]: loads = torch.bincount( route_ids.reshape(-1), minlength=experts, ).to(torch.int64) weighted = torch.zeros(experts, dtype=torch.float64) weighted.scatter_add_( 0, route_ids.reshape(-1), route_weights.reshape(-1).to(torch.float64), ) total = loads.sum().item() probabilities = loads.to(torch.float64) / max(total, 1) nonzero = probabilities[probabilities > 0] entropy = ( -(nonzero * nonzero.log()).sum().item() / math.log(experts) if len(nonzero) else 0.0 ) mean_load = loads.double().mean() cv = ( loads.double().std(unbiased=False) / mean_load if mean_load > 0 else torch.tensor(0.0) ) return { "tokens": route_ids.shape[0], "top_k": route_ids.shape[1], "decisions": total, "ordered_route_sha256": canonical_hash(route_ids.tolist()), "route_weight_sha256": tensor_hash(route_weights), "loads": loads.tolist(), "weighted_loads": weighted.tolist(), "unique_experts": int((loads > 0).sum().item()), "load_cv": cv.item(), "normalized_load_entropy": entropy, "max_load_share": probabilities.max().item(), } def normalized_load(loads: torch.Tensor) -> torch.Tensor: value = loads.to(torch.float64) return value / value.sum().clamp_min(1e-12) def compare_routes( left_ids: torch.Tensor, left_weights: torch.Tensor, right_ids: torch.Tensor, right_weights: torch.Tensor, experts: int, ) -> dict[str, Any]: if left_ids.shape != right_ids.shape: raise RuntimeError( f"route comparison shape mismatch: {left_ids.shape} != " f"{right_ids.shape}" ) ordered_exact = (left_ids == right_ids).all(dim=-1) left_sorted = left_ids.sort(dim=-1).values right_sorted = right_ids.sort(dim=-1).values set_exact = (left_sorted == right_sorted).all(dim=-1) left_hot = F.one_hot( left_ids, num_classes=experts, ).sum(dim=1).bool() right_hot = F.one_hot( right_ids, num_classes=experts, ).sum(dim=1).bool() intersection = (left_hot & right_hot).sum(dim=-1) union = (left_hot | right_hot).sum(dim=-1).clamp_min(1) left_token_weighted = torch.zeros( left_ids.shape[0], experts, dtype=torch.float64, ) right_token_weighted = torch.zeros_like(left_token_weighted) left_token_weighted.scatter_add_( 1, left_ids, left_weights.to(torch.float64), ) right_token_weighted.scatter_add_( 1, right_ids, right_weights.to(torch.float64), ) token_weighted_tv = 0.5 * ( left_token_weighted - right_token_weighted ).abs().sum(dim=-1) left_load = torch.bincount( left_ids.reshape(-1), minlength=experts, ) right_load = torch.bincount( right_ids.reshape(-1), minlength=experts, ) load_tv = 0.5 * ( normalized_load(left_load) - normalized_load(right_load) ).abs().sum() return { "tokens": left_ids.shape[0], "ordered_topk_exact_tokens": ordered_exact.sum().item(), "set_exact_tokens": set_exact.sum().item(), "mean_set_jaccard": ( intersection.to(torch.float64) / union.to(torch.float64) ).mean().item(), "mean_token_weighted_tv": token_weighted_tv.mean().item(), "aggregate_load_tv": load_tv.item(), } class TraceCapture: def __init__(self, experts: int) -> None: self.experts = experts self.positions: dict[str, dict[str, list[int]]] = {} self.batch_tokens = 0 self.hidden: dict[str, Any] = {} self.routes: dict[str, Any] = {} def start( self, positions: dict[str, dict[str, list[int]]], batch_tokens: int, ) -> None: self.positions = positions self.batch_tokens = batch_tokens self.hidden = {} self.routes = {} def hidden_hook(self, stage: str): def hook( _module: torch.nn.Module, _inputs: tuple[Any, ...], output: Any, ) -> None: tensor = output[0] if isinstance(output, tuple) else output if tensor.ndim != 3: raise RuntimeError( f"{stage} hidden rank {tensor.ndim}, expected 3" ) cpu = tensor.detach().cpu() condition_values: dict[str, dict[str, torch.Tensor]] = {} condition_rows = {} for row, condition in enumerate(special.CONDITIONS): condition_values[condition] = {} condition_rows[condition] = {} for scope in SCOPES: positions = self.positions[condition][scope] value = cpu[row, positions, :].contiguous() condition_values[condition][scope] = value condition_rows[condition][scope] = hidden_stats(value) comparisons = {} for name, (left, right) in EDGE_PAIRS.items(): comparisons[name] = compare_hidden( condition_values[left]["target_content"], condition_values[right]["target_content"], ) self.hidden[stage] = { "conditions": condition_rows, "target_comparisons": comparisons, } return hook def gate_hook(self, layer: int): def hook( _module: torch.nn.Module, _inputs: tuple[Any, ...], output: Any, ) -> None: route_ids, route_weights, _aux = output top_k = route_ids.shape[-1] ids = route_ids.detach().reshape( len(special.CONDITIONS), self.batch_tokens, top_k, ).cpu() weights = route_weights.detach().reshape( len(special.CONDITIONS), self.batch_tokens, top_k, ).cpu() values: dict[ str, dict[str, tuple[torch.Tensor, torch.Tensor]], ] = {} condition_rows = {} for row, condition in enumerate(special.CONDITIONS): values[condition] = {} condition_rows[condition] = {} for scope in SCOPES: positions = self.positions[condition][scope] scoped_ids = ids[row, positions, :].contiguous() scoped_weights = weights[ row, positions, :, ].contiguous() values[condition][scope] = ( scoped_ids, scoped_weights, ) condition_rows[condition][scope] = ( route_distribution( scoped_ids, scoped_weights, self.experts, ) ) comparisons = {} for name, (left, right) in EDGE_PAIRS.items(): left_ids, left_weights = values[left][ "target_content" ] right_ids, right_weights = values[right][ "target_content" ] comparisons[name] = compare_routes( left_ids, left_weights, right_ids, right_weights, self.experts, ) self.routes[f"layer_{layer:02d}"] = { "layer": layer, "conditions": condition_rows, "target_comparisons": comparisons, } return hook def aggregate_trace(sources: list[dict[str, Any]]) -> dict[str, Any]: hidden_stages = [ "embedding", *(f"layer_{index:02d}" for index in range(27)), "final_norm", ] router_layers = [ f"layer_{index:02d}" for index in range(1, 27) ] hidden = {} for stage in hidden_stages: hidden[stage] = {} for edge in EDGE_PAIRS: rows = [ source["hidden_stages"][stage][ "target_comparisons" ][edge] for source in sources ] tokens = sum(row["tokens"] for row in rows) hidden[stage][edge] = { "sources": len(rows), "tokens": tokens, "exact_hidden_rows": sum( row["exact_hidden_rows"] for row in rows ), "token_weighted_mean_cosine_similarity": sum( row["mean_cosine_similarity"] * row["tokens"] for row in rows ) / tokens, "source_mean_relative_l2": mean( row["mean_relative_l2"] for row in rows ), "max_abs_delta": max( row["max_abs_delta"] for row in rows ), } routes = {} for layer in router_layers: routes[layer] = {} for edge in EDGE_PAIRS: rows = [ source["router_layers"][layer][ "target_comparisons" ][edge] for source in sources ] tokens = sum(row["tokens"] for row in rows) routes[layer][edge] = { "sources": len(rows), "tokens": tokens, "ordered_topk_exact_tokens": sum( row["ordered_topk_exact_tokens"] for row in rows ), "set_exact_tokens": sum( row["set_exact_tokens"] for row in rows ), "token_weighted_mean_set_jaccard": sum( row["mean_set_jaccard"] * row["tokens"] for row in rows ) / tokens, "source_mean_token_weighted_tv": mean( row["mean_token_weighted_tv"] for row in rows ), "source_mean_aggregate_load_tv": mean( row["aggregate_load_tv"] for row in rows ), } target_tokens = sum( source["content_tokens"] for source in sources ) return { "sources": len(sources), "conditions_per_source": len(special.CONDITIONS), "decoder_layers_executed": 27, "moe_gates_executed": 26, "hidden_stages": len(hidden_stages), "target_content_tokens_per_condition": target_tokens, "target_route_decisions": ( target_tokens * len(special.CONDITIONS) * 26 * 6 ), "hidden": hidden, "routes": routes, } def main() -> None: args = parse_args() validate_args(args) captured_at = args.captured_at or datetime.now( timezone.utc ).isoformat() revision_contract = behavior.download_revision_contract( args.artifact_dir ) special.install_control_contract() tokenizer = AutoTokenizer.from_pretrained( args.artifact_dir, trust_remote_code=True, local_files_only=True, use_fast=True, ) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id if special.special_family_token_ids(tokenizer) != { "eos": 100001, "bos": 100000, "x": 87, "period": 13, }: raise RuntimeError("pinned tokenizer ID contract changed") source_rows, reference = behavior.selected_sources(args, tokenizer) for source in source_rows: source["variants"] = { condition: special.prior.render_boundary_variant( tokenizer, source["content"], condition, ) for condition in special.CONDITIONS } content_sequences = { tuple( source["variants"][condition]["token_ids"][position] for position in source["variants"][condition][ "content_positions" ] ) for condition in special.CONDITIONS } if len(content_sequences) != 1: raise RuntimeError( f"{source['id']} target token sequence is not exact " "across conditions" ) source["content_token_ids_sha256"] = canonical_hash( next(iter(content_sequences)) ) if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() load_started = time.perf_counter() _, official_modeling = special.base.load_official_modules( args.artifact_dir ) model = official_modeling.DeepseekV2ForCausalLM.from_pretrained( args.artifact_dir, local_files_only=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto", max_memory={ 0: args.gpu_memory, "cpu": args.cpu_memory, }, ) model.eval() if torch.cuda.is_available(): torch.cuda.synchronize() load_seconds = time.perf_counter() - load_started load_peak_cuda = ( torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None ) input_device = model.get_input_embeddings().weight.device capture = TraceCapture(model.config.n_routed_experts) handles = [ model.model.embed_tokens.register_forward_hook( capture.hidden_hook("embedding") ), model.model.norm.register_forward_hook( capture.hidden_hook("final_norm") ), ] for index, layer in enumerate(model.model.layers): handles.append( layer.register_forward_hook( capture.hidden_hook(f"layer_{index:02d}") ) ) if hasattr(layer.mlp, "gate"): handles.append( layer.mlp.gate.register_forward_hook( capture.gate_hook(index) ) ) traced_sources = [] try: for source in source_rows: unpadded = [ source["variants"][condition]["token_ids"] for condition in special.CONDITIONS ] prompt_lengths = [len(row) for row in unpadded] batch_tokens = max(prompt_lengths) padded = [ [int(tokenizer.pad_token_id)] * (batch_tokens - len(row)) + row for row in unpadded ] masks = [ [0] * (batch_tokens - len(row)) + [1] * len(row) for row in unpadded ] positions = condition_positions( source["variants"], batch_tokens, ) capture.start(positions, batch_tokens) input_ids = torch.tensor( padded, dtype=torch.long, device=input_device, ) attention_mask = torch.tensor( masks, dtype=torch.long, device=input_device, ) if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() started = time.perf_counter() with torch.inference_mode(): output = model.model( input_ids=input_ids, attention_mask=attention_mask, use_cache=False, output_attentions=False, output_hidden_states=False, return_dict=True, ) if torch.cuda.is_available(): torch.cuda.synchronize() elapsed = time.perf_counter() - started peak = ( torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None ) if output.last_hidden_state.shape[:2] != input_ids.shape: raise RuntimeError( "final hidden batch/sequence shape does not match input" ) expected_hidden = { "embedding", *(f"layer_{index:02d}" for index in range(27)), "final_norm", } expected_routes = { f"layer_{index:02d}" for index in range(1, 27) } if set(capture.hidden) != expected_hidden: raise RuntimeError( f"hidden stages mismatch: {set(capture.hidden)}" ) if set(capture.routes) != expected_routes: raise RuntimeError( f"router layers mismatch: {set(capture.routes)}" ) content_tokens = len( positions[special.CONDITIONS[0]][ "target_content" ] ) traced_sources.append( { key: value for key, value in source.items() if key not in {"content", "variants"} } | { "prompt_tokens_by_condition": { condition: prompt_lengths[index] for index, condition in enumerate( special.CONDITIONS ) }, "batch_prompt_tokens": batch_tokens, "content_tokens": content_tokens, "content_token_ids_sha256": source[ "content_token_ids_sha256" ], "boundary_crossing_tokens_by_condition": { condition: source["variants"][condition][ "boundary_crossing_tokens" ] for condition in special.CONDITIONS }, "forward_seconds": elapsed, "peak_cuda_memory_allocated_bytes": peak, "hidden_stages": capture.hidden, "router_layers": capture.routes, } ) finally: for handle in handles: handle.remove() device_map = getattr(model, "hf_device_map", {}) peaks = [ source["peak_cuda_memory_allocated_bytes"] for source in traced_sources if source["peak_cuda_memory_allocated_bytes"] is not None ] result = { "schema_version": 1, "captured_at": captured_at, "model": { "repo": "deepseek-ai/DeepSeek-V2-Lite-Chat", "revision": behavior.MODEL_REVISION, "checkpoint_identity": "SFT Chat", "architecture": type(model).__name__, "dtype": "bfloat16", "download_revision_contract": revision_contract, "files": { name: { "bytes": (args.artifact_dir / name).stat().st_size, "sha256": behavior.sha256( args.artifact_dir / name ), } for name in behavior.MODEL_FILES }, }, "tokenizer_contract": { "length": len(tokenizer), "vocab_size": tokenizer.vocab_size, "all_special_ids": tokenizer.all_special_ids, "pad_token_id": tokenizer.pad_token_id, "conditions": special.FACTORS, }, "source_contract": { **reference, "domains": args.domains, "per_domain": args.per_domain, "sources": len(source_rows), "full_source_text_used": True, "target_scope": ( "byte-identical interior content tokens; tokens whose " "offset crosses the content boundary are excluded" ), }, "trace_contract": { "conditions": list(special.CONDITIONS), "scopes": list(SCOPES), "edge_pairs": EDGE_PAIRS, "decoder_layers": 27, "moe_gate_layers": list(range(1, 27)), "experts": model.config.n_routed_experts, "top_k": model.config.num_experts_per_tok, "use_cache": False, "raw_hidden_tensors_saved": False, "raw_route_ids_saved": False, "load_vectors_saved": True, "hashes_cover_exact_bfloat16_hidden_and_route_values": True, }, "execution": { "python": platform.python_version(), "torch": torch.__version__, "transformers": transformers.__version__, "accelerate": accelerate.__version__, "safetensors": safetensors.__version__, "platform": platform.platform(), "gpu": ( torch.cuda.get_device_name(0) if torch.cuda.is_available() else None ), "gpu_memory_limit": args.gpu_memory, "cpu_memory_limit": args.cpu_memory, "pytorch_cuda_alloc_conf": os.environ.get( "PYTORCH_CUDA_ALLOC_CONF" ), "input_device": str(input_device), "device_map": device_map, "load_seconds": load_seconds, "forward_seconds": sum( source["forward_seconds"] for source in traced_sources ), "load_peak_cuda_memory_allocated_bytes": load_peak_cuda, "peak_cuda_memory_allocated_bytes": max( [load_peak_cuda, *peaks] ), "process_max_rss_kib": resource.getrusage( resource.RUSAGE_SELF ).ru_maxrss, }, "sources": traced_sources, "summary": aggregate_trace(traced_sources), "claim_boundary": [ "The trace observes one fixed SFT Chat checkpoint, not the Base checkpoint.", "Target comparisons exclude boundary-crossing tokens and require exact token IDs.", "Router and hidden associations do not establish mediation or capability causality.", "CPU in hf_device_map is offload residency, not proof of CPU matrix execution.", "Prompt-only use_cache=False traces are not generation-time KV or serving traces.", ], } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text( json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) payload = args.output.read_bytes() print( json.dumps( { "output": str(args.output), "sha256": hashlib.sha256(payload).hexdigest(), "bytes": len(payload), "sources": len(traced_sources), "decoder_layers": 27, "moe_gates": 26, "target_route_decisions": result["summary"][ "target_route_decisions" ], "load_seconds": load_seconds, "forward_seconds": result["execution"][ "forward_seconds" ], "peak_cuda_memory_allocated_bytes": result[ "execution" ]["peak_cuda_memory_allocated_bytes"], }, ensure_ascii=False, indent=2, ) ) if __name__ == "__main__": main()