class_docstr: Basic MoE layer.
class_docstr_tok: ["Basic","MoE","layer","."]
code: def __init__(self, config: TransformerConfig):
      super().__init__(config=config)
  
      self.config: TransformerConfig = config
  
      self.router = get_router_linear_layer(self.config)
      self.add_bias = config.add_bias_linear
      self.sequence_parallel = config.sequence_parallel
      self.route_algo = sinkhorn
      self.router_activation = torch.sigmoid
      self.expert_parallel_size = (
          parallel_state.get_expert_model_parallel_world_size()
      )
  
      assert (
          self.config.num_moe_experts % self.expert_parallel_size == 0
      )
      self.num_local_experts = (
          self.config.num_moe_experts // self.expert_parallel_size
      )
      local_expert_indices_offset = (
          parallel_state.get_expert_model_parallel_rank()
          * self.num_local_experts
      )
      self.local_expert_indices = [
          local_expert_indices_offset + i
          for i in range(self.num_local_experts)
      ]
      self.k = 1
  
file_path: megatron-lm/megatron/core/transformer/moe/base_moe_layer.py
imports: from abc import ABC, abstractmethod
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 1 to create a new BaseMoELayer object from inherited base classes: ABC, MegatronModule with example: obj = BaseMoELayer(config)