diff --git a/README.md b/README.md index fbf0eb0..cb14dd9 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ burn inference and training of (baby) [dragon hatchling](https://arxiv.org/abs/2 ## features +- [x] mixture-of-expert routing - [x] training benchmarks and reporting - [ ] adaptive tool discovery - [ ] conditional (deep) gating @@ -21,7 +22,6 @@ burn inference and training of (baby) [dragon hatchling](https://arxiv.org/abs/2 - [ ] episodic memory - [ ] fused kernels - [ ] hierarchical, memory-aware recurrent state -- [ ] mixture-of-expert routing - [ ] multi-modal architecture - [ ] neuromorphic backend - [ ] rl reasoning training diff --git a/bench/inference.rs b/bench/inference.rs index 20d101a..e3e0068 100644 --- a/bench/inference.rs +++ b/bench/inference.rs @@ -142,11 +142,11 @@ fn log_theoretical_profile(config: &BDHConfig, cfg: &InferenceConfig) { } fn compute_latent_per_head(config: &BDHConfig) -> usize { - (config.mlp_internal_dim_multiplier * config.n_embd) / config.n_head + config.latent_per_head() } fn compute_latent_total(config: &BDHConfig) -> usize { - compute_latent_per_head(config) * config.n_head + config.latent_total() } fn estimated_query_tensor_bytes(config: &BDHConfig, cfg: &InferenceConfig) -> u128 { diff --git a/bench/train.rs b/bench/train.rs index d56d66f..7dea2b7 100644 --- a/bench/train.rs +++ b/bench/train.rs @@ -148,11 +148,11 @@ fn log_theoretical_profile(config: &BDHConfig, cfg: &TrainConfig) { } fn compute_latent_per_head(config: &BDHConfig) -> usize { - (config.mlp_internal_dim_multiplier * config.n_embd) / config.n_head + config.latent_per_head() } fn compute_latent_total(config: &BDHConfig) -> usize { - compute_latent_per_head(config) * config.n_head + config.latent_total() } criterion_group!(benches, training_step_bench); diff --git a/config/base.toml b/config/base.toml index 6d949dc..4b67e32 100644 --- a/config/base.toml +++ b/config/base.toml @@ -28,6 +28,7 @@ n_layer = 6 n_embd = 256 n_head = 4 mlp_internal_dim_multiplier = 128 +experts = 2 dropout = 0.1 fused_kernels = true use_alibi = true diff --git a/src/bin/infer.rs b/src/bin/infer.rs index 7e73b73..026c336 100644 --- a/src/bin/infer.rs +++ b/src/bin/infer.rs @@ -110,6 +110,9 @@ fn build_model_config(overrides: &ModelOverrides) -> BDHConfig { if let Some(n_head) = overrides.n_head { model_config.n_head = n_head; } + if let Some(experts) = overrides.experts { + model_config.n_expert = experts.max(1); + } if let Some(multiplier) = overrides.mlp_internal_dim_multiplier { model_config.mlp_internal_dim_multiplier = multiplier; } diff --git a/src/config/mod.rs b/src/config/mod.rs index 4c99d24..1fb1bd6 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -88,6 +88,7 @@ pub struct ModelOverrides { pub n_layer: Option, pub n_embd: Option, pub n_head: Option, + pub experts: Option, pub mlp_internal_dim_multiplier: Option, pub dropout: Option, pub fused_kernels: Option, @@ -224,6 +225,7 @@ mod tests { "n_layer = 6", "n_embd = 256", "n_head = 4", + "experts = 2", "mlp_internal_dim_multiplier = 128", "dropout = 0.1", "fused_kernels = false", @@ -278,6 +280,7 @@ mod tests { assert_eq!(config.model.n_layer, Some(6)); assert_eq!(config.model.n_embd, Some(320)); assert_eq!(config.model.n_head, Some(4)); + assert_eq!(config.model.experts, Some(2)); assert_eq!(config.model.mlp_internal_dim_multiplier, Some(128)); assert_eq!(config.model.dropout, Some(0.1)); assert_eq!(config.model.fused_kernels, Some(true)); diff --git a/src/model/bdh.rs b/src/model/bdh.rs index a3c0d59..0d2da77 100644 --- a/src/model/bdh.rs +++ b/src/model/bdh.rs @@ -10,6 +10,7 @@ use crate::kernel::{BlockPattern1d, relu_lowrank}; use super::attention::Attention; use super::config::{BDHConfig, FusedKernelConfig}; +use super::router::Router; const LAYER_NORM_EPS: f32 = 1e-5; @@ -19,11 +20,13 @@ pub struct BDH { n_embd: usize, n_head: usize, mlp_internal_dim_multiplier: usize, + n_expert: usize, vocab_size: usize, kernel: FusedKernelConfig, embed: Embedding, dropout: Dropout, attention: Attention, + router: Router, encoder: Param>, encoder_v: Param>, decoder: Param>, @@ -43,6 +46,13 @@ impl BDH { device, &config.fused_kernels, ); + let router = Router::new( + config.n_expert, + config.n_head, + config.n_embd, + latent_per_head, + device, + ); let weight_init = |shape: [usize; 2]| { Tensor::::random(shape, TensorDistribution::Normal(0.0, 0.02), device) @@ -68,11 +78,13 @@ impl BDH { n_embd: config.n_embd, n_head: config.n_head, mlp_internal_dim_multiplier: config.mlp_internal_dim_multiplier, + n_expert: config.n_expert, vocab_size: config.vocab_size, kernel: config.fused_kernels, embed, dropout, attention, + router, encoder, encoder_v, decoder, @@ -131,6 +143,7 @@ impl BDH { activation::relu(y_latent) }; let xy_sparse = x_sparse * y_sparse; + let xy_sparse = self.router.route(state.clone(), xy_sparse); let xy_sparse = self.dropout.forward(xy_sparse); let mixed = xy_sparse.swap_dims(1, 2); diff --git a/src/model/mod.rs b/src/model/mod.rs index bbed98d..a20ce45 100644 --- a/src/model/mod.rs +++ b/src/model/mod.rs @@ -2,6 +2,7 @@ mod attention; mod bdh; mod config; mod loss; +mod router; pub use bdh::BDH; pub use config::{BDHConfig, FusedKernelConfig}; diff --git a/src/model/router.rs b/src/model/router.rs index e69de29..cf49ab9 100644 --- a/src/model/router.rs +++ b/src/model/router.rs @@ -0,0 +1,79 @@ +use burn::module::{Module, Param}; +use burn::tensor::backend::Backend; +use burn::tensor::{Distribution as TensorDistribution, Tensor, activation}; + +#[derive(Module, Debug)] +pub struct Router { + experts: usize, + n_head: usize, + latent_per_head: usize, + latent_per_expert: usize, + weight: Param>, + bias: Param>, +} + +impl Router { + pub fn new( + experts: usize, + n_head: usize, + n_embd: usize, + latent_per_head: usize, + device: &B::Device, + ) -> Self { + assert!(experts >= 1, "router requires at least one expert"); + assert!( + latent_per_head % experts == 0, + "latent size {latent_per_head} must be divisible by experts {experts}" + ); + + let weight = Tensor::::random( + [n_head, n_embd, experts], + TensorDistribution::Normal(0.0, 0.02), + device, + ); + let bias = Tensor::::zeros([n_head, experts], device); + + Self { + experts, + n_head, + latent_per_head, + latent_per_expert: latent_per_head / experts, + weight: Param::from_tensor(weight), + bias: Param::from_tensor(bias), + } + } + + pub fn route(&self, gating_input: Tensor, activations: Tensor) -> Tensor { + if self.experts == 1 { + return activations; + } + + let weight = self.weight.val().unsqueeze_dim::<4>(0); + let mut logits = gating_input.matmul(weight); + let bias = self.bias.val().reshape([1, self.n_head, 1, self.experts]); + logits = logits + bias; + + let routing = activation::softmax(logits, 3); + + let [batch, heads, time, latent] = activations.shape().dims(); + debug_assert_eq!(heads, self.n_head); + debug_assert_eq!(latent, self.latent_per_head); + + let routed = activations.reshape([ + batch, + self.n_head, + time, + self.experts, + self.latent_per_expert, + ]); + + let routing = routing.unsqueeze_dim::<5>(4); + let weighted = routing * routed; + + weighted.reshape([batch, self.n_head, time, self.latent_per_head]) + } + + pub fn experts(&self) -> usize { + self.experts + } +} diff --git a/src/train.rs b/src/train.rs index 4e3a8f7..a3afb36 100644 --- a/src/train.rs +++ b/src/train.rs @@ -468,6 +468,9 @@ fn build_model_config(overrides: &ModelOverrides) -> BDHConfig { if let Some(n_head) = overrides.n_head { model_config.n_head = n_head; } + if let Some(experts) = overrides.experts { + model_config.n_expert = experts.max(1); + } if let Some(multiplier) = overrides.mlp_internal_dim_multiplier { model_config.mlp_internal_dim_multiplier = multiplier; }