feat: trace DeepSeek V2-Lite real routes

This commit is contained in:
wuyang
2026-07-29 14:18:45 +08:00
parent a97653af69
commit e86420556f
14 changed files with 25897 additions and 34 deletions
+56
View File
@@ -0,0 +1,56 @@
# DeepSeek real-weight execution probes
The probes in this directory use official DeepSeek artifacts and keep their
scope deliberately narrower than a full-model benchmark.
## DeepSeek-V2-Lite truncated trace
`v2_lite_trace.py` executes layers 0–6 from the official BF16 checkpoint. Those
seven layers are fully contained in shard 1; layer 7 is split across shards 1
and 2 and is therefore outside the default evidence boundary.
Pinned model:
```text
deepseek-ai/DeepSeek-V2-Lite@604d5664dddd88a0433dbae533b7fe9472482de0
```
Required Python stack:
```text
torch==2.11.0+cu128
transformers==4.41.2
safetensors==0.8.0
```
The 2024 remote code does not import under Transformers 5.5 because
`is_torch_fx_available` was removed. The probe imports the official files as a
read-only local package; it does not patch the model source.
Download the metadata, tokenizer, remote code, index, and first shard with the
Hugging Face CLI, then run:
```bash
python experiments/deepseek/v2_lite_trace.py \
--artifact-dir /path/to/deepseek-v2-lite \
--output src/data/deepseek-v2-lite-trace.json
```
The result contains:
- real tokenizer pieces and model-derived hidden states;
- the actual `[B,T,576]` MLA compressed projection at each executed layer;
- the expanded key/value tensors stored by the Hugging Face eager cache;
- token-level top-6 routed expert IDs and weights for six MoE layers;
- per-layer and per-prompt expert-load summaries;
- explicit boundaries against global load, expert semantics, training traces,
full-model generation, and production serving claims.
Run the probe twice and compare deterministic evidence while excluding timing:
```bash
python experiments/deepseek/compare_v2_lite_traces.py \
--first /path/to/trace-1.json \
--second /path/to/trace-2.json \
--output src/data/deepseek-v2-lite-trace-repro.json
```
@@ -0,0 +1,110 @@
#!/usr/bin/env python3
"""Compare deterministic evidence from two DeepSeek-V2-Lite trace runs."""
from __future__ import annotations
import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--first", type=Path, required=True)
parser.add_argument("--second", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--captured-at", default=None)
return parser.parse_args()
def read(path: Path) -> dict[str, Any]:
return json.loads(path.read_text())
def sha256(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def main() -> None:
args = parse_args()
first = read(args.first)
second = read(args.second)
checks: list[dict[str, Any]] = []
def check(name: str, left: Any, right: Any) -> None:
checks.append({"name": name, "exact": left == right})
check("provenance", first["provenance"], second["provenance"])
check("configuration", first["configuration"], second["configuration"])
check("prompt_tokenization", first["prompts"], second["prompts"])
check(
"initial_hidden_hash",
first["execution"]["initial_hidden"]["sha256_fp32"],
second["execution"]["initial_hidden"]["sha256_fp32"],
)
check(
"final_hidden_hash",
first["execution"]["final_hidden"]["sha256_fp32"],
second["execution"]["final_hidden"]["sha256_fp32"],
)
for left, right in zip(
first["execution"]["layers"],
second["execution"]["layers"],
strict=True,
):
layer = left["layer"]
check(
f"layer_{layer}_hidden_hash",
left["hidden_after"]["sha256_fp32"],
right["hidden_after"]["sha256_fp32"],
)
check(f"layer_{layer}_mla_shapes", left["mla"], right["mla"])
if "routing" in left:
check(
f"layer_{layer}_aggregate_load",
left["routing"]["aggregate_load"],
right["routing"]["aggregate_load"],
)
check(
f"layer_{layer}_token_routes",
[
prompt["token_routes"]
for prompt in left["routing"]["per_prompt"]
],
[
prompt["token_routes"]
for prompt in right["routing"]["per_prompt"]
],
)
captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
result = {
"schema_version": 1,
"captured_at": captured_at,
"first": {
"captured_at": first["captured_at"],
"sha256": sha256(args.first),
},
"second": {
"captured_at": second["captured_at"],
"sha256": sha256(args.second),
},
"timing_compared": False,
"checks": checks,
"exact_checks": sum(item["exact"] for item in checks),
"total_checks": len(checks),
"all_exact": all(item["exact"] for item in checks),
}
if not result["all_exact"]:
failed = [item["name"] for item in checks if not item["exact"]]
raise AssertionError(f"trace mismatch: {failed}")
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
+588
View File
@@ -0,0 +1,588 @@
#!/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()