[Research] Custom Deepseek Routing

[kylesayrs]

Purpose

• Provide research the ability to calibration the deepseekv2 model as if it was in moe training mode. This means that the model performs a forward pass with all experts instead of just the top k.
  • Using the moe_top_k_activation option, the user can choose between using all experts to compute outputs (matching train-time activations) or just the top-k (matching inference-time activations)

class DeepseekV2MoE(nn.Module):
    def forward(self, hidden_states):
        ...
        if True:  #if self.training:
            hidden_states = hidden_states.repeat_interleave(
                self.num_experts_per_tok, dim=0
            )
            y = torch.empty_like(hidden_states)
            for i, expert in enumerate(self.experts):
                y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            y = y.to(hidden_states.dtype).view(*orig_shape)
            y = AddAuxiliaryLoss.apply(y, aux_loss)
        # TRACING: Give option to calibrate with top_k experts, as if in inference time
        #else:
        if self.config.moe_top_k_activation:
            y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
        if self.config.n_shared_experts is not None:
            y = y + self.shared_experts(identity)
        return y

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.