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