code: def __init__(self, d_model, n_expert, world_size, node_rank):
      global cr_cnt
      self.rep_lim = llims[cr_cnt % len(llims)]  # 4 - cr_cnt % 4
      cr_cnt += 1
      super().__init__(d_model, n_expert, world_size, top_k=2)
      self.ne_per_node = nw_per_node * n_expert
      self.ogn_ratio = 0.14
      self.node_rank = node_rank
  
      mask = [0] * world_size * n_expert
      for i in range(n_expert * world_size):
          if i // self.ne_per_node == self.node_rank:
              mask[i] = 1
      self.mask = torch.Tensor(mask).bool()
      self.stored_models = None
      self.policy_fn = None
  
file_path: thu-pacman-faster-moe/fmoe/gates/hir_gate.py
imports: from .naive_gate import NaiveGate
  import sys
  import torch
  import torch.distributed
  import torch.nn.functional
  from .utils import limit_by_capacity
  import fmoe_cuda
  from fmoe.functions import count_by_gate
label_desc: Use image node_id 1 to create a new HirGate object from inherited base classes: NaiveGate with example: obj = HirGate(d_model, n_expert, world_size, node_rank)