feat: trace DeepSeek V2-Lite real routes
This commit is contained in:
wuyang
@@ -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()
|
||||
Reference in New Issue
Block a user