handle JSON serialization

scripts/evalexperts.py

@@ -89,61 +89,61 @@ 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."""
-        # 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': torch.zeros(num_regular_experts, dtype=torch.long),
-                'small_expert_counts': torch.zeros(num_small_experts, dtype=torch.long) if num_small_experts > 0 else None,
-                'regular_expert_load': torch.zeros(num_regular_experts, dtype=torch.float),
-                'small_expert_load': torch.zeros(num_small_experts, dtype=torch.float) 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
+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()
         
-        # Update layer stats
-        layer_stats['total_tokens'] += num_tokens
+        layer_stats['regular_expert_counts'][expert_idx] += count
+        layer_stats['regular_expert_load'][expert_idx] += load
         
-        # Track regular experts
-        for expert_idx in range(num_regular_experts):
-            mask = (topk_experts_flat == expert_idx)
+        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['regular_expert_counts'][expert_idx] += count
-            layer_stats['regular_expert_load'][expert_idx] += load
+            layer_stats['small_expert_counts'][expert_idx] += count
+            layer_stats['small_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
-    
+            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 = {
@@ -495,15 +495,7 @@ def run_evaluation(args) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
 def save_results(results: Dict[str, Any], expert_stats: Dict[str, Any], args) -> str:
     """
-    Save evaluation results and expert statistics to file.
-    
-    Args:
-        results: Evaluation results
-        expert_stats: Expert usage statistics
-        args: Parsed command line arguments
-        
-    Returns:
-        str: Path to saved results file
+    Save evaluation results and expert statistics to file with proper serialization.
     """
     os.makedirs(args.output_dir, exist_ok=True)
     
@@ -514,10 +506,7 @@ def save_results(results: Dict[str, Any], expert_stats: Dict[str, Any], args) ->
         if len(args.tasks) > 3:
             tasks_str += f"_and_{len(args.tasks)-3}_more"
         
-        if args.model_type == "custom":
-            filename = f"{model_name}_custom_{tasks_str}_results_with_expert_stats.json"
-        else:
-            filename = f"{model_name}_transformers_{tasks_str}_results_with_expert_stats.json"
+        filename = f"{model_name}_{args.model_type}_{tasks_str}_results.json"
     else:
         filename = args.output_filename
         
@@ -534,7 +523,7 @@ def save_results(results: Dict[str, Any], expert_stats: Dict[str, Any], args) ->
         "num_fewshot": args.num_fewshot,
         "batch_size": args.batch_size,
         "device": args.device,
-        "dtype": args.dtype,
+        "dtype": str(args.dtype),  # Convert dtype to string
         "limit": args.limit,
     }
     
@@ -542,15 +531,31 @@ def save_results(results: Dict[str, Any], expert_stats: Dict[str, Any], args) ->
     if args.model_type == "custom":
         metadata["routing_type"] = "top-k (default)"
     
-    combined_results = {
+    def convert_for_json(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, torch.dtype)):
+            return str(obj) if isinstance(obj, torch.dtype) else obj.tolist()
+        elif isinstance(obj, (dict, list, tuple, str, int, float, bool, type(None))):
+            return obj
+        else:
+            return str(obj)
+    
+    # Convert all data to JSON-serializable format
+    serializable_results = {
         "metadata": metadata,
         "task_results": results,
-        "expert_statistics": expert_stats
+        "expert_statistics": {
+            k: convert_for_json(v) for k, v in expert_stats.items()
+        }
     }
     
     # Save to file
     with open(output_path, 'w') as f:
-        json.dump(combined_results, f, indent=2)
+        json.dump(serializable_results, f, indent=2)
     
     logger.info(f"Results saved to {output_path}")
     return output_path