push debugging info for evalexperts

scripts/evalexperts.py

@@ -25,7 +25,7 @@ from lm_eval.models.huggingface import HFLM
 
 # Set up logging
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.DEBUG,  # Changed from INFO to DEBUG
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
@@ -90,15 +90,22 @@ class ExpertTrackingHFLM(HFLM):
         return expert_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):
-        """Update expert usage statistics."""
+                        topk_probs: torch.Tensor, num_regular_experts: int,
+                        num_small_experts: int, batch_size: int, seq_len: int):
+        """Update expert usage statistics with debug logging."""
+        # Debug: Print input parameters
+        logger.debug(f"\n{'='*40}")
+        logger.debug(f"Updating stats for layer {layer_idx}")
+        logger.debug(f"Input shapes - experts: {topk_experts.shape}, probs: {topk_probs.shape}")
+        logger.debug(f"Num experts - regular: {num_regular_experts}, small: {num_small_experts}")
+        
         # Flatten the batch and sequence dimensions
         topk_experts_flat = topk_experts.view(-1, topk_experts.size(-1))
         topk_probs_flat = topk_probs.view(-1, topk_probs.size(-1))
         
         # Initialize layer stats if not present
         if layer_idx not in self.expert_stats['layer_stats']:
+            logger.debug(f"Initializing new layer stats with {num_regular_experts} regular and {num_small_experts} small experts")
             self.expert_stats['layer_stats'][layer_idx] = {
                 'total_tokens': 0,
                 'regular_expert_counts': [0] * num_regular_experts,
@@ -110,40 +117,60 @@ class ExpertTrackingHFLM(HFLM):
         layer_stats = self.expert_stats['layer_stats'][layer_idx]
         num_tokens = topk_experts_flat.size(0)
         
-        # Update global stats
-        self.expert_stats['total_tokens'] += num_tokens
-        
-        # Update layer stats
-        layer_stats['total_tokens'] += num_tokens
+        # Debug: Print current layer stats structure
+        logger.debug(f"Current layer stats structure: {layer_stats.keys()}")
+        if layer_stats['small_expert_counts'] is None:
+            logger.debug("Small expert counts is None - no small experts initialized")
+        else:
+            logger.debug(f"Small expert counts length: {len(layer_stats['small_expert_counts'])}")
         
         # Track regular experts
+        regular_expert_used = False
         for expert_idx in range(num_regular_experts):
             mask = (topk_experts_flat == expert_idx)
             count = mask.sum().item()
-            load = topk_probs_flat[mask].sum().item()
-            
-            layer_stats['regular_expert_counts'][expert_idx] += count
-            layer_stats['regular_expert_load'][expert_idx] += load
-            
-            if expert_idx not in self.expert_stats['regular_expert_usage']:
-                self.expert_stats['regular_expert_usage'][expert_idx] = 0
-            self.expert_stats['regular_expert_usage'][expert_idx] += count
+            if count > 0:
+                regular_expert_used = True
+                layer_stats['regular_expert_counts'][expert_idx] += count
+                layer_stats['regular_expert_load'][expert_idx] += topk_probs_flat[mask].sum().item()
+                
+                if expert_idx not in self.expert_stats['regular_expert_usage']:
+                    self.expert_stats['regular_expert_usage'][expert_idx] = 0
+                self.expert_stats['regular_expert_usage'][expert_idx] += count
+        
+        # Debug: Regular expert usage
+        logger.debug(f"Regular experts used this batch: {regular_expert_used}")
         
         # Track small experts if they exist
         if num_small_experts > 0:
+            small_expert_used = False
             for expert_idx in range(num_small_experts):
                 small_expert_num = expert_idx + num_regular_experts
                 mask = (topk_experts_flat == small_expert_num)
                 count = mask.sum().item()
-                load = topk_probs_flat[mask].sum().item()
                 
-                layer_stats['small_expert_counts'][expert_idx] += count
-                layer_stats['small_expert_load'][expert_idx] += load
-                
-                if expert_idx not in self.expert_stats['small_expert_usage']:
-                    self.expert_stats['small_expert_usage'][expert_idx] = 0
-                self.expert_stats['small_expert_usage'][expert_idx] += count
-    
+                if count > 0:
+                    small_expert_used = True
+                    layer_stats['small_expert_counts'][expert_idx] += count
+                    layer_stats['small_expert_load'][expert_idx] += topk_probs_flat[mask].sum().item()
+                    
+                    if expert_idx not in self.expert_stats['small_expert_usage']:
+                        self.expert_stats['small_expert_usage'][expert_idx] = 0
+                    self.expert_stats['small_expert_usage'][expert_idx] += count
+            
+            # Debug: Small expert usage
+            logger.debug(f"Small experts used this batch: {small_expert_used}")
+            if not small_expert_used:
+                logger.debug(f"Top-k experts sample: {topk_experts_flat[:5].tolist()}")
+                logger.debug(f"Num regular experts: {num_regular_experts}, looking for experts >= this number")
+        else:
+            logger.debug("No small experts configured for this layer")
+        
+        # Update token counts
+        self.expert_stats['total_tokens'] += num_tokens
+        layer_stats['total_tokens'] += num_tokens
+        logger.debug(f"Updated token counts - layer: {layer_stats['total_tokens']}, total: {self.expert_stats['total_tokens']}")
+
     def get_expert_stats(self) -> Dict[str, Any]:
         """Return expert usage statistics in a serializable format."""
         def convert(obj):