change back to topk

myolmoe/modeling_myolmoe.py

@@ -462,8 +462,11 @@ class OlmoeSparseMoeBlock(nn.Module):
         hidden_states = hidden_states.view(-1, hidden_dim)
         router_logits = self.gate(hidden_states)
         routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
-        selected_experts = torch.multinomial(routing_weights, self.top_k, replacement=False)
-        routing_weights = routing_weights.gather(1, selected_experts)
+        routing_weights, selected_experts = torch.topk(
+            routing_weights, self.top_k, dim=-1
+        )
+        # selected_experts = torch.multinomial(routing_weights, self.top_k, replacement=False)
+        # routing_weights = routing_weights.gather(1, selected_experts)
 
         if self.norm_topk_prob:
             routing_weights /= routing_weights.sum(dim=-1, keepdim=True)