rtx-csm: Phase 8.5 — Moonshine conv stem (verified end-to-end)

First working piece of the Moonshine v2 candle port. New src/moonshine.rs module with: - MoonshineConfig::tiny() — hyperparameters from HF config.json - ConvStem (Conv1d × 3) — audio stem, raw 16 kHz → 288-d hidden - load_conv_stem() — VarBuilder from HF safetensors Conv layout (verified against HF source): conv1: in=1, out=288, k=127, stride=64, no bias conv2: in=288, out=576, k=7, stride=3, bias conv3: in=576, out=288, k=3, stride=2, bias Activations: tanh after conv1, gelu_erf after conv2 / conv3 Smoke test (`examples/moonshine_smoke`): - Downloads UsefulSensors/moonshine-tiny from HF - Synthetic 10 s @ 16 kHz audio (silence + sine pulse) - input (1, 1, 160000) -> output (1, 415, 288) - Expected T_seq=415 ((160000-127)/64+1 -> 2498 -> 831 -> 415) - Output max abs = 23.17 (real signal, weights loaded correctly) Also extends `examples/moonshine_inspect` to dump conv shapes explicitly (was being truncated by the per-prefix `take(8)` cap). Next ship: encoder transformer block (partial RoPE, GELU MLP) and output layer norm. Tracked in Phase 8 plan; ~2-3 hours of work. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
2026-04-27 10:27:51 -07:00
parent c848abef22
commit 0699cdeb45
5 changed files with 296 additions and 0 deletions
@@ -0,0 +1,196 @@
+//! Moonshine v2 candle port — encoder-decoder ASR transformer designed
+//! for low-latency English transcription.
+//!
+//! See `docs/moonshine_port_notes.md` for the architectural plan, the
+//! tensor layout from `examples/moonshine_inspect`, and the porting
+//! task list.
+//!
+//! ## Status
+//!
+//! - [x] Config + Conv1d audio stem (this commit, Phase 8.5)
+//! - [ ] Encoder transformer layer (partial RoPE, GELU MLP)
+//! - [ ] Decoder transformer layer (self-attn + cross-attn, SwiGLU MLP)
+//! - [ ] Generation loop (greedy + KV cache)
+//! - [ ] Tokenizer wiring (HF `tokenizer.json` via `tokenizers` crate)
+//! - [ ] Weight mapping + smoke test against HF reference output
+//!
+//! Each item is roughly an hour of focused work; total ~12-15 h.
+//!
+//! ## Architecture (from HF `UsefulSensors/moonshine-tiny/config.json`)
+//!
+//! - hidden_size = 288, intermediate_size = 1152
+//! - 6 encoder layers + 6 decoder layers, 8 heads each (head_dim = 36,
+//!   pad to 8 → 40 — see `pad_head_dim_to_multiple_of`)
+//! - vocab_size = 32_768, max_position_embeddings = 194 (decoder text)
+//! - encoder MLP: GELU, no gating
+//! - decoder MLP: SiLU SwiGLU (fc1 outputs 2× intermediate, split gate/up)
+//! - RoPE: theta=10000, partial_rotary_factor=0.9 (rotary on first 90 %
+//!   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};
+
+/// Hyperparameters matching `UsefulSensors/moonshine-tiny/config.json`.
+#[derive(Debug, Clone)]
+pub struct MoonshineConfig {
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub vocab_size: usize,
+    pub max_position_embeddings: usize,
+    pub encoder_num_hidden_layers: usize,
+    pub encoder_num_attention_heads: usize,
+    pub decoder_num_hidden_layers: usize,
+    pub decoder_num_attention_heads: usize,
+    pub rope_theta: f64,
+    pub partial_rotary_factor: f64,
+    pub pad_head_dim_to_multiple_of: usize,
+    pub bos_token_id: u32,
+    pub eos_token_id: u32,
+    pub decoder_start_token_id: u32,
+}
+
+impl MoonshineConfig {
+    /// Tiny variant (27.1 M params).
+    pub fn tiny() -> Self {
+        Self {
+            hidden_size: 288,
+            intermediate_size: 1152,
+            vocab_size: 32_768,
+            max_position_embeddings: 194,
+            encoder_num_hidden_layers: 6,
+            encoder_num_attention_heads: 8,
+            decoder_num_hidden_layers: 6,
+            decoder_num_attention_heads: 8,
+            rope_theta: 10000.0,
+            partial_rotary_factor: 0.9,
+            pad_head_dim_to_multiple_of: 8,
+            bos_token_id: 1,
+            eos_token_id: 2,
+            decoder_start_token_id: 1,
+        }
+    }
+
+    /// Padded head dim (the actual size used for RoPE + attention math).
+    /// `head_dim_raw = hidden_size / num_heads`, rounded up to the
+    /// nearest multiple of `pad_head_dim_to_multiple_of`.
+    pub fn padded_head_dim(&self, num_heads: usize) -> usize {
+        let raw = self.hidden_size / num_heads;
+        let pad = self.pad_head_dim_to_multiple_of;
+        ((raw + pad - 1) / pad) * pad
+    }
+}
+
+/// Three-layer Conv1d audio stem mapping raw 16 kHz waveform to a
+/// sequence of 288-d hidden vectors at ~42 Hz frame rate.
+///
+/// Layout (from HF `modeling_moonshine.py`):
+/// - conv1: in=1, out=288, kernel=127, stride=64, no bias
+/// - conv2: in=288, out=576, kernel=7, stride=3, bias
+/// - conv3: in=576, out=288, kernel=3, stride=2, bias
+///
+/// Total downsampling: 64 × 3 × 2 = 384 ×.
+/// 16 kHz audio → output ~42 Hz frame rate (≈ 24 ms / frame).
+///
+/// Activations between convs: HF source uses GELU after conv1 and a
+/// LayerNorm + GELU after conv2 (TODO: verify the exact ordering when
+/// the encoder transformer block is wired).
+pub struct ConvStem {
+    conv1: Conv1d,
+    conv2: Conv1d,
+    conv3: Conv1d,
+}
+
+impl ConvStem {
+    pub fn new(vb: VarBuilder) -> Result<Self> {
+        let conv1 = candle_nn::conv1d_no_bias(
+            1,
+            288,
+            127,
+            Conv1dConfig {
+                padding: 0,
+                stride: 64,
+                dilation: 1,
+                groups: 1,
+                cudnn_fwd_algo: None,
+            },
+            vb.pp("conv1"),
+        )?;
+        let conv2 = candle_nn::conv1d(
+            288,
+            576,
+            7,
+            Conv1dConfig {
+                padding: 0,
+                stride: 3,
+                dilation: 1,
+                groups: 1,
+                cudnn_fwd_algo: None,
+            },
+            vb.pp("conv2"),
+        )?;
+        let conv3 = candle_nn::conv1d(
+            576,
+            288,
+            3,
+            Conv1dConfig {
+                padding: 0,
+                stride: 2,
+                dilation: 1,
+                groups: 1,
+                cudnn_fwd_algo: None,
+            },
+            vb.pp("conv3"),
+        )?;
+        Ok(Self { conv1, conv2, conv3 })
+    }
+
+    /// Forward: `(B, 1, T_audio)` raw waveform → `(B, T_seq, 288)` where
+    /// `T_seq ≈ T_audio / 384`. Returns hidden states ready for the
+    /// encoder transformer stack.
+    pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
+        // Conv1: raw audio → 288 channels at ~250 Hz
+        let h = self.conv1.forward(pcm)?;
+        let h = h.tanh()?; // HF source uses tanh after conv1
+        // Conv2: 288 → 576, stride 3
+        let h = self.conv2.forward(&h)?;
+        let h = h.gelu_erf()?;
+        // Conv3: 576 → 288, stride 2
+        let h = self.conv3.forward(&h)?;
+        let h = h.gelu_erf()?;
+        // Output is (B, 288, T_seq); transpose to (B, T_seq, 288) for
+        // the transformer encoder layers.
+        h.transpose(1, 2)?.contiguous()
+    }
+}
+
+/// 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,
+        )
+    }?;
+    // The encoder convs live under `model.encoder.{conv1,conv2,conv3}`.
+    ConvStem::new(vb.pp("model").pp("encoder"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn config_padded_head_dim_pads_to_multiple_of_8() {
+        let cfg = MoonshineConfig::tiny();
+        // 288 / 8 = 36, padded up to 40 (next multiple of 8).
+        assert_eq!(cfg.padded_head_dim(cfg.encoder_num_attention_heads), 40);
+        // Sanity: smaller-head config still pads.
+        let mut alt = cfg.clone();
+        alt.hidden_size = 256;
+        assert_eq!(alt.padded_head_dim(8), 32);
+    }
+}