try patch hook

myolmoe/config.json

@@ -30,9 +30,5 @@
   "torch_dtype": "float32",
   "transformers_version": "4.52.4",
   "use_cache": true,
-  "vocab_size": 50304,
-  "small_expert_intermediate_ratio": 16,
-  "small_expert_count": 64,
-  "small_expert_sparsity_coef": 0.1,
-  "max_small_expert_count": 64
+  "vocab_size": 50304
 }

scripts/evalexperts.py

@@ -57,38 +57,31 @@ class ExpertTrackingHFLM(HFLM):
                     self._make_expert_hook(layer_idx)
                 )
     
-    def _make_expert_hook(self, layer_idx: int):
-        """Create a forward hook for tracking expert usage in a specific layer."""
-        def expert_hook(module, input, output):
-            if not hasattr(module, 'gate') or not hasattr(module, 'experts'):
-                return
-                
-            hidden_states, router_logits = input[0], output[1]
-            batch_size, seq_len, hidden_dim = hidden_states.shape
-            
-            # Get routing probabilities
-            routing_probs = torch.softmax(router_logits, dim=-1)
-            
-            # Get top-k experts
-            topk_probs, topk_experts = torch.topk(
-                routing_probs, 
-                k=module.top_k, 
-                dim=-1
-            )
-            
-            # Update statistics
-            self.update_expert_stats(
-                layer_idx=layer_idx,
-                topk_experts=topk_experts,
-                topk_probs=topk_probs,
-                num_regular_experts=module.num_experts,
-                num_small_experts=module.num_small_experts if hasattr(module, 'num_small_experts') else 0,
-                batch_size=batch_size,
-                seq_len=seq_len
-            )
-            
-        return expert_hook
-    
+    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)
+
+            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 expert_id < num_small_experts:
+                        expert_stats[layer_idx]['small'] = expert_stats[layer_idx].get('small', 0) + 1
+
+        return hook
+
     def update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor,
                         topk_probs: torch.Tensor, num_regular_experts: int,
                         num_small_experts: int, batch_size: int, seq_len: int):