class_docstr: Initialize an MoE layer.
  
  Arguments:
      hidden_size (int): the hidden dimension of the model, importantly this is also the input and output dimension.
      expert (torch.nn.Module): the torch module that defines the expert (e.g., MLP, torch.linear).
      num_experts (int, optional): default=1, the total number of experts per layer.
      ep_size (int, optional): default=1, number of ranks in the expert parallel world or group.
      k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.
      capacity_factor (float, optional): default=1.0, the capacity of the expert at training time.
      eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time.
      min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor.
      use_residual (bool, optional): default=False, make this MoE layer a Residual MoE (https://arxiv.org/abs/2201.05596) layer.
      noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample' or 'None'.
      drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity).
      use_rts (bool, optional): default=True, whether to use Random Token Selection.
      use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed).
      enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts
class_docstr_tok: ["Initialize","an","MoE","layer",".","Arguments",":","hidden_size","(","int",")",":","the","hidden","dimension","of","the","model",",","importantly","this","is","also","the","input","and","output","dimension",".","expert","(","torch.nn.Module",")",":","the","torch","module","that","defines","the","expert","(","e.g.",",","MLP",",","torch.linear",")",".","num_experts","(","int",",","optional",")",":","default=1",",","the","total","number","of","experts","per","layer",".","ep_size","(","int",",","optional",")",":","default=1",",","number","of","ranks","in","the","expert","parallel","world","or","group",".","k","(","int",",","optional",")",":","default=1",",","top-k","gating","value",",","only","supports","k=1","or","k=2",".","capacity_factor","(","float",",","optional",")",":","default=1.0",",","the","capacity","of","the","expert","at","training","time",".","eval_capacity_factor","(","float",",","optional",")",":","default=1.0",",","the","capacity","of","the","expert","at","eval","time",".","min_capacity","(","int",",","optional",")",":","default=4",",","the","minimum","capacity","per","expert","regardless","of","the","capacity_factor",".","use_residual","(","bool",",","optional",")",":","default=False",",","make","this","MoE","layer","a","Residual","MoE","(","https",":","\/\/arxiv.org\/abs\/2201.05596",")","layer",".","noisy_gate_policy","(","str",",","optional",")",":","default=None",",","noisy","gate","policy",",","valid","options","are","'Jitter","'",",","'RSample","'","or","'None","'",".","drop_tokens","(","bool",",","optional",")",":","default=True",",","whether","to","drop","tokens","-","(","setting","to","False","is","equivalent","to","infinite","capacity",")",".","use_rts","(","bool",",","optional",")",":","default=True",",","whether","to","use","Random","Token","Selection",".","use_tutel","(","bool",",","optional",")",":","default=False",",","whether","to","use","Tutel","optimizations","(","if","installed",")",".","enable_expert_tensor_parallelism","(","bool",",","optional",")",":","default=False",",","whether","to","use","tensor","parallelism","for","experts"]
code: def __init__(
      self,
      hidden_size,
      expert,
      num_experts=1,
      ep_size=1,
      k=1,
      capacity_factor=1.0,
      eval_capacity_factor=1.0,
      min_capacity=4,
      use_residual=False,
      noisy_gate_policy: typing.Optional[str] = None,
      drop_tokens: bool = True,
      use_rts=True,
      use_tutel: bool = False,
      enable_expert_tensor_parallelism: bool = False,
  ):
      super(MoE, self).__init__()
  
      self.use_residual = use_residual
      self.enable_expert_tensor_parallelism = (
          enable_expert_tensor_parallelism
      )
      assert (
          num_experts % ep_size == 0
      ), f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})"
      self.ep_size = ep_size
      self.expert_group_name = f"ep_size_{self.ep_size}"
      self.num_experts = num_experts
      self.num_local_experts = num_experts // self.ep_size
  
      log_dist(
          f"Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}",
          [0],
      )
  
      assert noisy_gate_policy is None or noisy_gate_policy in [
          "None",
          "Jitter",
          "RSample",
      ], (
          "Unsupported noisy_gate_policy: " + noisy_gate_policy
      )
  
      experts = Experts(
          expert, self.num_local_experts, self.expert_group_name
      )
      self.deepspeed_moe = MOELayer(
          TopKGate(
              hidden_size,
              num_experts,
              k,
              capacity_factor,
              eval_capacity_factor,
              min_capacity,
              noisy_gate_policy,
              drop_tokens,
              use_rts,
          ),
          experts,
          self.expert_group_name,
          self.ep_size,
          self.num_local_experts,
          use_tutel=use_tutel,
      )
      if self.use_residual:
          self.mlp = expert
          # coefficient is used for weighted sum of the output of expert and mlp
          self.coefficient = torch.nn.Linear(hidden_size, 2)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/moe/layer.py
imports: import torch
  from deepspeed.utils import log_dist
  from deepspeed.utils import groups
  from .sharded_moe import MOELayer, TopKGate
  from .experts import Experts
  import typing
label_desc: Use image node_id 1 to create a new MoE object from inherited base classes: torch.nn with example: obj = MoE(hidden_size, expert, num_experts, ep_size, k, capacity_factor, eval_capacity_factor, min_capacity, use_residual, noisy_gate_policy, drop_tokens, use_rts, use_tutel, enable_expert_tensor_parallelism)