Upload SmalLmForCausalLM

README.md

@@ -2,6 +2,7 @@
 library_name: transformers
 tags:
 - generated_from_trainer
+- smallm
 model-index:
 - name: smallm_70_rope
   results: []

config.json

@@ -4,6 +4,10 @@
   ],
   "attention_bias": false,
   "attention_dropout": 0.1,
+  "auto_map": {
+    "AutoConfig": "config.SmalLmConfig",
+    "AutoModelForCausalLM": "model.SmalLmForCausalLM"
+  },
   "balancing_coef": 0.0001,
   "bos_token_id": 1,
   "embedding_dropout": 0.0,
@@ -39,7 +43,7 @@
   "sliding_window_attention": true,
   "sliding_window_context": 1024,
   "sliding_window_period": 4,
-  "static_residual": false,
+  "static_residual": true,
   "token_experts": 3,
   "torch_dtype": "float32",
   "transformers_version": "4.50.3",

config.py

@@ -56,14 +56,12 @@ class SmalLmConfig(PretrainedConfig):
         eos_token_id: int = 0,
         pad_token_id: int = 0,
         static_residual: bool = False,
-        moe_type: str = "default",
         **kwargs,
     ):
         if positional_bias_type not in ["alibi", "rope"]:
             raise ValueError(
                 f"positional_bias_type must be 'alibi' or 'rope', got {positional_bias_type}"
             )
-        self.moe_type = moe_type
         self.static_residual = not static_residual
         self.no_moe_layers = no_moe_layers
         self.moe_bias = moe_bias

model.py

@@ -10,7 +10,7 @@ from transformers.modeling_outputs import (
 from .config import SmalLmConfig
 from typing import Optional
 import logging
-from einops import rearrange, repeat
+from einops import rearrange
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from einops._torch_specific import allow_ops_in_compiled_graph
 
@@ -127,68 +127,6 @@ class MoE(nn.Module):
         return (out + shared_out).view(shape)
 
 
-class ComboMoe(nn.Module):
-    def __init__(self, config: SmalLmConfig, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.config = config
-        self.shared_experts = SwiGLU(
-            config.hidden_size,
-            config.shared_experts * config.expert_size,
-            config.moe_bias,
-        )
-        self.input_router = Router(config)
-        self.middle_router = Router(config)
-        self.out_router = Router(config)
-        self.routed_experts = nn.ModuleList(
-            [
-                nn.Linear(config.hidden_size, config.expert_size, bias=config.moe_bias)
-                for _ in range(config.routed_experts)
-            ]
-        )
-        self.middle_routed_experts = nn.ModuleList(
-            [
-                nn.Linear(config.expert_size, config.hidden_size, bias=config.moe_bias)
-                for _ in range(config.routed_experts)
-            ]
-        )
-        self.out_routed_experts = nn.ModuleList(
-            [
-                nn.Linear(config.expert_size, config.hidden_size, bias=config.moe_bias)
-                for _ in range(config.routed_experts)
-            ]
-        )
-        self.offset = config.routed_experts
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        shape = x.size()
-        x = x.view(-1, self.config.hidden_size)
-        iexpert_idx, iexpert_weights, icounts = self.input_router(x)
-        iout = torch.zeros((*x.shape[:-1], self.config.expert_size), device=x.device)
-        for i, expert in enumerate(self.routed_experts[: self.offset]):
-            if icounts[i] == 0:
-                continue
-            idx, pos = torch.where(iexpert_idx == i)
-            iout[idx] += expert(x[idx]) * iexpert_weights[idx, pos, None]
-
-        mexpert_idx, mexpert_weights, mcounts = self.middle_router(x)
-        for i, expert in enumerate(self.middle_routed_experts):
-            if mcounts[i] == 0:
-                continue
-            idx, pos = torch.where(mexpert_idx == i)
-            iout[idx] *= F.silu(expert(x[idx]) * mexpert_weights[idx, pos, None])
-
-        out = torch.zeros_like(x)
-        oexpert_idx, oexpert_weights, ocounts = self.out_router(iout)
-        for i, expert in enumerate(self.out_routed_experts):
-            if ocounts[i] == 0:
-                continue
-            idx, pos = torch.where(oexpert_idx == i)
-            out[idx] += expert(iout[idx]) * oexpert_weights[idx, pos, None]
-
-        shared_out = self.shared_experts(x)
-        return (out + shared_out).view(shape)
-
-
 def build_alibi_bias(config: SmalLmConfig) -> torch.Tensor:
     """Build ALiBi for specified number of heads:
 
@@ -471,9 +409,8 @@ class Block(nn.Module):
         self.dropout1 = nn.Dropout(config.layer_dropout)
         self.dropout2 = nn.Dropout(config.layer_dropout)
         self.attention = CausalSelfAttention(config, layer_idx)
-        moe_class = MoE if config.moe_type == "default" else ComboMoe
         self.mlp = (
-            moe_class(config)
+            MoE(config)
             if (
                 config.use_moe
                 and layer_idx % config.moe_period == 0