docstr: Search for the best swap combination to balance the data within the specified tolerance.
  And return the balanced data and the swap list. The swap list is used to record the swap.
  The swap list is a list of tuples. Each tuple is a swap operation.
  
  Args:
      data (List): expert load list.
          E.g. [[9.2, 8.3], [2.3, 10.0], [6.1, 7.2], [5.3, 3.2]]
          This means there are 4 devices and each devices has 2 experts.
          The value is the load of the expert.
      tolerance (float): tolerance for balance.
      beam_width (int): beam width for beam search.
      group_swap_factor (float): group swap factor for group swap loss.
          The bigger it is, the less times a group will be swapped.
      return_swapped_data (bool): whether to return the swapped data.
  
  Returns:
      Tuple: (balanced data, swap list).
          The swap list is a list of tuples. Each tuple is a swap operation.
          E.g. [(0, 0, 1, 0), (...), (...)]. The first tuple means
          the first expert of the first device is swapped with the first expert
          of the second device.
docstr_tok: ["Search","for","the","best","swap","combination","to","balance","the","data","within","the","specified","tolerance",".","And","return","the","balanced","data","and","the","swap","list",".","The","swap","list","is","used","to","record","the","swap",".","The","swap","list","is","a","list","of","tuples",".","Each","tuple","is","a","swap","operation",".","Args",":","data","(","List",")",":","expert","load","list",".","E.g",".","[","[","9.2",",","8.3","]",",","[","2.3",",","10.0","]",",","[","6.1",",","7.2","]",",","[","5.3",",","3.2","]","]","This","means","there","are","4","devices","and","each","devices","has","2","experts",".","The","value","is","the","load","of","the","expert",".","tolerance","(","float",")",":","tolerance","for","balance",".","beam_width","(","int",")",":","beam","width","for","beam","search",".","group_swap_factor","(","float",")",":","group","swap","factor","for","group","swap","loss",".","The","bigger","it","is",",","the","less","times","a","group","will","be","swapped",".","return_swapped_data","(","bool",")",":","whether","to","return","the","swapped","data",".","Returns",":","Tuple",":","(","balanced","data",",","swap","list",")",".","The","swap","list","is","a","list","of","tuples",".","Each","tuple","is","a","swap","operation",".","E.g",".","[","(","0",",","0",",","1",",","0",")",",","(","...",")",",","(","...",")","]",".","The","first","tuple","means","the","first","expert","of","the","first","device","is","swapped","with","the","first","expert","of","the","second","device","."]
code: def _search_balance(
      self,
      data: List,
      tolerance: Optional[float] = 0.1,
      beam_width: Optional[int] = 8,
      group_swap_factor: Optional[float] = 0.4,
      return_swapped_data: Optional[bool] = False,
  ) -> Tuple[List, List]:
      """
      Search for the best swap combination to balance the data within the specified tolerance.
      And return the balanced data and the swap list. The swap list is used to record the swap.
      The swap list is a list of tuples. Each tuple is a swap operation.
  
      Args:
          data (List): expert load list.
              E.g. [[9.2, 8.3], [2.3, 10.0], [6.1, 7.2], [5.3, 3.2]]
              This means there are 4 devices and each devices has 2 experts.
              The value is the load of the expert.
          tolerance (float): tolerance for balance.
          beam_width (int): beam width for beam search.
          group_swap_factor (float): group swap factor for group swap loss.
              The bigger it is, the less times a group will be swapped.
          return_swapped_data (bool): whether to return the swapped data.
  
      Returns:
          Tuple: (balanced data, swap list).
              The swap list is a list of tuples. Each tuple is a swap operation.
              E.g. [(0, 0, 1, 0), (...), (...)]. The first tuple means
              the first expert of the first device is swapped with the first expert
              of the second device.
      """
      norm_data = self._normalize_data(data)
      avg = sum(
          sum(sublist) / len(sublist) for sublist in norm_data
      ) / len(norm_data)
      results = [(norm_data, 0, [])]
      stop_flag = False
  
      while stop_flag == False:
          new_results = []
          best_score = results[0][1]
          for i in range(len(results)):
              new_results.extend(
                  self._beam_search(
                      results[i], beam_width, avg, group_swap_factor
                  )
              )
          if len(new_results) == 0:
              stop_flag = True
              break
          new_results.sort(key=lambda x: x[1], reverse=True)
          new_best_score = new_results[0][1]
          if new_best_score == best_score:
              stop_flag = True
              break
          new_results = new_results[:beam_width]
          results = new_results
          for i in results:
              if self._check_convergence(results[0][0], avg, tolerance):
                  stop_flag = True
                  break
  
      swap_list = results[0][2]
      if return_swapped_data:
          out = deepcopy(data)
          for swap in swap_list:
              self._swap_data(out, *swap)
          return out, swap_list
      else:
          return swap_list
  
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 11 for calling the LoadBalancer obj's underlying member method code with example usage: obj._search_balance(data, tolerance, beam_width, group_swap_factor, return_swapped_data) and returns: out,  swap_list, swap_list