376 lines
11 KiB
Rust
376 lines
11 KiB
Rust
//! WASM Model Loading and Inference
|
|
//!
|
|
//! Model loading and forward pass for WebAssembly inference.
|
|
|
|
use crate::tensor::WasmKvCache;
|
|
use serde::{Deserialize, Serialize};
|
|
use wasm_bindgen::prelude::*;
|
|
|
|
/// Model configuration
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelConfig {
|
|
/// Vocabulary size
|
|
pub vocab_size: usize,
|
|
/// Hidden dimension
|
|
pub hidden_size: usize,
|
|
/// Number of layers
|
|
pub num_layers: usize,
|
|
/// Number of attention heads
|
|
pub num_heads: usize,
|
|
/// Head dimension
|
|
pub head_dim: usize,
|
|
/// Intermediate size (FFN)
|
|
pub intermediate_size: usize,
|
|
/// Maximum sequence length
|
|
pub max_seq_len: usize,
|
|
/// RMS norm epsilon
|
|
pub rms_norm_eps: f32,
|
|
/// Whether model is quantized
|
|
pub quantized: bool,
|
|
}
|
|
|
|
impl Default for ModelConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
vocab_size: 32000,
|
|
hidden_size: 4096,
|
|
num_layers: 32,
|
|
num_heads: 32,
|
|
head_dim: 128,
|
|
intermediate_size: 11008,
|
|
max_seq_len: 2048,
|
|
rms_norm_eps: 1e-6,
|
|
quantized: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl ModelConfig {
|
|
/// Create a small model config for testing
|
|
#[wasm_bindgen]
|
|
pub fn tiny() -> Self {
|
|
Self {
|
|
vocab_size: 32000,
|
|
hidden_size: 512,
|
|
num_layers: 4,
|
|
num_heads: 8,
|
|
head_dim: 64,
|
|
intermediate_size: 1376,
|
|
max_seq_len: 512,
|
|
rms_norm_eps: 1e-6,
|
|
quantized: false,
|
|
}
|
|
}
|
|
|
|
/// Create a 1B-like model config
|
|
#[wasm_bindgen]
|
|
pub fn small() -> Self {
|
|
Self {
|
|
vocab_size: 32000,
|
|
hidden_size: 2048,
|
|
num_layers: 16,
|
|
num_heads: 16,
|
|
head_dim: 128,
|
|
intermediate_size: 5504,
|
|
max_seq_len: 2048,
|
|
rms_norm_eps: 1e-6,
|
|
quantized: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Layer weights
|
|
#[derive(Debug, Clone)]
|
|
struct LayerWeights {
|
|
/// Attention QKV projection
|
|
attn_qkv: Vec<f32>,
|
|
/// Attention output projection
|
|
attn_out: Vec<f32>,
|
|
/// FFN gate projection
|
|
ffn_gate: Vec<f32>,
|
|
/// FFN up projection
|
|
ffn_up: Vec<f32>,
|
|
/// FFN down projection
|
|
ffn_down: Vec<f32>,
|
|
/// Pre-attention norm
|
|
attn_norm: Vec<f32>,
|
|
/// Pre-FFN norm
|
|
ffn_norm: Vec<f32>,
|
|
}
|
|
|
|
/// WASM model for inference
|
|
#[wasm_bindgen]
|
|
pub struct WasmModel {
|
|
/// Model configuration
|
|
config: ModelConfig,
|
|
/// Embedding weights
|
|
embeddings: Vec<f32>,
|
|
/// Layer weights
|
|
layers: Vec<LayerWeights>,
|
|
/// Final norm
|
|
final_norm: Vec<f32>,
|
|
/// Output projection (lm_head)
|
|
lm_head: Vec<f32>,
|
|
/// Whether model is loaded
|
|
loaded: bool,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl WasmModel {
|
|
/// Create a new uninitialized model
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn new(config: ModelConfig) -> Self {
|
|
Self {
|
|
config,
|
|
embeddings: Vec::new(),
|
|
layers: Vec::new(),
|
|
final_norm: Vec::new(),
|
|
lm_head: Vec::new(),
|
|
loaded: false,
|
|
}
|
|
}
|
|
|
|
/// Load model from bytes
|
|
pub fn from_bytes(bytes: &[u8], quantized: bool) -> Result<Self, JsError> {
|
|
// Parse model format (simplified - would need actual format parsing)
|
|
// For now, create a dummy model for structure validation
|
|
|
|
if bytes.len() < 16 {
|
|
return Err(JsError::new("Invalid model data: too short"));
|
|
}
|
|
|
|
// Read config from header (simplified)
|
|
let config = ModelConfig {
|
|
quantized,
|
|
..ModelConfig::tiny()
|
|
};
|
|
|
|
let mut model = Self::new(config.clone());
|
|
|
|
// Initialize with random weights for now
|
|
// In production, parse actual weights from bytes
|
|
model.init_random_weights();
|
|
model.loaded = true;
|
|
|
|
Ok(model)
|
|
}
|
|
|
|
/// Initialize with random weights (for testing)
|
|
fn init_random_weights(&mut self) {
|
|
let c = &self.config;
|
|
|
|
// Embeddings
|
|
self.embeddings = (0..c.vocab_size * c.hidden_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect();
|
|
|
|
// Layers
|
|
self.layers = (0..c.num_layers)
|
|
.map(|_| LayerWeights {
|
|
attn_qkv: (0..c.hidden_size * c.hidden_size * 3)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect(),
|
|
attn_out: (0..c.hidden_size * c.hidden_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect(),
|
|
ffn_gate: (0..c.hidden_size * c.intermediate_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect(),
|
|
ffn_up: (0..c.hidden_size * c.intermediate_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect(),
|
|
ffn_down: (0..c.intermediate_size * c.hidden_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect(),
|
|
attn_norm: vec![1.0; c.hidden_size],
|
|
ffn_norm: vec![1.0; c.hidden_size],
|
|
})
|
|
.collect();
|
|
|
|
// Final norm
|
|
self.final_norm = vec![1.0; c.hidden_size];
|
|
|
|
// LM head
|
|
self.lm_head = (0..c.hidden_size * c.vocab_size)
|
|
.map(|_| (rand::random::<f32>() - 0.5) * 0.02)
|
|
.collect();
|
|
}
|
|
|
|
/// Check if model is loaded
|
|
#[wasm_bindgen]
|
|
pub fn is_loaded(&self) -> bool {
|
|
self.loaded
|
|
}
|
|
|
|
/// Get model config
|
|
#[wasm_bindgen]
|
|
pub fn config(&self) -> ModelConfig {
|
|
self.config.clone()
|
|
}
|
|
|
|
/// Memory usage in bytes
|
|
#[wasm_bindgen]
|
|
pub fn memory_usage(&self) -> usize {
|
|
let mut total = 0;
|
|
total += self.embeddings.len() * 4;
|
|
for layer in &self.layers {
|
|
total += layer.attn_qkv.len() * 4;
|
|
total += layer.attn_out.len() * 4;
|
|
total += layer.ffn_gate.len() * 4;
|
|
total += layer.ffn_up.len() * 4;
|
|
total += layer.ffn_down.len() * 4;
|
|
total += layer.attn_norm.len() * 4;
|
|
total += layer.ffn_norm.len() * 4;
|
|
}
|
|
total += self.final_norm.len() * 4;
|
|
total += self.lm_head.len() * 4;
|
|
total
|
|
}
|
|
|
|
/// Get parameter count
|
|
#[wasm_bindgen]
|
|
pub fn param_count(&self) -> usize {
|
|
let c = &self.config;
|
|
|
|
let embeddings = c.vocab_size * c.hidden_size;
|
|
let per_layer = c.hidden_size * c.hidden_size * 3 // QKV
|
|
+ c.hidden_size * c.hidden_size // O
|
|
+ c.hidden_size * c.intermediate_size * 3 // gate, up, down
|
|
+ c.hidden_size * 2; // norms
|
|
let lm_head = c.hidden_size * c.vocab_size;
|
|
|
|
embeddings + per_layer * c.num_layers + c.hidden_size + lm_head
|
|
}
|
|
}
|
|
|
|
// Internal methods (not exposed to WASM)
|
|
impl WasmModel {
|
|
/// Forward pass
|
|
pub fn forward(
|
|
&self,
|
|
input_ids: &[u32],
|
|
_kv_cache: Option<&mut WasmKvCache>,
|
|
) -> Result<Vec<f32>, JsError> {
|
|
if !self.loaded {
|
|
return Err(JsError::new("Model not loaded"));
|
|
}
|
|
|
|
let seq_len = input_ids.len();
|
|
let c = &self.config;
|
|
|
|
// Embed tokens
|
|
let mut hidden = vec![0.0; seq_len * c.hidden_size];
|
|
for (i, &token_id) in input_ids.iter().enumerate() {
|
|
if token_id as usize >= c.vocab_size {
|
|
return Err(JsError::new(&format!("Token ID {} out of range", token_id)));
|
|
}
|
|
let start = token_id as usize * c.hidden_size;
|
|
let end = start + c.hidden_size;
|
|
hidden[i * c.hidden_size..(i + 1) * c.hidden_size]
|
|
.copy_from_slice(&self.embeddings[start..end]);
|
|
}
|
|
|
|
// Process layers (simplified)
|
|
for layer in &self.layers {
|
|
// RMS Norm + Attention (simplified)
|
|
hidden = self.rms_norm(&hidden, &layer.attn_norm);
|
|
hidden = self.attention(&hidden);
|
|
|
|
// RMS Norm + FFN (simplified)
|
|
hidden = self.rms_norm(&hidden, &layer.ffn_norm);
|
|
hidden = self.ffn(&hidden, layer);
|
|
}
|
|
|
|
// Final norm
|
|
hidden = self.rms_norm(&hidden, &self.final_norm);
|
|
|
|
// Compute logits for last token only
|
|
let last_hidden = &hidden[(seq_len - 1) * c.hidden_size..seq_len * c.hidden_size];
|
|
let logits = self.compute_logits(last_hidden);
|
|
|
|
Ok(logits)
|
|
}
|
|
|
|
fn rms_norm(&self, x: &[f32], weight: &[f32]) -> Vec<f32> {
|
|
let hidden_size = self.config.hidden_size;
|
|
let num_tokens = x.len() / hidden_size;
|
|
let mut result = vec![0.0; x.len()];
|
|
|
|
for t in 0..num_tokens {
|
|
let start = t * hidden_size;
|
|
let end = start + hidden_size;
|
|
let slice = &x[start..end];
|
|
|
|
// Compute RMS
|
|
let rms: f32 = (slice.iter().map(|v| v * v).sum::<f32>() / hidden_size as f32
|
|
+ self.config.rms_norm_eps)
|
|
.sqrt();
|
|
|
|
// Normalize and apply weight
|
|
for (i, &v) in slice.iter().enumerate() {
|
|
result[start + i] = (v / rms) * weight[i];
|
|
}
|
|
}
|
|
|
|
result
|
|
}
|
|
|
|
fn attention(&self, x: &[f32]) -> Vec<f32> {
|
|
// Simplified attention - in production would use actual attention mechanism
|
|
// For now, just return input scaled (placeholder)
|
|
x.iter().map(|v| v * 0.9).collect()
|
|
}
|
|
|
|
fn ffn(&self, x: &[f32], _layer: &LayerWeights) -> Vec<f32> {
|
|
// Simplified FFN: SiLU(x @ gate) * (x @ up) @ down
|
|
// For now, just apply activation (placeholder)
|
|
x.iter()
|
|
.map(|v| {
|
|
let silu = v / (1.0 + (-v).exp());
|
|
silu * 0.9
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn compute_logits(&self, hidden: &[f32]) -> Vec<f32> {
|
|
let c = &self.config;
|
|
let mut logits = vec![0.0; c.vocab_size];
|
|
|
|
// Matrix multiply: hidden @ lm_head^T
|
|
for v in 0..c.vocab_size {
|
|
let mut sum = 0.0;
|
|
for h in 0..c.hidden_size {
|
|
sum += hidden[h] * self.lm_head[v * c.hidden_size + h];
|
|
}
|
|
logits[v] = sum;
|
|
}
|
|
|
|
logits
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_config() {
|
|
let config = ModelConfig::tiny();
|
|
assert_eq!(config.num_layers, 4);
|
|
assert_eq!(config.hidden_size, 512);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_memory() {
|
|
let config = ModelConfig::tiny();
|
|
let mut model = WasmModel::new(config);
|
|
model.init_random_weights();
|
|
model.loaded = true;
|
|
|
|
assert!(model.memory_usage() > 0);
|
|
assert!(model.param_count() > 0);
|
|
}
|
|
}
|