code: def __init__(self):
      super().__init__()
  
      self._moe_layer = tutel_moe.moe_layer(
          gate_type={
              "type": "top",
              "k": top_value,
              "fp32_gate": args.fp32_gate,
          },
          experts={
              "type": "ffn",
              "count_per_node": num_local_experts,
              "hidden_size_per_expert": hidden_size,
              "activation_fn": lambda x: F.relu(x),
          },
          model_dim=model_dim,
          scan_expert_func=lambda name, param: setattr(
              param, "skip_allreduce", True
          ),
          seeds=(1, dist_rank + 1, 1),
          a2a_ffn_overlap_degree=a2a_ffn_overlap_degree,
      )
  
      # Summary of different parameter types: gate, local_experts
      local_count = sum(
          [
              torch.numel(param)
              for name, param in self._moe_layer.get_parameter_iterator(
                  param_type="local_experts"
              )
          ]
      )
      shared_count = sum(
          [
              torch.numel(param)
              for name, param in self._moe_layer.get_parameter_iterator(
                  param_type="gate"
              )
          ]
      )
      dist_print(
          "[Statistics] param count for MoE local_experts = %s, param count for MoE gate = %s.\n"
          % (local_count, shared_count)
      )
  
file_path: tutel/tutel/examples/helloworld_ddp.py
imports: import os
  import torch
  import torch.optim
  import torch.nn.functional
  import torch.distributed
  from torch import nn
  import argparse
  from tutel import system
  from tutel import moe
label_desc: Use image node_id 1 to create a new ExampleModel object from inherited base classes: torch.nn with example: obj = ExampleModel()