focused key dump

dump_checkpoint_keys.py

@@ -1,34 +1,32 @@
 #!/usr/bin/env python3
-"""Dump checkpoint key names and shapes."""
+"""Dump checkpoint key names for specific layers only."""
 import json
 from pathlib import Path
+from safetensors.torch import load_file
 
 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:
+    with open(index_path) as f:
+        weight_map = json.load(f).get("weight_map", {})
+    
+    # Only show layer 0, 2, and 60 keys (non-expert)
+    for li in [0, 2, 60]:
+        prefix = f"model.layers.{li}."
+        keys = sorted(k for k in weight_map if k.startswith(prefix))
+        # Filter out individual expert weights
+        filtered = [k for k in keys if '.experts.' not in k]
+        print(f"\n=== Layer {li} keys ({len(filtered)} non-expert) ===")
+        for k in filtered:
             print(f"  {k}")
-    else:
-        print("No index file found, listing directory...")
-        for f in sorted(cdir.glob("*.safetensors"))[:3]:
-            print(f"  {f.name}")
+    
+    # Non-layer keys (short list)
+    other_keys = sorted(k for k in weight_map if not k.startswith("model.layers."))
+    print(f"\n=== Non-layer keys ({len(other_keys)}) ===")
+    for k in other_keys:
+        print(f"  {k}")
 
 if __name__ == "__main__":
     main()