Dynamic MOE Results and Analysis

common_utils.py

@@ -24,6 +24,9 @@ def get_topk_experts_from_json(path, top_k=5, mode="most", criterion=None, out_p
     experts_to_prune = {}
     for k, v in sorted_experts.items():
         k_int = int(k)
+        if top_k == 0:
+            experts_to_prune[k_int] = []
+            continue
         v_int = [int(x) for x in v]
         if mode == "most":
             selected = v_int[:top_k]
@@ -42,4 +45,4 @@ def get_topk_experts_from_json(path, top_k=5, mode="most", criterion=None, out_p
           (f" -> saved to '{out_path}'" if out_path else ""))
     print(f"Layers: {num_layers}, total experts selected: {total_experts}")
 
-    return experts_to_prune
+    return experts_to_prune

evaluation.py

@@ -18,8 +18,9 @@ def parse_args():
     parser.add_argument('--use_pruned_model', action='store_true', help='Whether to use pruned model')
     parser.add_argument('--pruned_metadata', type=str, default=None, help='Path to pruned expert metadata JSON')
     parser.add_argument('--mode', type=str, choices=['least', 'mode'], default='least', help='Strategy for selecting experts to prune')
-    parser.add_argument('--pruning_method', type=str, choices=['mask', 'zero'], default='zero', help='Method to use for pruning experts')
+    parser.add_argument('--pruning_method', type=str, choices=['mask', 'zero', 'dynamic'], default='zero', help='Method to use for pruning experts')
     parser.add_argument('--k', type=int, default=5, help='Number of experts to prune per layer')
+    parser.add_argument('--dynamic_routing_threshold', type=float, default=0.8, help='Cumulative probability threshold for dynamic routing')
     parser.add_argument('--device', type=str, default='cuda', help='Device for model')
     parser.add_argument('--output_file', type=str, default=None, help='File to save results JSON')
 
@@ -39,7 +40,12 @@ def main():
             top_k=args.k,
             mode=args.mode
         )
-        apply_pruning(model.model, experts_to_prune, mode=args.pruning_method)
+        apply_pruning(
+            model.model, 
+            experts_to_prune, 
+            mode=args.pruning_method, 
+            dynamic_routing_threshold=args.dynamic_routing_threshold
+        )
 
     # Prepare arguments for simple_evaluate
     eval_kwargs = dict(
@@ -55,12 +61,29 @@ def main():
     results = simple_evaluate(**eval_kwargs)
     print(results)
 
+    expert_activation_report = {}
+    if args.pruning_method == "dynamic":
+        total_avg = 0
+        num_layers = 0
+        for i, layer in enumerate(model.model.model.layers):
+            moe_block = layer.mlp
+            if hasattr(moe_block, "num_activated_experts_log") and moe_block.num_activated_experts_log:
+                avg_experts = sum(moe_block.num_activated_experts_log) / len(moe_block.num_activated_experts_log)
+                expert_activation_report[f"layer_{i}"] = f"{avg_experts:.2f}"
+                total_avg += avg_experts
+                num_layers += 1
+        if num_layers > 0:
+            expert_activation_report["overall_average"] = f"{total_avg / num_layers:.2f}"
+
     if args.output_file:
         output_file = args.output_file
         output_data = {
             "config": vars(args),
             "results": results["results"]
         }
+        if expert_activation_report:
+            output_data["expert_activation_report"] = expert_activation_report
+
         with open(output_file, "w") as f:
             json.dump(output_data, f, indent=4)
         print(f"Results and config saved to {output_file}")

model_utils.py

@@ -105,24 +105,106 @@ def patched_forward_zeroed_experts(self, hidden_states: torch.Tensor) -> torch.T
     final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
     return final_hidden_states, router_logits
 
-def apply_pruning(model, experts_to_prune, mode="zero"):
+def patched_forward_dynamic_routing(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    """
+    Patched forward method for dynamic routing based on cumulative probability threshold.
+    Selects a variable number of experts up to self.top_k.
+    """
+    batch_size, sequence_length, hidden_dim = hidden_states.shape
+    hidden_states = hidden_states.view(-1, hidden_dim)
+    router_logits = self.gate(hidden_states)
+
+    pruned_experts = getattr(self, "pruned_experts", None)
+    if pruned_experts is not None and len(pruned_experts) > 0:
+        idxs = torch.tensor(pruned_experts, device=router_logits.device, dtype=torch.long)
+        router_logits[:, idxs] = float('-inf')
+
+    routing_probs = F.softmax(router_logits, dim=1, dtype=torch.float)
+    # self.top_k is used as the max number of experts to consider
+    routing_weights, selected_experts = torch.topk(routing_probs, self.top_k, dim=-1)
+
+    # Dynamic routing based on cumulative probability
+    # The threshold is set on the MoE block during the patching process.
+    threshold = self.dynamic_routing_threshold
+    cumulative_weights = torch.cumsum(routing_weights, dim=-1)
+
+    # Create a mask to select experts until the cumulative probability exceeds the threshold
+    shifted_cumsum = torch.zeros_like(cumulative_weights)
+    shifted_cumsum[..., 1:] = cumulative_weights[..., :-1]
+    selection_mask = shifted_cumsum < threshold
+
+    # Apply the mask to the routing weights
+    routing_weights = routing_weights * selection_mask
+
+    # Track the average number of experts activated in this forward pass
+    if hasattr(self, "num_activated_experts_log"):
+        num_selected_experts_per_token = selection_mask.sum(dim=-1).float()
+        avg_experts_per_batch = num_selected_experts_per_token.mean().item()
+        self.num_activated_experts_log.append(avg_experts_per_batch)
+
+    if self.norm_topk_prob:
+        # Normalize routing weights for the selected experts
+        # Add a small epsilon for numerical stability to prevent division by zero.
+        routing_weights_sum = routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights / (routing_weights_sum + 1e-6)
+
+    routing_weights = routing_weights.to(hidden_states.dtype)
+
+    final_hidden_states = torch.zeros(
+        (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+    )
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+    # Apply the dynamic selection mask to the expert_mask
+    selection_mask_broadcast = selection_mask.permute(1, 0).unsqueeze(0)
+    expert_mask = expert_mask * selection_mask_broadcast
+
+    expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+    for expert_idx in expert_hit:
+        expert_layer = self.experts[expert_idx]
+        idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+
+        current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+        current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+        final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+
+    shared_expert_output = self.shared_expert(hidden_states)
+    shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
+
+    final_hidden_states = final_hidden_states + shared_expert_output
+
+    final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+    return final_hidden_states, router_logits
+
+def apply_pruning(model, experts_to_prune, mode="zero", dynamic_routing_threshold=0.8):
     """
     Monkey patch each OlmoeSparseMoeBlock forward and assign pruned experts per layer.
 
     Args:
         model: The model to prune.
         experts_to_prune: A dictionary where keys are layer indices and values are lists of expert indices to prune.
-        mode: "mask" (mask logits) or "zero" (zero out expert outputs)
+        mode: "mask" (mask logits), "zero" (zero out expert outputs), or "dynamic" (dynamic routing).
+        dynamic_routing_threshold (float): Cumulative probability threshold for dynamic routing.
     """
     if mode == "mask":
         patch_fn = patched_forward_masked_experts
     elif mode == "zero":
         patch_fn = patched_forward_zeroed_experts
+    elif mode == "dynamic":
+        patch_fn = patched_forward_dynamic_routing
     else:
-        raise ValueError(f"Unknown mode: {mode}. Use 'mask' or 'zero'.")
+        raise ValueError(f"Unknown mode: {mode}. Use 'mask', 'zero', or 'dynamic'.")
 
     for layer_idx, layer in enumerate(model.model.layers):
         moe_block = layer.mlp  
         if hasattr(moe_block, "gate") and hasattr(moe_block, "experts"):
+            moe_block.layer_id = layer_idx
+            print(f"INFO: Patching MoE layer {layer_idx} with mode '{mode}'.")
             moe_block.pruned_experts = experts_to_prune.get(layer_idx, [])
+            if mode == "dynamic":
+                moe_block.num_activated_experts_log = []
+                print(f"INFO: Setting dynamic_routing_threshold to {dynamic_routing_threshold} for layer {layer_idx}.")
+                moe_block.dynamic_routing_threshold = dynamic_routing_threshold
             moe_block.forward = patch_fn.__get__(moe_block, moe_block.__class__)