code: def backward(ctx, grad_out):
      (
          pos_s,
          pos_g,
          local_expert_count,
          global_expert_count,
          stored_models,
          _1,
          _2,
      ) = ctx.saved_tensors
      (
          fwd_batch_size,
          inp_batch_size,
          num_expert,
          world_size,
      ) = ctx.moe_args
  
      def _expert_backward(grad_y, grad_x, expert_idx, store_idx):
          y = ctx.gobs[store_idx]
          x = ctx.gibs[store_idx]
          torch.autograd.backward([y], [grad_y])
          grad_x.copy_(x.grad)
  
      experts = ctx.experts
  
      def stash_fn(store_idx, expert_idx):
          expert_utils.stash_expert_params(
              experts, ctx.shadows[store_idx], expert_idx
          )
  
      pop_fn = lambda idx: expert_utils.pop_expert_params(experts, idx)
  
      def collect_fn(store_idx, root, expert_idx):
          grad = ctx.shadows[store_idx]
          expert_utils.collect_expert_grads(experts, grad, expert_idx)
          fmoe_native.reduce_grad(grad, root, ctx.expert_size)
  
      set_grad_fn = (
          lambda store_idx, expert_idx: expert_utils.set_grads(
              experts, ctx.shadows[store_idx], expert_idx
          )
      )
  
      grad_out_buf = _local_scatter(grad_out.contiguous(), pos_g)
      grad_in_buf = fmoe_native.smart_sch_backward(
          grad_out_buf,
          local_expert_count,
          global_expert_count,
          stored_models,
          pos_s.shape[0],
          fwd_batch_size,
          world_size,
          _expert_backward,
          stash_fn,
          pop_fn,
          collect_fn,
          set_grad_fn,
      )
      grad_in = _local_gather(grad_in_buf, pos_s, inp_batch_size)
  
      return (
          None,
          None,
          grad_in,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
      )
  
file_path: fastmoe/fmoe/fastermoe/schedule.py
imports: import torch
  from torch.autograd.function import Function
  from fmoe.functions import prepare_forward, ensure_comm
  from fmoe.functions import _local_scatter, _local_gather
  import fmoe_cuda
  from fmoe.fastermoe import expert_utils
  from .shadow_policy import get_shadow_policy
label_desc: Use image node_id 2 for calling the MoEForward obj's underlying member method code with example usage: obj.backward(ctx, grad_out) and returns: None,  None,  grad_in,  None,  None,  None,  None,  None,  None,  None,  None,  None