code: def _test_faster_schedule(d_model, batch_size, n_expert):
      _ensure_initialized()
      rank = dist.get_rank()
      world_size = dist.get_world_size()
  
      x1 = torch.rand(batch_size, d_model).cuda()
      x1.requires_grad = True
      x2 = x1.data.clone()
      x2.requires_grad = True
      topk_idx = torch.randint(
          0, world_size * n_expert, (batch_size, 2)
      ).cuda()
      m1s = [
          torch.nn.Linear(d_model, d_model).cuda()
          for _ in range(n_expert)
      ]
      m2s = [
          torch.nn.Linear(d_model, d_model).cuda()
          for _ in range(n_expert)
      ]
      with torch.no_grad():
          for m1, m2 in zip(m1s, m2s):
              m2.weight.copy_(m1.weight)
              m2.bias.copy_(m1.bias)
  
      def ef1(x, fec, eidx):
          return m1s[eidx](x)
  
      def ef2(x, fec):
          o = 0
          ys = []
          for m, i in zip(m2s, fec):
              if i > 0:
                  ys.append(m(x[o : o + i]))
                  o += i
          y = torch.cat(ys)
          return y
  
      ensure_comm(x1, None)
      y1 = smart_fwd(
          x1, topk_idx, ef1, n_expert, world_size, experts=m1s
      )
      y1.sum().backward()
  
      y2 = naive_fwd(
          x2, topk_idx, ef2, n_expert, world_size, experts=m2s
      )
      y2.sum().backward()
      _assert_numerical(
          ["out", "grad_in"], [y1, x1.grad], [y2, x2.grad], rank
      )
      for i in range(n_expert):
          _assert_numerical(
              [f"grad_bias_{i}", f"grad_weight_{i}"],
              [m1s[i].bias.grad, m1s[i].weight.grad],
              [m2s[i].bias.grad, m2s[i].weight.grad],
              rank,
          )
  
file_path: fastmoe/tests/test_faster_schedule.py
imports: import pytest
  import os
  import sys
  import json
  import math
  import torch
  import torch.distributed
  import torch.nn.functional
  from fmoe.functions import ensure_comm
  from test_ddp import _ensure_initialized, _run_distributed
  from test_numerical import _assert_numerical
  from fmoe.fastermoe.schedule import _fmoe_general_global_forward
  from fmoe.layers import _fmoe_general_global_forward
label_desc: Use image node_id 2 for calling a global function with example usage: _test_faster_schedule(d_model, batch_size, n_expert) and returns: m1s, y