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