class_docstr: Top-1 Mixture of Experts Layer with Grouped GEMM. Routes input to one of N MLP "experts"
  Curently supports Sinkhorn based expert routing.
class_docstr_tok: ["Top-1","Mixture","of","Experts","Layer","with","Grouped","GEMM",".","Routes","input","to","one","of","N","MLP","``","experts","''","Curently","supports","Sinkhorn","based","expert","routing","."]
code: def __init__(self, config: TransformerConfig):
      super().__init__(config=config)
      self.config: TransformerConfig = config
  
      gg.assert_grouped_gemm_is_available()
      assert (
          config.add_bias_linear == False
      ), "bias in the expert layer is not supported in Grouped GEMM yet, please set '--disable-bias-linear' instead."
  
      self.expert_parallel = config.expert_model_parallel_size > 1
      if self.config.gated_linear_unit:
  
          def glu(x):
              x = torch.chunk(x, 2, dim=-1)
              return self.config.activation_func(x[0]) * x[1]
  
          self.activation_func = glu
      else:
          self.activation_func = self.config.activation_func
  
      # How many feature each rank holds for fc1 and fc2, respectively.
      tp_size = parallel_state.get_tensor_model_parallel_world_size()
      fc1_output_size = (
          self.config.ffn_hidden_size * self.num_local_experts
      )
      if config.gated_linear_unit:
          # Project to 4h. If using swiglu double the output width,
          # see https://arxiv.org/pdf/2002.05202.pdf
          fc1_output_size *= 2
      fc1_output_size_per_partition = divide(fc1_output_size, tp_size)
  
      fc2_input_size = (
          self.config.ffn_hidden_size * self.num_local_experts
      )
      fc2_input_size_per_partition = divide(fc2_input_size, tp_size)
  
      # Note: The current kernel implementations of grouped_gemm
      # does not support transposition with CUTLASS grouped GEMM
      # (https://github.com/fanshiqing/grouped_gemm/blob/main/csrc/grouped_gemm.cu#L355-L358)
      # and as a result we avoid allocate the transpose of weights.
      # Initialize weight.
      if config.use_cpu_initialization:
          self.weight1 = Parameter(
              torch.empty(
                  self.config.hidden_size,
                  fc1_output_size_per_partition,
                  dtype=config.params_dtype,
              )
          )
          self.weight2 = Parameter(
              torch.empty(
                  fc2_input_size_per_partition,
                  self.config.hidden_size,
                  dtype=config.params_dtype,
              )
          )
          if config.perform_initialization:
              _initialize_affine_weight_cpu(
                  self.weight1,
                  self.config.hidden_size,
                  fc1_output_size,
                  fc1_output_size_per_partition,
                  partition_dim=1,
                  init_method=config.init_method,
                  params_dtype=config.params_dtype,
              )
              _initialize_affine_weight_cpu(
                  self.weight2,
                  fc2_input_size,
                  self.config.hidden_size,
                  fc2_input_size_per_partition,
                  partition_dim=0,
                  init_method=config.output_layer_init_method,
                  params_dtype=config.params_dtype,
              )
      else:
          self.weight1 = Parameter(
              torch.empty(
                  self.config.hidden_size,
                  fc1_output_size_per_partition,
                  device=torch.cuda.current_device(),
                  dtype=config.params_dtype,
              )
          )
          self.weight2 = Parameter(
              torch.empty(
                  fc2_input_size_per_partition,
                  self.config.hidden_size,
                  device=torch.cuda.current_device(),
                  dtype=config.params_dtype,
              )
          )
          if config.perform_initialization:
              _initialize_affine_weight_gpu(
                  self.weight1,
                  config.init_method,
                  partition_dim=1,
                  expert_parallel=self.expert_parallel,
              )
              _initialize_affine_weight_gpu(
                  self.weight2,
                  config.output_layer_init_method,
                  partition_dim=0,
                  expert_parallel=self.expert_parallel,
              )
      setattr(self.weight1, "allreduce", not self.expert_parallel)
      setattr(self.weight2, "allreduce", not self.expert_parallel)
  
file_path: megatron-lm/megatron/core/transformer/moe/grouped_mlp.py
imports: import torch
  from torch.nn.parameter import Parameter
  from megatron.core import parallel_state
  from megatron.core.tensor_parallel.layers import _initialize_affine_weight_cpu, _initialize_affine_weight_gpu
  from megatron.core.tensor_parallel.utils import divide
  from megatron.core.transformer.moe import grouped_gemm_util
  from megatron.core.transformer.transformer_config import TransformerConfig
  from .base_moe_layer import BaseMoELayer
label_desc: Use image node_id 1 to create a new GroupedMLP object from inherited base classes: BaseMoELayer with example: obj = GroupedMLP(config)