1use crate::model_config;
9use super::traits::*;
10use anyhow::Result;
11use serde::{Serialize, Deserialize};
12
13model_config!(MPTConfig {
14 vocab_size: usize = 50432,
15 hidden_size: usize = 4096,
16 intermediate_size: usize = 16384,
17 num_hidden_layers: usize = 32,
18 num_attention_heads: usize = 32,
19 num_key_value_heads: usize = 32,
20 hidden_act: String = "gelu".to_string(),
21 max_position_embeddings: usize = 2048,
22 initializer_range: f32 = 0.02,
23 layer_norm_eps: f32 = 1e-5,
24 use_cache: bool = true,
25 pad_token_id: i64 = 0,
26 bos_token_id: i64 = 0,
27 eos_token_id: i64 = 0,
28 tie_word_embeddings: bool = true,
29 alibi: bool = true,
30 no_bias: bool = true,
31});
32
33impl MPTConfig {
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 ..Default::default()
44 }
45 }
46}
47
48pub struct MPTModelV2 {
49 config: MPTConfig,
50 device: Device,
51 wte: Tensor,
52 blocks: Vec<MPTBlock>,
53 norm_f: Tensor,
54}
55
56pub struct MPTBlock {
57 attn: MPTAttention,
58 ffn: MPTFFN,
59 norm_1: Tensor,
60 norm_2: Tensor,
61}
62
63pub struct MPTAttention {
64 wqkv: Tensor,
65 out_proj: Tensor,
66 num_heads: usize,
67 head_dim: usize,
68 scale: f32,
69 alibi: bool,
70}
71
72pub struct MPTFFN {
73 up_proj: Tensor,
74 down_proj: Tensor,
75}
76
77impl Model for MPTModelV2 {
78 type Config = MPTConfig;
79
80 fn new(config: MPTConfig) -> Result<Self> {
81 let device = Device::CPU;
82 let wte = ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?;
83 let norm_f = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
84
85 let mut blocks = Vec::with_capacity(config.num_hidden_layers);
86 for _ in 0..config.num_hidden_layers {
87 blocks.push(MPTBlock::new(&config, &device)?);
88 }
89
90 Ok(Self { config, device, wte, blocks, norm_f })
91 }
92
93 fn from_weights(config: MPTConfig, weights: ModelWeights) -> Result<Self> {
94 let mut model = Self::new(config)?;
95 if let Some(w) = weights.get("transformer.wte.weight") { model.wte = w.clone(); }
96 if let Some(w) = weights.get("transformer.norm_f.weight") { model.norm_f = w.clone(); }
97 for (i, block) in model.blocks.iter_mut().enumerate() { block.load_weights(&weights, i)?; }
98 Ok(model)
99 }
100
101 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
102 match inputs {
103 ModelInputs::Text { input_ids, .. } => {
104 let mut hidden = ops_fn::embedding(input_ids, &self.wte)?;
105 let seq_len = input_ids.shape()[1];
106
107 for block in &self.blocks {
108 hidden = block.forward(&hidden, seq_len)?;
109 }
110
111 hidden = ops_fn::layer_norm(&hidden, &self.norm_f, None, self.config.layer_norm_eps)?;
112
113 let wte_candle = self.wte.to_candle()?;
115 let hidden_candle = hidden.to_candle()?.contiguous()?;
116 let batch = hidden_candle.dims()[0];
117 let seq = hidden_candle.dims()[1];
118 let hidden_size = hidden_candle.dims()[2];
119 let flat = hidden_candle.reshape(&[batch * seq, hidden_size])?;
120 let logits_flat = flat.matmul(&wte_candle.t()?)?;
121 let logits_candle = logits_flat.reshape(&[batch, seq, self.config.vocab_size])?;
122 let logits = Tensor::from_candle(logits_candle);
123
124 Ok(ModelOutputs::Logits { logits, hidden_states: None })
125 }
126 _ => Err(anyhow::anyhow!("MPT only supports text inputs")),
127 }
128 }
129
130 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
131 use crate::tokenizer::Tokenizer;
132 use rand::Rng;
133 let tokenizer = Tokenizer::new();
134 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
135 for _ in 0..config.max_new_tokens {
136 let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
137 let input = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;
138 let outputs = self.forward(&ModelInputs::text(input))?;
139 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
140 let logits_candle = logits.to_candle()?;
141 let last = logits_candle.narrow(1, logits_candle.dims()[1] - 1, 1)?.squeeze(1)?.squeeze(0)?;
142 let logits_vec: Vec<f32> = last.to_vec1()?;
143 let next = if config.do_sample && config.temperature > 0.0 {
144 let scaled: Vec<f32> = logits_vec.iter().map(|&x| x / config.temperature).collect();
145 let max_v = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
146 let exp_sum: f32 = scaled.iter().map(|&x| (x - max_v).exp()).sum();
147 let probs: Vec<f32> = scaled.iter().map(|&x| (x - max_v).exp() / exp_sum).collect();
148 let mut rng = rand::thread_rng();
149 let r: f32 = rng.gen();
150 let mut cum = 0.0;
151 let mut s = 0u32;
152 for (i, &p) in probs.iter().enumerate() { cum += p; if r <= cum { s = i as u32; break; } }
153 s
154 } else {
155 logits_vec.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()).map(|(i, _)| i as u32).unwrap_or(0)
156 };
157 if next == config.eos_token_id { break; }
158 tokens.push(next);
159 }
160 Ok(tokenizer.decode(&tokens))
161 }
162
163 fn config(&self) -> &Self::Config { &self.config }
164 fn memory_requirements(&self) -> MemoryRequirements {
165 let p = self.config.vocab_size * self.config.hidden_size + self.config.num_hidden_layers * 8 * self.config.hidden_size.pow(2);
166 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 }
167 }
168 fn to_device(&mut self, device: &Device) -> Result<()> {
169 self.wte = self.wte.to_device(device)?;
170 self.norm_f = self.norm_f.to_device(device)?;
171 for b in &mut self.blocks { b.to_device(device)?; }
172 self.device = device.clone();
173 Ok(())
174 }
175}
176
177impl MPTBlock {
178 fn new(config: &MPTConfig, device: &Device) -> Result<Self> {
179 Ok(Self {
180 attn: MPTAttention::new(config, device)?,
181 ffn: MPTFFN::new(config, device)?,
182 norm_1: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
183 norm_2: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
184 })
185 }
186
187 fn forward(&self, hidden_states: &Tensor, seq_len: usize) -> Result<Tensor> {
188 let residual = hidden_states.clone();
189 let h = ops_fn::layer_norm(hidden_states, &self.norm_1, None, 1e-5)?;
190 let attn_out = self.attn.forward(&h, seq_len)?;
191 let h = ops_fn::add(&residual, &attn_out)?;
192
193 let residual = h.clone();
194 let h = ops_fn::layer_norm(&h, &self.norm_2, None, 1e-5)?;
195 let ffn_out = self.ffn.forward(&h)?;
196 ops_fn::add(&residual, &ffn_out)
197 }
198
199 fn load_weights(&mut self, weights: &ModelWeights, idx: usize) -> Result<()> {
200 let p = format!("transformer.blocks.{}", idx);
201 if let Some(w) = weights.get(&format!("{}.attn.Wqkv.weight", p)) { self.attn.wqkv = ops_fn::transpose(w)?; }
202 if let Some(w) = weights.get(&format!("{}.attn.out_proj.weight", p)) { self.attn.out_proj = ops_fn::transpose(w)?; }
203 if let Some(w) = weights.get(&format!("{}.ffn.up_proj.weight", p)) { self.ffn.up_proj = ops_fn::transpose(w)?; }
204 if let Some(w) = weights.get(&format!("{}.ffn.down_proj.weight", p)) { self.ffn.down_proj = ops_fn::transpose(w)?; }
205 if let Some(w) = weights.get(&format!("{}.norm_1.weight", p)) { self.norm_1 = w.clone(); }
206 if let Some(w) = weights.get(&format!("{}.norm_2.weight", p)) { self.norm_2 = w.clone(); }
207 Ok(())
208 }
209
210 fn to_device(&mut self, device: &Device) -> Result<()> {
211 self.attn.to_device(device)?;
212 self.ffn.to_device(device)?;
213 self.norm_1 = self.norm_1.to_device(device)?;
214 self.norm_2 = self.norm_2.to_device(device)?;
215 Ok(())
216 }
217}
218
219fn build_alibi_bias(num_heads: usize, seq_len: usize, device: &candle_core::Device) -> Result<candle_core::Tensor> {
220 let closest_power_of_2 = 2usize.pow((num_heads as f64).log2().floor() as u32);
221 let base = 2f32.powf(-(2f32.powf(-((closest_power_of_2 as f32).log2() - 3.0))));
222 let mut slopes = Vec::with_capacity(num_heads);
223 for i in 0..num_heads { slopes.push(base.powf((i + 1) as f32)); }
224
225 let mut bias_data = vec![0.0f32; num_heads * seq_len * seq_len];
226 for h in 0..num_heads {
227 for i in 0..seq_len {
228 for j in 0..seq_len {
229 if j <= i {
230 bias_data[h * seq_len * seq_len + i * seq_len + j] = slopes[h] * (j as i32 - i as i32) as f32;
231 } else {
232 bias_data[h * seq_len * seq_len + i * seq_len + j] = f32::NEG_INFINITY;
233 }
234 }
235 }
236 }
237 Ok(candle_core::Tensor::from_vec(bias_data, &[1, num_heads, seq_len, seq_len], device)?)
238}
239
240impl MPTAttention {
241 fn new(config: &MPTConfig, device: &Device) -> Result<Self> {
242 let head_dim = config.hidden_size / config.num_attention_heads;
243 Ok(Self {
244 wqkv: ops_fn::zeros(&[config.hidden_size, 3 * config.hidden_size], DataType::Float32, device)?,
245 out_proj: ops_fn::zeros(&[config.hidden_size, config.hidden_size], DataType::Float32, device)?,
246 num_heads: config.num_attention_heads,
247 head_dim,
248 scale: 1.0 / (head_dim as f32).sqrt(),
249 alibi: config.alibi,
250 })
251 }
252
253 fn forward(&self, hidden_states: &Tensor, seq_len: usize) -> Result<Tensor> {
254 let shape = hidden_states.shape();
255 let (batch, seq, hidden_size) = (shape[0], shape[1], shape[2]);
256
257 let qkv = ops_fn::matmul(hidden_states, &self.wqkv)?.to_candle()?;
258 let q = qkv.narrow(2, 0, hidden_size)?.reshape(&[batch, seq, self.num_heads, self.head_dim])?.transpose(1, 2)?;
259 let k = qkv.narrow(2, hidden_size, hidden_size)?.reshape(&[batch, seq, self.num_heads, self.head_dim])?.transpose(1, 2)?;
260 let v = qkv.narrow(2, 2 * hidden_size, hidden_size)?.reshape(&[batch, seq, self.num_heads, self.head_dim])?.transpose(1, 2)?;
261
262 let q = q.contiguous()?;
263 let k_t = k.transpose(2, 3)?.contiguous()?;
264 let scores = (q.matmul(&k_t)? * (self.scale as f64))?;
265
266 let device = scores.device();
267 let scores = if self.alibi {
268 let alibi_bias = build_alibi_bias(self.num_heads, seq_len, device)?;
269 scores.broadcast_add(&alibi_bias)?
270 } else {
271 let mask = {
272 let mut m = vec![0.0f32; seq * seq];
273 for i in 0..seq { for j in (i+1)..seq { m[i*seq+j] = f32::NEG_INFINITY; } }
274 candle_core::Tensor::from_vec(m, &[1, 1, seq, seq], device)?
275 };
276 scores.broadcast_add(&mask)?
277 };
278
279 let v = v.contiguous()?;
280 let attn = candle_nn::ops::softmax_last_dim(&scores)?.matmul(&v)?;
281 let out = attn.transpose(1, 2)?.reshape(&[batch, seq, hidden_size])?;
282 ops_fn::matmul(&Tensor::from_candle(out), &self.out_proj)
283 }
284
285 fn to_device(&mut self, device: &Device) -> Result<()> {
286 self.wqkv = self.wqkv.to_device(device)?;
287 self.out_proj = self.out_proj.to_device(device)?;
288 Ok(())
289 }
290}
291
292impl MPTFFN {
293 fn new(config: &MPTConfig, device: &Device) -> Result<Self> {
294 Ok(Self {
295 up_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
296 down_proj: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
297 })
298 }
299 fn forward(&self, x: &Tensor) -> Result<Tensor> {
300 let h = ops_fn::matmul(x, &self.up_proj)?;
301 let h = ops_fn::gelu(&h)?;
302 ops_fn::matmul(&h, &self.down_proj)
303 }
304 fn to_device(&mut self, device: &Device) -> Result<()> {
305 self.up_proj = self.up_proj.to_device(device)?;
306 self.down_proj = self.down_proj.to_device(device)?;
307 Ok(())
308 }
309}
310
311#[cfg(test)]
312mod tests {
313 use super::*;
314 #[test]
315 fn test_mpt_creation() {
316 let config = MPTConfig { vocab_size: 1000, hidden_size: 128, intermediate_size: 512, num_hidden_layers: 2, num_attention_heads: 4, num_key_value_heads: 4, ..Default::default() };
317 let model = MPTModelV2::new(config).unwrap();
318 assert_eq!(model.config().vocab_size(), 1000);
319 }
320 #[test]
321 fn test_mpt_forward() {
322 let config = MPTConfig { vocab_size: 100, hidden_size: 64, intermediate_size: 256, num_hidden_layers: 1, num_attention_heads: 4, num_key_value_heads: 4, ..Default::default() };
323 let model = MPTModelV2::new(config).unwrap();
324 let inputs = ModelInputs::text(ops_fn::zeros(&[2, 8], DataType::Int64, &Device::CPU).unwrap());
325 match model.forward(&inputs).unwrap() { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[2, 8, 100]), _ => panic!() }
326 }
327}