diff --git a/probe_keys.py b/probe_keys.py new file mode 100644 index 00000000..ebaffae2 --- /dev/null +++ b/probe_keys.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Probe actual checkpoint weight keys and shapes for DSV4-Pro NVFP4. +Focus on layer 0 (HCA) and layer 2 (CSA) non-expert weights. +""" +import json, os +from safetensors.torch import load_file +from pathlib import Path + +CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +cdir = Path(CHECKPOINT_DIR) + +with open(cdir / "model.safetensors.index.json") as f: + wmap = json.load(f)["weight_map"] + +# Find all keys for layer 0 and 2, excluding expert weights +for li in [0, 2]: + pfx = f"model.layers.{li}." + keys = sorted([k for k in wmap if k.startswith(pfx) and 'experts.' not in k and 'shared_expert' not in k]) + print(f"\n=== Layer {li} non-expert keys ({len(keys)}) ===") + for k in keys: + print(f" {k}") + +# Global keys +global_keys = sorted([k for k in wmap if not k.startswith("model.layers.")]) +print(f"\n=== Global keys ({len(global_keys)}) ===") +for k in global_keys: + print(f" {k}") + +# Load actual shapes for layer 0 and 2 +for li in [0, 2]: + pfx = f"model.layers.{li}." + keys = [k for k in wmap if k.startswith(pfx) and 'experts.' not in k and 'shared_expert' not in k] + shards = set(wmap[k] for k in keys) + print(f"\n=== Layer {li} shapes ===") + for sn in sorted(shards): + if not (cdir / sn).exists(): continue + w = load_file(str(cdir / sn)) + for k in sorted(w.keys()): + if k.startswith(pfx) and 'experts.' not in k and 'shared_expert' not in k: + print(f" {k}: {tuple(w[k].shape)} {w[k].dtype}") + +# Expert key pattern for layer 0 (just the first few) +pfx0 = f"model.layers.0.mlp." +expert_keys = sorted([k for k in wmap if k.startswith(pfx0) and 'experts.' in k])[:12] +print(f"\n=== Layer 0 expert keys (sample) ===") +for k in expert_keys: + print(f" {k}") + +# Shared expert keys for layer 0 +shared_keys = sorted([k for k in wmap if k.startswith(pfx0) and 'shared_expert' in k]) +print(f"\n=== Layer 0 shared expert keys ===") +for k in shared_keys: + print(f" {k}") + +# Router/gate keys for layer 0 +gate_keys = sorted([k for k in wmap if k.startswith(pfx0) and ('gate' in k or 'e_score' in k or 'tid2eid' in k or 'router' in k)]) +print(f"\n=== Layer 0 gate/router keys ===") +for k in gate_keys: + print(f" {k}")