76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Probe indexer and compressor weight shapes from the checkpoint.
|
|
This tells us the ACTUAL dimensions, not what we assume.
|
|
Run via: fire_b200_test probe_indexer_shapes.py
|
|
"""
|
|
import json, sys
|
|
from pathlib import Path
|
|
from safetensors.torch import load_file
|
|
|
|
CHECKPOINT = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
|
|
|
def main():
|
|
cdir = Path(CHECKPOINT)
|
|
with open(cdir / "config.json") as f:
|
|
cfg = json.load(f)
|
|
|
|
n_layers = cfg["num_hidden_layers"]
|
|
n_ih = cfg.get("index_n_heads", 64)
|
|
ihd = cfg.get("index_head_dim", 128)
|
|
hd = cfg["head_dim"]
|
|
cr = cfg.get("compress_ratios", [128] * n_layers)
|
|
|
|
print(f"Config: n_ih={n_ih}, ihd={ihd}, hd={hd}")
|
|
print(f"n_ih * ihd = {n_ih * ihd}")
|
|
print(f"2 * ihd = {2 * ihd}")
|
|
print(f"2 * hd = {2 * hd}")
|
|
print(f"Compress ratios: first5={cr[:5]}")
|
|
print()
|
|
|
|
# Load weight map to find indexer weights
|
|
idx_file = cdir / "model.safetensors.index.json"
|
|
if idx_file.exists():
|
|
with open(idx_file) as f:
|
|
wmap = json.load(f).get("weight_map", {})
|
|
|
|
# Find indexer/compressor weights for layer 2 (first CSA layer)
|
|
for li in [0, 1, 2, 3]:
|
|
pfx = f"model.layers.{li}.self_attn"
|
|
print(f"\n=== Layer {li} (ratio={cr[li] if li < len(cr) else '?'}) ===")
|
|
for k in sorted(wmap.keys()):
|
|
if k.startswith(pfx) and ('compressor' in k or 'indexer' in k or 'q_b_proj' in k or 'kv_proj' in k or 'gate_proj' in k):
|
|
shard = cdir / wmap[k]
|
|
print(f" {k} -> shard {wmap[k]}")
|
|
else:
|
|
print("No index file, loading all weights...")
|
|
|
|
# Actually load some weights and print shapes
|
|
# Just load the first shard to get shapes
|
|
print("\n=== Loading weight shapes ===")
|
|
all_w = {}
|
|
if idx_file.exists():
|
|
shards = set(wmap.values())
|
|
for sn in sorted(shards):
|
|
sf = cdir / sn
|
|
if sf.exists():
|
|
w = load_file(str(sf))
|
|
# Only print relevant keys
|
|
for k, v in w.items():
|
|
if ('compressor' in k or 'indexer' in k) and 'layers.2' in k:
|
|
print(f" {k}: shape={list(v.shape)} dtype={v.dtype}")
|
|
del w
|
|
|
|
# Also check q_b_proj for layer 2
|
|
print("\n=== Layer 2 attention projection shapes ===")
|
|
for sn in sorted(shards):
|
|
sf = cdir / sn
|
|
if sf.exists():
|
|
w = load_file(str(sf))
|
|
for k, v in w.items():
|
|
if 'layers.2.self_attn' in k and ('q_b' in k or 'kv_proj' in k or 'gate_proj' in k):
|
|
print(f" {k}: shape={list(v.shape)} dtype={v.dtype}")
|
|
del w
|
|
|
|
if __name__ == "__main__":
|
|
main()
|