MoE

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `OLMoE`: Configurations & modeling for training Mixture-of-Experts models.
 - Added ability to try loading latest checkpoint from save folder using `--try_load_latest_save`.
 - Added support for flash attention and gradient checkpointing to `hf_olmo`.
 

README.md

@@ -221,3 +221,15 @@ See [Debugging](https://github.com/allenai/OLMo/blob/main/docs/NOTES.md#debuggin
   journal={arXiv preprint},
 }
 ```
+
+```bibtex
+@misc{muennighoff2024olmoeopenmixtureofexpertslanguage,
+      title={OLMoE: Open Mixture-of-Experts Language Models}, 
+      author={Niklas Muennighoff and Luca Soldaini and Dirk Groeneveld and Kyle Lo and Jacob Morrison and Sewon Min and Weijia Shi and Pete Walsh and Oyvind Tafjord and Nathan Lambert and Yuling Gu and Shane Arora and Akshita Bhagia and Dustin Schwenk and David Wadden and Alexander Wettig and Binyuan Hui and Tim Dettmers and Douwe Kiela and Ali Farhadi and Noah A. Smith and Pang Wei Koh and Amanpreet Singh and Hannaneh Hajishirzi},
+      year={2024},
+      eprint={2409.02060},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2409.02060}, 
+}
+```

olmo/config.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from omegaconf.errors import OmegaConfBaseException
@@ -198,6 +199,11 @@ class BlockType(StrEnum):
     implementations of operations like attention to imitate the behavior of Llama.
     """
 
+    moe = "moe"
+    """
+    A block for OLMoE-style Mixture-of-Experts models.
+    """
+
 
 class InitFnType(StrEnum):
     mitchell = "mitchell"
@@ -457,6 +463,56 @@ class ModelConfig(BaseConfig):
     See :data:`TrainConfig.precision` instead.
     """
 
+    moe_num_experts: Optional[int] = 8
+    """
+    The number of experts to use in the MoE block.
+    """
+
+    moe_top_k: Optional[int] = 2
+    """
+    The number of experts to select for each token.
+    """
+
+    moe_mlp_impl: Optional[str] = "sparse"
+    """
+    Choose "grouped" for grouped GEMM installable via `pip install git+https://[email protected]/tgale96/grouped_gemm.git@66c7195e35e8c4f22fa6a014037ef511bfa397cb`.
+    """
+
+    moe_log_expert_assignment: Optional[bool] = True
+    """
+    Whether to log the expert assignment.
+    """
+
+    moe_shared_expert: Optional[bool] = False
+    """
+    Whether to have an always-used expert like in [DeepSeekMoE](https://arxiv.org/abs/2401.06066).
+    """
+
+    moe_lbl_in_fp32: Optional[bool] = False
+    """
+    Whether to perform load balancing in FP32.
+    """
+
+    moe_loss_weight: Optional[float] = 0.1
+    """
+    The weight to use for the MoE load balancing loss.
+    """
+
+    moe_zloss_weight: Optional[float] = None
+    """
+    Weight for MoE router z-loss where None means no router z-loss. 0.001 is a common value.
+    """
+
+    moe_dropless: Optional[bool] = True
+    """
+    Whether to use [dMoE](https://arxiv.org/abs/2211.15841).
+    """
+
+    moe_capacity_factor: Optional[float] = 1.25
+    """
+    The capacity factor to use in the MoE block. Only applies if not using dMoE.
+    """
+
     scale_emb_init: bool = False
     """
     If ``True``, embeddings are scaled up by ``sqrt(d_model)`` during initialization.
@@ -1283,3 +1339,41 @@ def update_legacy_settings(cls, config: D) -> D:
                 new_config.optimizer = OptimizerConfig.update_legacy_settings(new_config.optimizer)
 
         return new_config
+
+
+def config_to_moe_args(config: ModelConfig) -> Dict[str, Any]:
+    from megablocks.layers.arguments import Arguments as MoEArgs
+
+    from .model import Activation
+
+    hidden_size = (
+        config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model
+    )
+    act = Activation.build(config)
+    kwargs = {
+        "activation_fn": F.silu if "swiglu" in config.activation_type.lower() else Activation.build(config),
+        "mlp_type": "glu" if "glu" in config.activation_type.lower() else "mlp",
+        "mlp_impl": config.moe_mlp_impl,
+        "hidden_size": config.d_model,
+        "ffn_hidden_size": int(act.output_multiplier * hidden_size),
+        "moe_num_experts": config.moe_num_experts,
+        "num_layers": config.n_layers,
+        # Handled by FSDP (https://github.com/databricks/megablocks/issues/57#issuecomment-1854594483)
+        "moe_weight_parallelism": False,
+        "moe_expert_model_parallelism": False,
+        "moe_top_k": config.moe_top_k,
+        "moe_capacity_factor": config.moe_capacity_factor,
+        "moe_loss_weight": config.moe_loss_weight,
+        "device": config.init_device,
+        # Handled by FSDP
+        "bf16": False,
+        "fp16": False,
+        "bias": config.include_bias,
+        "return_bias": False,
+        "shared_expert": config.moe_shared_expert,
+        "moe_lbl_in_fp32": config.moe_lbl_in_fp32,
+    }
+    if config.moe_zloss_weight:
+        kwargs["moe_zloss_weight"] = config.moe_zloss_weight
+
+    return MoEArgs(**kwargs)

olmo/initialization.py

@@ -13,9 +13,15 @@ def init_normal(
     # weights
     if init_cutoff_factor is not None:
         cutoff_value = init_cutoff_factor * std
-        nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
+        if hasattr(module, "weight"):
+            nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
+        else:
+            nn.init.trunc_normal_(module, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
     else:
-        nn.init.normal_(module.weight, mean=0.0, std=std)
+        if hasattr(module, "weight"):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        else:
+            nn.init.normal_(module, mean=0.0, std=std)
 
     # biases
     if isinstance(module, nn.Linear) and module.bias is not None:

olmo/model.py

@@ -448,14 +448,15 @@ def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
             config.d_model, config.d_model, bias=config.include_bias, device=config.init_device
         )
 
-        # Feed-forward output projection.
-        self.ff_out = nn.Linear(
-            int(self.act.output_multiplier * self.hidden_size),
-            config.d_model,
-            bias=config.include_bias,
-            device=config.init_device,
-        )
-        self.ff_out._is_residual = True  # type: ignore
+        if self.config.block_type != BlockType.moe:
+            # Feed-forward output projection.
+            self.ff_out = nn.Linear(
+                int(self.act.output_multiplier * self.hidden_size),
+                config.d_model,
+                bias=config.include_bias,
+                device=config.init_device,
+            )
+            self.ff_out._is_residual = True  # type: ignore
 
         # Rotary embeddings.
         if self.config.rope:
@@ -664,10 +665,164 @@ def build(cls, layer_id: int, config: ModelConfig, cache: BufferCache) -> OLMoBl
             return OLMoSequentialBlock(layer_id, config, cache)
         elif config.block_type == BlockType.llama:
             return OLMoLlamaBlock(layer_id, config, cache)
+        elif config.block_type == BlockType.moe:
+            return OLMoEBlock(layer_id, config, cache)
         else:
             raise NotImplementedError(f"Unknown block type: '{config.block_type}'")
 
 
+class OLMoEBlock(OLMoBlock):
+    """
+    This is a transformer MoE block where the output is computed as ``MoE(LN(x + Attention(LN(x))))``
+    (plus another skip connection).
+    """
+
+    def __init__(self, layer_id: int, config: ModelConfig, cache: BufferCache):
+        try:
+            from megablocks.layers.dmoe import dMoE
+            from megablocks.layers.moe import MoE
+        except ImportError:
+            raise ImportError(
+                "To train MoEs, run `pip install git+https://github.com/Muennighoff/megablocks.git@olmoe`"
+            )
+        from .config import config_to_moe_args
+
+        super().__init__(layer_id, config, cache)
+
+        self.moe_args = config_to_moe_args(config)
+        self.ffn = dMoE(self.moe_args) if self.config.moe_dropless else MoE(self.moe_args)
+
+        self.attn_norm = LayerNorm.build(config)
+        self.ff_norm = LayerNorm.build(config)
+
+        # Attention input projection. Projects x -> (q, k, v)
+        head_dim = config.d_model // config.n_heads
+        self.fused_dims = (
+            config.d_model,
+            config.effective_n_kv_heads * head_dim,
+            config.effective_n_kv_heads * head_dim,
+        )
+        self.att_proj = nn.Linear(
+            config.d_model, sum(self.fused_dims), bias=config.include_bias, device=config.init_device
+        )
+
+    def reset_parameters(self):
+        if self.k_norm is not None:
+            self.k_norm.reset_parameters()
+        if self.q_norm is not None:
+            self.q_norm.reset_parameters()
+
+        if self.config.init_fn == InitFnType.normal:
+            attn_out_std = ff_out_std = in_std = self.config.init_std
+            cutoff_factor = self.config.init_cutoff_factor
+        elif self.config.init_fn == InitFnType.mitchell:
+            in_std = 1 / math.sqrt(self.config.d_model)
+            attn_out_std = 1 / (math.sqrt(2 * self.config.d_model * (self.layer_id + 1)))
+            ff_out_std = 1 / (math.sqrt(2 * self.ff_out.in_features * (self.layer_id + 1)))
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        elif self.config.init_fn == InitFnType.full_megatron:
+            in_std = self.config.init_std
+            attn_out_std = ff_out_std = self.config.init_std / math.sqrt(2.0 * self.config.n_layers)
+            cutoff_factor = self.config.init_cutoff_factor or 3.0
+        else:
+            raise NotImplementedError(self.config.init_fn)
+
+        init_normal(self.att_proj, std=in_std, init_cutoff_factor=cutoff_factor)
+        init_normal(self.attn_out, std=attn_out_std, init_cutoff_factor=cutoff_factor)
+        self.attn_norm.reset_parameters()
+        self.ff_norm.reset_parameters()
+        init_normal(self.ffn.experts.mlp.w1, std=in_std, init_cutoff_factor=cutoff_factor)
+        init_normal(self.ffn.experts.mlp.w2, std=ff_out_std, init_cutoff_factor=cutoff_factor)
+        if hasattr(self.ffn.experts.mlp, "v1"):
+            init_normal(self.ffn.experts.mlp.v1, std=in_std, init_cutoff_factor=cutoff_factor)
+        if self.ffn.experts.bias is not None:
+            torch.nn.init.zeros_(self.ffn.experts.bias)
+        init_normal(self.ffn.router.layer, std=in_std, init_cutoff_factor=cutoff_factor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_bias: Optional[torch.Tensor] = None,
+        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        max_doc_len: Optional[int] = None,
+        cu_doc_lens: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        # Get query, key, value projections.
+        # shape:
+        #  - for regular attn q, k, v: (batch_size, seq_len, d_model)
+        #  - for multi-query attn q: (batch_size, seq_len, d_model)
+        #                      k, v: (batch_size, seq_len, d_model // n_heads)
+        #  - for group query attn q: (batch_size, seq_len, d_model)
+        #                      k, v: (batch_size, seq_len, d_model // n_kv_heads)
+        if not self.config.norm_after:
+            if self._activation_checkpoint_fn is not None:
+                qkv = self.att_proj(self._activation_checkpoint_fn(self.attn_norm, x))
+            else:
+                qkv = self.att_proj(self.attn_norm(x))
+        else:
+            qkv = self.att_proj(x)
+
+        if self.config.clip_qkv is not None:
+            qkv.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
+
+        q, k, v = qkv.split(self.fused_dims, dim=-1)
+
+        # Get attention scores.
+        if self._activation_checkpoint_fn is not None:
+            att, cache = self._activation_checkpoint_fn(  # type: ignore
+                self.attention,
+                q,
+                k,
+                v,
+                attention_bias,
+                layer_past=layer_past,
+                use_cache=use_cache,
+                max_doc_len=max_doc_len,
+                cu_doc_lens=cu_doc_lens,
+            )
+        else:
+            att, cache = self.attention(
+                q,
+                k,
+                v,
+                attention_bias,
+                layer_past=layer_past,
+                use_cache=use_cache,
+                max_doc_len=max_doc_len,
+                cu_doc_lens=cu_doc_lens,
+            )
+
+        if self.config.norm_after:
+            if self._activation_checkpoint_fn is not None:
+                att = self._activation_checkpoint_fn(self.attn_norm, att)
+            else:
+                att = self.attn_norm(att)
+
+        # Add attention scores.
+        # shape: (B, T, C)
+        x = x + self.dropout(att)
+
+        # Add feed-forward projection.
+        # shape: (batch_size, seq_len, d_model)
+        og_x = x
+
+        if self.config.norm_after:
+            x = self.ffn(x)
+            if self._activation_checkpoint_fn is not None:
+                x = self._activation_checkpoint_fn(self.ff_norm, x)  # type: ignore
+            else:
+                x = self.ff_norm(x)
+            return og_x + self.dropout(x), cache
+        else:
+            if self._activation_checkpoint_fn is not None:
+                x = self._activation_checkpoint_fn(self.ff_norm, x)  # type: ignore
+            else:
+                x = self.ff_norm(x)
+            # Activation checkpointing for the MoE FFN is not supported
+            return og_x + self.dropout(self.ffn(x)), cache
+
+
 class OLMoSequentialBlock(OLMoBlock):
     """
     This is a typical transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))``
@@ -1552,7 +1707,7 @@ def fsdp_wrap_fn(module, recurse: bool = True, nonwrapped_numel: int = 0):
         else:
             raise NotImplementedError(wrap_strategy)
 
-    def num_params(self, include_embedding: bool = True) -> int:
+    def num_params(self, include_embedding: bool = True, include_inactive_params: bool = True) -> int:
         """
         Get the total number of parameters.
         """
@@ -1562,6 +1717,15 @@ def num_params(self, include_embedding: bool = True) -> int:
                 lambda np: ".wte." not in np[0] and ".wpe." not in np[0],
                 params,
             )
+        if not include_inactive_params:
+            # Need to reduce blocks to the number of experts that are selected
+            # If not dropless 'transformer.blocks.0.ffn.experts.mlp.w1' has shape (total_experts, in_dim, out_dim)
+            # change to 'transformer.blocks.0.ffn.experts.mlp.w1' with shape (selected_experts, in_dim, out_dim)
+            # If dropless, the total_experts & out_dim are combined into one dimension
+            idx = self.config.moe_top_k
+            if self.config.moe_dropless:
+                idx *= self.transformer.blocks[1].moe_args.ffn_hidden_size
+            params = [(np[0], np[1][:idx]) if "experts.mlp" in np[0] else np for np in params]  # type: ignore
         return sum(p.numel() for _, p in params)
 
     @property

olmo/optim.py

@@ -15,6 +15,13 @@
 from .config import OptimizerType, SchedulerConfig, SchedulerType, TrainConfig
 from .torch_util import get_default_device, is_distributed
 
+try:
+    from megablocks.layers.mlp import MLP, SparseMLP
+
+    megablocks_available = True
+except ImportError:
+    megablocks_available = False
+
 __all__ = [
     "Optimizer",
     "LionW",
@@ -858,6 +865,8 @@ def get_param_groups(cfg: TrainConfig, model: nn.Module) -> List[Dict[str, Any]]
                     decay.add(fpn)
                 else:
                     no_decay.add(fpn)
+            elif megablocks_available and pn.endswith(("w1", "w2", "v1")) and isinstance(m, (MLP, SparseMLP)):
+                decay.add(fpn)
 
     # Validate that we've considered every parameter
     inter_params = decay & no_decay