[moe] update train script (hpcaitech#4959)

* update

* update ckpt

* update train

* update train

colossalai/moe/checkpoint.py

@@ -48,14 +48,15 @@ def pre_load_model(self, model: nn.Module, state_dict: dict) -> dict:
  """
  for name, param in state_dict.items():
  if ".experts." in name:
- model_param = dict(model.named_parameters())[name]
- if is_moe_tensor(model_param):
- ep_rank = get_ep_rank(model_param)
- ep_size = get_ep_size(model_param)
- expert_num = param.shape[0] // ep_size
- assert param.shape[0] % ep_size == 0
- param = param[ep_rank * expert_num:(ep_rank + 1) * expert_num]
- state_dict[name] = param
+ if name in dict(model.named_parameters()):
+ model_param = dict(model.named_parameters())[name]
+ if is_moe_tensor(model_param):
+ ep_rank = get_ep_rank(model_param)
+ ep_size = get_ep_size(model_param)
+ expert_num = param.shape[0] // ep_size
+ assert param.shape[0] % ep_size == 0
+ param = param[ep_rank * expert_num:(ep_rank + 1) * expert_num]
+ state_dict[name] = param
  dist.barrier()
  return state_dict
 

examples/language/openmoe/benchmark/benchmark_cai.py

@@ -4,7 +4,7 @@
 import torch
 import torch.distributed as dist
 from huggingface_hub import snapshot_download
-from model.modeling_openmoe import OpenMoeForCausalLM
+from model.modeling_openmoe import OpenMoeForCausalLM, set_openmoe_args
 from model.openmoe_policy import OpenMoeForCausalLMPolicy
 from torch.utils.data import Dataset
 from tqdm import tqdm
@@ -19,7 +19,7 @@
 from colossalai.cluster import DistCoordinator
 from colossalai.moe.layers import apply_load_balance
 from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import set_moe_args, skip_init
+from colossalai.moe.utils import skip_init
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.utils import get_current_device
 
@@ -218,28 +218,12 @@ def main():
  # Build OpenMoe model
  repo_name = "hpcaitech/openmoe-" + args.model_name
  config = LlamaConfig.from_pretrained(repo_name)
- moe_args = {
- "num_experts": config.num_experts,
- "moe_layer_interval": config.moe_layer_interval,
- "router_topk": 2,
- "router_capacity_factor_train": 1.25,
- "router_capacity_factor_eval": 2.0,
- "router_min_capacity": 4,
- "router_noisy_policy": None,
- "router_drop_tks": True,
- "router_aux_loss_factor": 0.01,
- "router_z_loss_factor": 0.01,
- "mlp_gated": True,
- "label_smoothing": 0.001,
- "z_loss_factor": 0.01,
- "enable_load_balance": args.load_balance,
- "load_balance_tolerance": 0.1,
- "load_balance_beam_width": 8,
- "load_balance_group_swap_factor": 0.4,
- "enable_kernel": args.use_kernel,
- "enable_comm_overlap": args.overlap_alltoall,
- }
- set_moe_args(config, moe_args)
+ set_openmoe_args(config,
+ num_experts=config.num_experts,
+ moe_layer_interval=config.moe_layer_interval,
+ enable_load_balance=args.load_balance,
+ enable_kernel=args.use_kernel,
+ enable_comm_overlap=args.overlap_alltoall)
  with skip_init():
  model = OpenMoeForCausalLM(config)
  coordinator.print_on_master(f"Finish init model with config:\n{config}")

examples/language/openmoe/infer.py

@@ -1,12 +1,10 @@
 from argparse import ArgumentParser
 
 import torch
-from model.modeling_openmoe import OpenMoeForCausalLM
+from model.modeling_openmoe import OpenMoeForCausalLM, set_openmoe_args
 from transformers import T5Tokenizer
 from transformers.models.llama import LlamaConfig
 
-from colossalai.moe.utils import set_moe_args
-
 
 def parse_args():
  parser = ArgumentParser()
@@ -15,59 +13,22 @@ def parse_args():
 
 
 def inference(args):
-
  tokenizer = T5Tokenizer.from_pretrained("google/umt5-small")
  if args.model == "test":
  config = LlamaConfig.from_pretrained("hpcaitech/openmoe-base")
- moe_args = {
- "num_experts": config.num_experts,
- "moe_layer_interval": config.moe_layer_interval,
- "router_topk": 2,
- "router_capacity_factor_train": 1.25,
- "router_capacity_factor_eval": 2.0,
- "router_min_capacity": 4,
- "router_noisy_policy": None,
- "router_drop_tks": True,
- "router_aux_loss_factor": 0.01,
- "router_z_loss_factor": 0.01,
- "mlp_gated": True,
- "label_smoothing": 0.001,
- "z_loss_factor": 0.01,
- "enable_load_balance": False,
- "load_balance_tolerance": 0.1,
- "load_balance_beam_width": 8,
- "load_balance_group_swap_factor": 0.4,
- "enable_kernel": False,
- "enable_comm_overlap": False,
- }
- set_moe_args(config, moe_args)
+ set_openmoe_args(config,
+ num_experts=config.num_experts,
+ moe_layer_interval=config.moe_layer_interval,
+ enable_kernel=True)
  model = OpenMoeForCausalLM(config)
  else:
  config = LlamaConfig.from_pretrained(f"hpcaitech/openmoe-{args.model}")
- moe_args = {
- "num_experts": config.num_experts,
- "moe_layer_interval": config.moe_layer_interval,
- "router_topk": 2,
- "router_capacity_factor_train": 1.25,
- "router_capacity_factor_eval": 2.0,
- "router_min_capacity": 4,
- "router_noisy_policy": None,
- "router_drop_tks": True,
- "router_aux_loss_factor": 0.01,
- "router_z_loss_factor": 0.01,
- "mlp_gated": True,
- "label_smoothing": 0.001,
- "z_loss_factor": 0.01,
- "enable_load_balance": False,
- "load_balance_tolerance": 0.1,
- "load_balance_beam_width": 8,
- "load_balance_group_swap_factor": 0.4,
- "enable_kernel": False,
- "enable_comm_overlap": False,
- }
- set_moe_args(config, moe_args)
+ set_openmoe_args(config,
+ num_experts=config.num_experts,
+ moe_layer_interval=config.moe_layer_interval,
+ enable_kernel=False)
  model = OpenMoeForCausalLM.from_pretrained(f"hpcaitech/openmoe-{args.model}", config=config)
- model = model.eval().half()
+ model = model.eval().bfloat16()
  model = model.to(torch.cuda.current_device())
 
  input_str = """```
@@ -86,7 +47,7 @@ def inference(args):
  # print("model config: ", model.config)
  input_ids = tokenizer("<pad>" + input_str, return_tensors="pt", add_special_tokens=False)
  input_ids = input_ids.input_ids.to(torch.cuda.current_device())
- generation_output = model.generate(input_ids, use_cache=True, do_sample=True, max_new_tokens=16)
+ generation_output = model.generate(input_ids, use_cache=True, do_sample=True, max_new_tokens=64)
  out = tokenizer.decode(generation_output[0], skip_special_tokens=False)
  print(f"output: \n{out}\n")
 

examples/language/openmoe/model/modeling_openmoe.py

@@ -39,7 +39,7 @@
 from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
 from colossalai.moe.layers import SparseMLP
 from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import get_activation
+from colossalai.moe.utils import get_activation, set_moe_args
 
 if HAS_TRITON:
  from colossalai.kernel.triton.llama_act_combine_kernel import LlamaActCombine
@@ -49,6 +49,78 @@
 _CONFIG_FOR_DOC = "LlamaConfig"
 
 
+def set_openmoe_args(
+ config: LlamaConfig,
+ num_experts: int,
+ moe_layer_interval: int,
+ router_topk: int = 2,
+ router_capacity_factor_train: float = 1.25,
+ router_capacity_factor_eval: float = 2.0,
+ router_min_capacity: int = 4,
+ router_noisy_policy: str = None,
+ router_drop_tks: bool = True,
+ router_aux_loss_factor: float = 0.01,
+ router_z_loss_factor: float = 0.01,
+ mlp_gated: bool = True,
+ label_smoothing: float = 0.001,
+ z_loss_factor: float = 0.01,
+ enable_load_balance: bool = False,
+ load_balance_tolerance: float = 0.1,
+ load_balance_beam_width: int = 8,
+ load_balance_group_swap_factor: float = 0.4,
+ enable_kernel: bool = False,
+ enable_comm_overlap: bool = False,
+) -> None:
+ """
+ MoE related arguments.
+ It inserts the MoE arguments into the Llama config.
+
+ Args:
+ config (LlamaConfig): Transformers Llama config.
+ num_experts (int, optional): Number of experts.
+ moe_layer_interval (int, optional): The interval moe layer.
+ router_topk (int, optional): Moe router top k. Defaults to 2.
+ router_capacity_factor_train (float, optional): Moe router max capacity for train. Defaults to 1.25.
+ router_capacity_factor_eval (float, optional): Moe router max capacity for eval. Defaults to 2.0.
+ router_min_capacity (int, optional): Moe router min capacity. Defaults to 4.
+ router_noisy_policy (str, optional): Moe router noisy policy. You can choose [Jitter, Gaussian, None]. Defaults to None.
+ router_drop_tks (bool, optional): Whether moe router drop tokens which exceed max capacity. Defaults to True.
+ router_aux_loss_factor (float, optional): Moe router aux loss. You can refer to STMoE for details. Defaults to 0.01.
+ router_z_loss_factor (float, optional): Moe router z loss. You can refer to STMoE for details. Defaults to 0.01.
+ mlp_gated (bool, optional): Use gate in mlp. Defaults to True.
+ label_smoothing (float, optional): Label smoothing. Defaults to 0.001.
+ z_loss_factor (float, optional): The final outputs' classification z loss factor. Defaults to 0.01.
+ enable_load_balance (bool, optional): Expert load balance. Defaults to False.
+ load_balance_tolerance (float, optional): Expert load balance search's difference tolerance. Defaults to 0.1.
+ load_balance_beam_width (int, optional): Expert load balance search's beam width. Defaults to 8.
+ load_balance_group_swap_factor (float, optional): Expert load balance group swap factor. Longer value encourages less swap. Defaults to 0.4.
+ enable_kernel (bool, optional): Use kernel optimization. Defaults to False.
+ enable_comm_overlap (bool, optional): Use communication overlap for MoE. Recommended to enable for muiti-node training. Defaults to False.
+ """
+ moe_args = dict(
+ num_experts=num_experts,
+ moe_layer_interval=moe_layer_interval,
+ router_topk=router_topk,
+ router_capacity_factor_train=router_capacity_factor_train,
+ router_capacity_factor_eval=router_capacity_factor_eval,
+ router_min_capacity=router_min_capacity,
+ router_noisy_policy=router_noisy_policy,
+ router_drop_tks=router_drop_tks,
+ router_aux_loss_factor=router_aux_loss_factor,
+ router_z_loss_factor=router_z_loss_factor,
+ mlp_gated=mlp_gated,
+ label_smoothing=label_smoothing,
+ z_loss_factor=z_loss_factor,
+ enable_load_balance=enable_load_balance,
+ load_balance_tolerance=load_balance_tolerance,
+ load_balance_beam_width=load_balance_beam_width,
+ load_balance_group_swap_factor=load_balance_group_swap_factor,
+ enable_kernel=enable_kernel,
+ enable_comm_overlap=enable_comm_overlap,
+ )
+ set_moe_args(config, moe_args)
+
+
 # Copied from transformers.models.bart.modeling_bart._make_causal_mask
 def _make_causal_mask(input_ids_shape: torch.Size,
  dtype: torch.dtype,
@@ -96,7 +168,7 @@ def generate_fixed_pos_embedding(features, length, min_timescale=1.0, max_timesc
  output_sin: a float32 Tensor with shape [length, features]
  output_cos: a float32 Tensor with shape [length, features]
  """
- fraction = torch.arange(0, features, 2, dtype=torch.float64).cuda() / features
+ fraction = torch.arange(0, features, 2, dtype=torch.float32).cuda() / features
  timescale = min_timescale * (max_timescale / min_timescale)**fraction
  rotational_frequency = 1. / timescale
 
@@ -231,7 +303,7 @@ def __init__(self, config: LlamaConfig):
  self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
  self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
  self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
- self.sin, self.cos = generate_fixed_pos_embedding(self.head_dim, self.max_position_embeddings, 1e4)
+ self.sin, self.cos = generate_fixed_pos_embedding(self.head_dim, self.max_position_embeddings, 1.0, 1e4)
 
  def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
  return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

examples/language/openmoe/test_ci.sh

@@ -7,7 +7,13 @@ python infer.py --model "test"
 torchrun --standalone --nproc_per_node 4 train.py \
  --num_epoch 1 \
  --model_name "test" \
- --plugin zero2_ep \
+ --plugin "ep" \
+ --batch_size 1
+
+torchrun --standalone --nproc_per_node 4 train.py \
+ --num_epoch 1 \
+ --model_name "test" \
+ --plugin "ep_zero" \
  --batch_size 1
 
 torchrun --standalone --nproc_per_node 4 train.py \