rtx-csm: Phase 8.6 — Moonshine encoder transformer block
Full encoder forward path: conv stem -> 6 transformer layers -> final
LayerNorm. Loads HF safetensors, runs end-to-end on Metal.
Components added to src/moonshine.rs:
RotaryCache partial RoPE (32 of 36 head_dim, theta=10000)
EncoderAttention MHA (8 heads, no bias), partial RoPE on q/k
EncoderMlp 288 -> 1152 -> 288 with bias, GELU(erf) activation
EncoderLayer Pre-LN attn + Pre-LN MLP (LayerNorm weight-only)
Encoder stem + 6 layers + final LayerNorm
load_encoder() VarBuilder convenience for the standalone smoke
Smoke test (`examples/moonshine_smoke`) verified end-to-end:
input (1, 1, 160000) -> output (1, 415, 288)
forward: 132 ms (10 s of audio at 0.013x realtime)
max abs: 6.67 (signal preserved, not zeros)
Implementation notes captured in the diff:
- candle Metal 4D batched matmul had shape-mismatch issues for our
(B, H, T, D) pattern. Switched to (B*H, T, D) 3D form which is
unambiguous and avoids the kernel bug.
- LayerNorm is weight-only (no bias tensors in safetensors); we
construct LayerNorm with a zeros bias to satisfy candle's API.
- rotary_dim = floor(head_dim * 0.9 / 2) * 2 = 32 (must be even).
The remaining 4 head_dim channels pass through unchanged via
`narrow + cat` on dim 3.
Numerical parity vs HF Python reference is NOT yet verified — that's
the next bounded chunk (Phase 8.7). Shape + signal correctness are
verified by the smoke test.
Next: decoder transformer block (self-attn + cross-attn + SwiGLU).
~3-4 h of focused work.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -28,8 +28,8 @@
|
||||
//! of head_dim)
|
||||
//! - bos = 1, eos = 2, pad = 2, decoder_start = 1
|
||||
|
||||
use candle_core::{Device, Module, Result, Tensor};
|
||||
use candle_nn::{Conv1d, Conv1dConfig, VarBuilder};
|
||||
use candle_core::{DType, Device, Module, Result, Tensor};
|
||||
use candle_nn::{Conv1d, Conv1dConfig, LayerNorm, Linear, VarBuilder};
|
||||
|
||||
/// Hyperparameters matching `UsefulSensors/moonshine-tiny/config.json`.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -167,18 +167,228 @@ impl ConvStem {
|
||||
/// Loader: open the HF safetensors and construct a `ConvStem`. Useful
|
||||
/// for the standalone Phase 8.5 smoke test.
|
||||
pub fn load_conv_stem(weights_path: &std::path::Path, device: &Device) -> Result<ConvStem> {
|
||||
// Open the safetensors directly via VarBuilder. F32 storage.
|
||||
let vb = unsafe {
|
||||
VarBuilder::from_mmaped_safetensors(
|
||||
&[weights_path],
|
||||
candle_core::DType::F32,
|
||||
device,
|
||||
)
|
||||
VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
|
||||
}?;
|
||||
// The encoder convs live under `model.encoder.{conv1,conv2,conv3}`.
|
||||
ConvStem::new(vb.pp("model").pp("encoder"))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// Phase 8.6 — encoder transformer block
|
||||
// ---------------------------------------------------------------------
|
||||
|
||||
/// Cosine/sine tables for partial RoPE.
|
||||
///
|
||||
/// Moonshine uses `partial_rotary_factor = 0.9` with `head_dim = 36`
|
||||
/// (288 / 8 heads). `rotary_dim = floor(36 * 0.9 / 2) * 2 = 32` (must
|
||||
/// be even — RoPE rotates pairs). The remaining 4 dims pass through
|
||||
/// unchanged. RoPE base is `theta = 10_000`.
|
||||
struct RotaryCache {
|
||||
cos: Tensor, // (max_seq, rotary_dim/2)
|
||||
sin: Tensor, // (max_seq, rotary_dim/2)
|
||||
rotary_dim: usize,
|
||||
}
|
||||
|
||||
impl RotaryCache {
|
||||
fn new(rotary_dim: usize, max_seq: usize, theta: f64, dtype: DType, dev: &Device) -> Result<Self> {
|
||||
assert!(rotary_dim % 2 == 0, "rotary_dim must be even");
|
||||
let inv_freq: Vec<f32> = (0..rotary_dim)
|
||||
.step_by(2)
|
||||
.map(|i| (1.0 / theta.powf(i as f64 / rotary_dim as f64)) as f32)
|
||||
.collect();
|
||||
let inv_freq = Tensor::new(inv_freq, dev)?;
|
||||
let positions = Tensor::arange(0u32, max_seq as u32, dev)?
|
||||
.to_dtype(DType::F32)?
|
||||
.reshape((max_seq, 1))?;
|
||||
let freqs = positions.matmul(&inv_freq.reshape((1, rotary_dim / 2))?)?;
|
||||
let cos = freqs.cos()?.to_dtype(dtype)?;
|
||||
let sin = freqs.sin()?.to_dtype(dtype)?;
|
||||
Ok(Self { cos, sin, rotary_dim })
|
||||
}
|
||||
|
||||
/// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`.
|
||||
/// Rotates the first `rotary_dim` channels; leaves the rest as-is.
|
||||
fn apply(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let (_b, _h, t, head_dim) = x.dims4()?;
|
||||
let cos = self.cos.narrow(0, 0, t)?;
|
||||
let sin = self.sin.narrow(0, 0, t)?;
|
||||
if head_dim == self.rotary_dim {
|
||||
// Full rotary — straightforward path.
|
||||
return candle_nn::rotary_emb::rope_i(x, &cos, &sin);
|
||||
}
|
||||
// Partial rotary: split into rotary head and pass-through.
|
||||
let rot = x.narrow(3, 0, self.rotary_dim)?.contiguous()?;
|
||||
let pass = x.narrow(3, self.rotary_dim, head_dim - self.rotary_dim)?;
|
||||
let rot = candle_nn::rotary_emb::rope_i(&rot, &cos, &sin)?;
|
||||
Tensor::cat(&[&rot, &pass], 3)
|
||||
}
|
||||
}
|
||||
|
||||
/// Self-attention block matching `model.encoder.layers.X.self_attn`.
|
||||
/// 8 heads × 36 head_dim, no bias (`attention_bias=false`), partial
|
||||
/// RoPE on first 32 head_dim channels.
|
||||
struct EncoderAttention {
|
||||
q_proj: Linear,
|
||||
k_proj: Linear,
|
||||
v_proj: Linear,
|
||||
o_proj: Linear,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
}
|
||||
|
||||
impl EncoderAttention {
|
||||
fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
|
||||
let h = cfg.hidden_size;
|
||||
let n_heads = cfg.encoder_num_attention_heads;
|
||||
let head_dim = h / n_heads;
|
||||
let q_proj = candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?;
|
||||
let k_proj = candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?;
|
||||
let v_proj = candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?;
|
||||
let o_proj = candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?;
|
||||
Ok(Self { q_proj, k_proj, v_proj, o_proj, n_heads, head_dim })
|
||||
}
|
||||
|
||||
fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
|
||||
let (b, t, _h) = xs.dims3()?;
|
||||
let q = self.q_proj.forward(xs)?;
|
||||
let k = self.k_proj.forward(xs)?;
|
||||
let v = self.v_proj.forward(xs)?;
|
||||
// (B, T, H) -> (B, H_heads, T, head_dim)
|
||||
let q = q.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?;
|
||||
let k = k.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?;
|
||||
let v = v.reshape((b, t, self.n_heads, self.head_dim))?.transpose(1, 2)?.contiguous()?;
|
||||
let q = rope.apply(&q)?;
|
||||
let k = rope.apply(&k)?;
|
||||
// Collapse (B, H, T, D) -> (B*H, T, D) for the matmul. candle's
|
||||
// 4D batched matmul on Metal had shape-mismatch issues for our
|
||||
// pattern; the 3D form is unambiguous.
|
||||
let bh = b * self.n_heads;
|
||||
let q3 = q.reshape((bh, t, self.head_dim))?;
|
||||
let k3 = k.reshape((bh, t, self.head_dim))?;
|
||||
let v3 = v.reshape((bh, t, self.head_dim))?;
|
||||
let scale = 1.0 / (self.head_dim as f64).sqrt();
|
||||
let scores = (q3.matmul(&k3.transpose(1, 2)?.contiguous()?)? * scale)?;
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
let out = probs.matmul(&v3)?;
|
||||
// (B*H, T, D) -> (B, H, T, D) -> (B, T, H*D)
|
||||
let out = out
|
||||
.reshape((b, self.n_heads, t, self.head_dim))?
|
||||
.transpose(1, 2)?
|
||||
.contiguous()?
|
||||
.reshape((b, t, self.n_heads * self.head_dim))?;
|
||||
self.o_proj.forward(&out)
|
||||
}
|
||||
}
|
||||
|
||||
/// FFN block matching `model.encoder.layers.X.mlp`.
|
||||
/// 288 → 1152 (with bias) → GELU → 1152 → 288 (with bias).
|
||||
struct EncoderMlp {
|
||||
fc1: Linear,
|
||||
fc2: Linear,
|
||||
}
|
||||
|
||||
impl EncoderMlp {
|
||||
fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
|
||||
let fc1 = candle_nn::linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?;
|
||||
let fc2 = candle_nn::linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?;
|
||||
Ok(Self { fc1, fc2 })
|
||||
}
|
||||
|
||||
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
||||
let h = self.fc1.forward(xs)?;
|
||||
let h = h.gelu_erf()?;
|
||||
self.fc2.forward(&h)
|
||||
}
|
||||
}
|
||||
|
||||
/// One encoder layer: Pre-LN attn + Pre-LN MLP, residuals as standard.
|
||||
/// Tensors are `LayerNorm` weight-only (no bias) per the inspector
|
||||
/// dump — implemented via `LayerNorm` with bias zeros.
|
||||
struct EncoderLayer {
|
||||
input_ln: LayerNorm,
|
||||
self_attn: EncoderAttention,
|
||||
post_attn_ln: LayerNorm,
|
||||
mlp: EncoderMlp,
|
||||
}
|
||||
|
||||
impl EncoderLayer {
|
||||
fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
|
||||
let h = cfg.hidden_size;
|
||||
let input_ln = layer_norm_weight_only(h, 1e-5, vb.pp("input_layernorm"))?;
|
||||
let self_attn = EncoderAttention::new(cfg, vb.pp("self_attn"))?;
|
||||
let post_attn_ln = layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?;
|
||||
let mlp = EncoderMlp::new(cfg, vb.pp("mlp"))?;
|
||||
Ok(Self { input_ln, self_attn, post_attn_ln, mlp })
|
||||
}
|
||||
|
||||
fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
|
||||
let h = self.input_ln.forward(xs)?;
|
||||
let h = self.self_attn.forward(&h, rope)?;
|
||||
let xs = (xs + h)?;
|
||||
let h = self.post_attn_ln.forward(&xs)?;
|
||||
let h = self.mlp.forward(&h)?;
|
||||
xs + h
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper: `LayerNorm` with weight only (zeros bias). Moonshine's
|
||||
/// `*_layernorm.weight [288]` tensors lack a corresponding `.bias`.
|
||||
fn layer_norm_weight_only(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
|
||||
let weight = vb.get((size,), "weight")?;
|
||||
let bias = Tensor::zeros((size,), weight.dtype(), weight.device())?;
|
||||
Ok(LayerNorm::new(weight, bias, eps))
|
||||
}
|
||||
|
||||
/// Full encoder: conv stem → 6 transformer layers → final layer norm.
|
||||
/// Output shape: `(B, T_seq, hidden=288)` ready to feed cross-attn in
|
||||
/// the decoder (Phase 8.7).
|
||||
pub struct Encoder {
|
||||
stem: ConvStem,
|
||||
layers: Vec<EncoderLayer>,
|
||||
final_ln: LayerNorm,
|
||||
rope: RotaryCache,
|
||||
}
|
||||
|
||||
impl Encoder {
|
||||
pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
|
||||
let stem = ConvStem::new(vb.pp("encoder"))?;
|
||||
let head_dim = cfg.hidden_size / cfg.encoder_num_attention_heads;
|
||||
let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
|
||||
// max_seq for the encoder is bounded by the input audio length /
|
||||
// total stride 384. For 30 s of 16 kHz audio that's ~1250.
|
||||
// Allocate a generous table.
|
||||
let rope = RotaryCache::new(rotary_dim, 4096, cfg.rope_theta, vb.dtype(), vb.device())?;
|
||||
let mut layers = Vec::with_capacity(cfg.encoder_num_hidden_layers);
|
||||
let layer_vb = vb.pp("encoder").pp("layers");
|
||||
for i in 0..cfg.encoder_num_hidden_layers {
|
||||
layers.push(EncoderLayer::new(cfg, layer_vb.pp(i))?);
|
||||
}
|
||||
let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("encoder").pp("layer_norm"))?;
|
||||
Ok(Self { stem, layers, final_ln, rope })
|
||||
}
|
||||
|
||||
pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
|
||||
let mut h = self.stem.forward(pcm)?;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h, &self.rope)?;
|
||||
}
|
||||
self.final_ln.forward(&h)
|
||||
}
|
||||
}
|
||||
|
||||
/// Loader: open the HF safetensors and construct an `Encoder` with
|
||||
/// weights at `model.*` (the canonical HF prefix).
|
||||
pub fn load_encoder(
|
||||
weights_path: &std::path::Path,
|
||||
device: &Device,
|
||||
cfg: &MoonshineConfig,
|
||||
) -> Result<Encoder> {
|
||||
let vb = unsafe {
|
||||
VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
|
||||
}?;
|
||||
Encoder::new(cfg, vb.pp("model"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
Reference in New Issue
Block a user