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