wuyang
589 lines
22 KiB
Python
589 lines
22 KiB
Python
#!/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()
|