Merge pull request #21 from huggingface/ed-fix-modeling

Update modelling to work with new checkpoints, exposes output_router_logits

Author: ArthurZucker

src/transformers/models/openai_moe/modeling_openai_moe.py

@@ -102,7 +102,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
                     _, token_idx = torch.where(expert_mask[expert_idx[0]])
                 current_state = hidden_states[token_idx]
                 gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
-                gate, up = gate_up.chunk(2, dim=-1)
+                gate, up = gate_up[..., ::2], gate_up[..., 1::2]
                 glu = gate * torch.sigmoid(gate * self.alpha)
                 gated_output = (up + 1) * glu
                 out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
@@ -113,7 +113,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
             hidden_states = hidden_states.repeat(num_experts, 1)
             hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
             gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]
-            gate, up = gate_up.chunk(2, dim=-1)
+            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
             glu = gate * torch.sigmoid(gate * self.alpha)
             next_states = torch.bmm(((up + 1) * glu), self.down_proj)
             next_states = next_states + self.down_proj_bias[..., None, :]
@@ -666,7 +666,9 @@ def forward(
         >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
         "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
         ```"""
-
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
         outputs: MoeModelOutputWithPast = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,

src/transformers/models/openai_moe/modular_openai_moe.py

@@ -96,7 +96,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
                 gate_up = (
                     current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
                 )  # (num_tokens, 2 * interm_dim)
-                gate, up = gate_up.chunk(2, dim=-1)  # (num_tokens, interm_dim)
+                gate, up = gate_up[..., ::2], gate_up[..., 1::2]
                 glu = gate * torch.sigmoid(gate * self.alpha)  # (num_tokens, interm_dim)
                 gated_output = (up + 1) * glu  # (num_tokens, interm_dim)
                 out = (
@@ -109,7 +109,7 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
             hidden_states = hidden_states.repeat(num_experts, 1)
             hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
             gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]
-            gate, up = gate_up.chunk(2, dim=-1)  # not supported for DTensors
+            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
             glu = gate * torch.sigmoid(gate * self.alpha)
             next_states = torch.bmm(((up + 1) * glu), self.down_proj)
             next_states = next_states + self.down_proj_bias[..., None, :]