class_docstr: Built on the paper and library Megablocks as described in
  https://arxiv.org/abs/2211.15841. This implementation is
  strictly equivalent to standard MoE with full capacity (no
  dropped tokens). It's faster since it formulates MoE operations
  in terms of block-sparse operations to accomodate imbalanced
  assignments of tokens to experts, whereas standard MoE either
  (1) drop tokens at the cost of reduced performance or (2) set
  capacity factor to number of experts and thus waste computation
  and memory on padding.
class_docstr_tok: ["Built","on","the","paper","and","library","Megablocks","as","described","in","https",":","\/\/arxiv.org\/abs\/2211.15841",".","This","implementation","is","strictly","equivalent","to","standard","MoE","with","full","capacity","(","no","dropped","tokens",")",".","It","'s","faster","since","it","formulates","MoE","operations","in","terms","of","block-sparse","operations","to","accomodate","imbalanced","assignments","of","tokens","to","experts",",","whereas","standard","MoE","either","(","1",")","drop","tokens","at","the","cost","of","reduced","performance","or","(","2",")","set","capacity","factor","to","number","of","experts","and","thus","waste","computation","and","memory","on","padding","."]
code: def __init__(self, prefix, config: MixtralConfig, weights):
      super().__init__()
      self.hidden_dim = config.hidden_size
      self.ffn_dim = (
          config.intermediate_size // weights.process_group.size()
      )
      self.num_experts = config.num_local_experts
      self.top_k = config.num_experts_per_tok
  
      act = config.hidden_act
      if "gelu" in act:
          self.act = lambda x: torch.nn.functional.gelu(
              x,
              approximate="tanh"
              if act in ["gelu_fast", "gelu_pytorch_tanh"]
              else "none",
          )
      elif "silu" in act:
          self.act = torch.nn.functional.silu
      else:
          self.act = ACT2FN[act]
  
      # gating
      self.gate = FastLinear.load(
          config, f"{prefix}.gate", weights, bias=False
      )
  
      # merged expert weights, all of size  (n_experts * ffn_dim, hidden_dim)
      self.w1 = _load_experts(
          config, f"{prefix}.experts", "w1", weights
      )
      self.w2 = _load_experts(
          config, f"{prefix}.experts", "w2", weights
      )
      self.w3 = _load_experts(
          config, f"{prefix}.experts", "w3", weights
      )
  
      self.offsets = None
      self.offsets_block_rows = 0
  
      self.process_group = weights.process_group
  
      # Calculate the number of bits needed to represent the expert indices
      # so that we can pass it to radix sort.
      self.sort_end_bit = max(
          int(np.ceil(np.log2(self.num_experts))), 1
      )
      self.blocking = 128
      self.quantize_scatter_num_bits = -1
  
imports: import torch
  import torch.distributed
  import numpy
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from loguru import logger
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import FastLinear, FastRMSNorm, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear
label_desc: Use image node_id 1 to create a new BlockSparseMoE object from inherited base classes: nn with example: obj = BlockSparseMoE(prefix, config, weights)