Upload 4 files

aux_losses.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def log_mean(x, dim):
+    return torch.logsumexp(x, dim=dim) - torch.log(
+        torch.tensor(x.shape[dim], dtype=torch.float32)
+    )
+
+
+def entropy_reg(logits: torch.Tensor, mean_over_batch: bool = True):
+    """Entropy regularization for the router."""
+
+    entropy_l = lambda l: -(l * l.exp()).sum(-1)
+    # softmax over experts
+    # logits: [batch_size * sequence_length, num_experts]
+    logprobs = F.log_softmax(logits, dim=-1)
+    if mean_over_batch:
+        # take mean probability over batch
+        logprobs = log_mean(logprobs, 0)
+
+    return -entropy_l(logprobs).mean()
+
+
+# two losses below are adapted from
+# https://github.com/google/flaxformer/blob/b725bd2a51d70e866d819c92de166fbf24425e6a/flaxformer/architectures/moe/routing.py
+def load_balancing_loss(logits: torch.Tensor, expert_indices: torch.Tensor) -> float:
+    """Computes auxiliary load balancing loss as in Switch Transformer.
+
+    See Switch Transformer (https://arxiv.org/abs/2101.03961). This function
+    implements the loss function presented in equations (4) - (6). It aims to
+    penalize those cases where the routing between experts is unbalanced.
+
+    Args:
+      logits: logits assigned to each expert per token. Shape:
+        <float32>[batch_size * sequence_length, num_experts].
+      expert_indices: <int>[batch_size * sequence_length, num_selected_experts]
+        indices identifying the top num_selected_experts for a given token.
+
+    Returns:
+      The auxiliary loss.
+    """
+    # num_token = batch_size * sequence_length
+    num_token, num_experts = logits.shape
+
+    # Shape: [batch_size * sequence_length, num_selected_experts, num_experts].
+    expert_mask = F.one_hot(expert_indices, num_experts)
+    # For a given token, determine if it was routed to a given expert.
+    # Shape: [batch_size * sequence_length, num_experts]
+    expert_mask, _ = torch.max(expert_mask, dim=-2)
+
+    # shape [num_experts]
+    tokens_per_expert = torch.mean(expert_mask, dim=0, dtype=torch.float32)
+
+    # compute router probability per expert in log space for numerical stability
+    logprobs = F.log_softmax(logits, dim=-1)
+    # take mean probability over batch
+    # shape [num_experts]
+    logprobs = log_mean(logprobs, dim=0)
+    router_prob_per_expert = torch.exp(logprobs)
+    return (
+        torch.mean(  # mean over experts
+            tokens_per_expert * router_prob_per_expert,
+            dtype=torch.float32,
+        )
+        * num_experts
+    )
+
+
+def router_z_loss(router_logits: torch.Tensor) -> float:
+    """Compute router z-loss.
+
+     The router z-loss was introduced in Designing Effective Sparse Expert Models
+     (https://arxiv.org/abs/2202.08906). It encourages router logits to remain
+     small in an effort to improve stability.
+
+    Args:
+      router_logits: <float>[batch_size * sequence_length, num_experts]
+        router logits
+
+    Returns:
+      Scalar router z-loss.
+    """
+    num_tokens, _ = router_logits.shape
+    log_z = torch.logsumexp(router_logits, dim=-1)
+    z_loss = log_z**2
+    return torch.sum(z_loss, dtype=torch.float32) / (num_tokens)

configuration.py

@@ -0,0 +1,51 @@
+from transformers import PretrainedConfig
+
+class MoEGPTConfig(PretrainedConfig):
+    model_type = "moegpt"
+
+    def __init__(
+        self,
+        vocab_size=50304,
+        n_embd=768,
+        n_layer=12,
+        n_head=12,
+        sequence_length=1024,
+        moe=False,
+        moe_routing="standard_gating",
+        moe_num_experts=4,
+        moe_num_experts_per_tok=2,
+        moe_softmax_order="softmax_topk",
+        moe_router_loss="load_balancing_z_loss",
+        moe_aux_loss_factor=0.01,
+        moe_z_loss_factor=1.0,
+        mlp_dim_exp_factor=1.0,
+        dropout=0.0,
+        bias=False,
+        architectures=["MoEGPTForCausalLM"],
+        auto_map={
+            "AutoConfig": "configuration.MoEGPTConfig",
+            "AutoModelForCausalLM": "modeling.MoEGPTForCausalLM",
+            "AutoTokenizer": "GPT2TokenizerFast"
+        },
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.sequence_length = sequence_length
+        self.moe = moe
+        self.moe_routing = moe_routing
+        self.moe_num_experts = moe_num_experts
+        self.moe_num_experts_per_tok = moe_num_experts_per_tok
+        self.moe_softmax_order = moe_softmax_order
+        self.moe_router_loss = moe_router_loss
+        self.moe_aux_loss_factor = moe_aux_loss_factor
+        self.moe_z_loss_factor = moe_z_loss_factor
+        self.mlp_dim_exp_factor = mlp_dim_exp_factor
+        self.dropout = dropout
+        self.bias = bias
+        self.architectures = architectures
+        self.auto_map = auto_map
+    

modeling.py

@@ -0,0 +1,481 @@
+from transformers import PreTrainedModel
+from configuration import MoEGPTConfig
+# importa anche MoE, MaskedMoE, TimeDependantMoE ecc.
+import math
+import inspect
+from typing import Optional, Dict, Any
+from dataclasses import dataclass
+import tiktoken
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from huggingface_hub import PyTorchModelHubMixin
+from transformers.utils import ModelOutput
+
+
+from .moe import (
+    #ExpertChoiceMoE,
+    MaskedMoE,
+    TimeDependantMoE,
+    MoE,
+)
+
+from .aux_losses import (
+    entropy_reg,
+    load_balancing_loss,
+    router_z_loss,
+)
+
+# class Output(ModelOutput):
+#     def __init__(self, logits, loss=None, aux_losses=None, router_logits=None):
+#         self.logits = logits
+#         self.loss = loss
+#         self.aux_losses = aux_losses
+#         self.router_logits = router_logits
+@dataclass
+class Output(ModelOutput):
+    logits: torch.FloatTensor = None
+    loss: Optional[torch.FloatTensor] = None
+    aux_losses: Optional[Dict[str, torch.FloatTensor]] = None
+    router_logits: Optional[torch.FloatTensor] = None
+
+    def __repr__(self):
+        return f"Output(logits={self.logits}, loss={self.loss}, aux_losses={self.aux_losses}, router_logits={self.router_logits})"
+
+class LayerNorm(nn.Module):
+    """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
+        if not self.flash:
+            print(
+                "WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0"
+            )
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer(
+                "bias",
+                torch.tril(
+                    torch.ones(config.sequence_length, config.sequence_length)
+                ).view(1, 1, config.sequence_length, config.sequence_length),
+            )
+
+    def forward(self, x):
+        # batch size, sequence length, embedding dimensionality (n_embd)
+        (
+            B,
+            T,
+            C,
+        ) = x.size()
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        # (B, T, nh, hs)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=self.dropout, is_causal=True
+            )
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = (
+            y.transpose(1, 2).contiguous().view(B, T, C)
+        )  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dim_exp_factor = int(config.mlp_dim_exp_factor * 4)
+
+        self.c_fc = nn.Linear(
+            config.n_embd, self.dim_exp_factor * config.n_embd, bias=config.bias
+        )
+        self.c_proj = nn.Linear(
+            self.dim_exp_factor * config.n_embd, config.n_embd, bias=config.bias
+        )
+        self.dropout = nn.Dropout(config.dropout)
+        self.activation = nn.GELU()
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.activation(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        # need to return same type as the MoE block, but in this case it's empty
+        return x, {}
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.moe_config = config.moe_routing
+        if config.moe:
+            if config.moe_routing == "standard_gating":
+                self.mlp = MoE(config, MLP)
+            elif config.moe_routing == "masked":
+                self.mlp = TimeDependantMoE(config, MLP)
+            #elif config.moe_routing == "expert_choice":
+            #    self.mlp = ExpertChoiceMoE(config, MLP)
+            else:
+                raise ValueError(f"Unknown routing: {config.routing}")
+        else:
+            self.mlp = MLP(config)
+
+    def forward(self, x, date, *args, **kwargs):
+        x = x + self.attn(self.ln_1(x, *args, **kwargs))
+        if self.moe_config  == "masked":
+            x_, logits_and_experts = self.mlp(self.ln_2(x, *args, **kwargs), date)
+        else:
+            x_, logits_and_experts = self.mlp(self.ln_2(x, *args, **kwargs))
+        x = x + x_
+        return x, logits_and_experts
+
+
+class MoEGPTForCausalLM(PreTrainedModel):
+    config_class = MoEGPTConfig
+    def __init__(self, config):
+        super().__init__(config)
+        assert config.vocab_size is not None
+        assert config.sequence_length is not None
+        self.config = config
+        self.tokenizer = tiktoken.get_encoding("gpt2")
+        self.base_model_prefix = "timoe"
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.vocab_size, config.n_embd),
+                wpe=nn.Embedding(config.sequence_length, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+                ln_f=LayerNorm(config.n_embd, bias=config.bias),
+            )
+        )
+
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = (
+            self.lm_head.weight
+        )  # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith("c_proj.weight"):
+                torch.nn.init.normal_(
+                    p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)
+                )
+            if pn.endswith("router.weight"):
+                # special scaled init to moe router?
+                with torch.no_grad():
+                    dim = 1 if config.moe_routing == "standard_gating" else 0
+                    std = p.std()
+                    p.div_(p.sum(dim=dim, keepdim=True))
+                    p.mul_(std / p.std())
+
+    def get_router_losses(self, logits, selected_experts, eval=False):
+        # logits: (b * seq_len, n_experts)
+        # selected_experts: (b * seq_len, topk)
+        if eval:  # eval mode, compute all losses
+            return {
+                "moe_entropy_loss": entropy_reg(logits),
+                "moe_aux_loss": load_balancing_loss(logits, selected_experts),
+                "moe_z_loss": router_z_loss(logits),
+            }
+        if self.config.moe_router_loss == "entropy":
+            return {
+                "moe_entropy_loss": entropy_reg(logits),
+            }
+        elif self.config.moe_router_loss == "load_balancing_only":
+            return {
+                "moe_aux_loss": load_balancing_loss(logits, selected_experts),
+            }
+        elif self.config.moe_router_loss == "load_balancing_z_loss":
+            return {
+                "moe_aux_loss": load_balancing_loss(logits, selected_experts),
+                "moe_z_loss": router_z_loss(logits),
+            }
+        return {}
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, date=None, targets=None, attention_mask=None, get_logits=True, moe=False):
+        device = idx.device
+        b, t = idx.size()
+        assert (
+            t <= self.config.sequence_length
+        ), f"Cannot forward sequence of length {t}, block size is only {self.config.sequence_length}"
+        # shape (1, t)
+        if date is None:
+            # set all the date to 6
+            date = torch.full((1, b), 6, dtype=torch.long, device=device).squeeze(0)
+        else: 
+            date = (date - 2013) // 2 + 1
+            date = torch.full((1, b), date, dtype=torch.long, device=device).squeeze(0)
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(
+            pos
+        )  # position embeddings of shape (1, t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+
+        # router logits is a list for each layer's routing, each of shape (b * seq_len, n_experts)
+        router_logits = []
+        # experts is a list for each layer's selected experts, shape (b * seq_len, topk)
+        experts = []
+
+        # forward pass through all the transformer blocks
+        for block in self.transformer.h:
+            x, logits_and_experts = block(x, date)
+            if len(logits_and_experts) > 0:
+                router_logits.append(logits_and_experts["router_logits"])
+                experts.append(logits_and_experts["selected_experts"])
+        x = self.transformer.ln_f(x)
+
+        # aux_losses is a dict with keys for different auxiliary losses
+        aux_losses = {}
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
+            )
+            if moe and (self.config.moe_routing == "standard_gating" or self.config.moe_routing == "masked"):
+                # calculate the router losses per layer
+                for logit, expert_choice in zip(router_logits, experts):
+                    router_losses = self.get_router_losses(
+                        logit, expert_choice, eval=not self.training
+                    )
+                    for k, v in router_losses.items():
+                        aux_losses[k] = aux_losses.get(k, 0.0) + v
+                        if self.training:
+                            loss += (
+                                v
+                                * getattr(self.config, k + "_factor")
+                                / self.config.n_layer
+                            )
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(
+                #x[:, [-1], :]
+                x
+            )  # note: using list [-1] to preserve the time dim
+            loss = None
+        logits = logits if get_logits else None
+        router_logits = (
+            torch.stack(router_logits, dim=0) if len(router_logits) > 0 else None
+        )
+        # return {
+        #     "logits": logits,
+        #     "loss": loss,
+        #     "aux_losses": aux_losses,
+        #     "router_logits": router_logits,
+        # }
+        return Output(logits = logits, loss = loss, aux_losses = aux_losses, router_logits = router_logits)
+
+    def crop_sequence_length(self, sequence_length):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert sequence_length <= self.config.sequence_length
+        self.config.sequence_length = sequence_length
+        self.transformer.wpe.weight = nn.Parameter(
+            self.transformer.wpe.weight[:sequence_length]
+        )
+        for block in self.transformer.h:
+            block.attn.bias = block.attn.bias[:, :, :sequence_length, :sequence_length]
+
+
+    def get_parameter_group_specs(self):
+        """
+        This long function is unfortunately doing something very simple and is being very defensive:
+        We are separating out all parameters of the model into two buckets: those that will experience
+        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
+        We are then returning the PyTorch optimizer object.
+        """
+
+        # separate out all parameters to those that will and won't experience regularizing weight decay
+        decay = set()
+        no_decay = set()
+        whitelist_weight_modules = (torch.nn.Linear,)
+
+        BLACKLIST_WEIGHT_MODULES = (
+            torch.nn.LayerNorm,
+            LayerNorm,
+            torch.nn.Embedding,
+        )
+
+        for mn, m in self.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = "%s.%s" % (mn, pn) if mn else pn  # full param name
+                # random note: because named_modules and named_parameters are recursive
+                # we will see the same tensors p many many times. but doing it this way
+                # allows us to know which parent module any tensor p belongs to...
+                if pn.endswith("bias"):
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                elif pn.endswith("weight") and isinstance(m, BLACKLIST_WEIGHT_MODULES):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+
+        # subtle: 'transformer.wte.weight' and 'lm_head.weight' are tied, so they
+        # will appear in the no_decay and decay sets respectively after the above.
+        # In addition, because named_parameters() doesn't return duplicates, it
+        # will only return the first occurence, key'd by 'transformer.wte.weight', below.
+        # so let's manually remove 'lm_head.weight' from decay set. This will include
+        # this tensor into optimization via transformer.wte.weight only, and not decayed.
+        decay.remove("lm_head.weight")
+
+        # validate that we considered every parameter
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        inter_params = decay & no_decay
+        union_params = decay | no_decay
+        assert (
+            len(inter_params) == 0
+        ), "parameters %s made it into both decay/no_decay sets!" % (str(inter_params),)
+        assert (
+            len(param_dict.keys() - union_params) == 0
+        ), "parameters %s were not separated into either decay/no_decay set!" % (
+            str(param_dict.keys() - union_params),
+        )
+
+        # create the pytorch optimizer object
+        return [
+            {"params": sorted(list(decay))},
+            {"params": sorted(list(no_decay)), "weight_decay": 0.0},
+        ]
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens, date = None, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        idx = input_ids
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at sequence_length
+            idx_cond = (
+                idx
+                if idx.size(1) <= self.config.sequence_length
+                else idx[:, -self.config.sequence_length :]
+            )
+            # forward the model to get the logits for the index in the sequence
+            logits = self(idx_cond, date, get_logits=True).logits
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("Inf")
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+            # check if we hit the end of the sequence
+            if idx_next.item() == self.tokenizer.eot_token:
+                break
+
+        return idx
+
+    @torch.no_grad()
+    def generate_from_string(self, in_str, max_new_tokens, date = None, temperature=1.0, top_k=None):
+        idx = (
+            torch.tensor(
+                self.tokenizer.encode(in_str, allowed_special={"<|endoftext|>"})
+            )
+            .view(1, -1)
+            .to(self.lm_head.weight.device)
+        )
+        out_idx = (
+            self.generate(idx, max_new_tokens, date, temperature, top_k)
+            .view(-1)
+            .to("cpu")
+            .numpy()
+        )
+        return self.tokenizer.decode(out_idx).split(in_str)[-1]
+    
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+        # reset the lm_head to use the new embeddings
+        # this is necessary because the lm_head is tied to the input embeddings
+        self.lm_head = nn.Linear(
+            self.config.n_embd, new_embeddings.weight.shape[0] , bias=False
+        )
+        #self.transformer.wte.weight = (
+        #    self.lm_head.weight
+        #)

moe.py

@@ -0,0 +1,145 @@
+"""
+Simple MoE routing implementations that replace the MLP block in a standard transformer.
+References:
+1) Mistral Source for Mixtral MoEs: 
+https://github.com/mistralai/mistral-src
+2) ST-MoE:
+https://arxiv.org/abs/2202.08906
+3) Our notepad of MoE resources: 
+https://docs.google.com/document/d/1NuQ5jr7V-Jv1ui7p4KrxO_JTz-7bpYcYMmh49EeJ-QA/edit?usp=sharing
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MoE(nn.Module):
+    """
+    Simplest MoE implementation with a linear router and softmax over experts.
+
+    Note that in this implementation, we simply loop over the experts and
+    aggregate the results. This is not the most efficient way to do it, but
+    it also avoids the large memory overhead _and_ has no token dropping
+    (because we do not need the capacity factor).
+    """
+
+    def __init__(self, config, mlp):
+        super().__init__()
+        assert config.moe_num_experts > 0
+        self.experts = nn.ModuleList(
+            [mlp(config=config) for _ in range(config.moe_num_experts)]
+        )
+        self.router = nn.Linear(config.n_embd, config.moe_num_experts, bias=False)
+        self.top_k = config.moe_num_experts_per_tok
+        self.softmax_order = config.moe_softmax_order
+
+    def forward(self, inputs: torch.Tensor):
+        # [batch_size * sequence_length, n_embd]
+        inputs_squashed = inputs.view(-1, inputs.shape[-1])
+        # [batch_size * sequence_length, num_experts]
+        router_logits = self.router(inputs_squashed)
+
+        # note that selected experts will be the same for all orders:
+        # softmax doesnt change top-k, but the weights are different
+        if self.softmax_order == "softmax_topk":
+            all_probs = F.softmax(router_logits, dim=1, dtype=torch.float32)
+            weights, selected_experts = torch.topk(all_probs, self.top_k)
+        elif self.softmax_order == "topk_softmax":
+            weights, selected_experts = torch.topk(router_logits, self.top_k)
+            weights = F.softmax(weights, dim=-1, dtype=torch.float32)
+        else:
+            raise ValueError(f"Unknown softmax_order: {self.softmax_order}")
+
+        results = torch.zeros_like(inputs_squashed)
+        # naive looping over experts
+        for i, expert in enumerate(self.experts):
+            batch_idx, nth_expert = torch.where(selected_experts == i)
+            output, _ = expert(inputs_squashed[batch_idx])
+            results[batch_idx] += weights[batch_idx, nth_expert, None] * output
+
+        # return results and router logits (for aux loss calculation later)
+        return results.view_as(inputs), {
+            "router_logits": router_logits,
+            "selected_experts": selected_experts,
+        }
+
+
+class ExpertChoiceMoE(nn.Module):
+    """
+    This is the MoE implementation that uses the expert choice method from
+    https://arxiv.org/pdf/2202.09368v2.pdf.
+
+    The main difference is that the router takes the softmax over the tokens, not the experts
+    (i.e. each expert chooses its top-k tokens, not the other way around).
+    For the same capacity factor, in theory, the same compute will be used as in standard top-k routing.
+    AFAICT, there is no way around the capacity factor (whereas the code above does not need it).
+    """
+
+    def __init__(self, config, mlp):
+        super().__init__()
+        assert config.moe_num_experts > 0
+        self.n_experts = config.moe_num_experts
+        self.experts = nn.ModuleList(
+            [mlp(config=config) for _ in range(config.moe_num_experts)]
+        )
+        self.router = nn.Linear(config.n_embd, config.moe_num_experts, bias=False)
+        self.capacity_factor = config.capacity_factor
+        self.softmax_order = config.moe_softmax_order
+        self.top_k = int(
+            self.capacity_factor
+            * config.batch_size
+            * config.sequence_length
+            / config.moe_num_experts
+        )
+
+    def forward(self, inputs: torch.Tensor):
+        # [batch_size * sequence_length, n_embd]
+        inputs_squashed = inputs.view(-1, inputs.shape[-1])
+        num_tokens = inputs_squashed.shape[0]
+        top_k = min(self.top_k, int(self.capacity_factor * num_tokens / self.n_experts))
+        # [batch_size * sequence_length, num_experts]
+        router_logits = self.router(inputs_squashed)
+
+        # note that selected experts will be the same for all orders:
+        # softmax doesnt change top-k, but the weights are different
+        if self.softmax_order == "softmax_topk":
+            all_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+            # weights and selected tokens: [num_experts, top_k]
+            # topk over tokens!
+            weights, selected_tokens = torch.topk(all_probs.T, top_k)
+        elif self.softmax_order == "topk_softmax":
+            # weights and selected tokens: [num_experts, top_k]
+            weights, selected_tokens = torch.topk(router_logits.T, top_k)
+            weights = F.softmax(weights, dim=-1, dtype=torch.float32)
+        else:
+            raise ValueError(f"Unknown softmax_order: {self.softmax_order}")
+
+        """ this is the full parallel version with einsum -- this can OOM quickly """
+        # [num_experts, top_k, num_tokens]
+        # P = F.one_hot(selected_tokens, num_tokens).type_as(inputs_squashed)
+        # # [num_experts, top_k, n_embd]
+        # x_in = torch.matmul(P, inputs_squashed)
+        # # [num_experts, num_tokens, n_embd]
+        # experts_out = torch.stack(
+        #     [expert(x)[0] for expert, x in zip(self.experts, x_in)], dim=0
+        # )
+        # results = torch.einsum("ijl,ij,ijd->ld", P, weights, experts_out)
+
+        """ this is the naive loop version """
+        # loop through experts because of memory growing too large
+        # when doing everything in parallel.
+        # also, more hackable :)
+        results = torch.zeros_like(inputs_squashed)
+        for i, expert in enumerate(self.experts):
+            # [top_k]
+            batch_idx = selected_tokens[i]
+            # [top_k, n_embd]
+            output, _ = expert(inputs_squashed[batch_idx])
+            results[batch_idx] += weights[i, :, None] * output
+
+        # return results and router logits (for aux loss calculation later)
+        return results.view_as(inputs), {
+            "router_logits": router_logits,
+            "selected_experts": selected_tokens,
+        }