class_docstr: Basic MoE layer.
class_docstr_tok: ["Basic","MoE","layer","."]
docstr: Reverse process of `token_permutation()` which permutes the ouput of local
  experts locallay and across expert parallel rank into the original order to
  produce the final output.
  
  Args:
      hidden_states: 2D tensor of shape [sum_tokens_of_all_local_experts, HiddenSize],
      ouput of local experts.
      indices: 2D tensor of the indices of `local_indices` (which holds the un-sorted expert
      indices of tokens that local expert can process) that give its sorted order along dim 0.
      global_local_map (optional): 2D tensor, a mask of mapping between global and local tokens where each
      element is True if it's between the local_expert_indices. Only useful
      when cross device token permutation is enabled and **AllGahter** is performed.
      bias: bias if self.add_bias is enabled.
  
  Returns:
      output_total: un-permuted updated hidden states output from all local experts
      with shape of [SeqLen/TP, MBS, HiddenSize]
      output_bias_total: un-permuted bias output from all local experts if
      self.add_bias is enabled.
docstr_tok: ["Reverse","process","of","`","token_permutation","(",")","`","which","permutes","the","ouput","of","local","experts","locallay","and","across","expert","parallel","rank","into","the","original","order","to","produce","the","final","output",".","Args",":","hidden_states",":","2D","tensor","of","shape","[","sum_tokens_of_all_local_experts",",","HiddenSize","]",",","ouput","of","local","experts",".","indices",":","2D","tensor","of","the","indices","of","`","local_indices","`","(","which","holds","the","un-sorted","expert","indices","of","tokens","that","local","expert","can","process",")","that","give","its","sorted","order","along","dim","0.","global_local_map","(","optional",")",":","2D","tensor",",","a","mask","of","mapping","between","global","and","local","tokens","where","each","element","is","True","if","it","'s","between","the","local_expert_indices",".","Only","useful","when","cross","device","token","permutation","is","enabled","and","*","*","AllGahter","*","*","is","performed",".","bias",":","bias","if","self.add_bias","is","enabled",".","Returns",":","output_total",":","un-permuted","updated","hidden","states","output","from","all","local","experts","with","shape","of","[","SeqLen\/TP",",","MBS",",","HiddenSize","]","output_bias_total",":","un-permuted","bias","output","from","all","local","experts","if","self.add_bias","is","enabled","."]
code: def token_unpermutation(
      self, hidden_states, indices, global_local_map=None, bias=None
  ):
      """Reverse process of `token_permutation()` which permutes the ouput of local
      experts locallay and across expert parallel rank into the original order to
      produce the final output.
  
      Args:
          hidden_states: 2D tensor of shape [sum_tokens_of_all_local_experts, HiddenSize],
          ouput of local experts.
          indices: 2D tensor of the indices of `local_indices` (which holds the un-sorted expert
          indices of tokens that local expert can process) that give its sorted order along dim 0.
          global_local_map (optional): 2D tensor, a mask of mapping between global and local tokens where each
          element is True if it's between the local_expert_indices. Only useful
          when cross device token permutation is enabled and **AllGahter** is performed.
          bias: bias if self.add_bias is enabled.
  
      Returns:
          output_total: un-permuted updated hidden states output from all local experts
          with shape of [SeqLen/TP, MBS, HiddenSize]
          output_bias_total: un-permuted bias output from all local experts if
          self.add_bias is enabled.
      """
      # Stage1: unpermute the tokens and bias locally respectively.
      unpermuted_local_hidden = torch.zeros_like(hidden_states)
      assert indices.shape == hidden_states.shape
      unpermuted_local_hidden = unpermuted_local_hidden.scatter(
          0, indices, hidden_states
      )
  
      # Scale the expert output prior to reduction and subsequent to local unpermutation if k > 1.
      if self.k > 1:
          unpermuted_local_hidden = (
              unpermuted_local_hidden * self.max_prob.view(-1, 1)
          )
  
      unpermuted_local_bias = None
      if self.add_bias:
          assert bias is not None
          unpermuted_local_bias = torch.zeros_like(hidden_states)
          assert indices.shape == bias.shape
          unpermuted_local_bias = unpermuted_local_bias.scatter(
              0, indices, bias
          )
          if self.k > 1:
              unpermuted_local_bias = (
                  unpermuted_local_bias * self.max_prob.view(-1, 1)
              )
  
      output_total = unpermuted_local_hidden
      output_bias_total = unpermuted_local_bias
  
      # Stage2: unpermute the tokens across expert parallel devices.
      if self.sequence_parallel or (self.expert_parallel_size > 1):
          assert (
              global_local_map is not None
          ), "global_local_map is necessary for `AllGather`."
          ep_group_size = (
              parallel_state.get_tensor_and_expert_parallel_world_size()
          )
          # hidden_shape: [SeqLen/TP, MBS, HiddenSize], glboal_num_tokens = SeqLen/TP*MBS*(TP*EP)
          global_num_tokens = (
              self.hidden_shape[0]
              * self.hidden_shape[1]
              * ep_group_size
          )
          global_hidden_shape = [
              global_num_tokens,
              hidden_states.shape[-1],
          ]
          unpermuted_global_hidden = torch.zeros(
              global_hidden_shape,
              dtype=hidden_states.dtype,
              device=torch.cuda.current_device(),
          )
          # Reshape global_local_map to be compatible with Tensor.scatter
          assert global_local_map.shape == unpermuted_local_hidden.shape
          unpermuted_global_hidden = (
              unpermuted_global_hidden.scatter_add(
                  0, global_local_map, unpermuted_local_hidden
              )
          )
          output_total = tensor_parallel.reduce_scatter_to_sequence_parallel_region_from_moe(
              unpermuted_global_hidden
          )
          if self.add_bias:
              # Unpermute the bias across expert parallel devices.
              unpermuted_global_bias = torch.zeros_like(
                  unpermuted_global_hidden
              )
              unpermuted_global_bias = (
                  unpermuted_global_bias.scatter_add(
                      0, global_local_map, unpermuted_local_bias
                  )
              )
              output_bias_total = tensor_parallel.reduce_scatter_to_sequence_parallel_region_from_moe(
                  unpermuted_global_bias
              )
              # bias is duplicated across tensor parallelism ranks;
              # reduce scatter reduces bias across tensor parallel_ranks
              output_bias_total = (
                  output_bias_total
                  / parallel_state.get_tensor_model_parallel_world_size()
              )
      if self.k == 1:
          output_total = output_total * self.max_prob.view(-1, 1)
      output_total = output_total.view(self.hidden_shape)
      if self.add_bias:
          assert output_bias_total is not None
          if self.k == 1:
              output_bias_total = (
                  output_bias_total * self.max_prob.view(-1, 1)
              )
          output_bias_total = output_bias_total.view(self.hidden_shape)
      else:
          output_bias_total = None
  
      return output_total, output_bias_total
  
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 4 for calling the BaseMoELayer obj's underlying member method code with example usage: obj.token_unpermutation(hidden_states, indices, global_local_map, bias) and returns: output_total,  output_bias_total