filter expert keys

dump_checkpoint_keys.py

@@ -1,10 +1,9 @@
 #!/usr/bin/env python3
-"""Dump checkpoint key names and shapes to help understand the model structure."""
+"""Dump checkpoint key names and shapes — attention and compressor keys only."""
 import json
 from pathlib import Path
-from safetensors.torch import load_file
 
-CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
+CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro"
 
 def main():
     cdir = Path(CHECKPOINT_DIR)
@@ -12,25 +11,16 @@ def main():
     if index_path.exists():
         with open(index_path) as f:
             weight_map = json.load(f).get("weight_map", {})
-        # Collect unique key prefixes for layer 0 and layer 2 (CSA)
-        for li in [0, 1, 2, 3, 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 ===")
-            for k in keys:
+            # Filter: show everything EXCEPT individual expert weights
+            filtered = [k for k in keys if '.experts.' not in k]
+            print(f"\n=== Layer {li} keys (non-expert) ===")
+            for k in filtered:
                 print(f"  {k}")
     else:
-        print("No index file found, loading first shard...")
-        shards = sorted(cdir.glob("model-*.safetensors"))
-        if shards:
-            data = load_file(str(shards[0]))
-            # Print layer 0 and 2 keys
-            for li in [0, 1, 2]:
-                prefix = f"model.layers.{li}."
-                keys = sorted(k for k in data if k.startswith(prefix))
-                print(f"\n=== Layer {li} keys (from {shards[0].name}) ===")
-                for k in keys:
-                    print(f"  {k}: {data[k].shape}")
+        print("No index file found!")
 
 if __name__ == "__main__":
     main()