Files
nvfp4-megamoe-kernel/dump_checkpoint_keys.py
2026-05-31 21:24:58 +00:00

35 lines
1.1 KiB
Python

#!/usr/bin/env python3
"""Dump checkpoint key names and shapes."""
import json
from pathlib import Path
CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro"
def main():
cdir = Path(CHECKPOINT_DIR)
index_path = cdir / "model.safetensors.index.json"
if index_path.exists():
with open(index_path) as f:
weight_map = json.load(f).get("weight_map", {})
# Get all unique prefixes for layer 0-3 and 59-60
for li in [0, 1, 2, 3, 59, 60]:
prefix = f"model.layers.{li}."
keys = sorted(k for k in weight_map if k.startswith(prefix))
print(f"\n=== Layer {li} keys ({len(keys)} total) ===")
for k in keys:
print(f" {k}")
# Also show non-layer keys
other_keys = sorted(k for k in weight_map if not k.startswith("model.layers."))
print(f"\n=== Non-layer keys ===")
for k in other_keys:
print(f" {k}")
else:
print("No index file found, listing directory...")
for f in sorted(cdir.glob("*.safetensors"))[:3]:
print(f" {f.name}")
if __name__ == "__main__":
main()