Mixture of Experts

.gitignore

@@ -140,3 +140,4 @@ out*
 tests/assets/*
 .vscode/
 checkpoints/
+experiments/

MOE.md

@@ -0,0 +1,72 @@
+# Mixture of Experts Language Models
+
+## Dependencies
+
+```
+pip install megablocks
+pip install xformers==0.0.22.post4
+```
+
+## Train MoE
+
+To train an MoE, add the `--moe-X` related arguments to the training command:
+
+```
+torchrun --nproc-per-node 8 -m open_lm.main \
+    --train-num-samples 1000000000 \
+    --workers 2 \
+    --dataset-manifest "s3://laion-west/rpj_tokenized_upsampled_eleutherai/manifest.jsonl" "s3://laion-west/2T_no_rpj_tokenized_upsampled_25k_shards/manifest.jsonl" \
+    --train-data-mix-weights 0.725 0.275 \
+    --precision amp_bfloat16 \
+    --batch-size 8 \
+    --log-every-n-steps 20 \
+    --grad-clip-norm 1 \
+    --lr 6e-4 \
+    --warmup 200 \
+    --model aphid_neox \
+    --wd 0.01 \
+    --beta2 0.95 \
+    --epochs 4 \
+    --report-to wandb \
+    --wandb-project-name moe \
+    --name test_moe \
+    --logs /fsx/home-$USER/experiments/moe \
+    --resume latest \
+    --seed 124 \
+    --data-key 'json' \
+    --accum-freq 4 \
+    --model-norm gain_only_layer_norm \
+    --fsdp --fsdp-amp \
+    --lr-cooldown-end 1e-5 \
+    --no-skip-tokens \
+    --accurate-total-tokens \
+    --moe-freq 2 \
+    --moe-num-experts 8 \
+    --moe-top-k 2 \
+    --moe-capacity-factor 1.25 --moe-loss-weight 0.1
+```
+
+The above command will add an MoE FFN layer to every other Transformer block. You can use an arbitrary number of experts; you are only limited by total RAM across all GPUs. 
+
+
+You can also add the `moe_expert_model_parallelism` which will distribute experts across different GPUs. However, if the #GPU is larger than #Expert, additional #GPU/#Expert tensor parallelism is applied. Currently this is not eval-friendly though, so I would not recommend using it yet.
+
+You can evaluate the MoE in the same way as dense models:
+
+```
+torchrun --nproc-per-node 8 -m open_lm.main \
+                --val-data "pipe:aws s3 cp s3://laion-west/lmdata/validation_data_tokenized/open_lm//shard_00000000.tar -" \
+                --workers 6 \
+                --precision amp_bfloat16 \
+                --batch-size 8 \
+                --log-every-n-steps 1 \
+                --model open_lm_41m \
+                --fsdp --fsdp-amp \
+                --moe-num-experts 64 --moe-freq 2 \
+                --data-key json \
+                --train-num-samples 1000000000 \
+                --model-norm gain_only_layer_norm \
+                --name $RANDOM \
+                --resume /fsx/home-suching/experiments/mix_wo/test8086/checkpoints/epoch_1.pt \
+                --logs /fsx/home-$USER/experiments/eval
+```

open_lm/main.py

@@ -399,7 +399,11 @@ def main(args):
     if args.hf_model is not None:
         model = create_wrapped_hf_model(args)
     else:
-        with torch.device("meta" if args.fsdp else args.device):
+        # with torch.device("meta" if args.fsdp else args.device):
+        #     model = create_model(args)
+        # if not args.fsdp:
+        #     model.reset_parameters()
+        with torch.device(args.device):
             model = create_model(args)
         if not args.fsdp:
             model.reset_parameters()

open_lm/model.py

@@ -1,5 +1,6 @@
 import math
 import json
+import logging
 import re
 from copy import deepcopy
 from pathlib import Path
@@ -19,6 +20,15 @@
 from open_lm.positional_embedding.rotary import RotaryWithCast
 from open_lm.positional_embedding.llama_rotary import LLaMARotaryWithCast
 
+
+# from open_lm.moe.mixture_of_experts import MoE
+try:
+    from megablocks.layers.moe import MoE
+    from megablocks.layers.arguments import Arguments as MoEArgs
+except ImportError:
+    logging.warning(f"Megablocks not installed. To train MoE, install with pip install megablocks.")
+
+
 try:  # optional import
     from mamba_ssm import MambaLMHeadModel
 except ImportError:
@@ -77,6 +87,13 @@ class Params:
     weight_tying: bool = False
     norm_type: nn.Module = nn.LayerNorm
     apply_qk_norm: bool = False
+    moe_loss_weight: float = 0.1
+    moe_capacity_factor: float = 1.25
+    moe_expert_model_parallelism: bool = False
+    moe_weight_parallelism: bool = False
+    moe_num_experts: int = 8
+    moe_top_k: int = 2
+    moe_freq: int = 0
     positional_embedding_type: str = "rotary"
     ffn_type: str = "swiglu"
 
@@ -237,10 +254,10 @@ def __init__(self, layer_id, args: Params):
         super().__init__()
         self.n_heads = args.n_heads
         self.dim = args.dim
+
         self.head_dim = args.dim // args.n_heads
         self.attention = CustomAttn(layer_id, args)
         self._ffn_type = args.ffn_type
-
         if args.ffn_type == "swiglu":
             # this follows llama / lit llama -- go to multiple of 256
             self.hidden_dim = 256 * ((int(2 * 4 * args.dim / 3) + 256 - 1) // 256)
@@ -251,6 +268,21 @@ def __init__(self, layer_id, args: Params):
             self._ff_w1 = nn.Linear(args.dim, self.hidden_dim, bias=False)
             self._ff_w2 = nn.Linear(self.hidden_dim, args.dim, bias=False)
             self.feed_forward = nn.Sequential(self._ff_w1, nn.GELU(approximate="none"), self._ff_w2)
+        elif args.ffn_type == "moe":
+            moe_args = MoEArgs(
+                hidden_size=args.dim,
+                ffn_hidden_size=args.dim * 4,
+                moe_num_experts=args.moe_num_experts,
+                moe_weight_parallelism=args.moe_weight_parallelism,
+                moe_expert_model_parallelism=args.moe_expert_model_parallelism,
+                moe_top_k=args.moe_top_k,
+                moe_capacity_factor=args.moe_capacity_factor,
+                moe_loss_weight=args.moe_loss_weight,
+                device=torch.cuda.current_device(),
+                bf16=False,
+                fp16=False,
+            )
+            self.feed_forward = MoE(moe_args)
 
         self.layer_id = layer_id
         self.attention_norm = args.norm_type(
@@ -289,7 +321,11 @@ def forward(self, x, past_key_value=None, use_cache=False):
             use_cache=use_cache,
         )
         h = x + h
-        out = h + self.feed_forward(self.ffn_norm(h))
+        if self._ffn_type == "moe":
+            ffn_out, _ = self.feed_forward(self.ffn_norm(h))
+        else:
+            ffn_out = self.feed_forward(self.ffn_norm(h))
+        out = h + ffn_out
         return out, past_key_value
 
 
@@ -298,8 +334,10 @@ def __init__(self, params):
         super().__init__()
         # for convenience we often share param names with llama
         self.params = params
+        self.dim = params.dim
         self.vocab_size = params.vocab_size
         self.n_layers = params.n_layers
+        self.moe_num_experts = params.moe_num_experts
         self.seq_len = params.seq_len
         self.post_embed_norm = (
             params.norm_type(
@@ -314,7 +352,12 @@ def __init__(self, params):
         self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
 
         self.layers = torch.nn.ModuleList()
+        ffn_type_ = params.ffn_type
         for layer_id in range(params.n_layers):
+            if params.moe_freq > 0 and layer_id % params.moe_freq == 0:
+                params.ffn_type = "moe"
+            else:
+                params.ffn_type = ffn_type_
             self.layers.append(Block(layer_id, params))
 
         # get class for normalization layers
@@ -405,6 +448,13 @@ def create_params(args):
             apply_qk_norm=cfg.get("qk_norm", args.qk_norm),
             positional_embedding_type=cfg.get("positional_embedding_type", args.positional_embedding_type),
             ffn_type=cfg.get("ffn_type", args.ffn_type),
+            moe_num_experts=cfg.get("moe_num_experts", args.moe_num_experts),
+            moe_loss_weight=cfg.get("moe_loss_weight", args.moe_loss_weight),
+            moe_expert_model_parallelism=cfg.get("moe_expert_model_parallelism", args.moe_expert_model_parallelism),
+            moe_weight_parallelism=cfg.get("moe_weight_parallelism", args.moe_weight_parallelism),
+            moe_capacity_factor=cfg.get("moe_capacity_factor", args.moe_capacity_factor),
+            moe_freq=cfg.get("moe_freq", args.moe_freq),
+            moe_top_k=cfg.get("moe_top_k", args.moe_top_k),
         )
 
 

open_lm/params.py

@@ -56,6 +56,50 @@ def add_model_args(parser):
         default="rotary",
         help="Type of positional embedding to use. This might be overridden by the model config.",
     )
+    parser.add_argument(
+        "--moe-freq",
+        type=int,
+        default=0,
+        help="if set > 0, we will add MoE layer to every moe_freq layer.",
+    )
+    parser.add_argument(
+        "--moe-num-experts",
+        type=int,
+        default=None,
+        help="Number of experts for MoE",
+    )
+
+    parser.add_argument(
+        "--moe-weight-parallelism",
+        action="store_true",
+        help="Add weight parallelism to MoE",
+    )
+
+    parser.add_argument(
+        "--moe-expert-model-parallelism",
+        action="store_true",
+        help="Add expert model parallelism to MoE",
+    )
+
+    parser.add_argument(
+        "--moe-capacity-factor",
+        type=float,
+        default=1.25,
+        help="MoE capacity factor",
+    )
+
+    parser.add_argument(
+        "--moe-loss-weight",
+        type=float,
+        default=0.1,
+        help="MoE loss weight",
+    )
+    parser.add_argument(
+        "--moe-top-k",
+        type=int,
+        default=2,
+        help="MoE top k experts",
+    )
 
 
 def check_replacement_type(replacement, original):

open_lm/train.py

@@ -11,6 +11,13 @@
 from torch.distributed.distributed_c10d import ReduceOp
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
+
+try:
+    from megablocks.layers.moe import batched_load_balancing_loss, clear_load_balancing_loss
+    from megablocks.layers.arguments import Arguments as MoEArgs
+except ImportError:
+    logging.warning(f"Megablocks not installed. To train MoE, install with pip install megablocks.")
+
 try:
     import wandb
 except ImportError:
@@ -155,6 +162,7 @@ def train_one_epoch(model, data, loss, epoch, step, optimizer, scaler, scheduler
     sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
 
     losses_m = AverageMeter()
+    load_balancing_losses_m = AverageMeter()
     batch_time_m = AverageMeter()
     data_time_m = AverageMeter()
 
@@ -228,11 +236,37 @@ def train_one_epoch(model, data, loss, epoch, step, optimizer, scaler, scheduler
                             * inputs_ii.shape[0]
                             / inputs.shape[0]
                         )
+                    if args.moe_freq > 0:
+                        moe_args = MoEArgs(
+                            hidden_size=model.dim,
+                            ffn_hidden_size=model.dim * 4,
+                            moe_num_experts=args.moe_num_experts,
+                            num_layers=model.n_layers // 2,
+                            moe_expert_model_parallelism=True,
+                            moe_top_k=args.moe_top_k,
+                            device=torch.distributed.get_rank(),
+                            moe_capacity_factor=args.moe_capacity_factor,
+                            moe_loss_weight=args.moe_loss_weight,
+                            fp16=False,
+                        )
+                        local_load_balancing_loss = batched_load_balancing_loss(moe_args)
+                        clear_load_balancing_loss()
+
+                        local_loss = local_load_balancing_loss + local_loss
+
                     backward(local_loss, scaler)
                 if ii == 0:
-                    total_loss = local_loss
+                    if args.moe_freq > 0:
+                        total_loss = local_loss - local_load_balancing_loss
+                        total_load_balancing_loss = local_load_balancing_loss
+                    else:
+                        total_loss = local_loss
                 else:
-                    total_loss += local_loss
+                    if args.moe_freq > 0:
+                        total_loss += local_loss - local_load_balancing_loss
+                        total_load_balancing_loss += local_load_balancing_loss
+                    else:
+                        total_loss += local_loss
 
         if scaler is not None:
             if args.grad_clip_norm is not None:
@@ -268,13 +302,18 @@ def train_one_epoch(model, data, loss, epoch, step, optimizer, scaler, scheduler
 
             # gathered_loss = [torch.zeros_like(total_loss) for _ in range(args.world_size)]
             # torch.distributed.all_gather(gathered_loss, total_loss)
+
             # losses_m.update(sum(gathered_loss).item() / args.world_size, batch_size * args.world_size)
             losses_m.update(global_loss_tensor.item(), batch_size)
+            if args.moe_freq > 0:
+                load_balancing_losses_m.update(total_load_balancing_loss.item(), batch_size)
             samples_per_second = inputs.numel() * args.world_size / batch_time_m.val
             samples_per_second_per_gpu = inputs.numel() / batch_time_m.val
+            loss_str = f"Loss: {losses_m.avg:.3f}"
+            loss_str += f" LB-Loss: {load_balancing_losses_m.avg:.3f}" if args.moe_freq > 0 else ""
             logging.info(
                 f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                f"Loss: {losses_m.avg:.3f} "
+                f"{loss_str} "
                 f"Data (t): {data_time_m.avg:.3f} "
                 f"Batch (t): {batch_time_m.avg:.3f}, {samples_per_second:#g}/s, {samples_per_second_per_gpu:#g}/s/gpu "
                 f"LR: {optimizer.param_groups[0]['lr']:5f} "
@@ -283,6 +322,7 @@ def train_one_epoch(model, data, loss, epoch, step, optimizer, scaler, scheduler
             # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
             log_data = {
                 "loss": losses_m.val,
+                "load_balancing_loss": load_balancing_losses_m.val,
                 "data_time": data_time_m.val,
                 "batch_time": batch_time_m.val,
                 "samples_per_second": samples_per_second,

scripts/generate.py

@@ -64,7 +64,6 @@ def main():
 
     add_model_args(parser)
     args = parser.parse_args()
-
     print("Loading model into the right classes...")
     open_lm = OpenLMforCausalLM(OpenLMConfig(create_params(args)))
 

tests/shared.py

@@ -73,6 +73,13 @@ def __init__(self, model, **kwargs):
         self.target_mask_left = None
         self.target_mask_individual = None
         self.ignore_parse_errors = False
+        self.moe_num_experts = None
+        self.moe_freq = 0
+        self.moe_weight_parallelism = False
+        self.moe_expert_model_parallelism = False
+        self.moe_capacity_factor = 1.25
+        self.moe_loss_weight = 0.1
+        self.moe_top_k = 2
 
         for k, v in kwargs.items():
             if hasattr(self, k):

tests/test_generate_kv_cache_time.py

@@ -33,6 +33,13 @@ def test_generate_kv_cache(wiki_page, context_len, max_gen_len):
             "qk_norm": False,
             "positional_embedding_type": "rotary",
             "ffn_type": "swiglu",
+            "moe_num_experts": None,
+            "moe_freq": 0,
+            "moe_weight_parallelism": False,
+            "moe_expert_model_parallelism": False,
+            "moe_capacity_factor": 1.25,
+            "moe_loss_weight": 0.1,
+            "moe_top_k": 2,
         }
     )