Repository: huggingface/text-embeddings-inference (Rust)

Issue:
Support for `nomic-ai/nomic-embed-text-v2-moe`
### Model description

TEI currently does not seem to support Nomic's latest embedding model: `nomic-ai/nomic-embed-text-v2-moe`. When I try to load the model using the container image `ghcr.io/huggingface/text-embeddings-inference:1.6.0`, I get this error:

```
tei  | 2025-02-21T12:03:12.972121Z  INFO text_embeddings_router: router/src/main.rs:175: Args { model_id: "nom**-**/*****-*****-****-**-moe", revision: None, tokenization_workers: None, dtype: None, pooling: None, max_concurrent_requests: 512, max_batch_tokens: 16384, max_batch_requests: None, max_client_batch_size: 32, auto_truncate: false, default_prompt_name: None, default_prompt: None, hf_api_token: None, hostname: "379cf6f93689", port: 80, uds_path: "/tmp/text-embeddings-inference-server", huggingface_hub_cache: Some("/data"), payload_limit: 2000000, api_key: None, json_output: false, otlp_endpoint: None, otlp_service_name: "text-embeddings-inference.server", cors_allow_origin: None }
tei  | 2025-02-21T12:03:12.972326Z  INFO hf_hub: /root/.cargo/registry/src/index.crates.io-6f17d22bba15001f/hf-hub-0.3.2/src/lib.rs:55: Token file not found "/root/.cache/huggingface/token"
tei  | 2025-02-21T12:03:13.061604Z  INFO download_artifacts: text_embeddings_core::download: core/src/download.rs:20: Starting download
tei  | 2025-02-21T12:03:13.061624Z  INFO download_artifacts:download_pool_config: text_embeddings_core::download: core/src/download.rs:53: Downloading `1_Pooling/config.json`
tei  | 2025-02-21T12:03:13.061695Z  INFO download_artifacts:download_new_st_config: text_embeddings_core::download: core/src/download.rs:77: Downloading `config_sentence_transformers.json`
tei  | 2025-02-21T12:03:13.061735Z  INFO download_artifacts: text_embeddings_core::download: core/src/download.rs:40: Downloading `config.json`
tei  | 2025-02-21T12:03:13.061771Z  INFO download_artifacts: text_embeddings_core::download: core/src/download.rs:43: Downloading `tokenizer.json`
tei  | 2025-02-21T12:03:13.061798Z  INFO download_artifacts: text_embeddings_core::download: core/src/download.rs:47: Model artifacts downloaded in 195.324µs
tei  | 2025-02-21T12:03:13.529301Z  INFO text_embeddings_router: router/src/lib.rs:188: Maximum number of tokens per request: 512
tei  | 2025-02-21T12:03:13.529350Z  INFO text_embeddings_core::tokenization: core/src/tokenization.rs:28: Starting 4 tokenization workers
tei  | 2025-02-21T12:03:14.590727Z  INFO text_embeddings_router: router/src/lib.rs:230: Starting model backend
tei  | 2025-02-21T12:03:14.590785Z  INFO text_embeddings_backend: backends/src/lib.rs:360: Downloading `model.safetensors`
tei  | 2025-02-21T12:03:14.590856Z  INFO text_embeddings_backend: backends/src/lib.rs:244: Model weights downloaded in 71.202µs
tei  | 2025-02-21T12:03:15.057394Z  INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:332: Starting FlashNomicBert model on Cuda(CudaDevice(DeviceId(1)))
tei  | 2025-02-21T12:03:15.057766Z ERROR text_embeddings_backend: backends/src/lib.rs:255: Could not start Candle backend: Could not start backend: config is not supported
tei  | Error: Could not create backend
tei  |
tei  | Caused by:
tei  |     Could not start backend: Could not start a suitable backend
tei exited with code 1
```


### Open source status

- [x] The model implementation is available
- [x] The model weights are available

### Provide useful links for the implementation

- https://huggingface.co/nomic-ai/nomic-embed-text-v2-moe



Fix (patch):
diff --git a/README.md b/README.md
index d30f5787..3d23a8cb 100644
--- a/README.md
+++ b/README.md
@@ -81,6 +81,7 @@ Below are some examples of the currently supported models:
 | 49        | 0.5B                | XLM-RoBERTa | [intfloat/multilingual-e5-large-instruct](https://hf.co/intfloat/multilingual-e5-large-instruct) |
 | N/A       | 0.4B                | Alibaba GTE | [Alibaba-NLP/gte-large-en-v1.5](https://hf.co/Alibaba-NLP/gte-large-en-v1.5)                     |
 | N/A       | 0.4B                | ModernBERT  | [answerdotai/ModernBERT-large](https://hf.co/answerdotai/ModernBERT-large) |
+| N/A       | 0.3B                | NomicBert   | [nomic-ai/nomic-embed-text-v2-moe](https://hf.co/nomic-ai/nomic-embed-text-v2-moe) |
 | N/A       | 0.1B                | NomicBert   | [nomic-ai/nomic-embed-text-v1](https://hf.co/nomic-ai/nomic-embed-text-v1)                       |
 | N/A       | 0.1B                | NomicBert   | [nomic-ai/nomic-embed-text-v1.5](https://hf.co/nomic-ai/nomic-embed-text-v1.5)                   |
 | N/A       | 0.1B                | JinaBERT    | [jinaai/jina-embeddings-v2-base-en](https://hf.co/jinaai/jina-embeddings-v2-base-en)             |
diff --git a/backends/candle/src/layers/linear.rs b/backends/candle/src/layers/linear.rs
index fc8af1dd..4c875a89 100644
--- a/backends/candle/src/layers/linear.rs
+++ b/backends/candle/src/layers/linear.rs
@@ -11,6 +11,16 @@ pub enum HiddenAct {
     Swiglu,
 }
 
+impl HiddenAct {
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        match self {
+            Self::Gelu => x.gelu(),
+            Self::Relu => x.relu(),
+            Self::Swiglu => candle_nn::ops::swiglu(x),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Linear {
     weight: Tensor,
diff --git a/backends/candle/src/models/flash_nomic.rs b/backends/candle/src/models/flash_nomic.rs
index 5683a86c..32cd31b6 100644
--- a/backends/candle/src/models/flash_nomic.rs
+++ b/backends/candle/src/models/flash_nomic.rs
@@ -1,6 +1,6 @@
 use crate::flash_attn::flash_attn_varlen;
 use crate::layers::{get_cos_sin, get_inv_freqs, LayerNorm, Linear};
-use crate::models::nomic::{NomicBertEmbeddings, NomicBertGatedMLP};
+use crate::models::nomic::{NomicBertEmbeddings, NomicMLP};
 use crate::models::{Model, NomicConfig};
 use candle::{DType, Device, IndexOp, Result, Tensor, D};
 use candle_nn::VarBuilder;
@@ -25,16 +25,25 @@ impl NomicAttention {
         let attention_head_size = config.n_embd / config.n_head;
         let hidden_size = config.n_embd;
 
-        let qkv_weight = vb.pp("Wqkv").get(
-            (3 * num_attention_heads * attention_head_size, hidden_size),
-            "weight",
-        )?;
-        let qkv_linear = Linear::new(qkv_weight, None, None);
+        let qkv_dim = 3 * num_attention_heads * attention_head_size;
+
+        let qkv_weight = vb.pp("Wqkv").get((qkv_dim, hidden_size), "weight")?;
+        let qkv_bias = if config.qkv_proj_bias {
+            Some(vb.pp("Wqkv").get((qkv_dim,), "bias")?)
+        } else {
+            None
+        };
+        let qkv_linear = Linear::new(qkv_weight, qkv_bias, None);
 
         let out_proj_weight = vb
             .pp("out_proj")
             .get((hidden_size, hidden_size), "weight")?;
-        let out_proj = Linear::new(out_proj_weight, None, None);
+        let out_proj_bias = if config.qkv_proj_bias {
+            Some(vb.pp("out_proj").get((hidden_size,), "bias")?)
+        } else {
+            None
+        };
+        let out_proj = Linear::new(out_proj_weight, out_proj_bias, None);
 
         let softmax_scale = (1. / (attention_head_size as f64).sqrt()) as f32;
 
@@ -93,7 +102,7 @@ impl NomicAttention {
 
 struct NomicBertBlock {
     attention: NomicAttention,
-    mlp: NomicBertGatedMLP,
+    mlp: NomicMLP,
     post_attention_layer_norm: LayerNorm,
     output_layer_norm: LayerNorm,
 
@@ -101,9 +110,10 @@ struct NomicBertBlock {
 }
 
 impl NomicBertBlock {
-    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+    pub fn load(vb: VarBuilder, index: usize, config: &NomicConfig) -> Result<Self> {
         let attention = NomicAttention::load(vb.pp("attn"), config)?;
-        let mlp = NomicBertGatedMLP::load(vb.pp("mlp"), config)?;
+
+        let mlp = NomicMLP::load(vb.pp("mlp"), index, config)?;
 
         let post_attention_layer_norm =
             LayerNorm::load(vb.pp("norm1"), config.n_embd, config.layer_norm_epsilon)?;
@@ -132,6 +142,7 @@ impl NomicBertBlock {
         let attn_output = self
             .attention
             .forward(&hidden_states, cu_seqlens, cos, sin, max_s)?;
+
         let hidden_states = self
             .post_attention_layer_norm
             .forward(&hidden_states, Some(&attn_output))?;
@@ -145,13 +156,14 @@ impl NomicBertBlock {
 
 struct NomicBertEncoder {
     layers: Vec<NomicBertBlock>,
+
     span: tracing::Span,
 }
 
 impl NomicBertEncoder {
     pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
         let layers = (0..config.n_layer)
-            .map(|index| NomicBertBlock::load(vb.pp(format!("layers.{index}")), config))
+            .map(|index| NomicBertBlock::load(vb.pp(format!("layers.{index}")), index, config))
             .collect::<Result<Vec<_>>>()?;
 
         let span = tracing::span!(tracing::Level::TRACE, "encoder");
@@ -170,7 +182,6 @@ impl NomicBertEncoder {
 
         let mut hidden_states = hidden_states.clone();
 
-        // Use a loop rather than a fold as it's easier to modify when adding debug/...
         for layer in self.layers.iter() {
             hidden_states = layer.forward(&hidden_states, cu_seqlens, cos, sin, max_s)?
         }
@@ -419,6 +430,7 @@ impl Model for FlashNomicBertModel {
     fn is_padded(&self) -> bool {
         false
     }
+
     fn embed(&self, batch: Batch) -> Result<(Option<Tensor>, Option<Tensor>)> {
         self.forward(batch)
     }
diff --git a/backends/candle/src/models/nomic.rs b/backends/candle/src/models/nomic.rs
index c860cfbc..8748db38 100644
--- a/backends/candle/src/models/nomic.rs
+++ b/backends/candle/src/models/nomic.rs
@@ -4,6 +4,7 @@ use crate::layers::{
 use crate::models::Model;
 use candle::{DType, Device, IndexOp, Module, Result, Tensor, D};
 use candle_nn::{Embedding, VarBuilder};
+use candle_transformers::models::deepseek2::{BincountOp, NonZeroOp, TopKLastDimOp, TopKOutput};
 use serde::Deserialize;
 use text_embeddings_backend_core::{Batch, ModelType, Pool};
 
@@ -20,6 +21,11 @@ pub struct NomicConfig {
     #[serde(default = "default_max_trained_positions")]
     pub max_trained_positions: usize,
 
+    pub moe_every_n_layers: Option<usize>,
+    pub moe_normalize_expert_weights: Option<bool>,
+    pub moe_top_k: Option<usize>,
+    pub num_experts: Option<usize>,
+
     pub n_embd: usize,
     pub n_head: usize,
     pub n_inner: usize,
@@ -42,12 +48,8 @@ impl NomicConfig {
     pub fn valid(&self) -> bool {
         !self.prenorm
             && self.rotary_emb_fraction == 1.0
-            && !self.qkv_proj_bias
             && !self.rotary_emb_interleaved
-            && !self.mlp_fc1_bias
-            && !self.mlp_fc2_bias
             && self.type_vocab_size > 0
-            && self.activation_function == HiddenAct::Swiglu
     }
 }
 
@@ -56,6 +58,7 @@ pub struct NomicBertEmbeddings {
     word_embeddings: Embedding,
     token_type_embeddings: Embedding,
     layer_norm: LayerNorm,
+
     span: tracing::Span,
 }
 
@@ -92,8 +95,8 @@ impl NomicBertEmbeddings {
 }
 
 pub struct NomicBertGatedMLP {
-    gate_up_proj: Linear,
-    down_proj: Linear,
+    fc1: Linear,
+    fc2: Linear,
 
     span: tracing::Span,
 }
@@ -102,29 +105,91 @@ impl NomicBertGatedMLP {
     pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
         let intermediate_size = config.n_inner;
 
-        let gate_proj_weight = vb
+        let fc11_weight = vb
+            .pp("fc11")
+            .get((intermediate_size, config.n_embd), "weight")?;
+        let fc12_weight = vb
             .pp("fc12")
             .get((intermediate_size, config.n_embd), "weight")?;
+        let fc1_weight = Tensor::cat(&[fc12_weight, fc11_weight], 0)?;
 
-        let up_proj_weight = vb
-            .pp("fc11")
-            .get((intermediate_size, config.n_embd), "weight")?;
+        let fc1_bias = if config.mlp_fc1_bias {
+            let fc11_bias = vb.pp("fc11").get((intermediate_size,), "bias")?;
+            let fc12_bias = vb.pp("fc12").get((intermediate_size,), "bias")?;
+            Some(Tensor::cat(&[fc12_bias, fc11_bias], 0)?)
+        } else {
+            None
+        };
 
-        let gate_up_proj_weight = Tensor::cat(&[&gate_proj_weight, &up_proj_weight], 0)?;
-        let gate_up_proj = Linear::new(
-            gate_up_proj_weight,
-            None,
+        let fc1 = Linear::new(
+            fc1_weight,
+            fc1_bias,
             Some(config.activation_function.clone()),
         );
 
-        let down_proj_weight = vb
+        let fc2_weight = vb
             .pp("fc2")
             .get((config.n_embd, intermediate_size), "weight")?;
-        let down_proj = Linear::new(down_proj_weight, None, None);
+        let fc2_bias = if config.mlp_fc2_bias {
+            Some(vb.pp("fc2").get((config.n_embd,), "bias")?)
+        } else {
+            None
+        };
+        let fc2 = Linear::new(fc2_weight, fc2_bias, None);
 
         Ok(Self {
-            gate_up_proj,
-            down_proj,
+            fc1,
+            fc2,
+            span: tracing::span!(tracing::Level::TRACE, "gated_mlp"),
+        })
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        let _enter = self.span.enter();
+
+        let gate_up_states = self.fc1.forward(hidden_states)?;
+        self.fc2.forward(&gate_up_states)
+    }
+}
+
+pub struct NomicBertMLP {
+    fc1: Linear,
+    fc2: Linear,
+
+    span: tracing::Span,
+}
+
+impl NomicBertMLP {
+    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+        let intermediate_size = config.n_inner;
+
+        let fc1_weight = vb
+            .pp("fc1")
+            .get((intermediate_size, config.n_embd), "weight")?;
+        let fc1_bias = if config.mlp_fc1_bias {
+            Some(vb.pp("fc1").get((intermediate_size,), "bias")?)
+        } else {
+            None
+        };
+        let fc1 = Linear::new(
+            fc1_weight,
+            fc1_bias,
+            Some(config.activation_function.clone()),
+        );
+
+        let fc2_weight = vb
+            .pp("fc2")
+            .get((config.n_embd, intermediate_size), "weight")?;
+        let fc2_bias = if config.mlp_fc2_bias {
+            Some(vb.pp("fc2").get((config.n_embd,), "bias")?)
+        } else {
+            None
+        };
+        let fc2 = Linear::new(fc2_weight, fc2_bias, None);
+
+        Ok(Self {
+            fc1,
+            fc2,
             span: tracing::span!(tracing::Level::TRACE, "mlp"),
         })
     }
@@ -132,8 +197,226 @@ impl NomicBertGatedMLP {
     pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
         let _enter = self.span.enter();
 
-        let gate_up_states = self.gate_up_proj.forward(hidden_states)?;
-        self.down_proj.forward(&gate_up_states)
+        let y = self.fc1.forward(hidden_states)?;
+        self.fc2.forward(&y)
+    }
+}
+
+pub struct NomicRouter {
+    layer: Linear,
+    moe_top_k: usize,
+
+    span: tracing::Span,
+}
+
+impl NomicRouter {
+    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+        let num_experts = config.num_experts.unwrap();
+        let moe_top_k = config.moe_top_k.unwrap();
+
+        let layer_weight = vb.pp("layer").get((num_experts, config.n_embd), "weight")?;
+        let layer = Linear::new(layer_weight, None, None);
+
+        Ok(Self {
+            layer,
+            moe_top_k,
+            span: tracing::span!(tracing::Level::TRACE, "router"),
+        })
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor) -> Result<(Tensor, Tensor)> {
+        let _enter = self.span.enter();
+
+        let weights = hidden_states.reshape(((), hidden_states.dim(D::Minus1)?))?;
+        let weights = self.layer.forward(&weights)?.to_dtype(DType::F32)?;
+        let weights = candle_nn::ops::softmax_last_dim(&weights)?;
+
+        let TopKOutput { values, indices } = weights.topk(self.moe_top_k)?;
+
+        let values = values.to_dtype(hidden_states.dtype())?;
+
+        Ok((values, indices))
+    }
+}
+
+pub struct NomicExpertMLP {
+    w1: Tensor,
+    w2: Tensor,
+    activation: HiddenAct,
+
+    span: tracing::Span,
+}
+
+impl NomicExpertMLP {
+    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+        let hidden_size = config.n_embd;
+        let ffn_hidden_size = config.n_inner;
+        let moe_num_experts = config.num_experts.unwrap();
+        let activation = config.activation_function.clone();
+
+        let w1 = vb
+            .get((moe_num_experts * ffn_hidden_size, hidden_size), "w1")?
+            .reshape((moe_num_experts, ffn_hidden_size, hidden_size))?;
+        let w2 = vb
+            .get((moe_num_experts * ffn_hidden_size, hidden_size), "w2")?
+            .reshape((moe_num_experts, ffn_hidden_size, hidden_size))?;
+
+        Ok(Self {
+            w1,
+            w2,
+            activation,
+            span: tracing::span!(tracing::Level::TRACE, "expert_mlp"),
+        })
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor, expert_idx: usize) -> Result<Tensor> {
+        let _enter = self.span.enter();
+
+        let expert_w1 = self.w1.narrow(0, expert_idx, 1)?.squeeze(0)?.t()?;
+        let expert_w2 = self.w2.narrow(0, expert_idx, 1)?.squeeze(0)?;
+
+        let hidden_states = hidden_states.broadcast_matmul(&expert_w1)?;
+        let hidden_states = self.activation.forward(&hidden_states)?;
+
+        hidden_states.broadcast_matmul(&expert_w2)
+    }
+}
+
+pub struct NomicExperts {
+    moe_num_experts: usize,
+    mlp: NomicExpertMLP,
+    bias: Tensor,
+
+    span: tracing::Span,
+}
+
+impl NomicExperts {
+    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+        let moe_num_experts = config.num_experts.unwrap();
+
+        let mlp = NomicExpertMLP::load(vb.pp("mlp"), config)?;
+
+        let bias = vb.get((config.n_embd,), "bias")?;
+
+        Ok(Self {
+            moe_num_experts,
+            mlp,
+            bias,
+            span: tracing::span!(tracing::Level::TRACE, "experts"),
+        })
+    }
+
+    pub fn forward(
+        &self,
+        hidden_states: &Tensor,
+        top_weights: &Tensor,
+        top_experts: &Tensor,
+    ) -> Result<Tensor> {
+        let _enter = self.span.enter();
+
+        let dims = hidden_states.dims();
+        let ndim = dims.len();
+
+        let (bs, seq_len, hidden_size) = match ndim {
+            3 => (dims[0], dims[1], dims[2]),
+            2 => (1, dims[0], dims[1]),
+            _ => unreachable!(),
+        };
+
+        let hidden_states = hidden_states.reshape(((), hidden_size))?;
+
+        let mut out = Tensor::zeros_like(&hidden_states)?;
+
+        let counts = top_experts
+            .flatten_all()?
+            .bincount(self.moe_num_experts as u32)?;
+
+        for (expert_idx, &count) in counts.iter().enumerate().take(self.moe_num_experts) {
+            if count == 0u32 {
+                continue;
+            }
+
+            let idx_top = top_experts.eq(expert_idx as f64)?.nonzero()?.t()?;
+            let idx = &idx_top.i(0)?.contiguous()?;
+            let top = &idx_top.i(1)?.contiguous()?;
+
+            let expert_out = self
+                .mlp
+                .forward(&hidden_states.index_select(idx, 0)?, expert_idx)?
+                .broadcast_mul(
+                    &top_weights
+                        .index_select(idx, 0)?
+                        .gather(&top.unsqueeze(1)?, 1)?
+                        .squeeze(1)?
+                        .unsqueeze(D::Minus1)?
+                        .to_dtype(hidden_states.dtype())?,
+                )?;
+
+            out = out.index_add(idx, &expert_out, 0)?;
+        }
+
+        if ndim == 3 {
+            out = out.reshape((bs, seq_len, hidden_size))?;
+        }
+
+        out.broadcast_add(&self.bias)
+    }
+}
+
+pub struct NomicMoELayer {
+    router: NomicRouter,
+    experts: NomicExperts,
+
+    span: tracing::Span,
+}
+
+impl NomicMoELayer {
+    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+        let router = NomicRouter::load(vb.pp("router"), config)?;
+        let experts = NomicExperts::load(vb.pp("experts"), config)?;
+
+        Ok(Self {
+            router,
+            experts,
+            span: tracing::span!(tracing::Level::TRACE, "moe"),
+        })
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        let _enter = self.span.enter();
+
+        let (top_weights, top_experts) = self.router.forward(hidden_states)?;
+
+        self.experts
+            .forward(hidden_states, &top_weights, &top_experts)
+    }
+}
+
+pub enum NomicMLP {
+    MoE(NomicMoELayer),
+    GatedMLP(NomicBertGatedMLP),
+    Mlp(NomicBertMLP),
+}
+
+impl NomicMLP {
+    pub fn load(vb: VarBuilder, index: usize, config: &NomicConfig) -> Result<Self> {
+        let use_moe = matches!(config.moe_every_n_layers, Some(n) if n > 0 && index % n == 1);
+
+        if use_moe {
+            Ok(Self::MoE(NomicMoELayer::load(vb, config)?))
+        } else if config.activation_function == HiddenAct::Gelu {
+            Ok(Self::Mlp(NomicBertMLP::load(vb, config)?))
+        } else {
+            Ok(Self::GatedMLP(NomicBertGatedMLP::load(vb, config)?))
+        }
+    }
+
+    pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        match self {
+            Self::MoE(layer) => layer.forward(hidden_states),
+            Self::GatedMLP(layer) => layer.forward(hidden_states),
+            Self::Mlp(layer) => layer.forward(hidden_states),
+        }
     }
 }
 
@@ -155,16 +438,25 @@ impl NomicAttention {
         let attention_head_size = config.n_embd / config.n_head;
         let hidden_size = config.n_embd;
 
-        let qkv_weight = vb.pp("Wqkv").get(
-            (3 * num_attention_heads * attention_head_size, hidden_size),
-            "weight",
-        )?;
-        let qkv_linear = Linear::new(qkv_weight, None, None);
+        let qkv_dim = 3 * num_attention_heads * attention_head_size;
+
+        let qkv_weight = vb.pp("Wqkv").get((qkv_dim, hidden_size), "weight")?;
+        let qkv_bias = if config.qkv_proj_bias {
+            Some(vb.pp("Wqkv").get((qkv_dim,), "bias")?)
+        } else {
+            None
+        };
+        let qkv_linear = Linear::new(qkv_weight, qkv_bias, None);
 
         let out_proj_weight = vb
             .pp("out_proj")
             .get((hidden_size, hidden_size), "weight")?;
-        let out_proj = Linear::new(out_proj_weight, None, None);
+        let out_proj_bias = if config.qkv_proj_bias {
+            Some(vb.pp("out_proj").get((hidden_size,), "bias")?)
+        } else {
+            None
+        };
+        let out_proj = Linear::new(out_proj_weight, out_proj_bias, None);
 
         let softmax_scale = 1. / (attention_head_size as f64).sqrt();
 
@@ -282,7 +574,7 @@ impl NomicAttention {
 
 struct NomicBertBlock {
     attention: NomicAttention,
-    mlp: NomicBertGatedMLP,
+    mlp: NomicMLP,
     post_attention_layer_norm: LayerNorm,
     output_layer_norm: LayerNorm,
 
@@ -290,9 +582,10 @@ struct NomicBertBlock {
 }
 
 impl NomicBertBlock {
-    pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
+    pub fn load(vb: VarBuilder, index: usize, config: &NomicConfig) -> Result<Self> {
         let attention = NomicAttention::load(vb.pp("attn"), config)?;
-        let mlp = NomicBertGatedMLP::load(vb.pp("mlp"), config)?;
+
+        let mlp = NomicMLP::load(vb.pp("mlp"), index, config)?;
 
         let post_attention_layer_norm =
             LayerNorm::load(vb.pp("norm1"), config.n_embd, config.layer_norm_epsilon)?;
@@ -320,6 +613,7 @@ impl NomicBertBlock {
         let attn_output = self
             .attention
             .forward(hidden_states, attention_bias, cos, sin)?;
+
         let hidden_states = self
             .post_attention_layer_norm
             .forward(hidden_states, Some(&attn_output))?;
@@ -333,13 +627,14 @@ impl NomicBertBlock {
 
 struct NomicBertEncoder {
     layers: Vec<NomicBertBlock>,
+
     span: tracing::Span,
 }
 
 impl NomicBertEncoder {
     pub fn load(vb: VarBuilder, config: &NomicConfig) -> Result<Self> {
         let layers = (0..config.n_layer)
-            .map(|index| NomicBertBlock::load(vb.pp(format!("layers.{index}")), config))
+            .map(|index| NomicBertBlock::load(vb.pp(format!("layers.{index}")), index, config))
             .collect::<Result<Vec<_>>>()?;
 
         let span = tracing::span!(tracing::Level::TRACE, "encoder");
@@ -357,7 +652,6 @@ impl NomicBertEncoder {
 
         let mut hidden_states = hidden_states.clone();
 
-        // Use a loop rather than a fold as it's easier to modify when adding debug/...
         for layer in self.layers.iter() {
             hidden_states = layer.forward(&hidden_states, attention_bias, cos, sin)?
         }
@@ -444,6 +738,7 @@ impl NomicBertModel {
             span: tracing::span!(tracing::Level::TRACE, "model"),
         })
     }
+
     pub fn forward(&self, batch: Batch) -> Result<(Option<Tensor>, Option<Tensor>)> {
         let _enter = self.span.enter();