inspect

scripts/evalexperts.py

@@ -57,30 +57,30 @@ class ExpertTrackingHFLM(HFLM):
                     self._make_expert_hook(layer_idx)
                 )
     
-    def _make_expert_hook(layer_idx, model):
-        def hook(module, input, output):
-            # Get expert routing data from output
-            if isinstance(output, tuple) and len(output) == 2:
-                hidden_states, routing_weights = output
-            else:
-                hidden_states = output
-                routing_weights = None
+def _make_expert_hook(layer_idx, model):
+    def hook(module, input, output):
+        # Get expert routing data from output
+        if isinstance(output, tuple) and len(output) == 2:
+            hidden_states, routing_weights = output
+        else:
+            hidden_states = output
+            routing_weights = None
 
-            # Always use the config value for num_small_experts
-            num_small_experts = getattr(model.config, 'small_expert_count', 0)
+        # Always use the config value for num_small_experts
+        num_small_experts = getattr(model.config, 'small_expert_count', 0)
 
-            expert_stats[layer_idx] = expert_stats.get(layer_idx, {})
-            expert_stats[layer_idx]['total'] = expert_stats[layer_idx].get('total', 0) + 1
+        expert_stats[layer_idx] = expert_stats.get(layer_idx, {})
+        expert_stats[layer_idx]['total'] = expert_stats[layer_idx].get('total', 0) + 1
 
-            if routing_weights is not None:
-                top_expert = routing_weights.argmax(dim=-1)
-                for expert_id in top_expert.view(-1).tolist():
-                    expert_stats[layer_idx][expert_id] = expert_stats[layer_idx].get(expert_id, 0) + 1
+        if routing_weights is not None:
+            top_expert = routing_weights.argmax(dim=-1)
+            for expert_id in top_expert.view(-1).tolist():
+                expert_stats[layer_idx][expert_id] = expert_stats[layer_idx].get(expert_id, 0) + 1
 
-                    if expert_id < num_small_experts:
-                        expert_stats[layer_idx]['small'] = expert_stats[layer_idx].get('small', 0) + 1
+                if expert_id < num_small_experts:
+                    expert_stats[layer_idx]['small'] = expert_stats[layer_idx].get('small', 0) + 1
 
-        return hook
+    return hook
 
     def update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor,
                         topk_probs: torch.Tensor, num_regular_experts: int,

scripts/inspectexperts.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+import argparse
+import os
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path to the fine-tuned checkpoint directory (e.g., ./checkpoints/checkpoint-16000)",
+    )
+    parser.add_argument(
+        "--custom_model_path",
+        type=str,
+        required=False,
+        help="(Optional) Path to the model implementation source if needed",
+    )
+    args = parser.parse_args()
+
+    print(f"Loading config from: {args.model_path}")
+    config = AutoConfig.from_pretrained(args.model_path)
+
+    if hasattr(config, "num_small_experts"):
+        num_small_experts = config.num_small_experts
+    else:
+        raise ValueError("The model config does not contain 'num_small_experts'.")
+
+    print(f"Number of small experts: {num_small_experts}")
+
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+    model.eval()
+
+    print("Inspecting small expert weights...")
+    total_params = 0
+    matched_params = 0
+    for name, param in model.named_parameters():
+        total_params += 1
+        if f"small_experts." in name:
+            matched_params += 1
+            print(f"[Matched] {name} - shape: {tuple(param.shape)}")
+    print(f"\nMatched {matched_params}/{total_params} parameters containing 'small_experts.'")
+
+    print("Done.")
+
+if __name__ == "__main__":
+    main()