code: def patch_forward_step(forward_step_func):
      r"""
      Patch model's forward_step_func to support balance loss
      """
  
      from megatron.mpu import is_pipeline_last_stage
      from megatron import get_args
  
      if not get_args().balance_strategy:
          return forward_step_func
  
      def forward_step_with_balance_loss(
          data_iterator, model, input_tensor="not given"
      ):
          args = get_args()
          if input_tensor == "not given":
              output = forward_step_func(data_iterator, model)
          else:
              output = forward_step_func(
                  data_iterator, model, input_tensor
              )
  
          if args.balance_strategy == "baseorig":
              loss, state_dict = output
          else:
              if (
                  not is_pipeline_last_stage()
                  or not args.balance_strategy
                  or args.balance_strategy in noloss_strategies
              ):
                  return output
              loss_name = args.balance_strategy + "_loss"
  
              while hasattr(model, "module"):
                  model = model.module
  
              loss_list = [
                  l.mlp.gate.get_loss(clear=False).view(1)
                  for l in model.language_model.transformer.layers
              ]
              (loss, state_dict), bal_loss = (
                  output,
                  torch.cat(loss_list).mean()
                  * args.balance_loss_weight,
              )
  
              # avarage across moe group
              moe_group = get_moe_group()
              world_size = torch.distributed.get_world_size(
                  group=moe_group
              )
              averaged_bal_loss = bal_loss.clone().detach()
              torch.distributed.all_reduce(
                  averaged_bal_loss, group=moe_group
              )
              averaged_bal_loss /= world_size
  
              loss += bal_loss
              state_dict[loss_name] = averaged_bal_loss
  
          return loss, state_dict
  
      return forward_step_with_balance_loss
  
file_path: thu-pacman-faster-moe/fmoe/megatron/balance.py
imports: import torch
  from fmoe.balance import reset_balance_profile
  from fmoe.balance import update_balance_profile
  from fmoe.utils import get_torch_default_comm
  from .distributed import get_moe_group
label_desc: Use image node_id 5 for calling a global function with example usage: patch_forward_step(forward_step_func) and returns: forward_step_with_balance_loss, forward_step_func, loss,  state_dict, output