robust get_expert_stats and remove repeat

scripts/evalexperts.py

@@ -146,158 +146,25 @@ class ExpertTrackingHFLM(HFLM):
     
     def get_expert_stats(self) -> Dict[str, Any]:
         """Return expert usage statistics in a serializable format."""
+        def convert(obj):
+            """Recursively convert objects to JSON-serializable formats."""
+            if isinstance(obj, (np.integer, np.floating)):
+                return int(obj) if isinstance(obj, np.integer) else float(obj)
+            elif isinstance(obj, np.ndarray):
+                return obj.tolist()
+            elif isinstance(obj, torch.Tensor):
+                return obj.cpu().numpy().tolist()
+            elif isinstance(obj, torch.dtype):
+                return str(obj)
+            elif isinstance(obj, (dict)):
+                return {k: convert(v) for k, v in obj.items()}
+            elif isinstance(obj, (list, tuple)):
+                return [convert(v) for v in obj]
+            else:
+                return obj
+
         stats = {
-            'total_tokens': self.expert_stats['total_tokens'],
-            'regular_expert_usage': {},
-            'small_expert_usage': {},
-            'layer_stats': {}
-        }
-        
-        # Convert regular expert usage
-        for expert_idx, count in self.expert_stats['regular_expert_usage'].items():
-            stats['regular_expert_usage'][expert_idx] = {
-                'count': count,
-                'percentage': count / (self.expert_stats['total_tokens'] * getattr(self.model.config, 'top_k', 1)) * 100
-            }
-        
-        # Convert small expert usage if they exist
-        if self.expert_stats['small_expert_usage']:
-            for expert_idx, count in self.expert_stats['small_expert_usage'].items():
-                stats['small_expert_usage'][expert_idx] = {
-                    'count': count,
-                    'percentage': count / (self.expert_stats['total_tokens'] * getattr(self.model.config, 'top_k', 1)) * 100
-                }
-        
-        # Convert layer stats
-        for layer_idx, layer_stat in self.expert_stats['layer_stats'].items():
-            stats['layer_stats'][layer_idx] = {
-                'total_tokens': layer_stat['total_tokens'],
-                'regular_expert_counts': layer_stat['regular_expert_counts'],
-                'regular_expert_load': layer_stat['regular_expert_load'],
-                'small_expert_counts': layer_stat['small_expert_counts'],
-                'small_expert_load': layer_stat['small_expert_load']
-            }
-        
-        return stats
-    
-    def print_expert_stats(self) -> None:
-        """Print expert usage statistics in a human-readable format."""
-        if not self.expert_stats['total_tokens']:
-            print("No expert usage statistics collected.")
-            return
-            
-        total_tokens = self.expert_stats['total_tokens']
-        top_k = getattr(self.model.config, 'top_k', 1)
-        total_expert_activations = total_tokens * top_k
-        
-        print("\n" + "="*80)
-        print("EXPERT USAGE STATISTICS")
-        print("="*80)
-        print(f"Total tokens processed: {total_tokens:,}")
-        print(f"Total expert activations (top-{top_k}): {total_expert_activations:,}")
-        print("\nOverall Expert Usage:")
-        
-        # Print regular experts
-        if self.expert_stats['regular_expert_usage']:
-            print("\nRegular Experts:")
-            for expert_idx, count in sorted(self.expert_stats['regular_expert_usage'].items()):
-                percentage = count / total_expert_activations * 100
-                print(f"  Expert {expert_idx}: {count:,} ({percentage:.2f}%)")
-        
-        # Print small experts if they exist
-        if self.expert_stats['small_expert_usage']:
-            print("\nSmall Experts:")
-            for expert_idx, count in sorted(self.expert_stats['small_expert_usage'].items()):
-                percentage = count / total_expert_activations * 100
-                print(f"  Small Expert {expert_idx}: {count:,} ({percentage:.2f}%)")
-        
-        # Print layer-wise statistics
-        print("\nLayer-wise Statistics:")
-        for layer_idx, layer_stat in self.expert_stats['layer_stats'].items():
-            print(f"\nLayer {layer_idx}:")
-            print(f"  Tokens processed: {layer_stat['total_tokens']:,}")
-            
-            # Regular experts
-            print("  Regular Experts:")
-            for expert_idx, (count, load) in enumerate(zip(
-                layer_stat['regular_expert_counts'], 
-                layer_stat['regular_expert_load']
-            )):
-                count_pct = count / (layer_stat['total_tokens'] * top_k) * 100
-                load_pct = load / layer_stat['total_tokens'] * 100
-                print(f"    Expert {expert_idx}: Count={count:,} ({count_pct:.2f}%), Load={load:.2f} ({load_pct:.2f}%)")
-            
-            # Small experts if they exist
-            if layer_stat['small_expert_counts'] is not None:
-                print("  Small Experts:")
-                for expert_idx, (count, load) in enumerate(zip(
-                    layer_stat['small_expert_counts'], 
-                    layer_stat['small_expert_load']
-                )):
-                    count_pct = count / (layer_stat['total_tokens'] * top_k) * 100
-                    load_pct = load / layer_stat['total_tokens'] * 100
-                    print(f"    Small Expert {expert_idx}: Count={count:,} ({count_pct:.2f}%), Load={load:.2f} ({load_pct:.2f}%)")
-        
-        print("="*80 + "\n")
-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 with serializable data types."""
-    # 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']:
-        self.expert_stats['layer_stats'][layer_idx] = {
-            'total_tokens': 0,
-            'regular_expert_counts': [0] * num_regular_experts,  # Use list instead of tensor
-            'small_expert_counts': [0] * num_small_experts if num_small_experts > 0 else None,
-            'regular_expert_load': [0.0] * num_regular_experts,
-            'small_expert_load': [0.0] * num_small_experts if num_small_experts > 0 else None
-        }
-    
-    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
-    
-    # Track regular experts
-    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
-    
-    # Track small experts if they exist
-    if num_small_experts > 0:
-        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
- 
-    def get_expert_stats(self) -> Dict[str, Any]:
-        """Return expert usage statistics in a serializable format."""
-        stats = {
-            'total_tokens': self.expert_stats['total_tokens'],
+            'total_tokens': convert(self.expert_stats['total_tokens']),
             'regular_expert_usage': {},
             'small_expert_usage': {},
             'layer_stats': {}
@@ -306,30 +173,30 @@ def _update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor,
         # Convert regular expert usage
         for expert_idx, count in self.expert_stats['regular_expert_usage'].items():
             stats['regular_expert_usage'][expert_idx] = {
-                'count': count,
-                'percentage': count / (self.expert_stats['total_tokens'] * self.model.config.top_k) * 100
+                'count': convert(count),
+                'percentage': convert(count / (self.expert_stats['total_tokens'] * getattr(self.model.config, 'top_k', 1)) * 100)
             }
         
         # Convert small expert usage if they exist
         if self.expert_stats['small_expert_usage']:
             for expert_idx, count in self.expert_stats['small_expert_usage'].items():
                 stats['small_expert_usage'][expert_idx] = {
-                    'count': count,
-                    'percentage': count / (self.expert_stats['total_tokens'] * self.model.config.top_k) * 100
+                    'count': convert(count),
+                    'percentage': convert(count / (self.expert_stats['total_tokens'] * getattr(self.model.config, 'top_k', 1)) * 100)
                 }
         
         # Convert layer stats
         for layer_idx, layer_stat in self.expert_stats['layer_stats'].items():
             stats['layer_stats'][layer_idx] = {
-                'total_tokens': layer_stat['total_tokens'],
-                'regular_expert_counts': layer_stat['regular_expert_counts'].tolist(),
-                'regular_expert_load': layer_stat['regular_expert_load'].tolist(),
-                'small_expert_counts': layer_stat['small_expert_counts'].tolist() if layer_stat['small_expert_counts'] is not None else None,
-                'small_expert_load': layer_stat['small_expert_load'].tolist() if layer_stat['small_expert_load'] is not None else None
+                'total_tokens': convert(layer_stat['total_tokens']),
+                'regular_expert_counts': convert(layer_stat['regular_expert_counts']),
+                'regular_expert_load': convert(layer_stat['regular_expert_load']),
+                'small_expert_counts': convert(layer_stat['small_expert_counts']),
+                'small_expert_load': convert(layer_stat['small_expert_load'])
             }
         
         return stats
-    
+
     def print_expert_stats(self) -> None:
         """Print expert usage statistics in a human-readable format."""
         if not self.expert_stats['total_tokens']:
@@ -390,7 +257,6 @@ def _update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor,
         
         print("="*80 + "\n")
 
-
 def parse_args():
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(