docstr: Swap moe param and optim.
  We use different strategies to swap expert and gate.
  For expert, we exchange the param and optim of the expert by p2p.
  For gate, we all gather the gate choose the part we want.
  
  Args:
      swap_list (List)
      optim (LowLevelZeroOptimizer)
docstr_tok: ["Swap","moe","param","and","optim",".","We","use","different","strategies","to","swap","expert","and","gate",".","For","expert",",","we","exchange","the","param","and","optim","of","the","expert","by","p2p",".","For","gate",",","we","all","gather","the","gate","choose","the","part","we","want",".","Args",":","swap_list","(","List",")","optim","(","LowLevelZeroOptimizer",")"]
code: def _swap_moe_param(
      self, swap_list: List, optim: LowLevelZeroOptimizer
  ) -> None:
      """
      Swap moe param and optim.
      We use different strategies to swap expert and gate.
      For expert, we exchange the param and optim of the expert by p2p.
      For gate, we all gather the gate choose the part we want.
  
      Args:
          swap_list (List)
          optim (LowLevelZeroOptimizer)
      """
      # get all experts weights
      local_rank = dist.get_rank(self.moe_ep_group)
      if self.experts.gated:
          weight_list = [self.experts.wi_up, self.experts.wi_gate]
      else:
          weight_list = [self.experts.wi]
      weight_list.append(self.experts.wo)
  
      # gate optim should be obtained first
      gate_shape = self.gate.shape
      # get master weight and optim
      master_gate_weight = optim._param_store.working_to_master_param[
          id(self.gate)
      ]
      gate_exp_avg = optim.optim.state[master_gate_weight]["exp_avg"]
      gate_exp_avg_sq = optim.optim.state[master_gate_weight][
          "exp_avg_sq"
      ]
      # gather
      global_master_gate_weight = self._gather_global_dp_group(
          master_gate_weight
      ).view(gate_shape)
      global_gate_exp_avg = self._gather_global_dp_group(
          gate_exp_avg
      ).view(gate_shape)
      global_gate_exp_avg_sq = self._gather_global_dp_group(
          gate_exp_avg_sq
      ).view(gate_shape)
      assert (
          self.gate.shape
          == global_master_gate_weight.shape
          == global_gate_exp_avg.shape
          == global_gate_exp_avg_sq.shape
      )
  
      for swap in swap_list:
          source_group, source_idx, target_group, target_idx = swap
          source_rank = self.moe_ep_ranks[source_group]
          target_rank = self.moe_ep_ranks[target_group]
          # exchange expert
          if local_rank in [source_group, target_group]:
              for weight in weight_list:
                  if local_rank == source_group:
                      self._swap_expert_param_and_optim(
                          weight,
                          source_idx,
                          self.moe_ep_group,
                          True,
                          target_rank,
                          optim,
                      )
                  elif local_rank == target_group:
                      self._swap_expert_param_and_optim(
                          weight,
                          target_idx,
                          self.moe_ep_group,
                          False,
                          source_rank,
                          optim,
                      )
          # exchange gate
          source_expert_pos = (
              source_group * self.local_expert_num + source_idx
          )
          target_expert_pos = (
              target_group * self.local_expert_num + target_idx
          )
          for gate in [
              self.gate,
              global_master_gate_weight,
              global_gate_exp_avg,
              global_gate_exp_avg_sq,
          ]:
              origin_source = (
                  gate.data[source_expert_pos].clone().detach()
              )
              origin_target = (
                  gate.data[target_expert_pos].clone().detach()
              )
              (
                  gate.data[source_expert_pos],
                  gate.data[target_expert_pos],
              ) = (
                  origin_target,
                  origin_source,
              )
  
      # update gate
      global_master_gate_weight = global_master_gate_weight.view(
          -1
      ).split(global_master_gate_weight.numel() // self.global_dp_size)[
          self.global_dp_rank
      ]
      master_gate_weight.data.copy_(global_master_gate_weight)
      global_gate_exp_avg = global_gate_exp_avg.view(-1).split(
          global_gate_exp_avg.numel() // self.global_dp_size
      )[self.global_dp_rank]
      gate_exp_avg.data.copy_(global_gate_exp_avg)
      global_gate_exp_avg_sq = global_gate_exp_avg_sq.view(-1).split(
          global_gate_exp_avg_sq.numel() // self.global_dp_size
      )[self.global_dp_rank]
      gate_exp_avg_sq.data.copy_(global_gate_exp_avg_sq)
  
imports: from copy import deepcopy
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor, nn
  from torch.distributed import ProcessGroup
  from colossalai.cluster import ProcessGroupMesh
  from colossalai.moe.experts import MLPExperts
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.zero.low_level import LowLevelZeroOptimizer
label_desc: Use image node_id 15 for calling the LoadBalancer obj's underlying member method code with example usage: obj._swap_moe_param(swap_list, optim) without return types