HFLM forward to loglikelihoods

scripts/evalexperts.py

@@ -169,12 +169,12 @@ def load_custom_model(args) -> Tuple[AutoModelForCausalLM, AutoTokenizer, HFLM]:
     
     # Create HFLM wrapper for evaluation
     hf_model = HFLM(
-        pretrained=model,
-        tokenizer=tokenizer,
+        pretrained=args.model_path,  # Pass model path as string
         device=args.device,
         batch_size=args.batch_size,
         max_batch_size=args.max_batch_size,
-        dtype=args.dtype
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code
     )
     
     logger.info("Custom model, tokenizer, and HFLM wrapper loaded successfully")
@@ -194,16 +194,20 @@ def track_expert_usage(model, input_ids: torch.Tensor) -> List[Dict[int, int]]:
     expert_usage = [{} for _ in range(model.config.num_hidden_layers)]
     
     def hook_fn(module, input, output, layer_idx):
-        # Assuming the module outputs selected expert indices
         if hasattr(module, 'selected_experts'):  # Hypothetical attribute
             selected_experts = module.selected_experts  # Shape: (batch_size, seq_len, top_k)
             for expert_idx in selected_experts.flatten().tolist():
                 expert_usage[layer_idx][expert_idx] = expert_usage[layer_idx].get(expert_idx, 0) + 1
+        elif hasattr(module, 'routing_weights'):  # Alternative: use routing weights
+            weights = module.routing_weights  # Shape: (batch_size, seq_len, num_experts)
+            top_k_indices = torch.topk(weights, k=model.config.top_k, dim=-1).indices
+            for expert_idx in top_k_indices.flatten().tolist():
+                expert_usage[layer_idx][expert_idx] = expert_usage[layer_idx].get(expert_idx, 0) + 1
     
     # Register hooks for each MoE layer
     hooks = []
     for i, layer in enumerate(model.transformer.layers):  # Adjust based on actual model structure
-        if hasattr(layer, 'moe'):  # Check if layer has MoE component
+        if hasattr(layer, 'moe'):
             hook = layer.moe.register_forward_hook(lambda m, inp, out: hook_fn(m, inp, out, i))
             hooks.append(hook)
     
@@ -237,19 +241,22 @@ def run_evaluation_with_tracking(model, hf_model, tokenizer, args) -> Tuple[Dict
     # Initialize expert usage tracking for each task
     task_expert_usage = {task: [] for task in args.tasks}
     
-    # Custom evaluation loop to track expert usage
-    def custom_forward(model, batch):
-        input_ids = batch["input_ids"].to(model.device)
-        # Track expert usage for this batch
-        batch_expert_usage = track_expert_usage(model, input_ids)
-        # Accumulate usage for the task
-        task_name = batch.get("task_name", args.tasks[0])  # Fallback to first task
-        task_expert_usage[task_name].append(batch_expert_usage)
-        return model(input_ids)
-    
-    # Override HFLM's forward method to include expert tracking
-    original_forward = hf_model.forward
-    hf_model.forward = lambda batch: custom_forward(model, batch)
+    # Custom batch processing to track expert usage
+    def custom_loglikelihood(self, requests):
+        from lm_eval.api.instance import Instance
+        res = []
+        for request in requests:
+            input_ids = tokenizer(request.arguments[0], return_tensors="pt").input_ids.to(model.device)
+            # Track expert usage
+            batch_expert_usage = track_expert_usage(model, input_ids)
+            task_expert_usage[request.task_name].append(batch_expert_usage)
+            # Original loglikelihood computation
+            res.append(self._loglikelihood([request]))
+        return [item for sublist in res for item in sublist]
+    
+    # Override HFLM's loglikelihood method
+    original_loglikelihood = hf_model.loglikelihood
+    hf_model.loglikelihood = custom_loglikelihood.__get__(hf_model, HFLM)
     
     # Run evaluation
     results = evaluator.simple_evaluate(
@@ -261,8 +268,8 @@ def run_evaluation_with_tracking(model, hf_model, tokenizer, args) -> Tuple[Dict
         max_batch_size=args.max_batch_size,
     )
     
-    # Restore original forward method
-    hf_model.forward = original_forward
+    # Restore original method
+    hf_model.loglikelihood = original_loglikelihood
     
     # Aggregate expert usage per task
     aggregated_usage = {}