Update moep_llama.py

moep_llama.py

@@ -19,6 +19,12 @@ from transformers.modeling_outputs import (
 @dataclass
 class BaseModelOutputWithAux(BaseModelOutput):
     aux_loss: Optional[torch.Tensor] = None
+    one_hot: Optional[torch.Tensor] = None
+
+@dataclass
+class CausalLMOutputWithAux(CausalLMOutput):
+    aux_loss: Optional[torch.Tensor] = None
+    one_hot: Optional[torch.Tensor] = None
 
 
 def create_causal_mask(seq_len: int, device=None, dtype=torch.float32):
@@ -183,7 +189,9 @@ class Top_k_Router(nn.Module):
         f_mean = one_hot.float().mean(dim=0) / float(self.k)
         aux_loss = (E * (p_mean * f_mean).sum())
 
-        return topk_idx, topk_prob, aux_loss
+        one_hot = one_hot.view(B, S, E)
+
+        return topk_idx, topk_prob, aux_loss, one_hot
 
 
 class MoE_Block(nn.Module):
@@ -201,7 +209,7 @@ class MoE_Block(nn.Module):
         T = B * S
         x_flat = x.view(T, D)
 
-        topk_idx, topk_prob, aux_loss = self.router(x_flat)  # [T,k], [T,k]
+        topk_idx, topk_prob, aux_loss, one_hot = self.router(x_flat)  # [T,k], [T,k]
         y_flat = x_flat.new_zeros((T, self.dim_out))
 
         for slot in range(self.k):
@@ -214,7 +222,7 @@ class MoE_Block(nn.Module):
                 y_e = self.experts[e](x_flat.index_select(0, idx))  # [Ne, dim_out]
                 y_flat.index_add_(0, idx, y_e * e_prob.index_select(0, idx))
 
-        return y_flat.view(B, S, self.dim_out), aux_loss
+        return y_flat.view(B, S, self.dim_out), aux_loss, one_hot
 
 
 class ParallelLayer(nn.Module):
@@ -233,13 +241,13 @@ class ParallelLayer(nn.Module):
         Y = torch.stack([y.view(T, D) for y in outs], dim=1)  # [T, n_blocks, D]
 
         x_flat = x.view(T, D)
-        topk_idx, topk_prob, aux_loss = self.router(x_flat)  # [T,k], [T,k]
+        topk_idx, topk_prob, aux_loss, one_hot = self.router(x_flat)  # [T,k], [T,k]
 
         gather_idx = topk_idx.unsqueeze(-1).expand(-1, -1, D)  # [T,k,D]
         Y_topk = torch.gather(Y, dim=1, index=gather_idx)      # [T,k,D]
         y_flat = (Y_topk * topk_prob.unsqueeze(-1)).sum(dim=1) # [T,D]
 
-        return y_flat.view(B, S, D), aux_loss
+        return y_flat.view(B, S, D), aux_loss, one_hot
 
 
 class MoEPllamaLMConfig(PretrainedConfig):
@@ -401,6 +409,7 @@ class MoEPllamaModel(PreTrainedModel):
         causal = create_causal_mask(S, device=device, dtype=mask_dtype)
       
         aux_total = None
+        all_one_hots = []
         
         freqs = self.freqs
 
@@ -408,26 +417,34 @@ class MoEPllamaModel(PreTrainedModel):
         x = self.layers[0](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)
 
         # MoE shrink
-        x, aux = self.layers[1](x)
+        x, aux, one_hot = self.layers[1](x)
  
         if aux is not None:
             aux_total = aux
 
+        if one_hot is not None:
+            all_one_hots.append(one_hot)
+
         freqs_p = self.freqs_p
         
 
         idx = 2
         for i in range(self.config.n_layers - 2):
-            x, aux = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs_p)
+            x, aux, one_hot = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs_p)
             if aux is not None:
                 aux_total += aux
+            if one_hot is not None:
+                all_one_hots.append(one_hot)
             idx += 1
 
         # MoE grow
-        x, aux = self.layers[idx](x)
+        x, aux, one_hot = self.layers[idx](x)
         
         if aux is not None:
                 aux_total += aux
+
+        if one_hot is not None:
+            all_one_hots.append(one_hot)
             
         idx += 1
 
@@ -445,6 +462,7 @@ class MoEPllamaModel(PreTrainedModel):
                 hidden_states=None,
                 attentions=None,
                 aux_loss=aux_total,
+                one_hot=all_one_hots,
             )
         else:
             return BaseModelOutput(
@@ -525,12 +543,22 @@ class MoEPllamaForCausalLM(PreTrainedModel):
         if not return_dict:
             return (loss, logits) if loss is not None else (logits,)
 
-        return CausalLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=base_out.hidden_states,
-            attentions=base_out.attentions,
-        )
+        if return_aux:
+            return CausalLMOutputWithAux(
+                loss=loss,
+                logits=logits,
+                hidden_states=base_out.hidden_states,
+                attentions=base_out.attentions,
+                aux_loss=getattr(base_out, "aux_loss", None),
+                one_hot=getattr(base_out, "one_hot", None),
+            )
+        else:
+            return CausalLMOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=base_out.hidden_states,
+                attentions=base_out.attentions,
+            )
 
 AutoConfig.register("moep_llama_lm", MoEPllamaLMConfig)
 AutoModelForCausalLM.register(MoEPllamaLMConfig, MoEPllamaForCausalLM)