Add MoE

Closes EleutherAI#479

configs/125M-moe.yml

@@ -0,0 +1,98 @@
+# GPT-2 pretraining setup
+{
+ # Have 2 experts per layer (every 2 layers by default)
+ "num_experts": 2,
+
+ # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages
+ # across the node boundaries )
+ "pipe_parallel_size": 1,
+ "model_parallel_size": 1,
+
+ # model settings
+ "num_layers": 12,
+ "hidden_size": 768,
+ "num_attention_heads": 12,
+ "seq_length": 2048,
+ "max_position_embeddings": 2048,
+ "norm": "layernorm",
+ "pos_emb": "rotary",
+ "no_weight_tying": true,
+ "gpt_j_residual": false,
+ "output_layer_parallelism": "column",
+
+ # these should provide some speedup but takes a while to build, set to true if desired
+ "scaled_upper_triang_masked_softmax_fusion": false,
+ "bias_gelu_fusion": false,
+ "rope_fusion": false,
+
+ # init methods
+ "init_method": "small_init",
+ "output_layer_init_method": "wang_init",
+
+
+ # optimizer settings
+ "optimizer": {
+ "type": "Adam",
+ "params": {
+ "lr": 0.0006,
+ "betas": [0.9, 0.95],
+ "eps": 1.0e-8,
+ }
+ },
+ "min_lr": 0.00006,
+
+ # for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training
+ "zero_optimization": {
+ "stage": 1,
+ "allgather_partitions": True,
+ "allgather_bucket_size": 500000000,
+ "overlap_comm": True,
+ "reduce_scatter": True,
+ "reduce_bucket_size": 500000000,
+ "contiguous_gradients": True,
+ },
+
+ # batch / data settings
+ "train_micro_batch_size_per_gpu": 4,
+ "data_impl": "mmap",
+
+ # activation checkpointing
+ "checkpoint_activations": true,
+ "checkpoint_num_layers": 1,
+ "partition_activations": true,
+ "synchronize_each_layer": true,
+
+ # regularization
+ "gradient_clipping": 1.0,
+ "weight_decay": 0.1,
+ "hidden_dropout": 0.0,
+ "attention_dropout": 0.0,
+
+ # precision settings
+ "fp16": {
+ "enabled": true,
+ "loss_scale": 0,
+ "loss_scale_window": 1000,
+ "hysteresis": 2,
+ "min_loss_scale": 1
+ },
+
+ # misc. training settings
+ "train_iters": 320000,
+ "lr_decay_iters": 320000,
+ "distributed_backend": "nccl",
+ "lr_decay_style": "cosine",
+ "warmup": 0.01,
+ "checkpoint_factor": 10000,
+ "eval_interval": 1000,
+ "eval_iters": 10,
+
+ # logging
+ "log_interval": 10,
+ "steps_per_print": 10,
+ "keep_last_n_checkpoints": 4,
+ "wall_clock_breakdown": true,
+
+ # networking
+ "hostfile": "/mock_path"
+}

megatron/model/transformer.py

@@ -43,6 +43,7 @@
  bias_dropout_add_fused_inference,
 )
 from megatron.model.utils import configure_sparse_attention
+from deepspeed.moe.layer import MoE
 
 # flags required to enable jit fusion kernels
 torch._C._jit_set_profiling_mode(False)
@@ -82,7 +83,13 @@ class ParallelMLP(nn.Module):
  """
 
  def __init__(
- self, neox_args, init_method, output_layer_init_method, parallel_output=False
+ self,
+ neox_args,
+ init_method,
+ output_layer_init_method,
+ parallel_output=False,
+ MOE=False,
+ MoE_mp_size=1,
  ):
  super().__init__()
 
@@ -104,6 +111,8 @@ def __init__(
  gather_output=False,
  init_method=init_method,
  skip_bias_add=True,
+ MOE=MOE,
+ MoE_mp_size=MoE_mp_size,
  )
  ff_dim_in = ff_dim // 2 if self.activation_type == "geglu" else ff_dim
  # Project back to h.
@@ -113,8 +122,10 @@ def __init__(
  output_size=neox_args.hidden_size,
  input_is_parallel=True,
  init_method=output_layer_init_method,
- skip_bias_add=True,
  parallel_output=parallel_output,
+ skip_bias_add=True,
+ MOE=MOE,
+ MoE_mp_size=MoE_mp_size,
  )
 
  def forward(self, hidden_states):
@@ -817,6 +828,44 @@ def __init__(
  else:
  raise KeyError(neox_args.mlp_type)
 
+ self.num_experts = neox_args.num_experts
+ args = neox_args
+ if self.num_experts <= 1:
+ self.mlp = ParallelMLP(
+ neox_args=neox_args,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ parallel_output=self.gpt_j_residual,
+ )
+ else:
+ from torch import distributed as dist
+
+ if self.num_experts > dist.get_world_size():
+ moe_mp_size = 1
+ else:
+ moe_mp_size = dist.get_world_size() // self.num_experts
+
+ self.mlp = MoE(
+ args.hidden_size,
+ ParallelMLP(
+ neox_args=neox_args,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ parallel_output=self.gpt_j_residual,
+ MOE=True,
+ MoE_mp_size=moe_mp_size,
+ ),
+ num_experts=self.num_experts,
+ ep_size=args.moe_expert_parallel_size,
+ k=args.topk,
+ use_residual=(args.mlp_type == "residual"),
+ capacity_factor=args.moe_train_capacity_factor,
+ eval_capacity_factor=args.moe_eval_capacity_factor,
+ min_capacity=args.moe_min_capacity,
+ drop_tokens=args.moe_token_dropping,
+ use_tutel=args.use_tutel,
+ )
+
  self.layer_past = None # used to cache k/v pairs in inference
 
  def _get_bias_dropout(self):
@@ -913,10 +962,19 @@ def forward(self, x, attention_mask, layer_past=None):
  )
 
  # output = x + mlp(ln2(x))
- mlp_output, mlp_bias = self.mlp(
- self.post_attention_layernorm(attention_output)
+ layernorm_output = self.post_attention_layernorm(attention_output)
+ moe_loss = torch.tensor(
+ 0.0, device=layernorm_output.device, dtype=layernorm_output.dtype
+ )
+ mlp_bias = torch.tensor(
+ 0.0, device=layernorm_output.device, dtype=layernorm_output.dtype
  )
 
+ if self.num_experts == 1:
+ mlp_output, mlp_bias = self.mlp(layernorm_output)
+ else:
+ mlp_output, moe_loss, _ = self.mlp(layernorm_output)
+
  with torch.enable_grad():
  if self.mlp_type == "llama":
  # No dropout either
@@ -930,7 +988,7 @@ def forward(self, x, attention_mask, layer_past=None):
  prob=self.hidden_dropout,
  )
 
- return output
+ return output, moe_loss
 
 
 class ParallelTransformerLayerPipe(ParallelTransformerLayer):
@@ -942,7 +1000,7 @@ def forward(self, args):
  ), "ParallelTransformerLayerPipe expects 2 arguments - hidden_states and attention_mask"
  hidden_states, attention_mask = args
  # we are returning just [hidden_states, mask]
- return super().forward(hidden_states, attention_mask), attention_mask
+ return super().forward(hidden_states, attention_mask)[0], attention_mask
 
 
 class ParallelLinearPipe(ParallelLinear):

megatron/model/utils.py

@@ -28,8 +28,12 @@ def get_params_for_weight_decay_optimization(module, neox_args):
  """Divide params into with-weight-decay and without-weight-decay groups.
  Layernorms and biases will have no weight decay but the rest will.
  """
- weight_decay_params = {"params": []}
- no_weight_decay_params = {"params": [], "weight_decay": 0.0}
+ weight_decay_params = {"params": [], "name": "weight_decay_params"}
+ no_weight_decay_params = {
+ "params": [],
+ "weight_decay": 0.0,
+ "name": "no_weight_decay_params",
+ }
  for module_ in module.modules():
  if any(
  [

megatron/mpu/__init__.py

@@ -31,6 +31,7 @@
 from .initialize import get_pipe_parallel_group
 from .initialize import get_pipe_parallel_rank
 from .initialize import get_pipe_parallel_world_size
+from .initialize import get_tensor_model_parallel_world_size
 from .initialize import get_io_parallel_group
 from .initialize import initialize_model_parallel
 from .initialize import model_parallel_is_initialized

megatron/mpu/initialize.py

@@ -266,6 +266,12 @@ def get_pipe_parallel_world_size():
  return torch.distributed.get_world_size(group=get_pipe_parallel_group())
 
 
+# Needed for MOE. True tensor parallelism todo.
+def get_tensor_model_parallel_world_size():
+ """Return world size for the tensor model parallel group."""
+ return torch.distributed.get_world_size()
+
+
 def destroy_model_parallel():
  """Set the groups to none."""
  global _MODEL_PARALLEL_GROUP

megatron/mpu/layers.py

@@ -413,6 +413,8 @@ def __init__(
  stride=1,
  keep_master_weight_for_test=False,
  skip_bias_add=False,
+ MOE=False,
+ MoE_mp_size=1,
  mup_rescale_parameters=False,
  ):
  super(ColumnParallelLinear, self).__init__()
@@ -422,7 +424,7 @@ def __init__(
  self.output_size = output_size
  self.gather_output = gather_output
  # Divide the weight matrix along the last dimension.
- world_size = get_model_parallel_world_size()
+ world_size = MoE_mp_size if MOE else get_model_parallel_world_size()
  self.output_size_per_partition = divide(output_size, world_size)
  self.skip_bias_add = skip_bias_add
  self.init_method = init_method
@@ -605,6 +607,8 @@ def __init__(
  stride=1,
  keep_master_weight_for_test=False,
  skip_bias_add=False,
+ MOE=False,
+ MoE_mp_size=1,
  parallel_output=False,
  mup_rescale_parameters=False,
  ):
@@ -615,7 +619,7 @@ def __init__(
  self.output_size = output_size
  self.input_is_parallel = input_is_parallel
  # Divide the weight matrix along the last dimension.
- world_size = get_model_parallel_world_size()
+ world_size = MoE_mp_size if MOE else get_model_parallel_world_size()
  self.input_size_per_partition = divide(input_size, world_size)
  self.skip_bias_add = skip_bias_add
  self.parallel_output = parallel_output

megatron/neox_arguments/arguments.py

@@ -1032,7 +1032,13 @@ def calculate_derived(self):
  # Update 'is pipe parallel' flag
  # if we set pipe_parallel_size to 0 or 1, GPT2ModelPipe.to_sequential() is called, and we run training with
  # the sequential model without the PipelineModule wrapper to avoid the overhead it incurs
- self.update_value("is_pipe_parallel", self.pipe_parallel_size >= 1)
+ self.update_value(
+ "is_pipe_parallel", self.pipe_parallel_size > 1 and self.num_experts == 1
+ )
+ assert not (
+ (self.is_pipe_parallel or self.pipe_parallel_size > 1)
+ and self.num_experts > 1
+ ), "MoE not supported with pipeline parallelism"
 
  # Attention config
  if self.attention_config is None:

megatron/neox_arguments/neox_args.py

@@ -83,6 +83,11 @@ class NeoXArgsParallelism(NeoXArgsTemplate):
  according to pipeline parallel size.
  """
 
+ expert_interval: int = 2
+ """
+ Have one MoE layer every expert_interval layers
+ """
+
 
 @dataclass
 class NeoXArgsModel(NeoXArgsTemplate):
@@ -1154,3 +1159,48 @@ class NeoXArgsTextgen(NeoXArgsTemplate):
  """
  Tasks to evaluate on using lm_eval_harness
  """
+
+ topk: int = 2
+ """
+ Activate top K experts in MoE
+ """
+
+ use_tutel: bool = False
+ """
+ Use Tutel optimizations in MoE
+ """
+
+ num_experts: int = 1
+ """
+ Number of MoE experts
+ """
+
+ moe_train_capacity_factor: float = 1.0
+ """
+ The capacity of the expert at train time
+ """
+
+ moe_eval_capacity_factor: float = 1.0
+ """
+ The capacity of the expert at eval time
+ """
+
+ moe_min_capacity: int = 4
+ """
+ The minimum capacity per expert regardless of the capacity_factor
+ """
+
+ moe_token_dropping: bool = True
+ """
+ Whether to drop tokens when exceeding capacity
+ """
+
+ create_moe_param_group: bool = True
+ """
+ Whether to create a separate parameter group for MoE parameters
+ """
+
+ moe_expert_parallel_size: int = 1
+ """
+ Number of parallel experts in MoE
+ """

megatron/training.py

@@ -478,6 +478,16 @@ def get_optimizer(model, neox_args):
  f'Configuring Optimizer type: {neox_args.optimizer_type} with params: {neox_args.optimizer["params"]}'
  )
 
+ if neox_args.create_moe_param_group:
+ from deepspeed.moe.utils import (
+ is_moe_param,
+ split_params_into_different_moe_groups_for_optimizer,
+ )
+
+ param_groups = split_params_into_different_moe_groups_for_optimizer(
+ param_groups
+ )
+
  # Add model parallel attribute if it is not set.
  for param_group in param_groups:
  for param in param_group["params"]: