code: def _test_faster_gate(
      d_model, batch_size, n_expert, gpu_per_node, frac
  ):
      _ensure_initialized()
      rank = dist.get_rank()
      node_rank = rank // gpu_per_node
  
      gate = FasterGate(
          d_model, n_expert, dist.get_world_size(), node_rank
      ).cuda()
      x = torch.rand(batch_size, d_model).cuda()
      topk_idx, topk_val = gate(x)
  
      cnto = 0
      idxs = topk_idx[:, 0].cpu().view(-1).numpy()
      for v in idxs:
          assert v != -1
          if v // n_expert // gpu_per_node != rank // gpu_per_node:
              cnto += 1
      assert cnto <= math.ceil(batch_size * frac)
  
file_path: fastmoe/tests/test_faster_gate.py
imports: import pytest
  import os
  import sys
  import json
  import math
  import torch
  import torch.distributed
  import torch.nn.functional
  from fmoe.gates.faster_gate import FasterGate
  from test_ddp import _ensure_initialized, _run_distributed
label_desc: Use image node_id 2 for calling a global function with example usage: _test_faster_gate(d_model, batch_size, n_expert, gpu_per_node, frac) without return types