Mixture of Experts (MoE)

Layer specification

class deepspeed.moe.layer.MoE(hidden_size, expert, num_experts=1, k=1, capacity_factor=1.0, eval_capacity_factor=1.0, min_capacity=4, noisy_gate_policy: Optional[str] = None)

Initialize an MoE layer.

Parameters:

• hidden_size (int) – the hidden dimension of the model, importantly this is also the input and output dimension.
• expert (torch.nn.Module) – the torch module that defines the expert (e.g., MLP, torch.linear).
• num_experts (int, optional) – default=1, the total number of experts per layer.
• k (int, optional) – default=1, top-k gating value, only supports k=1 or k=2.
• capacity_factor (float, optional) – default=1.0, the capacity of the expert at training time.
• eval_capacity_factor (float, optional) – default=1.0, the capacity of the expert at eval time.
• min_capacity (int, optional) – default=4, the minimum capacity per expert regardless of the capacity_factor.
• noisy_gate_policy (str, optional) – default=None, noisy gate policy, valid options are ‘Jitter’, ‘RSample’ or ‘None’.

forward(hidden_states, used_token=None)

MoE forward

Parameters:

• hidden_states (Tensor) – input to the layer
• used_token (Tensor, optional) – default: None, mask only used tokens

Returns:

A tuple including output, gate loss, and expert count.

• output (Tensor): output of the model
• l_aux (Tensor): gate loss value
• exp_counts (int): expert count

Groups initialization

deepspeed.utils.groups.initialize(ep_size=1, mpu=None)

Process groups initialization supporting expert (E), data (D), and model (M) parallelism. DeepSpeed considers the following scenarios w.r.t. process group creation.

• S1: There is no expert parallelism or model parallelism, only data (D):

  model = my_model(args)
  engine = deepspeed.initialize(model) # initialize groups without mpu
  

• S2: There is expert parallelism but no model parallelism (E+D):

  deepspeed.utils.groups.initialize(ep_size) # groups will be initialized here
  model = my_model(args)
  engine = deepspeed.initialize(model)
  

• S3: There is model parallelism but no expert parallelism (M):

  mpu.init() # client initializes it's model parallel unit
  model = my_model(args)
  engine = deepspeed.initialize(model, mpu=mpu) # init w. mpu but ep_size = dp_world_size
  

• S4: There is model, data, and expert parallelism (E+D+M):

  mpu.init() # client initializes it's model parallel unit
  deepspeed.utils.groups.initialize(ep_size, mpu) # initialize expert groups wrt mpu
  model = my_model(args)
  engine = deepspeed.initialize(model, mpu=mpu) # passing mpu is optional in this case
  

Parameters:

• ep_size (int, optional) – default=1, expert parallel size
• mpu (module, optional) – default=None, model parallel unit (e.g., from Megatron) that describes model/data parallel ranks.