1use crate::model_config;
9use super::traits::*;
10use anyhow::Result;
11use serde::{Serialize, Deserialize};
12
13model_config!(DbrxConfig {
14 vocab_size: usize = 100352,
15 hidden_size: usize = 6144,
16 intermediate_size: usize = 10752,
17 num_hidden_layers: usize = 40,
18 num_attention_heads: usize = 48,
19 num_key_value_heads: usize = 8,
20 hidden_act: String = "silu".to_string(),
21 max_position_embeddings: usize = 32768,
22 rms_norm_eps: f32 = 1e-5,
23 use_cache: bool = true,
24 pad_token_id: i64 = 0,
25 bos_token_id: i64 = 1,
26 eos_token_id: i64 = 2,
27 tie_word_embeddings: bool = false,
28 rope_theta: f32 = 500000.0,
29 num_experts: usize = 16,
30 num_experts_per_tok: usize = 4,
31});
32
33impl DbrxConfig {
34 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
35 Self {
36 vocab_size: gguf.vocab_size,
37 hidden_size: gguf.hidden_size,
38 intermediate_size: gguf.intermediate_size,
39 num_hidden_layers: gguf.num_hidden_layers,
40 num_attention_heads: gguf.num_attention_heads,
41 num_key_value_heads: gguf.num_key_value_heads,
42 max_position_embeddings: gguf.max_position_embeddings,
43 rope_theta: gguf.rope_theta,
44 ..Default::default()
45 }
46 }
47}
48
49pub struct DbrxModelV2 {
50 config: DbrxConfig,
51 device: Device,
52 embed_tokens: Tensor,
53 layers: Vec<DbrxLayer>,
54 norm: Tensor,
55 lm_head: Tensor,
56}
57
58pub struct DbrxLayer {
59 self_attn: DbrxAttention,
60 moe: DbrxMoE,
61 input_layernorm: Tensor,
62 post_attention_layernorm: Tensor,
63}
64
65pub struct DbrxAttention {
66 qkv_proj: Tensor,
67 o_proj: Tensor,
68 num_heads: usize,
69 num_kv_heads: usize,
70 head_dim: usize,
71 scale: f32,
72}
73
74pub struct DbrxMoE {
75 router: Tensor,
76 experts: Vec<DbrxExpert>,
77 num_experts_per_tok: usize,
78}
79
80pub struct DbrxExpert {
81 gate_proj: Tensor,
82 up_proj: Tensor,
83 down_proj: Tensor,
84}
85
86fn apply_rope_dbrx(
87 q: &candle_core::Tensor,
88 k: &candle_core::Tensor,
89 seq_len: usize,
90 head_dim: usize,
91 rope_theta: f32,
92) -> Result<(candle_core::Tensor, candle_core::Tensor)> {
93 let device = q.device();
94 let half_dim = head_dim / 2;
95 let inv_freq: Vec<f32> = (0..half_dim)
96 .map(|i| 1.0 / rope_theta.powf((2 * i) as f32 / head_dim as f32))
97 .collect();
98
99 let positions: Vec<f32> = (0..seq_len).map(|p| p as f32).collect();
100 let mut angles = Vec::with_capacity(seq_len * half_dim);
101 for pos in &positions {
102 for freq in &inv_freq {
103 angles.push(pos * freq);
104 }
105 }
106
107 let angles_tensor = candle_core::Tensor::from_vec(angles, &[seq_len, half_dim], device)?;
108 let cos = angles_tensor.cos()?.unsqueeze(0)?.unsqueeze(0)?;
109 let sin = angles_tensor.sin()?.unsqueeze(0)?.unsqueeze(0)?;
110
111 let q_half1 = q.narrow(3, 0, half_dim)?;
112 let q_half2 = q.narrow(3, half_dim, half_dim)?;
113 let k_half1 = k.narrow(3, 0, half_dim)?;
114 let k_half2 = k.narrow(3, half_dim, half_dim)?;
115
116 let q_rot1 = (q_half1.broadcast_mul(&cos)? - q_half2.broadcast_mul(&sin)?)?;
117 let q_rot2 = (q_half1.broadcast_mul(&sin)? + q_half2.broadcast_mul(&cos)?)?;
118 let k_rot1 = (k_half1.broadcast_mul(&cos)? - k_half2.broadcast_mul(&sin)?)?;
119 let k_rot2 = (k_half1.broadcast_mul(&sin)? + k_half2.broadcast_mul(&cos)?)?;
120
121 Ok((
122 candle_core::Tensor::cat(&[&q_rot1, &q_rot2], 3)?,
123 candle_core::Tensor::cat(&[&k_rot1, &k_rot2], 3)?
124 ))
125}
126
127impl Model for DbrxModelV2 {
128 type Config = DbrxConfig;
129
130 fn new(config: DbrxConfig) -> Result<Self> {
131 let device = Device::CPU;
132 let embed_tokens = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?;
133 let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
134 let lm_head = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?;
135
136 let mut layers = Vec::with_capacity(config.num_hidden_layers);
137 for _ in 0..config.num_hidden_layers {
138 layers.push(DbrxLayer::new(&config, &device)?);
139 }
140
141 Ok(Self { config, device, embed_tokens, layers, norm, lm_head })
142 }
143
144 fn from_weights(config: DbrxConfig, weights: ModelWeights) -> Result<Self> {
145 let mut model = Self::new(config)?;
146 if let Some(w) = weights.get("model.embed_tokens.weight") { model.embed_tokens = w.clone(); }
147 if let Some(w) = weights.get("model.norm.weight") { model.norm = w.clone(); }
148 if let Some(w) = weights.get("lm_head.weight") { model.lm_head = w.clone(); }
149 for (i, layer) in model.layers.iter_mut().enumerate() { layer.load_weights(&weights, i)?; }
150 Ok(model)
151 }
152
153 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
154 match inputs {
155 ModelInputs::Text { input_ids, .. } => {
156 let seq_len = input_ids.shape()[1];
157 let mut hidden = ops_fn::embedding(input_ids, &self.embed_tokens)?;
158
159 for layer in &self.layers {
160 hidden = layer.forward(&hidden, seq_len, self.config.rope_theta)?;
161 }
162
163 hidden = ops_fn::rms_norm(&hidden, &self.norm, self.config.rms_norm_eps)?;
164 let logits = ops_fn::matmul(&hidden, &ops_fn::transpose(&self.lm_head)?)?;
165
166 Ok(ModelOutputs::Logits { logits, hidden_states: None })
167 }
168 _ => Err(anyhow::anyhow!("DBRX only supports text inputs")),
169 }
170 }
171
172 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
173 use crate::tokenizer::Tokenizer;
174 use rand::Rng;
175 let tokenizer = Tokenizer::new();
176 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
177 for _ in 0..config.max_new_tokens {
178 let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
179 let input = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;
180 let outputs = self.forward(&ModelInputs::text(input))?;
181 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
182 let logits_candle = logits.to_candle()?;
183 let last = logits_candle.narrow(1, logits_candle.dims()[1] - 1, 1)?.squeeze(1)?.squeeze(0)?;
184 let logits_vec: Vec<f32> = last.to_vec1()?;
185 let next = if config.do_sample && config.temperature > 0.0 {
186 let scaled: Vec<f32> = logits_vec.iter().map(|&x| x / config.temperature).collect();
187 let max_v = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
188 let exp_sum: f32 = scaled.iter().map(|&x| (x - max_v).exp()).sum();
189 let probs: Vec<f32> = scaled.iter().map(|&x| (x - max_v).exp() / exp_sum).collect();
190 let mut rng = rand::thread_rng();
191 let r: f32 = rng.gen();
192 let mut cum = 0.0;
193 let mut s = 0u32;
194 for (i, &p) in probs.iter().enumerate() { cum += p; if r <= cum { s = i as u32; break; } }
195 s
196 } else {
197 logits_vec.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()).map(|(i, _)| i as u32).unwrap_or(0)
198 };
199 if next == config.eos_token_id { break; }
200 tokens.push(next);
201 }
202 Ok(tokenizer.decode(&tokens))
203 }
204
205 fn config(&self) -> &Self::Config { &self.config }
206 fn memory_requirements(&self) -> MemoryRequirements {
207 let p = self.config.vocab_size * self.config.hidden_size + self.config.num_hidden_layers * 8 * self.config.hidden_size.pow(2);
208 MemoryRequirements { gpu_memory: p * 4, cpu_memory: p, kv_cache_memory: 2 * self.config.num_hidden_layers * self.config.max_position_embeddings * self.config.hidden_size * 4, peak_memory: p * 5 }
209 }
210 fn to_device(&mut self, device: &Device) -> Result<()> {
211 self.embed_tokens = self.embed_tokens.to_device(device)?;
212 self.norm = self.norm.to_device(device)?;
213 self.lm_head = self.lm_head.to_device(device)?;
214 for layer in &mut self.layers { layer.to_device(device)?; }
215 self.device = device.clone();
216 Ok(())
217 }
218}
219
220impl DbrxLayer {
221 fn new(config: &DbrxConfig, device: &Device) -> Result<Self> {
222 Ok(Self {
223 self_attn: DbrxAttention::new(config, device)?,
224 moe: DbrxMoE::new(config, device)?,
225 input_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
226 post_attention_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
227 })
228 }
229
230 fn forward(&self, hidden_states: &Tensor, seq_len: usize, rope_theta: f32) -> Result<Tensor> {
231 let residual = hidden_states.clone();
232 let h = ops_fn::rms_norm(hidden_states, &self.input_layernorm, 1e-5)?;
233 let attn_out = self.self_attn.forward(&h, seq_len, rope_theta)?;
234 let h = ops_fn::add(&residual, &attn_out)?;
235
236 let residual = h.clone();
237 let h = ops_fn::rms_norm(&h, &self.post_attention_layernorm, 1e-5)?;
238 let moe_out = self.moe.forward(&h)?;
239 ops_fn::add(&residual, &moe_out)
240 }
241
242 fn load_weights(&mut self, weights: &ModelWeights, idx: usize) -> Result<()> {
243 let p = format!("model.layers.{}", idx);
244 if let Some(w) = weights.get(&format!("{}.self_attn.Wqkv.weight", p)) { self.self_attn.qkv_proj = ops_fn::transpose(w)?; }
245 if let Some(w) = weights.get(&format!("{}.self_attn.out_proj.weight", p)) { self.self_attn.o_proj = ops_fn::transpose(w)?; }
246 if let Some(w) = weights.get(&format!("{}.input_layernorm.weight", p)) { self.input_layernorm = w.clone(); }
247 if let Some(w) = weights.get(&format!("{}.post_attention_layernorm.weight", p)) { self.post_attention_layernorm = w.clone(); }
248 Ok(())
249 }
250
251 fn to_device(&mut self, device: &Device) -> Result<()> {
252 self.self_attn.to_device(device)?;
253 self.moe.to_device(device)?;
254 self.input_layernorm = self.input_layernorm.to_device(device)?;
255 self.post_attention_layernorm = self.post_attention_layernorm.to_device(device)?;
256 Ok(())
257 }
258}
259
260impl DbrxAttention {
261 fn new(config: &DbrxConfig, device: &Device) -> Result<Self> {
262 let head_dim = config.hidden_size / config.num_attention_heads;
263 let qkv_size = config.num_attention_heads * head_dim + 2 * config.num_key_value_heads * head_dim;
264 Ok(Self {
265 qkv_proj: ops_fn::zeros(&[config.hidden_size, qkv_size], DataType::Float32, device)?,
266 o_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
267 num_heads: config.num_attention_heads,
268 num_kv_heads: config.num_key_value_heads,
269 head_dim,
270 scale: 1.0 / (head_dim as f32).sqrt(),
271 })
272 }
273
274 fn forward(&self, hidden_states: &Tensor, seq_len: usize, rope_theta: f32) -> Result<Tensor> {
275 let shape = hidden_states.shape();
276 let (batch, _seq, hidden_size) = (shape[0], shape[1], shape[2]);
277
278 let qkv = ops_fn::matmul(hidden_states, &self.qkv_proj)?.to_candle()?;
279
280 let q_size = self.num_heads * self.head_dim;
281 let kv_size = self.num_kv_heads * self.head_dim;
282
283 let q = qkv.narrow(2, 0, q_size)?.reshape(&[batch, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
284 let k = qkv.narrow(2, q_size, kv_size)?.reshape(&[batch, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
285 let v = qkv.narrow(2, q_size + kv_size, kv_size)?.reshape(&[batch, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
286
287 let (q, k) = apply_rope_dbrx(&q, &k, seq_len, self.head_dim, rope_theta)?;
288
289 let num_groups = self.num_heads / self.num_kv_heads;
290 let (k, v) = if num_groups > 1 {
291 let k_exp = k.unsqueeze(2)?.broadcast_as(&[batch, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch, self.num_heads, seq_len, self.head_dim])?;
292 let v_exp = v.unsqueeze(2)?.broadcast_as(&[batch, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch, self.num_heads, seq_len, self.head_dim])?;
293 (k_exp, v_exp)
294 } else {
295 (k, v)
296 };
297
298 let q = q.contiguous()?;
299 let k_t = k.transpose(2, 3)?.contiguous()?;
300 let scores = (q.matmul(&k_t)? * (self.scale as f64))?;
301
302 let device = scores.device();
303 let mask = {
304 let mut m = vec![0.0f32; seq_len * seq_len];
305 for i in 0..seq_len { for j in (i+1)..seq_len { m[i*seq_len+j] = f32::NEG_INFINITY; } }
306 candle_core::Tensor::from_vec(m, &[1, 1, seq_len, seq_len], device)?
307 };
308 let scores = scores.broadcast_add(&mask)?;
309
310 let v = v.contiguous()?;
311 let attn = candle_nn::ops::softmax_last_dim(&scores)?.matmul(&v)?;
312 let out = attn.transpose(1, 2)?.reshape(&[batch, seq_len, hidden_size])?;
313 ops_fn::matmul(&Tensor::from_candle(out), &self.o_proj)
314 }
315
316 fn to_device(&mut self, device: &Device) -> Result<()> {
317 self.qkv_proj = self.qkv_proj.to_device(device)?;
318 self.o_proj = self.o_proj.to_device(device)?;
319 Ok(())
320 }
321}
322
323impl DbrxMoE {
324 fn new(config: &DbrxConfig, device: &Device) -> Result<Self> {
325 let router = ops_fn::zeros(&[config.hidden_size, config.num_experts], DataType::Float32, device)?;
326 let mut experts = Vec::with_capacity(config.num_experts);
327 for _ in 0..config.num_experts {
328 experts.push(DbrxExpert::new(config, device)?);
329 }
330 Ok(Self { router, experts, num_experts_per_tok: config.num_experts_per_tok })
331 }
332
333 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
334 let shape = hidden_states.shape();
335 let (batch_size, seq_len, hidden_size) = (shape[0], shape[1], shape[2]);
336 let num_tokens = batch_size * seq_len;
337 let k = self.num_experts_per_tok;
338
339 let flat_hidden = hidden_states.reshape(&[num_tokens, hidden_size])?;
340 let router_logits = ops_fn::matmul(&flat_hidden, &self.router)?;
341
342 let (topk_weights, topk_indices) = ops_fn::topk(&router_logits, k, -1)?;
343 let routing_weights = ops_fn::softmax(&topk_weights, -1)?;
344
345 let all_indices: Vec<i64> = topk_indices.to_candle()?.flatten_all()?.to_vec1()?;
346 let all_weights: Vec<f32> = routing_weights.to_candle()?.flatten_all()?.to_vec1()?;
347 let flat_hidden_candle = flat_hidden.to_candle()?;
348
349 let mut output_data = vec![0.0f32; num_tokens * hidden_size];
350
351 for tok_idx in 0..num_tokens {
352 let token_hidden = flat_hidden_candle.get(tok_idx)?;
353 let token_tensor = Tensor::from_candle(token_hidden.unsqueeze(0)?);
354
355 let start = tok_idx * k;
356 let indices = &all_indices[start..start + k];
357 let weights = &all_weights[start..start + k];
358
359 let mut token_output = ops_fn::zeros(&[1, hidden_size], hidden_states.dtype(), hidden_states.device())?;
360
361 for (i, &expert_idx) in indices.iter().enumerate() {
362 let expert = &self.experts[expert_idx as usize];
363 let expert_output = expert.forward(&token_tensor)?;
364 let scaled_output = ops_fn::scale(&expert_output, weights[i])?;
365 token_output = ops_fn::add(&token_output, &scaled_output)?;
366 }
367
368 let token_data: Vec<f32> = token_output.to_candle()?.flatten_all()?.to_vec1()?;
369 for (i, &v) in token_data.iter().enumerate() {
370 output_data[tok_idx * hidden_size + i] = v;
371 }
372 }
373
374 let output = Tensor::from_f32_slice(&output_data, &[num_tokens, hidden_size], hidden_states.device())?;
375 output.reshape(&[batch_size, seq_len, hidden_size])
376 }
377
378 fn to_device(&mut self, device: &Device) -> Result<()> {
379 self.router = self.router.to_device(device)?;
380 for expert in &mut self.experts { expert.to_device(device)?; }
381 Ok(())
382 }
383}
384
385impl DbrxExpert {
386 fn new(config: &DbrxConfig, device: &Device) -> Result<Self> {
387 Ok(Self {
388 gate_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
389 up_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
390 down_proj: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
391 })
392 }
393 fn forward(&self, x: &Tensor) -> Result<Tensor> {
394 let gate = ops_fn::matmul(x, &self.gate_proj)?;
395 let up = ops_fn::matmul(x, &self.up_proj)?;
396 let h = ops_fn::mul(&ops_fn::silu(&gate)?, &up)?;
397 ops_fn::matmul(&h, &self.down_proj)
398 }
399 fn to_device(&mut self, device: &Device) -> Result<()> {
400 self.gate_proj = self.gate_proj.to_device(device)?;
401 self.up_proj = self.up_proj.to_device(device)?;
402 self.down_proj = self.down_proj.to_device(device)?;
403 Ok(())
404 }
405}
406
407#[cfg(test)]
408mod tests {
409 use super::*;
410 #[test]
411 fn test_dbrx_creation() {
412 let config = DbrxConfig { vocab_size: 1000, hidden_size: 64, intermediate_size: 256, num_hidden_layers: 2, num_attention_heads: 4, num_key_value_heads: 2, num_experts: 4, num_experts_per_tok: 2, ..Default::default() };
413 let model = DbrxModelV2::new(config).unwrap();
414 assert_eq!(model.config().vocab_size(), 1000);
415 }
416 #[test]
417 fn test_dbrx_forward() {
418 let config = DbrxConfig { vocab_size: 100, hidden_size: 64, intermediate_size: 256, num_hidden_layers: 1, num_attention_heads: 4, num_key_value_heads: 2, num_experts: 4, num_experts_per_tok: 2, ..Default::default() };
419 let model = DbrxModelV2::new(config).unwrap();
420 let inputs = ModelInputs::text(ops_fn::zeros(&[1, 4], DataType::Int64, &Device::CPU).unwrap());
421 match model.forward(&inputs).unwrap() { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[1, 4, 100]), _ => panic!() }
422 }
423}