fix: updates

modeling_hf_nomic_bert.py

@@ -41,14 +41,22 @@ from transformers.modeling_outputs import (
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
 from transformers.utils.hub import cached_file, get_checkpoint_shard_files
 
-from .configuration_hf_nomic_bert import NomicBertConfig
+from configuration_hf_nomic_bert import NomicBertConfig
+logger = logging.getLogger(__name__)
 
 try:
     from torch.nn.functional import scaled_dot_product_attention
 except ImportError:
+    logger.warning("scaled_dot_product_attention not available, using torch.matmul instead")
     scaled_dot_product_attention = None
 
-logger = logging.getLogger(__name__)
+try:
+    from megablocks.layers import dmoe
+    from megablocks.layers.arguments import Arguments
+except ImportError:
+    logger.warning("!!!!!!!!!!!!megablocks not available, using torch.matmul instead")
+    dmoe = None
+
 
 
 # adapted from flash attention, added safe serialization option for hf models
@@ -1083,6 +1091,171 @@ class NomciBertGatedMLP(nn.Module):
         y = self.fc2(y)
         return y if not self.return_residual else (y, x)
 
+class NomicRouter(nn.Module):
+    def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int,
+                 moe_jitter_eps: Optional[float] = None,
+                 moe_normalize_expert_weights: Optional[float] = None,
+                 uniform_expert_assignment: bool = False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.moe_num_experts = moe_num_experts
+        self.moe_top_k = moe_top_k
+        self.moe_jitter_eps = moe_jitter_eps
+        self.moe_normalize_expert_weights = moe_normalize_expert_weights
+        self.uniform_expert_assignment = uniform_expert_assignment
+
+        self.layer = nn.Linear(self.hidden_size,
+                               self.moe_num_experts,
+                               bias=False)
+
+    def jitter(self, x: torch.Tensor) -> torch.Tensor:
+        if self.moe_jitter_eps is None:
+            raise RuntimeError('The router does not have moe_jitter_eps set.')
+        low = 1.0 - self.moe_jitter_eps
+        high = 1.0 + self.moe_jitter_eps
+        noise = torch.rand(x.size(), dtype=x.dtype, device=x.device)
+        return low + noise * (high - low)
+
+    def forward(
+            self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
+        if self.training and self.moe_jitter_eps is not None:
+            x = x * self.jitter(x)
+
+        weights = self.layer(x.view(-1,
+                                    x.shape[-1])).softmax(dim=-1,
+                                                          dtype=torch.float32)
+        top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)
+
+        if self.moe_normalize_expert_weights:
+            top_weights = top_weights / torch.norm(
+                top_weights,
+                p=self.moe_normalize_expert_weights,
+                dim=-1,
+                keepdim=True)
+
+        if self.uniform_expert_assignment:
+            with torch.no_grad():
+                uniform_tensor = torch.arange(
+                    0,
+                    top_experts.numel(),
+                    device=top_experts.device,
+                    dtype=top_experts.dtype) % self.moe_num_experts
+                top_experts = uniform_tensor.reshape(top_experts.shape)
+                # Note, weights and top_weights are not changed
+
+        weights = weights.to(x.dtype)
+        top_weights = top_weights.to(x.dtype)
+        return weights, top_weights, top_experts  # type: ignore
+
+        
+class NomicExpertMLP(nn.Module):
+
+    def __init__(self, hidden_size: int, ffn_hidden_size: int,
+                 moe_num_experts: int, ffn_act_fn: dict):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.ffn_hidden_size = ffn_hidden_size
+        self.moe_num_experts = moe_num_experts
+
+        self.w1 = nn.Parameter(
+            torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
+        self.w2 = nn.Parameter(
+            torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
+        self.activation_fn = ffn_act_fn
+
+    def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:
+        expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size,
+                                 self.hidden_size)[expert_idx]
+        expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size,
+                                 self.hidden_size)[expert_idx]
+
+        x1 = x.matmul(expert_w1.t())
+        act_out = self.activation_fn(x1)
+        x2 = act_out.matmul(expert_w2)
+        return x2
+
+class NomicExperts(nn.Module):
+    def __init__(self, config, hidden_size: int, ffn_hidden_size: int,
+                 moe_num_experts: int):
+        super().__init__()
+        self.moe_num_experts = moe_num_experts
+        activation = (
+            F.sigmoid
+            if config.activation_function == "glu"
+            else (F.silu if config.activation_function == "swiglu" else F.gelu)
+        )
+        self.mlp = NomicExpertMLP(
+           hidden_size=config.n_embd,
+           ffn_hidden_size=config.n_inner,
+           moe_num_experts=moe_num_experts,
+           ffn_act_fn=activation,
+        )
+        self.bias = nn.Parameter(torch.zeros(config.n_embd))
+
+    def forward(self, x: torch.Tensor, weights: torch.Tensor,
+                top_weights: torch.Tensor,
+                top_experts: torch.LongTensor) -> torch.Tensor:
+        bsz, q_len, hidden_size = x.shape
+        x = x.view(-1, hidden_size)
+        out = torch.zeros_like(x)
+
+        expert_mask = nn.functional.one_hot(
+            top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
+        for expert_idx in range(0, self.moe_num_experts):
+            topk_idx, token_idx = torch.where(expert_mask[expert_idx])
+            if token_idx.shape[0] == 0:
+                continue
+
+            token_list = token_idx.tolist()
+            topk_list = topk_idx.tolist()
+
+            expert_tokens = x[None, token_list].reshape(-1, hidden_size)
+            expert_out = self.mlp(
+                expert_tokens, expert_idx) * top_weights[token_list, topk_list,
+                                                         None]
+
+            out.index_add_(0, token_idx, expert_out)
+
+        out = out.reshape(bsz, q_len, hidden_size)
+        return out + self.bias
+
+
+class NomicMoELayer(nn.Module):
+
+    def __init__(self, config: NomicBertConfig):
+        super().__init__()
+
+        self.router = NomicRouter(
+            config.n_embd,
+            moe_num_experts=config.num_experts,
+            moe_top_k=config.moe_top_k,
+        )
+
+        self.experts = NomicExperts(
+            config,
+            hidden_size=config.n_embd,
+            ffn_hidden_size=config.n_inner,
+            moe_num_experts=config.num_experts,
+        )
+
+
+    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        batch_size, seq_len, hidden_dim = x.shape
+        if attention_mask is not None:
+            valid_indices = attention_mask.bool().view(-1)
+            x_valid = x.view(-1, hidden_dim)[valid_indices]
+
+        weights, top_weights, top_experts = self.router(x)
+        out = self.experts(x, weights, top_weights, top_experts)
+
+        if attention_mask is not None:
+            full_out = torch.zeros(batch_size * seq_len, hidden_dim, dtype=out.dtype, device=out.device)
+            full_out[valid_indices] = out
+            out = full_out.view(batch_size, seq_len, hidden_dim)
+
+        return out
+
 
 def rotate_half(x, interleaved=False):
     if not interleaved:
@@ -1431,6 +1604,7 @@ class NomicBertBlock(NomicBertPreTrainedModel):
     def __init__(
         self,
         config,
+        moe=False,
     ):
         super().__init__(config=config)
         self.prenorm = config.prenorm
@@ -1442,25 +1616,46 @@ class NomicBertBlock(NomicBertPreTrainedModel):
             if config.activation_function == "glu"
             else (F.silu if config.activation_function == "swiglu" else F.gelu)
         )
-        if config.activation_function in ["glu", "swiglu", "geglu"]:
-            self.mlp = NomciBertGatedMLP(
-                config.n_embd,
-                hidden_features=config.n_inner,
-                bias1=config.mlp_fc1_bias,
-                bias2=config.mlp_fc2_bias,
-                activation=activation,
-                fused_bias_fc=config.fused_bias_fc,
-                norm_layer=getattr(config, "norm_mlp", False),
-            )
+        if moe:
+            if dmoe is not None:
+                megablocks_args = Arguments(
+                    moe_num_experts=config.num_experts,
+                    moe_top_k=config.moe_top_k,
+                    hidden_size=config.n_embd,
+                    ffn_hidden_size=config.n_inner,
+                    num_layers=config.n_layer,
+                    moe_normalize_expert_weights=config.moe_normalize_expert_weights,
+                    activation_fn=activation,
+                    mlp_type="glu" if config.activation_function == "swiglu" else "mlp",
+                    fp16=True,
+                    bf16=False,
+                    return_bias=False,
+                )
+                self.mlp = dmoe.dMoE(megablocks_args)
+            else:
+                self.mlp = NomicMoELayer(
+                    config 
+                )
         else:
-            self.mlp = NomicBertMLP(
-                config.n_embd,
-                hidden_features=config.n_inner,
-                bias1=config.mlp_fc1_bias,
-                bias2=config.mlp_fc2_bias,
-                activation=activation,
-                fused_bias_fc=config.fused_bias_fc,
-            )
+            if config.activation_function in ["glu", "swiglu", "geglu"]:
+                self.mlp = NomciBertGatedMLP(
+                    config.n_embd,
+                    hidden_features=config.n_inner,
+                    bias1=config.mlp_fc1_bias,
+                    bias2=config.mlp_fc2_bias,
+                    activation=activation,
+                    fused_bias_fc=config.fused_bias_fc,
+                    norm_layer=getattr(config, "norm_mlp", False),
+                )
+            else:
+                self.mlp = NomicBertMLP(
+                    config.n_embd,
+                    hidden_features=config.n_inner,
+                    bias1=config.mlp_fc1_bias,
+                    bias2=config.mlp_fc2_bias,
+                    activation=activation,
+                    fused_bias_fc=config.fused_bias_fc,
+                )
 
         self.dropout1 = nn.Dropout(config.resid_pdrop)
         self.norm1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
@@ -1530,7 +1725,11 @@ class NomicBertBlock(NomicBertPreTrainedModel):
 class NomicBertEncoder(nn.Module):
     def __init__(self, config: GPT2Config):
         super().__init__()
-        self.layers = nn.ModuleList([NomicBertBlock(config) for _ in range(config.n_layer)])
+        if getattr(config, "moe_every_n_layers", 0) > 0:
+            every_n = config.moe_every_n_layers
+            self.layers = nn.ModuleList([NomicBertBlock(config, moe=i%every_n == 1) for i in range(config.n_layer)])
+        else:
+            self.layers = nn.ModuleList([NomicBertBlock(config, moe=False) for _ in range(config.n_layer)])
         self.gradient_checkpointing = False
         self.config = config