#!/usr/bin/env python3 """Trace real DeepSeek-V2-Lite MLA states and MoE routes from official weights. The script deliberately executes only the contiguous layers fully contained in the first official safetensors shard (layers 0 through 6). It loads one decoder layer at a time, so a 32 GB workstation GPU can produce model-derived hidden states without downloading or materializing the complete 15.7B-parameter model. No model source is patched. The pinned remote-code files are imported as a read-only local package, and every decoder layer is the official class. """ from __future__ import annotations import argparse import gc import hashlib import importlib.util import json import math import platform import statistics import subprocess import sys import time import types from datetime import datetime, timezone from itertools import combinations from pathlib import Path from typing import Any import torch import torch.nn.functional as F from safetensors import safe_open from transformers import AutoTokenizer, DynamicCache DEFAULT_PROMPTS = [ { "id": "zh_explanation", "label": "中文解释", "text": "用通俗的语言解释，为什么稀疏专家模型可以拥有很多参数，但每个 token 只使用其中一小部分。", }, { "id": "en_architecture", "label": "English architecture", "text": "Explain how a compressed key-value latent changes the memory cost of autoregressive decoding.", }, { "id": "code", "label": "Python code", "text": "Write a Python function that returns the first repeated element in a list and explain its complexity.", }, { "id": "math", "label": "数学推理", "text": "若正数 x 满足 x 加上它的倒数等于 3，求 x 的平方加上倒数的平方。", }, ] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--artifact-dir", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--layers", type=int, default=7) parser.add_argument("--device", default="cuda") parser.add_argument("--captured-at", default=None) return parser.parse_args() def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for block in iter(lambda: handle.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def tensor_sha256(tensor: torch.Tensor) -> str: value = tensor.detach().float().contiguous().cpu().numpy() return hashlib.sha256(value.tobytes()).hexdigest() def load_official_modules(root: Path) -> tuple[Any, Any]: """Import relative official files without Transformers' remote-code scanner.""" package_name = "deepseek_v2_lite_official" package = types.ModuleType(package_name) package.__path__ = [str(root)] sys.modules[package_name] = package loaded = {} for leaf in ("configuration_deepseek", "modeling_deepseek"): name = f"{package_name}.{leaf}" spec = importlib.util.spec_from_file_location(name, root / f"{leaf}.py") if spec is None or spec.loader is None: raise RuntimeError(f"cannot load official module: {leaf}") module = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) loaded[leaf] = module return loaded["configuration_deepseek"], loaded["modeling_deepseek"] def percentile(values: torch.Tensor, q: float) -> float: return torch.quantile(values.float(), q).item() def hidden_stats(hidden: torch.Tensor, mask: torch.Tensor) -> dict[str, Any]: valid = hidden[mask.bool()] rms = valid.float().square().mean(dim=-1).sqrt() norms = valid.float().norm(dim=-1) return { "valid_tokens": valid.shape[0], "hidden_rms": { "mean": rms.mean().item(), "p10": percentile(rms, 0.10), "p50": percentile(rms, 0.50), "p90": percentile(rms, 0.90), }, "l2_norm": { "mean": norms.mean().item(), "p10": percentile(norms, 0.10), "p50": percentile(norms, 0.50), "p90": percentile(norms, 0.90), }, "finite": bool(torch.isfinite(valid).all()), "sha256_fp32": tensor_sha256(valid), } def load_metrics(load: torch.Tensor) -> dict[str, Any]: values = load.float() mean = values.mean() ordered = values.sort().values count = values.numel() indices = torch.arange(1, count + 1, dtype=torch.float32) denominator = count * ordered.sum() gini = ( ((2 * indices - count - 1) * ordered).sum() / denominator if denominator else torch.tensor(0.0) ) probabilities = values / values.sum().clamp_min(1) nonzero = probabilities[probabilities > 0] entropy = -(nonzero * nonzero.log()).sum() return { "routes": int(values.sum()), "used_experts": int((values > 0).sum()), "zero_experts": int((values == 0).sum()), "mean": mean.item(), "std": values.std().item(), "cv": (values.std() / mean).item() if mean else 0.0, "gini": gini.item(), "entropy_nats": entropy.item(), "effective_experts": math.exp(entropy.item()), "min": values.min().item(), "p50": percentile(values, 0.50), "p90": percentile(values, 0.90), "max": values.max().item(), } def route_trace( tokenizer: Any, input_ids: torch.Tensor, attention_mask: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, prompts: list[dict[str, str]], expert_count: int, ) -> dict[str, Any]: batch, sequence, top_k = topk_ids.shape valid_routes = attention_mask.bool().unsqueeze(-1).expand_as(topk_ids) aggregate_load = torch.bincount( topk_ids[valid_routes].cpu(), minlength=expert_count ) per_prompt = [] used_sets: dict[str, set[int]] = {} for batch_index, prompt in enumerate(prompts): length = int(attention_mask[batch_index].sum()) ids = input_ids[batch_index, :length].tolist() tokens = tokenizer.convert_ids_to_tokens(ids) routes = topk_ids[batch_index, :length].cpu() weights = topk_weights[batch_index, :length].float().cpu() load = torch.bincount(routes.flatten(), minlength=expert_count) used_sets[prompt["id"]] = set(torch.nonzero(load, as_tuple=False).flatten().tolist()) token_rows = [] for position, (token_id, token, experts, scores) in enumerate( zip(ids, tokens, routes.tolist(), weights.tolist(), strict=True) ): ordered = sorted( zip(experts, scores, strict=True), key=lambda pair: pair[1], reverse=True, ) token_rows.append( { "position": position, "token_id": token_id, "token": token, "decoded_piece": tokenizer.decode( [token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False, ), "experts_by_weight": [ {"expert": expert, "weight": weight} for expert, weight in ordered ], "selected_weight_sum": sum(scores), "top1_top2_margin": ordered[0][1] - ordered[1][1], } ) top_experts = sorted( enumerate(load.tolist()), key=lambda pair: pair[1], reverse=True )[:8] per_prompt.append( { "id": prompt["id"], "label": prompt["label"], "tokens": length, "load": load.tolist(), "metrics": load_metrics(load), "top_experts": [ {"expert": expert, "routes": routes} for expert, routes in top_experts if routes ], "token_routes": token_rows, } ) jaccard = [] for left, right in combinations(prompts, 2): a = used_sets[left["id"]] b = used_sets[right["id"]] jaccard.append( { "left": left["id"], "right": right["id"], "used_expert_jaccard": len(a & b) / len(a | b), "shared_experts": len(a & b), "union_experts": len(a | b), } ) valid_weights = topk_weights[valid_routes].float().cpu().view(-1, top_k) weight_sums = valid_weights.sum(dim=-1) ordered_weights = valid_weights.sort(dim=-1, descending=True).values top_experts = sorted( enumerate(aggregate_load.tolist()), key=lambda pair: pair[1], reverse=True )[:12] return { "aggregate_load": aggregate_load.tolist(), "aggregate_metrics": load_metrics(aggregate_load), "top_experts": [ {"expert": expert, "routes": routes} for expert, routes in top_experts if routes ], "selected_weight_sum": { "mean": weight_sums.mean().item(), "p10": percentile(weight_sums, 0.10), "p50": percentile(weight_sums, 0.50), "p90": percentile(weight_sums, 0.90), }, "top1_top2_margin": { "mean": (ordered_weights[:, 0] - ordered_weights[:, 1]).mean().item(), "p50": percentile(ordered_weights[:, 0] - ordered_weights[:, 1], 0.50), "p90": percentile(ordered_weights[:, 0] - ordered_weights[:, 1], 0.90), }, "per_prompt": per_prompt, "prompt_pair_jaccard": jaccard, } def nvidia_smi() -> dict[str, str]: fields = [ "name", "driver_version", "memory.total", "power.limit", "clocks.max.sm", ] output = subprocess.check_output( [ "nvidia-smi", f"--query-gpu={','.join(fields)}", "--format=csv,noheader,nounits", ], text=True, ).strip() return dict(zip(fields, [item.strip() for item in output.split(",")], strict=True)) def main() -> None: args = parse_args() root = args.artifact_dir.resolve() shard = root / "model-00001-of-000004.safetensors" required = [ root / "config.json", root / "configuration_deepseek.py", root / "modeling_deepseek.py", root / "model.safetensors.index.json", root / "tokenizer.json", shard, ] missing = [str(path) for path in required if not path.exists()] if missing: raise FileNotFoundError(f"missing official artifacts: {missing}") if args.device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable") if not 1 <= args.layers <= 7: raise ValueError("shard 1 fully contains only layers 0 through 6") configuration, modeling = load_official_modules(root) config = configuration.DeepseekV2Config.from_pretrained(root) config._attn_implementation = "eager" tokenizer = AutoTokenizer.from_pretrained( root, trust_remote_code=True, local_files_only=True, ) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" encoded = tokenizer( [prompt["text"] for prompt in DEFAULT_PROMPTS], add_special_tokens=True, padding=True, return_tensors="pt", ) input_ids = encoded.input_ids attention_mask_cpu = encoded.attention_mask device = torch.device(args.device) start = time.perf_counter() with safe_open(shard, framework="pt", device="cpu") as handle: embedding = handle.get_tensor("model.embed_tokens.weight") hidden = F.embedding(input_ids, embedding).to(device) embedding_ms = (time.perf_counter() - start) * 1000 attention_mask = attention_mask_cpu.to(device) sequence = input_ids.shape[1] position_ids = torch.arange(sequence, device=device).unsqueeze(0) causal_mask = modeling._prepare_4d_causal_attention_mask( attention_mask, input_ids.shape, hidden, 0, ) cache = DynamicCache() layer_results = [] initial_hidden = hidden_stats(hidden, attention_mask) for layer_index in range(args.layers): prefix = f"model.layers.{layer_index}." load_start = time.perf_counter() with safe_open(shard, framework="pt", device="cpu") as handle: state = { key[len(prefix) :]: handle.get_tensor(key) for key in handle.keys() if key.startswith(prefix) } state_numel = sum(value.numel() for value in state.values()) state_bytes = sum(value.numel() * value.element_size() for value in state.values()) with torch.device("meta"): layer = modeling.DeepseekV2DecoderLayer(config, layer_index) layer.to_empty(device="cpu") layer.load_state_dict(state, strict=True, assign=True) del state load_ms = (time.perf_counter() - load_start) * 1000 torch.cuda.reset_peak_memory_stats(device) if device.type == "cuda" else None transfer_start = time.perf_counter() layer = layer.to(device=device, dtype=torch.bfloat16).eval() if device.type == "cuda": torch.cuda.synchronize(device) transfer_ms = (time.perf_counter() - transfer_start) * 1000 captures: dict[str, torch.Tensor] = {} def capture_kv(_module: Any, _inputs: Any, output: torch.Tensor) -> None: captures["compressed_kv"] = output.detach() handles = [ layer.self_attn.kv_a_proj_with_mqa.register_forward_hook(capture_kv) ] if layer_index > 0: def capture_gate(_module: Any, _inputs: Any, output: Any) -> None: captures["topk_ids"] = output[0].detach() captures["topk_weights"] = output[1].detach() handles.append(layer.mlp.gate.register_forward_hook(capture_gate)) before = hidden_stats(hidden, attention_mask) if device.type == "cuda": torch.cuda.synchronize(device) forward_start = time.perf_counter() with torch.inference_mode(): outputs = layer( hidden, attention_mask=causal_mask, position_ids=position_ids, past_key_value=cache, use_cache=True, ) hidden = outputs[0] if device.type == "cuda": torch.cuda.synchronize(device) forward_ms = (time.perf_counter() - forward_start) * 1000 after = hidden_stats(hidden, attention_mask) for handle in handles: handle.remove() compressed = captures["compressed_kv"] latent = compressed[..., : config.kv_lora_rank] rope_key = compressed[..., config.kv_lora_rank :] key_cache = cache.key_cache[layer_index] value_cache = cache.value_cache[layer_index] result: dict[str, Any] = { "layer": layer_index, "ffn": "dense" if layer_index == 0 else "moe", "state_tensors": len(layer.state_dict()), "state_numel": state_numel, "state_bytes": state_bytes, "timing_ms": { "load_cpu": load_ms, "transfer_to_device": transfer_ms, "forward": forward_ms, }, "peak_device_allocated_mib": ( torch.cuda.max_memory_allocated(device) / 2**20 if device.type == "cuda" else None ), "hidden_before": before, "hidden_after": after, "mla": { "compressed_projection_shape": list(compressed.shape), "latent_shape": list(latent.shape), "rope_key_shape": list(rope_key.shape), "latent_rms": latent.float().square().mean().sqrt().item(), "rope_key_rms": rope_key.float().square().mean().sqrt().item(), "eager_key_cache_shape": list(key_cache.shape), "eager_value_cache_shape": list(value_cache.shape), "eager_cache_bytes": key_cache.numel() * key_cache.element_size() + value_cache.numel() * value_cache.element_size(), }, } if layer_index > 0: topk_ids = captures["topk_ids"].view( input_ids.shape[0], input_ids.shape[1], -1 ) topk_weights = captures["topk_weights"].view( input_ids.shape[0], input_ids.shape[1], -1 ) result["routing"] = route_trace( tokenizer, input_ids, attention_mask_cpu, topk_ids.cpu(), topk_weights.cpu(), DEFAULT_PROMPTS, config.n_routed_experts, ) layer_results.append(result) del layer, compressed, latent, rope_key, key_cache, value_cache captures.clear() gc.collect() if device.type == "cuda": torch.cuda.empty_cache() valid_token_count = int(attention_mask_cpu.sum()) latent_elements = config.kv_lora_rank + config.qk_rope_head_dim eager_elements = config.num_attention_heads * ( config.qk_nope_head_dim + config.qk_rope_head_dim + config.v_head_dim ) captured_at = args.captured_at or datetime.now(timezone.utc).isoformat() index = json.loads((root / "model.safetensors.index.json").read_text()) result = { "schema_version": 1, "captured_at": captured_at, "evidence_identity": "X / official BF16 weights and tokenizer, local truncated forward", "boundary": { "model": "DeepSeek-V2-Lite base", "executed_layers": list(range(args.layers)), "total_model_layers": config.num_hidden_layers, "full_model_generation": False, "training_or_global_expert_load": False, "expert_semantics_inferred": False, "prompt_sample": "four authored prompts; descriptive trace, not population estimate", "cache_note": "HF eager materializes expanded K/V; latent-cache arithmetic is derived from official dimensions", }, "provenance": { "huggingface_model": "deepseek-ai/DeepSeek-V2-Lite", "huggingface_revision": "604d5664dddd88a0433dbae533b7fe9472482de0", "sha256": { "config": sha256(root / "config.json"), "modeling_code": sha256(root / "modeling_deepseek.py"), "tokenizer": sha256(root / "tokenizer.json"), "index": sha256(root / "model.safetensors.index.json"), "shard_1": sha256(shard), }, "checkpoint_tensor_bytes": index["metadata"]["total_size"], "shard_1_bytes": shard.stat().st_size, }, "environment": { "python": platform.python_version(), "platform": platform.platform(), "libc": list(platform.libc_ver()), "torch": torch.__version__, "torch_cuda": torch.version.cuda, "transformers": __import__("transformers").__version__, "safetensors": __import__("safetensors").__version__, "device": str(device), "nvidia_smi": nvidia_smi() if device.type == "cuda" else None, "matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32 if device.type == "cuda" else None, }, "configuration": { "total_parameters_reported": "15.7B", "activated_parameters_reported": "2.4B", "layers": config.num_hidden_layers, "hidden": config.hidden_size, "attention_heads": config.num_attention_heads, "qk_nope_head_dim": config.qk_nope_head_dim, "qk_rope_head_dim": config.qk_rope_head_dim, "v_head_dim": config.v_head_dim, "kv_lora_rank": config.kv_lora_rank, "routed_experts": config.n_routed_experts, "active_routed_experts": config.num_experts_per_tok, "shared_experts": config.n_shared_experts, "expert_intermediate": config.moe_intermediate_size, "first_dense_layers": config.first_k_dense_replace, "router_scoring": config.scoring_func, "router_topk_method": config.topk_method, "normalize_selected_weights": config.norm_topk_prob, }, "prompts": [ { **prompt, "token_ids": input_ids[index, : int(attention_mask_cpu[index].sum())].tolist(), "tokens": tokenizer.convert_ids_to_tokens( input_ids[index, : int(attention_mask_cpu[index].sum())].tolist() ), } for index, prompt in enumerate(DEFAULT_PROMPTS) ], "execution": { "batch": input_ids.shape[0], "padded_sequence": input_ids.shape[1], "valid_tokens": valid_token_count, "embedding_ms": embedding_ms, "initial_hidden": initial_hidden, "layers": layer_results, "final_hidden": hidden_stats(hidden, attention_mask), }, "cache_accounting": { "dtype": "BF16", "latent_elements_per_token_layer": latent_elements, "latent_bytes_per_token_layer": latent_elements * 2, "hf_eager_elements_per_token_layer": eager_elements, "hf_eager_bytes_per_token_layer": eager_elements * 2, "eager_over_latent_ratio": eager_elements / latent_elements, "latent_reduction_vs_eager": 1 - latent_elements / eager_elements, "components": { "latent_content": config.kv_lora_rank, "rope_key": config.qk_rope_head_dim, "expanded_key": config.num_attention_heads * (config.qk_nope_head_dim + config.qk_rope_head_dim), "expanded_value": config.num_attention_heads * config.v_head_dim, }, }, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n") print(json.dumps(result, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()