#!/usr/bin/env python3 """Run a fixed, multi-domain DeepSeek-V2-Lite routing corpus. The probe executes the official embedding and decoder layers 0 through 6 from the first official BF16 safetensors shard. Layers 1 through 6 are MoE layers. Their real top-6 routed-expert IDs are aggregated at prompt level before any statistics are computed. Sampling is deterministic and source-addressable. Confidence intervals use a prompt-level, within-domain bootstrap; tokens from the same prompt are never treated as independent observations. """ from __future__ import annotations import argparse import gc import gzip import hashlib import importlib.util import json import math import platform import re import subprocess import sys import types from datetime import datetime, timezone from itertools import combinations from pathlib import Path from typing import Any import numpy as np import pyarrow.parquet as pq import torch import torch.nn.functional as F from safetensors import safe_open from transformers import AutoTokenizer DOMAIN_ORDER = ("english", "chinese", "code", "math") DOMAIN_LABELS = { "english": "English encyclopedia", "chinese": "中文新闻", "code": "Python code", "math": "Grade-school math", } MIN_TOKENS = { "english": 24, "chinese": 8, "code": 24, "math": 16, } SAMPLE_SALT = "llm-atlas-deepseek-routing-v1" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--artifact-dir", 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("--per-domain", type=int, default=32) parser.add_argument("--max-tokens", type=int, default=96) parser.add_argument("--batch-size", type=int, default=16) parser.add_argument("--layers", type=int, default=7) parser.add_argument("--bootstrap", type=int, default=2000) parser.add_argument("--seed", type=int, default=20260729) 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 text_sha256(value: str) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest() def git_revision(path: Path) -> str | None: current = path.resolve() for candidate in (current, *current.parents): if (candidate / ".git").exists(): return subprocess.check_output( ["git", "-C", str(candidate), "rev-parse", "HEAD"], text=True, ).strip() return None def scoped_seed(seed: int, scope: str) -> int: payload = f"{seed}:{scope}".encode() return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") def load_official_modules(root: Path) -> tuple[Any, Any]: package_name = "deepseek_v2_lite_official_routing_corpus" 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 read_jsonl(path: Path) -> list[dict[str, Any]]: with path.open(encoding="utf-8") as handle: return [json.loads(line) for line in handle if line.strip()] def load_candidates(args: argparse.Namespace) -> dict[str, list[dict[str, str]]]: with gzip.open(args.human_eval, "rt", encoding="utf-8") as handle: human_eval = [json.loads(line) for line in handle if line.strip()] gsm8k = read_jsonl(args.gsm8k) tnews = read_jsonl(args.tnews) wikitext = pq.read_table(args.wikitext, columns=["text"]).column("text").to_pylist() return { "english": [ { "id": f"wikitext2/raw-validation/{index:04d}", "text": str(text).strip(), } for index, text in enumerate(wikitext) if str(text).strip() and not re.fullmatch(r"=+\s*.*?\s*=+", str(text).strip()) ], "chinese": [ { "id": f"tnews/test/{row['id']}", "text": str(row["sentence"]).strip(), } for row in tnews if str(row.get("sentence", "")).strip() ], "code": [ { "id": str(row["task_id"]), "text": str(row["prompt"]).rstrip(), } for row in human_eval if str(row.get("prompt", "")).strip() ], "math": [ { "id": f"gsm8k/test/{index:04d}", "text": str(row["question"]).strip(), } for index, row in enumerate(gsm8k) if str(row.get("question", "")).strip() ], } def select_corpus( tokenizer: Any, candidates: dict[str, list[dict[str, str]]], per_domain: int, max_tokens: int, ) -> tuple[list[dict[str, Any]], dict[str, dict[str, int]]]: selected: list[dict[str, Any]] = [] counts: dict[str, dict[str, int]] = {} for domain in DOMAIN_ORDER: rows = candidates[domain] texts = [row["text"] for row in rows] encoded: list[list[int]] = [] for start in range(0, len(texts), 512): result = tokenizer( texts[start : start + 512], add_special_tokens=True, truncation=True, max_length=max_tokens, padding=False, ) encoded.extend(result.input_ids) eligible = [] for row, token_ids in zip(rows, encoded, strict=True): if len(token_ids) < MIN_TOKENS[domain]: continue rank = hashlib.sha256( f"{SAMPLE_SALT}|{domain}|{row['id']}".encode() ).hexdigest() eligible.append( { "id": row["id"], "domain": domain, "label": DOMAIN_LABELS[domain], "text": row["text"], "text_sha256": text_sha256(row["text"]), "characters": len(row["text"]), "token_ids": token_ids, "tokens": len(token_ids), "selection_rank": rank, } ) eligible.sort(key=lambda row: (row["selection_rank"], row["id"])) if len(eligible) < per_domain: raise RuntimeError( f"{domain} has only {len(eligible)} eligible prompts; need {per_domain}" ) domain_selection = eligible[:per_domain] for within_domain_index, row in enumerate(domain_selection): row["within_domain_index"] = within_domain_index selected.extend(domain_selection) counts[domain] = { "candidate_records_after_text_filter": len(rows), "eligible_records": len(eligible), "selected_records": len(domain_selection), "valid_tokens": sum(row["tokens"] for row in domain_selection), } return selected, counts def make_batches( samples: list[dict[str, Any]], batch_size: int, pad_token_id: int, ) -> list[dict[str, Any]]: ordered = sorted( samples, key=lambda row: (row["tokens"], row["domain"], row["id"]), ) batches = [] for start in range(0, len(ordered), batch_size): rows = ordered[start : start + batch_size] sequence = max(row["tokens"] for row in rows) input_ids = torch.full( (len(rows), sequence), pad_token_id, dtype=torch.long, ) attention_mask = torch.zeros((len(rows), sequence), dtype=torch.long) for index, row in enumerate(rows): length = row["tokens"] input_ids[index, :length] = torch.tensor(row["token_ids"]) attention_mask[index, :length] = 1 batches.append( { "samples": rows, "input_ids": input_ids, "attention_mask": attention_mask, "padded_sequence": sequence, } ) return batches def distribution(loads: np.ndarray, mode: str) -> np.ndarray: if mode == "token_weighted": values = loads.sum(axis=0, dtype=np.float64) return values / values.sum() if mode == "prompt_balanced": prompt_distributions = loads / loads.sum(axis=1, keepdims=True) values = prompt_distributions.mean(axis=0, dtype=np.float64) return values / values.sum() raise ValueError(mode) def metric_vector(distributions: np.ndarray) -> dict[str, np.ndarray]: values = np.atleast_2d(distributions).astype(np.float64, copy=False) expert_count = values.shape[1] means = values.mean(axis=1) cv = values.std(axis=1) / means ordered = np.sort(values, axis=1) indices = np.arange(1, expert_count + 1, dtype=np.float64) gini = ( ((2 * indices - expert_count - 1) * ordered).sum(axis=1) / (expert_count * ordered.sum(axis=1)) ) log_values = np.zeros_like(values) np.log(values, out=log_values, where=values > 0) entropy = -(values * log_values).sum(axis=1) return { "cv": cv, "gini": gini, "effective_experts": np.exp(entropy), "top_expert_share": values.max(axis=1), "used_experts": (values > 0).sum(axis=1).astype(np.float64), } def interval(values: np.ndarray) -> list[float]: low, high = np.quantile(values, [0.025, 0.975], axis=0) if np.ndim(low) == 0: return [float(low), float(high)] return np.stack([low, high], axis=-1).tolist() def bootstrap_domain( loads: np.ndarray, mode: str, replicates: int, seed: int, scope: str, ) -> dict[str, Any]: point_distribution = distribution(loads, mode) prompt_count = loads.shape[0] rng = np.random.default_rng(scoped_seed(seed, scope)) sampled = rng.integers( 0, prompt_count, size=(replicates, prompt_count), endpoint=False, ) if mode == "token_weighted": bootstrap_loads = loads[sampled].sum(axis=1, dtype=np.float64) bootstrap_distributions = bootstrap_loads / bootstrap_loads.sum( axis=1, keepdims=True ) else: prompt_distributions = loads / loads.sum(axis=1, keepdims=True) bootstrap_distributions = prompt_distributions[sampled].mean(axis=1) bootstrap_distributions /= bootstrap_distributions.sum(axis=1, keepdims=True) point_metrics = metric_vector(point_distribution) bootstrap_metrics = metric_vector(bootstrap_distributions) metrics = { name: { "point": float(point_metrics[name][0]), "ci95": interval(bootstrap_metrics[name]), } for name in point_metrics } order = np.argsort(point_distribution)[::-1][:8] return { "distribution": point_distribution.tolist(), "expert_share_ci95": interval(bootstrap_distributions), "metrics": metrics, "top_experts": [ { "expert": int(expert), "share": float(point_distribution[expert]), "ci95": interval(bootstrap_distributions[:, expert]), } for expert in order ], } def js_divergence(left: np.ndarray, right: np.ndarray) -> np.ndarray: left = np.atleast_2d(left).astype(np.float64, copy=False) right = np.atleast_2d(right).astype(np.float64, copy=False) midpoint = (left + right) / 2 left_log_ratio = np.zeros_like(left) right_log_ratio = np.zeros_like(right) left_ratio = np.ones_like(left) right_ratio = np.ones_like(right) np.divide(left, midpoint, out=left_ratio, where=left > 0) np.divide(right, midpoint, out=right_ratio, where=right > 0) np.log(left_ratio, out=left_log_ratio, where=left > 0) np.log(right_ratio, out=right_log_ratio, where=right > 0) return 0.5 * ( (left * left_log_ratio).sum(axis=1) + (right * right_log_ratio).sum(axis=1) ) def bootstrap_pair( left: np.ndarray, right: np.ndarray, mode: str, replicates: int, seed: int, scope: str, ) -> dict[str, Any]: left_point = distribution(left, mode) right_point = distribution(right, mode) rng = np.random.default_rng(scoped_seed(seed, scope)) left_indices = rng.integers( 0, left.shape[0], size=(replicates, left.shape[0]), endpoint=False ) right_indices = rng.integers( 0, right.shape[0], size=(replicates, right.shape[0]), endpoint=False ) if mode == "token_weighted": left_boot = left[left_indices].sum(axis=1, dtype=np.float64) right_boot = right[right_indices].sum(axis=1, dtype=np.float64) left_boot /= left_boot.sum(axis=1, keepdims=True) right_boot /= right_boot.sum(axis=1, keepdims=True) else: left_prompt = left / left.sum(axis=1, keepdims=True) right_prompt = right / right.sum(axis=1, keepdims=True) left_boot = left_prompt[left_indices].mean(axis=1) right_boot = right_prompt[right_indices].mean(axis=1) left_boot /= left_boot.sum(axis=1, keepdims=True) right_boot /= right_boot.sum(axis=1, keepdims=True) bootstrap_jsd = js_divergence(left_boot, right_boot) return { "point": float(js_divergence(left_point, right_point)[0]), "ci95": interval(bootstrap_jsd), "unit": "nats", "upper_bound": math.log(2), } def layer_statistics( prompt_rows: list[dict[str, Any]], replicates: int, seed: int, layer_index: int, ) -> dict[str, Any]: loads_by_domain = { domain: np.asarray( [row["load"] for row in prompt_rows if row["domain"] == domain], dtype=np.int64, ) for domain in DOMAIN_ORDER } modes: dict[str, Any] = {} for mode in ("token_weighted", "prompt_balanced"): domains = { domain: bootstrap_domain( loads_by_domain[domain], mode, replicates, seed, f"layer={layer_index}|mode={mode}|domain={domain}", ) for domain in DOMAIN_ORDER } pairs = [] for left, right in combinations(DOMAIN_ORDER, 2): pairs.append( { "left": left, "right": right, "js_divergence": bootstrap_pair( loads_by_domain[left], loads_by_domain[right], mode, replicates, seed, f"layer={layer_index}|mode={mode}|pair={left}:{right}", ), } ) modes[mode] = {"domains": domains, "pairs": pairs} return modes def gpu_identity(device: torch.device) -> dict[str, Any] | None: if device.type != "cuda": return None index = device.index or 0 properties = torch.cuda.get_device_properties(index) driver = subprocess.check_output( [ "nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits", f"--id={index}", ], text=True, ).strip() return { "name": properties.name, "total_memory_bytes": properties.total_memory, "compute_capability": list(torch.cuda.get_device_capability(index)), "driver": driver, } 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, args.human_eval, args.gsm8k, args.tnews, args.tnews_archive, args.wikitext, ] missing = [str(path) for path in required if not path.exists()] if missing: raise FileNotFoundError(f"missing artifacts: {missing}") if args.device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable") if not 2 <= args.layers <= 7: raise ValueError("need layer 0 plus at least one MoE layer; shard ends at layer 6") if args.per_domain < 2: raise ValueError("per-domain sample must be at least two") if args.bootstrap < 100: raise ValueError("bootstrap replicates must be at least 100") torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) torch.backends.cuda.matmul.allow_tf32 = False 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" candidates = load_candidates(args) samples, corpus_counts = select_corpus( tokenizer, candidates, args.per_domain, args.max_tokens, ) batches = make_batches(samples, args.batch_size, tokenizer.pad_token_id) device = torch.device(args.device) with safe_open(shard, framework="pt", device="cpu") as handle: embedding = handle.get_tensor("model.embed_tokens.weight") for batch in batches: batch["hidden"] = F.embedding(batch["input_ids"], embedding) del embedding layer_results = [] for layer_index in range(args.layers): prefix = f"model.layers.{layer_index}." 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 layer = layer.to(device=device, dtype=torch.bfloat16).eval() prompt_rows = [] next_hidden = [] for batch in batches: hidden = batch["hidden"].to(device) attention_mask = batch["attention_mask"].to(device) sequence = batch["padded_sequence"] position_ids = torch.arange(sequence, device=device).unsqueeze(0) causal_mask = modeling._prepare_4d_causal_attention_mask( attention_mask, batch["input_ids"].shape, hidden, 0, ) captures: dict[str, torch.Tensor] = {} hook = None if layer_index > 0: def capture_gate(_module: Any, _inputs: Any, output: Any) -> None: captures["topk_ids"] = output[0].detach().cpu() hook = layer.mlp.gate.register_forward_hook(capture_gate) with torch.inference_mode(): outputs = layer( hidden, attention_mask=causal_mask, position_ids=position_ids, past_key_value=None, use_cache=False, ) next_hidden.append(outputs[0].cpu()) if hook is not None: hook.remove() topk_ids = captures["topk_ids"].view( len(batch["samples"]), sequence, config.num_experts_per_tok, ) for sample_index, sample in enumerate(batch["samples"]): length = sample["tokens"] load = torch.bincount( topk_ids[sample_index, :length].flatten(), minlength=config.n_routed_experts, ) prompt_rows.append( { "id": sample["id"], "domain": sample["domain"], "tokens": length, "routes": int(load.sum()), "load": load.tolist(), } ) del hidden, attention_mask, causal_mask, outputs for batch, hidden in zip(batches, next_hidden, strict=True): batch["hidden"] = hidden result: dict[str, Any] = { "layer": layer_index, "ffn": "dense" if layer_index == 0 else "moe", "state_numel": state_numel, "state_bytes": state_bytes, } if layer_index > 0: prompt_rows.sort( key=lambda row: ( DOMAIN_ORDER.index(row["domain"]), row["id"], ) ) result["valid_tokens"] = sum(row["tokens"] for row in prompt_rows) result["routes"] = sum(row["routes"] for row in prompt_rows) result["prompts"] = prompt_rows result["statistics"] = layer_statistics( prompt_rows, args.bootstrap, args.seed, layer_index, ) layer_results.append(result) del layer, next_hidden gc.collect() if device.type == "cuda": torch.cuda.empty_cache() captured_at = args.captured_at or datetime.now(timezone.utc).isoformat() index = json.loads((root / "model.safetensors.index.json").read_text()) selected_identity = [ { "id": sample["id"], "domain": sample["domain"], "within_domain_index": sample["within_domain_index"], "selection_rank": sample["selection_rank"], "text_sha256": sample["text_sha256"], "characters": sample["characters"], "tokens": sample["tokens"], } for sample in samples ] result = { "schema_version": 1, "captured_at": captured_at, "evidence_identity": "X / official BF16 weights, fixed public corpus, local truncated forward", "boundary": { "model": "DeepSeek-V2-Lite base", "executed_layers": list(range(args.layers)), "measured_moe_layers": list(range(1, args.layers)), "total_model_layers": config.num_hidden_layers, "full_model_generation": False, "training_or_online_load": False, "expert_semantics_inferred": False, "population": ( f"{len(samples)} fixed public prompts across four domains; " "not training data, online traffic, or a task-performance benchmark" ), "code_execution": False, "answers_used": False, }, "provenance": { "model": { "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, }, "corpora": { "english": { "name": "WikiText-2 raw validation", "url": "https://huggingface.co/datasets/Salesforce/wikitext", "revision": "b08601e04326c79dfdd32d625aee71d232d685c3", "file_sha256": sha256(args.wikitext), "field_used": "text", }, "chinese": { "name": "CLUE TNEWS public test", "url": "https://github.com/CLUEbenchmark/CLUE", "download_url": "https://storage.googleapis.com/cluebenchmark/tasks/tnews_public.zip", "revision": "9e61ddd3659ddb57ed82b4d0ba0a8613dfb55a2e", "archive_sha256": sha256(args.tnews_archive), "file_sha256": sha256(args.tnews), "field_used": "sentence", }, "code": { "name": "OpenAI HumanEval", "url": "https://github.com/openai/human-eval", "revision": git_revision(args.human_eval), "file_sha256": sha256(args.human_eval), "field_used": "prompt", }, "math": { "name": "OpenAI GSM8K test", "url": "https://github.com/openai/grade-school-math", "revision": git_revision(args.gsm8k), "file_sha256": sha256(args.gsm8k), "field_used": "question", }, }, }, "environment": { "python": platform.python_version(), "platform": platform.platform(), "torch": torch.__version__, "torch_cuda": torch.version.cuda, "transformers": __import__("transformers").__version__, "safetensors": __import__("safetensors").__version__, "numpy": np.__version__, "pyarrow": __import__("pyarrow").__version__, "device": str(device), "gpu": gpu_identity(device), "matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32, }, "configuration": { "layers": config.num_hidden_layers, "hidden": config.hidden_size, "routed_experts": config.n_routed_experts, "active_routed_experts": config.num_experts_per_tok, "shared_experts": config.n_shared_experts, "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, }, "corpus_contract": { "domains": list(DOMAIN_ORDER), "domain_labels": DOMAIN_LABELS, "sample_salt": SAMPLE_SALT, "selection": "ascending SHA256(salt|domain|source_id), then source_id", "per_domain": args.per_domain, "max_tokens": args.max_tokens, "minimum_tokens": MIN_TOKENS, "special_tokens": True, "truncation": "right", "counts": corpus_counts, "selected": selected_identity, }, "inference_contract": { "batch_size": args.batch_size, "batches": len(batches), "batch_order": "ascending token count, then domain, then source id", "attention": "official eager", "dtype": "BF16", "total_prompts": len(samples), "valid_tokens": sum(sample["tokens"] for sample in samples), "routes_per_moe_layer": sum(sample["tokens"] for sample in samples) * config.num_experts_per_tok, }, "statistical_contract": { "resampling_unit": "prompt", "strata": "domain", "replicates": args.bootstrap, "seed": args.seed, "interval": "95% percentile bootstrap", "modes": { "token_weighted": "sum route counts, so longer prompts contribute more", "prompt_balanced": "normalize each prompt first, then give every prompt equal weight", }, "metrics": { "cv": "population standard deviation across 64 expert shares divided by their mean", "gini": "Gini coefficient across 64 expert shares", "effective_experts": "exp(Shannon entropy) in nats", "top_expert_share": "largest routed-expert share", "used_experts": "experts with nonzero share in the resample", "js_divergence": "Jensen-Shannon divergence in nats; bounded by ln(2)", }, "multiple_comparison_correction": False, }, "layers": layer_results, } 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( { "output": str(args.output), "prompts": len(samples), "tokens": result["inference_contract"]["valid_tokens"], "routes_per_moe_layer": result["inference_contract"][ "routes_per_moe_layer" ], "measured_moe_layers": args.layers - 1, }, ensure_ascii=False, ) ) if __name__ == "__main__": main()