ruvllm_esp32/
model.rs

1//! Model definition and loading for ESP32
2//!
3//! Supports tiny transformer models with INT8 quantization.
4
5use crate::quantized::{QuantParams, QuantizationType};
6use heapless::Vec as HVec;
7use serde::{Deserialize, Serialize};
8
9/// Maximum number of transformer layers
10pub const MAX_LAYERS: usize = 2;
11/// Maximum embedding table size (vocab * embed_dim bytes)
12pub const MAX_EMBEDDING_SIZE: usize = 32 * 1024; // 32KB
13/// Maximum weight size per layer
14pub const MAX_LAYER_SIZE: usize = 16 * 1024; // 16KB
15
16/// Model configuration
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ModelConfig {
19    /// Vocabulary size
20    pub vocab_size: usize,
21    /// Embedding dimension
22    pub embed_dim: usize,
23    /// Hidden dimension in FFN
24    pub hidden_dim: usize,
25    /// Number of transformer layers
26    pub num_layers: usize,
27    /// Number of attention heads
28    pub num_heads: usize,
29    /// Maximum sequence length
30    pub max_seq_len: usize,
31    /// Quantization type
32    pub quant_type: QuantizationType,
33}
34
35impl Default for ModelConfig {
36    fn default() -> Self {
37        // Tiny model suitable for ESP32
38        Self {
39            vocab_size: 256,
40            embed_dim: 32,
41            hidden_dim: 64,
42            num_layers: 1,
43            num_heads: 2,
44            max_seq_len: 16,
45            quant_type: QuantizationType::Int8,
46        }
47    }
48}
49
50impl ModelConfig {
51    /// Validate configuration fits ESP32 constraints
52    pub fn validate(&self, variant: crate::Esp32Variant) -> crate::Result<()> {
53        let model_size = self.estimate_size();
54        let max_ram = variant.max_model_ram();
55
56        if model_size > max_ram {
57            return Err(crate::Error::ModelTooLarge {
58                required: model_size,
59                available: max_ram,
60            });
61        }
62
63        if self.embed_dim % self.num_heads != 0 {
64            return Err(crate::Error::InvalidModel(
65                "embed_dim must be divisible by num_heads"
66            ));
67        }
68
69        if self.num_layers > MAX_LAYERS {
70            return Err(crate::Error::InvalidModel("Too many layers"));
71        }
72
73        Ok(())
74    }
75
76    /// Estimate total model size in bytes
77    pub fn estimate_size(&self) -> usize {
78        let bytes_per_weight = match self.quant_type {
79            QuantizationType::Int8 => 1,
80            QuantizationType::Int4 => 1, // 2 weights per byte
81            QuantizationType::Binary => 1, // 8 weights per byte
82            QuantizationType::Fixed16 => 2,
83        };
84
85        let divisor = match self.quant_type {
86            QuantizationType::Int4 => 2,
87            QuantizationType::Binary => 8,
88            _ => 1,
89        };
90
91        // Embedding table
92        let embed_size = (self.vocab_size * self.embed_dim * bytes_per_weight) / divisor;
93
94        // Per-layer weights
95        let qkv_size = 3 * self.embed_dim * self.embed_dim * bytes_per_weight / divisor;
96        let ffn_size = 3 * self.embed_dim * self.hidden_dim * bytes_per_weight / divisor;
97        let layer_size = qkv_size + ffn_size;
98
99        // Output projection
100        let output_size = (self.vocab_size * self.embed_dim * bytes_per_weight) / divisor;
101
102        embed_size + (layer_size * self.num_layers) + output_size
103    }
104
105    /// Get recommended config for variant
106    pub fn for_variant(variant: crate::Esp32Variant) -> Self {
107        match variant {
108            crate::Esp32Variant::Esp32 | crate::Esp32Variant::Esp32S3 => {
109                // ~300KB available, use larger model (but fits in stack)
110                Self {
111                    vocab_size: 256,
112                    embed_dim: 64,
113                    hidden_dim: 128,
114                    num_layers: 2,
115                    num_heads: 4,
116                    max_seq_len: 32,
117                    quant_type: QuantizationType::Int8,
118                }
119            }
120            crate::Esp32Variant::Esp32S2 => {
121                // ~120KB available, use smaller model
122                Self {
123                    vocab_size: 128,
124                    embed_dim: 32,
125                    hidden_dim: 64,
126                    num_layers: 1,
127                    num_heads: 2,
128                    max_seq_len: 16,
129                    quant_type: QuantizationType::Int8,
130                }
131            }
132            crate::Esp32Variant::Esp32C3 | crate::Esp32Variant::Esp32C6 => {
133                // ~200KB available
134                Self {
135                    vocab_size: 256,
136                    embed_dim: 48,
137                    hidden_dim: 96,
138                    num_layers: 2,
139                    num_heads: 3,
140                    max_seq_len: 24,
141                    quant_type: QuantizationType::Int8,
142                }
143            }
144        }
145    }
146}
147
148/// Layer weights for a single transformer layer
149#[derive(Clone)]
150pub struct LayerWeights {
151    /// Query projection weights [embed_dim, embed_dim]
152    pub wq: HVec<i8, MAX_LAYER_SIZE>,
153    /// Key projection weights
154    pub wk: HVec<i8, MAX_LAYER_SIZE>,
155    /// Value projection weights
156    pub wv: HVec<i8, MAX_LAYER_SIZE>,
157    /// Output projection weights
158    pub wo: HVec<i8, MAX_LAYER_SIZE>,
159
160    /// FFN up projection [embed_dim, hidden_dim]
161    pub w_up: HVec<i8, MAX_LAYER_SIZE>,
162    /// FFN gate projection
163    pub w_gate: HVec<i8, MAX_LAYER_SIZE>,
164    /// FFN down projection [hidden_dim, embed_dim]
165    pub w_down: HVec<i8, MAX_LAYER_SIZE>,
166
167    /// Quantization params
168    pub q_params: QuantParams,
169    pub k_params: QuantParams,
170    pub v_params: QuantParams,
171    pub o_params: QuantParams,
172    pub up_params: QuantParams,
173    pub gate_params: QuantParams,
174    pub down_params: QuantParams,
175}
176
177impl Default for LayerWeights {
178    fn default() -> Self {
179        Self {
180            wq: HVec::new(),
181            wk: HVec::new(),
182            wv: HVec::new(),
183            wo: HVec::new(),
184            w_up: HVec::new(),
185            w_gate: HVec::new(),
186            w_down: HVec::new(),
187            q_params: QuantParams::default(),
188            k_params: QuantParams::default(),
189            v_params: QuantParams::default(),
190            o_params: QuantParams::default(),
191            up_params: QuantParams::default(),
192            gate_params: QuantParams::default(),
193            down_params: QuantParams::default(),
194        }
195    }
196}
197
198impl LayerWeights {
199    /// Initialize with random weights (for testing)
200    pub fn random(config: &ModelConfig, seed: u32) -> crate::Result<Self> {
201        let mut layer = Self::default();
202
203        let embed_dim = config.embed_dim;
204        let hidden_dim = config.hidden_dim;
205
206        // Simple LCG random number generator
207        let mut rng_state = seed;
208        let mut next_rand = || {
209            rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
210            // Get value in range 0-127, then map to -64 to 63
211            (((rng_state >> 16) & 0x7F) as i16 - 64) as i8
212        };
213
214        // QKV projections [embed_dim, embed_dim]
215        let qkv_size = embed_dim * embed_dim;
216        for _ in 0..qkv_size {
217            layer.wq.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
218            layer.wk.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
219            layer.wv.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
220            layer.wo.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
221        }
222
223        // FFN projections
224        let up_size = embed_dim * hidden_dim;
225        for _ in 0..up_size {
226            layer.w_up.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
227            layer.w_gate.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
228        }
229
230        let down_size = hidden_dim * embed_dim;
231        for _ in 0..down_size {
232            layer.w_down.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
233        }
234
235        // Initialize quant params with reasonable defaults
236        let scale = 1.0 / 64.0; // For weights in range [-64, 63]
237        layer.q_params = QuantParams { scale, zero_point: 0.0, min_val: -1.0, max_val: 1.0 };
238        layer.k_params = layer.q_params;
239        layer.v_params = layer.q_params;
240        layer.o_params = layer.q_params;
241        layer.up_params = layer.q_params;
242        layer.gate_params = layer.q_params;
243        layer.down_params = layer.q_params;
244
245        Ok(layer)
246    }
247
248    /// Memory size of this layer
249    pub fn memory_size(&self) -> usize {
250        self.wq.len() + self.wk.len() + self.wv.len() + self.wo.len()
251            + self.w_up.len() + self.w_gate.len() + self.w_down.len()
252    }
253}
254
255/// Complete tiny model
256pub struct TinyModel {
257    /// Model configuration
258    pub config: ModelConfig,
259    /// Embedding table [vocab_size, embed_dim]
260    pub embedding_table: HVec<i8, MAX_EMBEDDING_SIZE>,
261    /// Transformer layers
262    pub layers: [LayerWeights; MAX_LAYERS],
263    /// Output projection [embed_dim, vocab_size]
264    pub output_proj: HVec<i8, MAX_EMBEDDING_SIZE>,
265    /// Input quantization params
266    pub input_params: QuantParams,
267    /// Output quantization params
268    pub output_params: QuantParams,
269}
270
271impl TinyModel {
272    /// Create a new model with random weights
273    pub fn new(config: ModelConfig) -> crate::Result<Self> {
274        config.validate(crate::Esp32Variant::Esp32)?;
275
276        let mut embedding_table = HVec::new();
277        let mut output_proj = HVec::new();
278
279        // Initialize embedding table
280        let embed_size = config.vocab_size * config.embed_dim;
281        let mut rng_state = 12345u32;
282        let mut next_rand = || {
283            rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
284            // Get value in range 0-255, then map to -128 to 127
285            (((rng_state >> 16) & 0xFF) as i16 - 128) as i8
286        };
287
288        for _ in 0..embed_size {
289            embedding_table.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
290        }
291
292        // Initialize output projection
293        for _ in 0..embed_size {
294            output_proj.push(next_rand()).map_err(|_| crate::Error::BufferOverflow)?;
295        }
296
297        // Initialize layers
298        let mut layers: [LayerWeights; MAX_LAYERS] = Default::default();
299        for i in 0..config.num_layers {
300            layers[i] = LayerWeights::random(&config, (i * 1000) as u32)?;
301        }
302
303        Ok(Self {
304            config,
305            embedding_table,
306            layers,
307            output_proj,
308            input_params: QuantParams::default(),
309            output_params: QuantParams::default(),
310        })
311    }
312
313    /// Total memory size of model
314    pub fn memory_size(&self) -> usize {
315        let mut size = self.embedding_table.len();
316        size += self.output_proj.len();
317        for i in 0..self.config.num_layers {
318            size += self.layers[i].memory_size();
319        }
320        size
321    }
322
323    /// Load model from bytes (e.g., from flash)
324    pub fn from_bytes(data: &[u8]) -> crate::Result<Self> {
325        // Parse header
326        if data.len() < 32 {
327            return Err(crate::Error::InvalidModel("Data too small"));
328        }
329
330        // Magic number check
331        if &data[0..4] != b"RUVM" {
332            return Err(crate::Error::InvalidModel("Invalid magic number"));
333        }
334
335        // Parse config from header
336        let vocab_size = u16::from_le_bytes([data[4], data[5]]) as usize;
337        let embed_dim = u16::from_le_bytes([data[6], data[7]]) as usize;
338        let hidden_dim = u16::from_le_bytes([data[8], data[9]]) as usize;
339        let num_layers = data[10] as usize;
340        let num_heads = data[11] as usize;
341        let max_seq_len = data[12] as usize;
342        let quant_type = match data[13] {
343            0 => QuantizationType::Int8,
344            1 => QuantizationType::Int4,
345            2 => QuantizationType::Binary,
346            3 => QuantizationType::Fixed16,
347            _ => return Err(crate::Error::InvalidModel("Unknown quantization type")),
348        };
349
350        let config = ModelConfig {
351            vocab_size,
352            embed_dim,
353            hidden_dim,
354            num_layers,
355            num_heads,
356            max_seq_len,
357            quant_type,
358        };
359
360        config.validate(crate::Esp32Variant::Esp32)?;
361
362        // For now, create random weights - real implementation would parse from data
363        Self::new(config)
364    }
365
366    /// Export model to bytes
367    pub fn to_bytes(&self) -> HVec<u8, 256> {
368        let mut header: HVec<u8, 256> = HVec::new();
369
370        // Magic number
371        let _ = header.extend_from_slice(b"RUVM");
372
373        // Config
374        let _ = header.extend_from_slice(&(self.config.vocab_size as u16).to_le_bytes());
375        let _ = header.extend_from_slice(&(self.config.embed_dim as u16).to_le_bytes());
376        let _ = header.extend_from_slice(&(self.config.hidden_dim as u16).to_le_bytes());
377        let _ = header.push(self.config.num_layers as u8);
378        let _ = header.push(self.config.num_heads as u8);
379        let _ = header.push(self.config.max_seq_len as u8);
380        let _ = header.push(match self.config.quant_type {
381            QuantizationType::Int8 => 0,
382            QuantizationType::Int4 => 1,
383            QuantizationType::Binary => 2,
384            QuantizationType::Fixed16 => 3,
385        });
386
387        // Padding to 32 bytes
388        while header.len() < 32 {
389            let _ = header.push(0);
390        }
391
392        header
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    #[test]
401    fn test_default_config() {
402        let config = ModelConfig::default();
403        assert!(config.validate(crate::Esp32Variant::Esp32S2).is_ok());
404
405        let size = config.estimate_size();
406        println!("Default model size: {} bytes ({:.1} KB)", size, size as f32 / 1024.0);
407        assert!(size < 50 * 1024); // < 50KB for testing
408    }
409
410    #[test]
411    fn test_variant_configs() {
412        for variant in [
413            crate::Esp32Variant::Esp32,
414            crate::Esp32Variant::Esp32S2,
415            crate::Esp32Variant::Esp32S3,
416            crate::Esp32Variant::Esp32C3,
417            crate::Esp32Variant::Esp32C6,
418        ] {
419            let config = ModelConfig::for_variant(variant);
420            assert!(config.validate(variant).is_ok());
421
422            let size = config.estimate_size();
423            println!("{:?}: {} bytes ({:.1} KB)", variant, size, size as f32 / 1024.0);
424        }
425    }
426
427    #[test]
428    fn test_model_creation() {
429        let config = ModelConfig::default();
430        let model = TinyModel::new(config).unwrap();
431
432        let size = model.memory_size();
433        println!("Actual model size: {} bytes ({:.1} KB)", size, size as f32 / 1024.0);
434    }
435
436    #[test]
437    fn test_serialization() {
438        let config = ModelConfig::default();
439        let model = TinyModel::new(config).unwrap();
440
441        let header = model.to_bytes();
442        assert_eq!(&header[0..4], b"RUVM");
443    }
444}
ruvllm_esp32/model.rs

ruvllm_esp32/
model.rs
