feat(batch13): EMA model weights, cross-layer weight sharing, schedule-free optimizer

- ModelEma: decay-weighted shadow weights with warmup ramp, bias correction,
  apply/restore swap for eval, and shadow_drift L2 diagnostic
- SharedLayerStack: FullSharing/GroupedSharing/AlternatingPairs strategies
  (ALBERT-style); memory_reduction_ratio(); LCG-seeded SharedFfnWeight
- ScheduleFreeOptimizer: Defazio 2024 z/x dual sequences, c_t cubic
  interpolation coefficient, Adam+SGD variants, weight decay

48 tests + 5 doctests

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/training/rtx-transformers/src/layers/mod.rs

@@ -67,6 +67,9 @@ pub mod ring_attention;
 
 // Token Merging (ToMe) — bipartite soft matching for ViT-style speedup
 pub mod token_merging;
+
+// Cross-layer parameter sharing (ALBERT-style)
+pub mod shared_layers;
 // pub mod mamba_integration;
 // pub mod mamba_cuda_kernels;
 // pub mod mamba_cuda_integration;
@@ -233,6 +236,7 @@ pub use token_merging::{
     apply_merge, apply_unmerge, bipartite_soft_matching, MergeResult, TokenMergingConfig,
     TokenMergingLayer,
 };
+pub use shared_layers::{SharedFfnWeight, SharedLayerConfig, SharedLayerStack, SharingStrategy};
 
 // TransformerConfig is defined above and available for import
 

crates/training/rtx-transformers/src/layers/shared_layers.rs

@@ -0,0 +1,645 @@
+//! Cross-layer parameter sharing (ALBERT-style).
+//!
+//! ALBERT (arXiv:1909.11942) reduces transformer memory by sharing weight tensors
+//! across multiple layers. Instead of N independent weight sets, you have G groups
+//! where each group's layers reuse the same weights. This gives N/G× parameter
+//! reduction while preserving depth for representational power.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::layers::shared_layers::{SharedLayerConfig, SharedLayerStack, SharingStrategy};
+//!
+//! // ALBERT-style: all 12 layers share one weight set
+//! let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, 768, 3072);
+//! let stack = SharedLayerStack::new(config);
+//!
+//! assert_eq!(stack.layer_group_mapping(), vec![0usize; 12]);
+//! assert_eq!(stack.num_unique_params(), stack.num_params_without_sharing() / 12);
+//! ```
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/// Strategy for grouping transformer layers into weight-sharing groups.
+///
+/// Controls how the `num_layers` virtual layers are partitioned into groups,
+/// where every layer within a group reuses the same weight tensors.
+#[derive(Debug, Clone, PartialEq)]
+pub enum SharingStrategy {
+    /// All layers share one set of weights (ALBERT-style full sharing).
+    ///
+    /// `num_groups = 1`. Maximum compression: N× fewer parameters.
+    FullSharing,
+
+    /// Layers divided into G groups; each group shares one set of weights.
+    ///
+    /// For example, 12 layers with 3 groups gives groups of 4 layers each,
+    /// resulting in a 4× parameter reduction.
+    GroupedSharing { num_groups: usize },
+
+    /// Even-indexed layers share one weight set; odd-indexed layers share another.
+    ///
+    /// Produces exactly 2 unique weight sets regardless of `num_layers`.
+    AlternatingPairs,
+}
+
+/// Configuration for cross-layer parameter sharing.
+///
+/// Describes the sharing topology without holding any weights. Pass this to
+/// [`SharedLayerStack::new`] to create the actual weight store.
+#[derive(Debug, Clone)]
+pub struct SharedLayerConfig {
+    /// Total number of virtual transformer layers (depth).
+    pub num_layers: usize,
+    /// How layers are partitioned into weight-sharing groups.
+    pub strategy: SharingStrategy,
+    /// Dimension of hidden (residual) representations.
+    pub hidden_size: usize,
+    /// Inner dimension of the feed-forward network.
+    pub ffn_size: usize,
+}
+
+impl SharedLayerConfig {
+    /// Create a new `SharedLayerConfig`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `num_layers == 0`, `hidden_size == 0`, or `ffn_size == 0`.
+    /// For `GroupedSharing`, also panics if `num_groups == 0` or
+    /// `num_groups > num_layers`.
+    #[must_use]
+    pub fn new(
+        num_layers: usize,
+        strategy: SharingStrategy,
+        hidden_size: usize,
+        ffn_size: usize,
+    ) -> Self {
+        assert!(num_layers > 0, "num_layers must be > 0");
+        assert!(hidden_size > 0, "hidden_size must be > 0");
+        assert!(ffn_size > 0, "ffn_size must be > 0");
+        if let SharingStrategy::GroupedSharing { num_groups } = strategy {
+            assert!(num_groups > 0, "num_groups must be > 0");
+            assert!(
+                num_groups <= num_layers,
+                "num_groups ({num_groups}) must be <= num_layers ({num_layers})"
+            );
+        }
+        Self {
+            num_layers,
+            strategy,
+            hidden_size,
+            ffn_size,
+        }
+    }
+
+    /// Number of unique weight sets that will be allocated.
+    ///
+    /// - `FullSharing` → 1
+    /// - `GroupedSharing { num_groups }` → `num_groups`
+    /// - `AlternatingPairs` → 2
+    #[must_use]
+    pub fn num_unique_weight_sets(&self) -> usize {
+        match &self.strategy {
+            SharingStrategy::FullSharing => 1,
+            SharingStrategy::GroupedSharing { num_groups } => *num_groups,
+            SharingStrategy::AlternatingPairs => 2,
+        }
+    }
+
+    /// Which weight-group index the layer at `layer_depth` should use.
+    ///
+    /// `layer_depth` is zero-based and must be `< num_layers`.
+    ///
+    /// # Panics
+    ///
+    /// Panics in debug mode if `layer_depth >= num_layers`.
+    #[must_use]
+    pub fn group_for_layer(&self, layer_depth: usize) -> usize {
+        debug_assert!(
+            layer_depth < self.num_layers,
+            "layer_depth {layer_depth} out of range (num_layers={})",
+            self.num_layers
+        );
+        match &self.strategy {
+            SharingStrategy::FullSharing => 0,
+            SharingStrategy::GroupedSharing { num_groups } => {
+                // Integer division: layers [0, layers_per_group) → group 0, etc.
+                // Last group absorbs any remainder layers.
+                let layers_per_group = self.num_layers / num_groups;
+                let group = layer_depth / layers_per_group;
+                // Clamp to avoid exceeding num_groups for remainder layers
+                group.min(num_groups - 1)
+            }
+            SharingStrategy::AlternatingPairs => layer_depth % 2,
+        }
+    }
+
+    /// Memory reduction ratio relative to no sharing: `num_unique / num_layers`.
+    ///
+    /// Values close to 0 indicate heavy compression; 1.0 means no reduction.
+    #[must_use]
+    pub fn memory_reduction_ratio(&self) -> f32 {
+        self.num_unique_weight_sets() as f32 / self.num_layers as f32
+    }
+
+    /// Approximate number of layers per group.
+    ///
+    /// For `FullSharing` returns `num_layers`. For `AlternatingPairs` returns
+    /// `num_layers / 2` (rounded down). For `GroupedSharing` returns
+    /// `num_layers / num_groups` (rounded down).
+    #[must_use]
+    pub fn layers_per_group(&self) -> usize {
+        match &self.strategy {
+            SharingStrategy::FullSharing => self.num_layers,
+            SharingStrategy::GroupedSharing { num_groups } => self.num_layers / num_groups,
+            SharingStrategy::AlternatingPairs => self.num_layers / 2,
+        }
+    }
+}
+
+// ============================================================================
+// FFN weight store
+// ============================================================================
+
+/// Simulated transformer FFN weights.
+///
+/// Stores W1, W2 and their biases for a two-layer feed-forward network:
+/// `ReLU(x @ W1 + b1) @ W2 + b2`.
+///
+/// Uses `Vec<f32>` so the type is fully testable without GPU infrastructure.
+#[derive(Debug, Clone)]
+pub struct SharedFfnWeight {
+    /// W1: \[hidden\_size × ffn\_size\] stored row-major.
+    pub w1: Vec<f32>,
+    /// W2: \[ffn\_size × hidden\_size\] stored row-major.
+    pub w2: Vec<f32>,
+    /// b1: \[ffn\_size\].
+    pub b1: Vec<f32>,
+    /// b2: \[hidden\_size\].
+    pub b2: Vec<f32>,
+}
+
+impl SharedFfnWeight {
+    /// Create a zero-initialised weight set.
+    #[must_use]
+    pub fn new(hidden_size: usize, ffn_size: usize) -> Self {
+        Self {
+            w1: vec![0.0_f32; hidden_size * ffn_size],
+            w2: vec![0.0_f32; ffn_size * hidden_size],
+            b1: vec![0.0_f32; ffn_size],
+            b2: vec![0.0_f32; hidden_size],
+        }
+    }
+
+    /// Create weights initialised with a simple LCG pseudo-random number generator.
+    ///
+    /// Values are scaled to `[-0.5, 0.5)` so the network is not saturated.
+    /// The LCG parameters are the classic Numerical Recipes constants.
+    #[must_use]
+    pub fn new_random(hidden_size: usize, ffn_size: usize, seed: u64) -> Self {
+        let mut state = seed;
+        let mut next = move || -> f32 {
+            // LCG: x_{n+1} = (a * x_n + c) mod 2^32
+            state = state
+                .wrapping_mul(1_664_525)
+                .wrapping_add(1_013_904_223)
+                & 0xFFFF_FFFF;
+            // Map to [-0.5, 0.5)
+            (state as f32 / u32::MAX as f32) - 0.5
+        };
+
+        let w1: Vec<f32> = (0..hidden_size * ffn_size).map(|_| next()).collect();
+        let w2: Vec<f32> = (0..ffn_size * hidden_size).map(|_| next()).collect();
+        let b1: Vec<f32> = (0..ffn_size).map(|_| next()).collect();
+        let b2: Vec<f32> = (0..hidden_size).map(|_| next()).collect();
+
+        Self { w1, w2, b1, b2 }
+    }
+
+    /// FFN forward pass: `ReLU(x @ W1 + b1) @ W2 + b2`.
+    ///
+    /// # Arguments
+    ///
+    /// * `x`       — Input slice of length `seq_len * hidden`.
+    /// * `seq_len` — Number of token positions.
+    /// * `hidden`  — Hidden (residual) dimension; must equal `self.w1.len() / ffn`.
+    /// * `ffn`     — FFN inner dimension; must equal `self.w2.len() / hidden`.
+    ///
+    /// # Returns
+    ///
+    /// Output slice of length `seq_len * hidden`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if dimension arguments are inconsistent with stored weight sizes.
+    #[must_use]
+    pub fn forward(&self, x: &[f32], seq_len: usize, hidden: usize, ffn: usize) -> Vec<f32> {
+        assert_eq!(
+            x.len(),
+            seq_len * hidden,
+            "input length mismatch: expected {}, got {}",
+            seq_len * hidden,
+            x.len()
+        );
+        assert_eq!(self.w1.len(), hidden * ffn, "w1 size mismatch");
+        assert_eq!(self.w2.len(), ffn * hidden, "w2 size mismatch");
+        assert_eq!(self.b1.len(), ffn, "b1 size mismatch");
+        assert_eq!(self.b2.len(), hidden, "b2 size mismatch");
+
+        // hidden1[t, j] = ReLU(sum_i x[t,i] * W1[i,j] + b1[j])
+        let mut hidden1 = vec![0.0_f32; seq_len * ffn];
+        for t in 0..seq_len {
+            for j in 0..ffn {
+                let mut acc = self.b1[j];
+                for i in 0..hidden {
+                    acc += x[t * hidden + i] * self.w1[i * ffn + j];
+                }
+                hidden1[t * ffn + j] = acc.max(0.0); // ReLU
+            }
+        }
+
+        // out[t, k] = sum_j hidden1[t,j] * W2[j,k] + b2[k]
+        let mut out = vec![0.0_f32; seq_len * hidden];
+        for t in 0..seq_len {
+            for k in 0..hidden {
+                let mut acc = self.b2[k];
+                for j in 0..ffn {
+                    acc += hidden1[t * ffn + j] * self.w2[j * hidden + k];
+                }
+                out[t * hidden + k] = acc;
+            }
+        }
+
+        out
+    }
+}
+
+// ============================================================================
+// SharedLayerStack
+// ============================================================================
+
+/// Manages shared weight sets and routes layers to the correct group.
+///
+/// Each of the `num_layers` virtual transformer layers receives a reference to
+/// one of the `num_unique_weight_sets` stored `SharedFfnWeight` instances,
+/// determined by the [`SharingStrategy`] in the [`SharedLayerConfig`].
+pub struct SharedLayerStack {
+    config: SharedLayerConfig,
+    /// One `SharedFfnWeight` per unique group (length == `num_unique_weight_sets`).
+    shared_weights: Vec<SharedFfnWeight>,
+}
+
+impl SharedLayerStack {
+    /// Build a new stack with zero-initialised shared weights.
+    #[must_use]
+    pub fn new(config: SharedLayerConfig) -> Self {
+        let n_sets = config.num_unique_weight_sets();
+        let shared_weights = (0..n_sets)
+            .map(|_| SharedFfnWeight::new(config.hidden_size, config.ffn_size))
+            .collect();
+        Self {
+            config,
+            shared_weights,
+        }
+    }
+
+    /// Borrow the weight set assigned to `layer_depth`.
+    ///
+    /// # Panics
+    ///
+    /// Panics (debug) if `layer_depth >= num_layers`.
+    #[must_use]
+    pub fn weights_for_layer(&self, layer_depth: usize) -> &SharedFfnWeight {
+        let group = self.config.group_for_layer(layer_depth);
+        &self.shared_weights[group]
+    }
+
+    /// Run a forward pass through all `num_layers` virtual layers sequentially.
+    ///
+    /// Each layer applies its assigned shared FFN weights. The output of one
+    /// layer is the input to the next.
+    ///
+    /// # Arguments
+    ///
+    /// * `x`       — Input of length `seq_len * hidden_size`.
+    /// * `seq_len` — Number of token positions.
+    ///
+    /// # Returns
+    ///
+    /// Final output after all layers, length `seq_len * hidden_size`.
+    #[must_use]
+    pub fn forward(&self, x: &[f32], seq_len: usize) -> Vec<f32> {
+        let hidden = self.config.hidden_size;
+        let ffn = self.config.ffn_size;
+        let mut current = x.to_vec();
+        for depth in 0..self.config.num_layers {
+            let w = self.weights_for_layer(depth);
+            current = w.forward(&current, seq_len, hidden, ffn);
+        }
+        current
+    }
+
+    /// Total number of unique weight parameters (across all groups).
+    ///
+    /// Each weight set contributes `w1 + w2 + b1 + b2` elements.
+    #[must_use]
+    pub fn num_unique_params(&self) -> usize {
+        let hidden = self.config.hidden_size;
+        let ffn = self.config.ffn_size;
+        // Per set: w1(hidden*ffn) + w2(ffn*hidden) + b1(ffn) + b2(hidden)
+        let params_per_set = hidden * ffn + ffn * hidden + ffn + hidden;
+        self.config.num_unique_weight_sets() * params_per_set
+    }
+
+    /// Total parameter count if every layer had its own independent weights.
+    #[must_use]
+    pub fn num_params_without_sharing(&self) -> usize {
+        let hidden = self.config.hidden_size;
+        let ffn = self.config.ffn_size;
+        let params_per_set = hidden * ffn + ffn * hidden + ffn + hidden;
+        self.config.num_layers * params_per_set
+    }
+
+    /// Actual memory reduction: `num_unique_params / num_params_without_sharing`.
+    ///
+    /// Equivalent to `config.memory_reduction_ratio()` but computed from
+    /// `num_unique_params` and `num_params_without_sharing` directly.
+    #[must_use]
+    pub fn actual_reduction_ratio(&self) -> f32 {
+        self.num_unique_params() as f32 / self.num_params_without_sharing() as f32
+    }
+
+    /// Returns a `Vec` of length `num_layers` where entry `i` is the group
+    /// index that layer `i` maps to.
+    #[must_use]
+    pub fn layer_group_mapping(&self) -> Vec<usize> {
+        (0..self.config.num_layers)
+            .map(|d| self.config.group_for_layer(d))
+            .collect()
+    }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ------------------------------------------------------------------
+    // SharingStrategy / SharedLayerConfig tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_full_sharing_all_layers_same_group() {
+        let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, 64, 256);
+        for depth in 0..12 {
+            assert_eq!(
+                config.group_for_layer(depth),
+                0,
+                "layer {depth} should map to group 0"
+            );
+        }
+    }
+
+    #[test]
+    fn test_grouped_sharing_correct_assignment() {
+        // 12 layers, 3 groups → layers 0-3 → group 0, 4-7 → group 1, 8-11 → group 2
+        let config = SharedLayerConfig::new(
+            12,
+            SharingStrategy::GroupedSharing { num_groups: 3 },
+            64,
+            256,
+        );
+        for depth in 0..4 {
+            assert_eq!(config.group_for_layer(depth), 0, "layer {depth} → group 0");
+        }
+        for depth in 4..8 {
+            assert_eq!(config.group_for_layer(depth), 1, "layer {depth} → group 1");
+        }
+        for depth in 8..12 {
+            assert_eq!(config.group_for_layer(depth), 2, "layer {depth} → group 2");
+        }
+    }
+
+    #[test]
+    fn test_alternating_even_odd_groups() {
+        let config = SharedLayerConfig::new(6, SharingStrategy::AlternatingPairs, 64, 256);
+        assert_eq!(config.group_for_layer(0), 0);
+        assert_eq!(config.group_for_layer(1), 1);
+        assert_eq!(config.group_for_layer(2), 0);
+        assert_eq!(config.group_for_layer(3), 1);
+        assert_eq!(config.group_for_layer(4), 0);
+        assert_eq!(config.group_for_layer(5), 1);
+    }
+
+    #[test]
+    fn test_num_unique_weight_sets_full() {
+        let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, 64, 256);
+        assert_eq!(config.num_unique_weight_sets(), 1);
+    }
+
+    #[test]
+    fn test_num_unique_weight_sets_grouped() {
+        let config = SharedLayerConfig::new(
+            12,
+            SharingStrategy::GroupedSharing { num_groups: 3 },
+            64,
+            256,
+        );
+        assert_eq!(config.num_unique_weight_sets(), 3);
+    }
+
+    #[test]
+    fn test_memory_reduction_full() {
+        // FullSharing with 12 layers → ratio = 1/12
+        let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, 64, 256);
+        let expected = 1.0_f32 / 12.0_f32;
+        let actual = config.memory_reduction_ratio();
+        assert!(
+            (actual - expected).abs() < 1e-6,
+            "expected {expected}, got {actual}"
+        );
+    }
+
+    #[test]
+    fn test_memory_reduction_grouped() {
+        // 3 groups, 12 layers → ratio = 3/12 = 0.25
+        let config = SharedLayerConfig::new(
+            12,
+            SharingStrategy::GroupedSharing { num_groups: 3 },
+            64,
+            256,
+        );
+        let expected = 3.0_f32 / 12.0_f32;
+        let actual = config.memory_reduction_ratio();
+        assert!(
+            (actual - expected).abs() < 1e-6,
+            "expected {expected}, got {actual}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // SharedFfnWeight tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_ffn_forward_output_shape() {
+        let w = SharedFfnWeight::new_random(32, 128, 42);
+        let seq_len = 5;
+        let hidden = 32;
+        let ffn = 128;
+        let x: Vec<f32> = (0..seq_len * hidden).map(|i| i as f32 * 0.01).collect();
+        let out = w.forward(&x, seq_len, hidden, ffn);
+        assert_eq!(out.len(), seq_len * hidden, "output shape mismatch");
+    }
+
+    #[test]
+    fn test_ffn_forward_zero_weights_zero_output() {
+        // Zero-initialised weights and zero bias → all-zero output regardless of input
+        let w = SharedFfnWeight::new(16, 64);
+        let x: Vec<f32> = vec![1.0_f32; 4 * 16]; // non-zero input
+        let out = w.forward(&x, 4, 16, 64);
+        for &v in &out {
+            assert_eq!(v, 0.0_f32, "expected zero output with zero weights");
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // SharedLayerStack tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_stack_full_sharing_same_result() {
+        // Deterministic: same input twice → same output
+        let config = SharedLayerConfig::new(4, SharingStrategy::FullSharing, 16, 64);
+        let stack = SharedLayerStack::new(config);
+        let x: Vec<f32> = (0..3 * 16).map(|i| i as f32 * 0.1).collect();
+
+        let out1 = stack.forward(&x, 3);
+        let out2 = stack.forward(&x, 3);
+
+        assert_eq!(out1, out2, "same input must produce same output");
+    }
+
+    #[test]
+    fn test_stack_num_unique_params() {
+        // FullSharing with 12 layers: 1 weight set, not 12
+        let hidden = 64_usize;
+        let ffn = 256_usize;
+        let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, hidden, ffn);
+        let stack = SharedLayerStack::new(config);
+
+        let expected_per_set = hidden * ffn + ffn * hidden + ffn + hidden;
+        assert_eq!(stack.num_unique_params(), expected_per_set);
+        assert_eq!(
+            stack.num_params_without_sharing(),
+            12 * expected_per_set,
+            "unshared count must be 12×"
+        );
+    }
+
+    #[test]
+    fn test_layer_group_mapping_length() {
+        let config = SharedLayerConfig::new(
+            8,
+            SharingStrategy::GroupedSharing { num_groups: 2 },
+            32,
+            128,
+        );
+        let stack = SharedLayerStack::new(config);
+        let mapping = stack.layer_group_mapping();
+        assert_eq!(mapping.len(), 8, "mapping must have exactly num_layers entries");
+    }
+
+    #[test]
+    fn test_actual_reduction_ratio() {
+        // FullSharing 12 layers: actual ratio must match theoretical
+        let config = SharedLayerConfig::new(12, SharingStrategy::FullSharing, 64, 256);
+        let theoretical = config.memory_reduction_ratio();
+        let stack = SharedLayerStack::new(config);
+        let actual = stack.actual_reduction_ratio();
+        assert!(
+            (actual - theoretical).abs() < 1e-6,
+            "actual {actual} vs theoretical {theoretical}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Additional edge-case / coverage tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_grouped_sharing_single_group_equals_full() {
+        // GroupedSharing { num_groups: 1 } behaves identically to FullSharing
+        let config_g = SharedLayerConfig::new(
+            6,
+            SharingStrategy::GroupedSharing { num_groups: 1 },
+            32,
+            128,
+        );
+        let config_f = SharedLayerConfig::new(6, SharingStrategy::FullSharing, 32, 128);
+        for d in 0..6 {
+            assert_eq!(
+                config_g.group_for_layer(d),
+                config_f.group_for_layer(d),
+                "layer {d}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_alternating_pairs_two_unique_sets() {
+        let config = SharedLayerConfig::new(10, SharingStrategy::AlternatingPairs, 32, 128);
+        assert_eq!(config.num_unique_weight_sets(), 2);
+        let stack = SharedLayerStack::new(config);
+        assert_eq!(stack.shared_weights.len(), 2);
+    }
+
+    #[test]
+    fn test_stack_weights_for_layer_full_sharing_identity() {
+        // In FullSharing every layer returns a pointer to the same weight set
+        let config = SharedLayerConfig::new(6, SharingStrategy::FullSharing, 16, 64);
+        let stack = SharedLayerStack::new(config);
+        let ptr0 = stack.weights_for_layer(0) as *const SharedFfnWeight;
+        for d in 1..6 {
+            let ptr = stack.weights_for_layer(d) as *const SharedFfnWeight;
+            assert_eq!(ptr, ptr0, "layer {d} must point to the same weight set");
+        }
+    }
+
+    #[test]
+    fn test_ffn_random_seed_determinism() {
+        // Same seed → identical weights
+        let w1 = SharedFfnWeight::new_random(16, 64, 123);
+        let w2 = SharedFfnWeight::new_random(16, 64, 123);
+        assert_eq!(w1.w1, w2.w1);
+        assert_eq!(w1.w2, w2.w2);
+        assert_eq!(w1.b1, w2.b1);
+        assert_eq!(w1.b2, w2.b2);
+    }
+
+    #[test]
+    fn test_grouped_sharing_remainder_layers_clamped_to_last_group() {
+        // 5 layers, 3 groups → layers_per_group=1, groups: 0,1,2,2,2
+        // (remainder folds into last group)
+        let config = SharedLayerConfig::new(
+            5,
+            SharingStrategy::GroupedSharing { num_groups: 3 },
+            16,
+            64,
+        );
+        // Layer 0 → group 0
+        assert_eq!(config.group_for_layer(0), 0);
+        // Layer 1 → group 1
+        assert_eq!(config.group_for_layer(1), 1);
+        // Layers 2,3,4 → at most group 2
+        for d in 2..5 {
+            let g = config.group_for_layer(d);
+            assert!(g < 3, "layer {d} group {g} must be < 3");
+        }
+    }
+}

crates/training/rtx-transformers/src/optimizers/mod.rs

@@ -43,6 +43,7 @@ pub use adam::AdamState;
 
 pub mod matrix_utils;
 pub mod galore;
+pub mod schedule_free;
 
 // Re-export tensor bridge traits for all optimizer modules
 pub use crate::tensor_bridge::{TensorBridge, TensorBridgeStatic, TensorCompat};
@@ -61,6 +62,7 @@ pub use crate::tensor_bridge::{TensorBridge, TensorBridgeStatic, TensorCompat};
 pub use adam::AdamOptimizer;
 pub use adamw::AdamWOptimizer;
 pub use galore::{GaLoreAdamW, GaLoreConfig, GaLoreParamState};
+pub use schedule_free::{ScheduleFreeConfig, ScheduleFreeOptimizer, ScheduleFreeState};
 pub use layer_lr_decay::{LayerLrDecayBuilder, LayerLrDecayConfig, LayerLrScheduler, ParamGroupLR};
 
 // ============================================================================

crates/training/rtx-transformers/src/optimizers/schedule_free.rs

@@ -0,0 +1,619 @@
+//! Schedule-Free optimizer (SGD and Adam variants).
+//!
+//! Eliminates the need for an external learning rate schedule by maintaining two
+//! parameter sequences:
+//!
+//! - `z` — the "worker" sequence updated by gradient descent steps
+//! - `x` — the weighted-average "evaluation" sequence used for gradient computation
+//!
+//! The interpolation coefficient at step `t` is:
+//!
+//! ```text
+//! c_t = 1 - (1 + β * t)^{-r}
+//! ```
+//!
+//! And the update rule is:
+//!
+//! ```text
+//! x_t = (1 - c_t) * z_t + c_t * x_{t-1}
+//! ```
+//!
+//! Gradients are always computed at `x` (the averaged sequence), which provides
+//! implicit momentum and eliminates the need for a decaying schedule.
+//!
+//! # References
+//!
+//! Defazio et al., "The Road Less Scheduled", arXiv:2405.15682 (2024)
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::optimizers::schedule_free::{ScheduleFreeConfig, ScheduleFreeOptimizer};
+//!
+//! let mut opt = ScheduleFreeOptimizer::adam(1e-3);
+//! let params = vec![1.0_f32, 2.0, 3.0];
+//! let grad = vec![0.1_f32, 0.0, -0.1];
+//!
+//! // Pass x (eval params) into step; on first call these equal the initial params.
+//! let x_new = opt.step("layer0.weight", &params, &grad);
+//! assert_eq!(x_new.len(), params.len());
+//! ```
+
+use std::collections::HashMap;
+
+// ============================================================================
+// Configuration
+// ============================================================================
+
+/// Configuration for the Schedule-Free optimizer.
+///
+/// All fields have sensible defaults via [`Default`]; the most common entry
+/// points are [`ScheduleFreeOptimizer::adam`] and [`ScheduleFreeOptimizer::sgd`].
+#[derive(Debug, Clone)]
+pub struct ScheduleFreeConfig {
+    /// Constant learning rate (no schedule needed).
+    pub lr: f32,
+    /// Momentum coefficient.
+    ///
+    /// In Adam mode this is the EMA decay for the second moment (`β₂` in
+    /// standard Adam notation). In SGD mode it is unused.
+    pub momentum: f32,
+    /// Averaging speed exponent `r` (default 0.6).
+    ///
+    /// Controls how fast the `c_t` interpolation coefficient reaches 1. Larger
+    /// values make `x` converge faster to `z`.
+    pub r: f32,
+    /// Averaging weight bias `β` (default 10.0).
+    ///
+    /// Shifts the step at which `c_t` starts to grow meaningfully.
+    pub beta: f32,
+    /// L2 weight-decay coefficient (default 0.0 — no decay).
+    pub weight_decay: f32,
+    /// Epsilon for Adam-mode numerical stability (default 1e-8).
+    pub epsilon: f32,
+    /// `true` → Adam-style second-moment normalisation; `false` → pure SGD.
+    pub use_adam: bool,
+}
+
+impl Default for ScheduleFreeConfig {
+    fn default() -> Self {
+        Self {
+            lr: 0.01,
+            momentum: 0.9,
+            r: 0.6,
+            beta: 10.0,
+            weight_decay: 0.0,
+            epsilon: 1e-8,
+            use_adam: true,
+        }
+    }
+}
+
+// ============================================================================
+// Per-parameter state
+// ============================================================================
+
+/// Optimizer state for a single named parameter tensor.
+#[derive(Debug, Clone)]
+pub struct ScheduleFreeState {
+    /// Inner "worker" parameters updated by gradient descent steps.
+    pub z: Vec<f32>,
+    /// Weighted-average "evaluation" parameters (used for gradient computation).
+    pub x: Vec<f32>,
+    /// Second-moment EMA (`v` in Adam notation). Empty in SGD mode.
+    pub v: Vec<f32>,
+    /// Number of update steps taken for this parameter (1-indexed during update).
+    pub step: usize,
+}
+
+impl ScheduleFreeState {
+    /// Initialize state from a parameter slice.
+    ///
+    /// Sets `z = x = params` and `v = 0`, which is correct for step 0 because
+    /// `c_0 = 0` so the first update degenerates to a plain gradient step on `z`.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use rtx_transformers::optimizers::schedule_free::ScheduleFreeState;
+    ///
+    /// let params = vec![1.0_f32, 2.0, 3.0];
+    /// let state = ScheduleFreeState::new(&params);
+    /// assert_eq!(state.z, params);
+    /// assert_eq!(state.x, params);
+    /// assert_eq!(state.v, vec![0.0_f32; 3]);
+    /// assert_eq!(state.step, 0);
+    /// ```
+    #[must_use]
+    pub fn new(params: &[f32]) -> Self {
+        Self {
+            z: params.to_vec(),
+            x: params.to_vec(),
+            v: vec![0.0_f32; params.len()],
+            step: 0,
+        }
+    }
+}
+
+// ============================================================================
+// Optimizer
+// ============================================================================
+
+/// Schedule-Free optimizer (SGD or Adam variant).
+///
+/// Maintains two sequences per parameter (`z` and `x`) and eliminates the need
+/// for an external learning-rate schedule.
+///
+/// # Usage
+///
+/// 1. Call [`step`](ScheduleFreeOptimizer::step) with the **current `x`** (evaluation
+///    parameters) and the gradient computed at that point.
+/// 2. The returned `Vec<f32>` is the new `x` — store it as your model's parameters.
+/// 3. For inference / loss evaluation always use `x`.
+pub struct ScheduleFreeOptimizer {
+    config: ScheduleFreeConfig,
+    /// Per-parameter state, keyed by the caller-supplied name.
+    states: HashMap<String, ScheduleFreeState>,
+    /// Total number of [`step`](ScheduleFreeOptimizer::step) calls across all parameters.
+    global_step: usize,
+}
+
+impl ScheduleFreeOptimizer {
+    /// Create a new optimizer with an explicit configuration.
+    #[must_use]
+    pub fn new(config: ScheduleFreeConfig) -> Self {
+        Self {
+            config,
+            states: HashMap::new(),
+            global_step: 0,
+        }
+    }
+
+    /// Convenience constructor: Schedule-Free Adam with recommended defaults.
+    ///
+    /// Uses `momentum = 0.999`, `r = 0.6`, `β = 10.0`, `ε = 1e-8`.
+    #[must_use]
+    pub fn adam(lr: f32) -> Self {
+        Self::new(ScheduleFreeConfig {
+            lr,
+            momentum: 0.999,
+            r: 0.6,
+            beta: 10.0,
+            weight_decay: 0.0,
+            epsilon: 1e-8,
+            use_adam: true,
+        })
+    }
+
+    /// Convenience constructor: Schedule-Free SGD with momentum.
+    ///
+    /// Sets `use_adam = false`; the `momentum` field is ignored in SGD mode
+    /// (gradients are applied without second-moment normalisation).
+    #[must_use]
+    pub fn sgd(lr: f32, momentum: f32) -> Self {
+        Self::new(ScheduleFreeConfig {
+            lr,
+            momentum,
+            r: 0.6,
+            beta: 10.0,
+            weight_decay: 0.0,
+            epsilon: 1e-8,
+            use_adam: false,
+        })
+    }
+
+    /// Compute the interpolation coefficient `c_t` at step `t`.
+    ///
+    /// ```text
+    /// c_t = 1 - (1 + β * t)^{-r}
+    /// ```
+    ///
+    /// At `t = 0` this returns `0.0` (no averaging yet). As `t → ∞` it
+    /// approaches `1.0`.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use rtx_transformers::optimizers::schedule_free::ScheduleFreeOptimizer;
+    ///
+    /// let opt = ScheduleFreeOptimizer::adam(1e-3);
+    /// assert_eq!(opt.compute_ct(0), 0.0);
+    /// let c1 = opt.compute_ct(1);
+    /// let c2 = opt.compute_ct(2);
+    /// assert!(c2 > c1, "c_t must be monotonically increasing");
+    /// assert!(opt.compute_ct(100_000) < 1.0);
+    /// ```
+    #[must_use]
+    pub fn compute_ct(&self, t: usize) -> f32 {
+        1.0 - (1.0 + self.config.beta * t as f32).powf(-self.config.r)
+    }
+
+    /// Run one update step for the named parameter.
+    ///
+    /// `params` must equal the `x` sequence from the previous step (or the
+    /// initial parameter values on the first call).  `grad` is the gradient
+    /// computed at that `x`.
+    ///
+    /// Internally:
+    /// 1. Initialises state on the first call (lazy init).
+    /// 2. Advances `step` for this parameter.
+    /// 3. Computes `c_t`.
+    /// 4. **Adam mode**: updates the second moment `v` then normalises the
+    ///    gradient; **SGD mode**: uses the raw gradient.
+    /// 5. Applies optional weight decay and advances `z`.
+    /// 6. Blends `z` and old `x` to produce new `x`.
+    ///
+    /// Returns the new `x` (evaluation parameters).
+    pub fn step(&mut self, name: &str, params: &[f32], grad: &[f32]) -> Vec<f32> {
+        assert_eq!(
+            params.len(),
+            grad.len(),
+            "params and grad must have equal length"
+        );
+
+        // Lazy state initialisation — first call uses `params` as the starting point.
+        if !self.states.contains_key(name) {
+            self.states
+                .insert(name.to_owned(), ScheduleFreeState::new(params));
+        }
+
+        let cfg = &self.config;
+        let state = self.states.get_mut(name).unwrap();
+
+        // Advance per-parameter step counter (1-indexed during this update).
+        state.step += 1;
+        let t = state.step;
+        let c_t = 1.0 - (1.0 + cfg.beta * t as f32).powf(-cfg.r);
+
+        let len = state.z.len();
+        let mut z_new = Vec::with_capacity(len);
+
+        if cfg.use_adam {
+            // Adam mode: update second moment then compute normalised gradient.
+            for i in 0..len {
+                // EMA of squared gradient (β is reused as the second-moment decay).
+                state.v[i] = cfg.momentum * state.v[i] + (1.0 - cfg.momentum) * grad[i] * grad[i];
+
+                let g_hat = grad[i] / (state.v[i].sqrt() + cfg.epsilon);
+                let z_i = (1.0 - cfg.weight_decay * cfg.lr) * state.z[i] - cfg.lr * g_hat;
+                z_new.push(z_i);
+            }
+        } else {
+            // SGD mode: apply gradient directly without second-moment normalisation.
+            for i in 0..len {
+                let z_i = (1.0 - cfg.weight_decay * cfg.lr) * state.z[i] - cfg.lr * grad[i];
+                z_new.push(z_i);
+            }
+        }
+
+        // Blend z and old x to obtain new x.
+        let old_x = state.x.clone();
+        let x_new: Vec<f32> = (0..len)
+            .map(|i| (1.0 - c_t) * z_new[i] + c_t * old_x[i])
+            .collect();
+
+        state.z = z_new;
+        state.x = x_new.clone();
+
+        self.global_step += 1;
+
+        x_new
+    }
+
+    /// Return the current evaluation parameters (`x` sequence) for `name`,
+    /// or `None` if this parameter has not yet been seen by [`step`](Self::step).
+    #[must_use]
+    pub fn eval_params(&self, name: &str) -> Option<&[f32]> {
+        self.states.get(name).map(|s| s.x.as_slice())
+    }
+
+    /// Return the current inner parameters (`z` sequence) for `name`,
+    /// or `None` if this parameter has not yet been seen.
+    #[must_use]
+    pub fn z_params(&self, name: &str) -> Option<&[f32]> {
+        self.states.get(name).map(|s| s.z.as_slice())
+    }
+
+    /// Total number of [`step`](Self::step) calls across all parameters.
+    #[must_use]
+    pub fn global_step(&self) -> usize {
+        self.global_step
+    }
+
+    /// Number of distinct parameter groups tracked.
+    #[must_use]
+    pub fn num_params(&self) -> usize {
+        self.states.len()
+    }
+
+    /// Reset all optimizer state (per-parameter state and global step counter).
+    pub fn reset(&mut self) {
+        self.states.clear();
+        self.global_step = 0;
+    }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Helper: absolute difference between two f32 slices.
+    fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+        a.iter()
+            .zip(b.iter())
+            .map(|(x, y)| (x - y).abs())
+            .fold(0.0_f32, f32::max)
+    }
+
+    // -----------------------------------------------------------------------
+    // ScheduleFreeState initialisation
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_state_initialized_z_equals_params() {
+        let params = vec![1.0_f32, -2.0, 0.5, 3.14];
+        let state = ScheduleFreeState::new(&params);
+        assert_eq!(state.z, params, "z must equal initial params");
+    }
+
+    #[test]
+    fn test_state_initialized_x_equals_params() {
+        let params = vec![0.1_f32, 0.2, 0.3];
+        let state = ScheduleFreeState::new(&params);
+        assert_eq!(state.x, params, "x must equal initial params");
+    }
+
+    // -----------------------------------------------------------------------
+    // compute_ct behaviour
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_compute_ct_zero_at_step_zero() {
+        let opt = ScheduleFreeOptimizer::adam(1e-3);
+        let c0 = opt.compute_ct(0);
+        assert_eq!(c0, 0.0, "c_0 must be exactly 0.0 (t=0 gives (1+β*0)^-r = 1)");
+    }
+
+    #[test]
+    fn test_compute_ct_increases_with_step() {
+        let opt = ScheduleFreeOptimizer::adam(1e-3);
+        let mut prev = opt.compute_ct(0);
+        for t in 1..=20 {
+            let curr = opt.compute_ct(t);
+            assert!(
+                curr > prev,
+                "c_t must be strictly increasing: c_{t} = {curr} <= c_{} = {prev}",
+                t - 1
+            );
+            prev = curr;
+        }
+    }
+
+    #[test]
+    fn test_compute_ct_approaches_one() {
+        let opt = ScheduleFreeOptimizer::adam(1e-3);
+        let c_large = opt.compute_ct(1_000_000);
+        assert!(
+            c_large < 1.0,
+            "c_t must remain strictly below 1.0: got {c_large}"
+        );
+        assert!(
+            c_large > 0.999,
+            "c_t at 1M steps must be very close to 1.0: got {c_large}"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // step() mechanics
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_step_creates_state_on_first_call() {
+        let mut opt = ScheduleFreeOptimizer::adam(1e-3);
+        assert_eq!(opt.num_params(), 0, "no state before any step");
+        let params = vec![1.0_f32, 2.0];
+        let grad = vec![0.1_f32, 0.1];
+        let _ = opt.step("weight", &params, &grad);
+        assert_eq!(opt.num_params(), 1, "state must be created after first step");
+    }
+
+    #[test]
+    fn test_step_returns_x_not_z() {
+        // After one step with a non-trivial gradient the returned vector must
+        // equal the internal x, not z.
+        let mut opt = ScheduleFreeOptimizer::adam(1e-2);
+        let params = vec![0.0_f32, 0.0, 0.0];
+        let grad = vec![1.0_f32, 1.0, 1.0];
+        let returned = opt.step("w", &params, &grad);
+
+        // Compare returned slice to the stored x.
+        let stored_x = opt.eval_params("w").expect("state must exist after step");
+        assert_eq!(
+            returned, stored_x,
+            "returned vector must equal stored x sequence"
+        );
+
+        // Additionally verify it is NOT identical to z.
+        let stored_z = opt.z_params("w").expect("z state must exist after step");
+        // After exactly one step c_1 = 1-(1+10)^{-0.6} > 0, so x ≠ z unless both are zero.
+        // With lr=0.01 and a large gradient, z will be non-zero and x will differ.
+        assert_ne!(
+            returned, stored_z,
+            "returned x must differ from z after interpolation"
+        );
+    }
+
+    #[test]
+    fn test_zero_gradient_no_update_to_z() {
+        let mut opt = ScheduleFreeOptimizer::new(ScheduleFreeConfig {
+            lr: 0.1,
+            weight_decay: 0.0,
+            use_adam: false,
+            ..Default::default()
+        });
+        let params = vec![3.0_f32, -1.0, 2.5];
+        let zero_grad = vec![0.0_f32; 3];
+
+        let _ = opt.step("w", &params, &zero_grad);
+        let z = opt.z_params("w").unwrap().to_vec();
+        // z must not move when gradient is zero and weight decay is zero.
+        assert!(
+            max_abs_diff(&z, &params) < 1e-6,
+            "z must not change with zero grad and no weight decay; z={z:?}, params={params:?}"
+        );
+    }
+
+    #[test]
+    fn test_weight_decay_shrinks_z() {
+        let mut opt = ScheduleFreeOptimizer::new(ScheduleFreeConfig {
+            lr: 0.1,
+            weight_decay: 0.1,
+            use_adam: false,
+            ..Default::default()
+        });
+        let params = vec![2.0_f32, 2.0, 2.0];
+        let zero_grad = vec![0.0_f32; 3];
+
+        let _ = opt.step("w", &params, &zero_grad);
+        let z = opt.z_params("w").unwrap().to_vec();
+
+        // With weight decay and zero grad: z_i = (1 - wd*lr) * params_i < params_i
+        for (z_i, p_i) in z.iter().zip(params.iter()) {
+            assert!(
+                z_i.abs() < p_i.abs(),
+                "|z| must shrink with weight decay; z={z_i}, param={p_i}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_sgd_mode_x_changes_after_step() {
+        let mut opt = ScheduleFreeOptimizer::sgd(0.1, 0.9);
+        let params = vec![1.0_f32, 1.0, 1.0];
+        let grad = vec![0.5_f32, 0.5, 0.5];
+
+        let x_new = opt.step("w", &params, &grad);
+
+        // z must have moved: z = params - lr*grad = 1 - 0.05 = 0.95
+        let expected_z = 1.0 - 0.1 * 0.5;
+        let z = opt.z_params("w").unwrap();
+        for &z_i in z {
+            assert!(
+                (z_i - expected_z).abs() < 1e-6,
+                "SGD z update incorrect: expected {expected_z}, got {z_i}"
+            );
+        }
+
+        // x must differ from initial params.
+        assert_ne!(
+            x_new, params,
+            "x must change from initial params after an SGD step"
+        );
+    }
+
+    #[test]
+    fn test_adam_mode_normalizes_large_gradient() {
+        // With a very large gradient, Adam normalises it via sqrt(v)+ε.
+        // The resulting z update should be much smaller than lr*grad.
+        let lr = 0.01;
+        let mut opt = ScheduleFreeOptimizer::adam(lr);
+        let params = vec![0.0_f32];
+        let large_grad = vec![1000.0_f32];
+
+        let _ = opt.step("w", &params, &large_grad);
+        let z = opt.z_params("w").unwrap();
+
+        // After one step: v = (1-momentum)*grad^2 = 0.001 * 1e6 = 1000
+        // g_hat = 1000 / (sqrt(1000) + 1e-8) ≈ 31.6
+        // z update ≈ -0.01 * 31.6 = -0.316  (much less than -0.01*1000=-10)
+        assert!(
+            z[0].abs() < 1.0,
+            "Adam must normalise large gradient; z[0] = {}",
+            z[0]
+        );
+        assert!(
+            z[0].abs() < lr * large_grad[0],
+            "Adam step must be smaller than raw lr*grad"
+        );
+    }
+
+    #[test]
+    fn test_global_step_increments() {
+        let mut opt = ScheduleFreeOptimizer::adam(1e-3);
+        let params = vec![0.0_f32];
+        let grad = vec![0.1_f32];
+
+        assert_eq!(opt.global_step(), 0);
+        let _ = opt.step("a", &params, &grad);
+        assert_eq!(opt.global_step(), 1);
+        let _ = opt.step("b", &params, &grad);
+        assert_eq!(opt.global_step(), 2);
+        let _ = opt.step("a", &params, &grad);
+        assert_eq!(opt.global_step(), 3);
+    }
+
+    #[test]
+    fn test_reset_clears_state() {
+        let mut opt = ScheduleFreeOptimizer::adam(1e-3);
+        let params = vec![1.0_f32, 2.0];
+        let grad = vec![0.1_f32, 0.1];
+
+        let _ = opt.step("weight", &params, &grad);
+        let _ = opt.step("bias", &params, &grad);
+        assert_eq!(opt.num_params(), 2, "expect 2 param groups before reset");
+        assert_eq!(opt.global_step(), 2);
+
+        opt.reset();
+
+        assert_eq!(opt.num_params(), 0, "num_params must be 0 after reset");
+        assert_eq!(opt.global_step(), 0, "global_step must be 0 after reset");
+        assert!(opt.eval_params("weight").is_none(), "state must be cleared");
+    }
+
+    // -----------------------------------------------------------------------
+    // Additional edge-case tests
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_multiple_params_tracked_independently() {
+        let mut opt = ScheduleFreeOptimizer::adam(1e-3);
+        let p1 = vec![1.0_f32, 2.0];
+        let p2 = vec![10.0_f32, 20.0];
+        let g = vec![0.1_f32, 0.1];
+
+        let _ = opt.step("layer0", &p1, &g);
+        let _ = opt.step("layer1", &p2, &g);
+
+        assert_eq!(opt.num_params(), 2);
+        // The two parameter groups must hold different z values.
+        let z0 = opt.z_params("layer0").unwrap().to_vec();
+        let z1 = opt.z_params("layer1").unwrap().to_vec();
+        assert_ne!(z0, z1, "independent params must have independent z state");
+    }
+
+    #[test]
+    fn test_convergence_towards_minimum() {
+        // Minimise f(x) = (x - 1)^2 via Schedule-Free Adam.
+        // Gradient: f'(x) = 2*(x - 1).
+        //
+        // The x sequence converges slower than z because it is a weighted
+        // average: we allow 500 steps and a tolerance of 0.2.
+        let mut opt = ScheduleFreeOptimizer::adam(0.1);
+        let mut x = vec![5.0_f32]; // Start far from optimum at x=1.
+
+        for _ in 0..500 {
+            let grad = vec![2.0 * (x[0] - 1.0)];
+            x = opt.step("x", &x, &grad);
+        }
+
+        assert!(
+            (x[0] - 1.0).abs() < 0.2,
+            "Schedule-Free Adam must converge near x=1.0; got x={}",
+            x[0]
+        );
+    }
+}

crates/training/rtx-transformers/src/training/mod.rs

@@ -1,5 +1,6 @@
 //! Training infrastructure for transformers
 
+pub mod model_ema;
 pub mod draft_distill;
 pub mod comprehensive_integration_test;
 pub mod end_to_end_training_example;
@@ -52,6 +53,7 @@ pub use draft_distill::{
     distill_loss_step, kl_divergence, log_softmax, softmax, token_acceptance_estimate,
 };
 pub use gradient_noise_scale::{GnsEstimate, GnsTracker, GradientNoiseScale};
+pub use model_ema::{ModelEma, ModelEmaConfig};
 
 /// Training state structure
 #[derive(Debug, Clone)]

crates/training/rtx-transformers/src/training/model_ema.rs

@@ -0,0 +1,570 @@
+//! Exponential Moving Average (EMA) of model parameters (Polyak averaging).
+//!
+//! During training, model weights fluctuate due to gradient noise. EMA maintains a
+//! shadow copy updated as:
+//!
+//! ```text
+//! shadow = decay * shadow + (1 - decay) * param
+//! ```
+//!
+//! At inference the shadow weights are used instead of live training weights, which
+//! significantly improves final model quality. This technique is employed in diffusion
+//! models, GANs, and increasingly in LLM fine-tuning.
+//!
+//! # Usage
+//!
+//! ```rust
+//! use rtx_transformers::training::{ModelEma, ModelEmaConfig};
+//!
+//! let config = ModelEmaConfig {
+//!     decay: 0.9999,
+//!     warmup_steps: 100,
+//!     use_bias_correction: false,
+//! };
+//! let mut ema = ModelEma::new(config);
+//!
+//! // After each optimizer step:
+//! let params = [("weight", &[1.0_f32, 2.0, 3.0][..])];
+//! ema.update(&params);
+//!
+//! // Before evaluation: swap in shadow weights.
+//! let shadow = ema.apply_shadow();
+//!
+//! // After evaluation: restore live training weights.
+//! let _backup = ema.restore();
+//! ```
+
+use std::collections::HashMap;
+
+/// Configuration for model EMA.
+#[derive(Debug, Clone)]
+pub struct ModelEmaConfig {
+    /// EMA decay rate (e.g. 0.9999 for slow averaging, 0.999 for faster).
+    pub decay: f32,
+    /// Number of warmup steps before EMA starts (shadow = current during warmup).
+    pub warmup_steps: usize,
+    /// Whether to use bias-corrected decay:
+    /// `effective_decay = decay * (1 - decay^step) / (1 - decay)`.
+    /// Starts near 0 and warms up to `decay`. Default: false.
+    pub use_bias_correction: bool,
+}
+
+impl Default for ModelEmaConfig {
+    fn default() -> Self {
+        Self {
+            decay: 0.9999,
+            warmup_steps: 100,
+            use_bias_correction: false,
+        }
+    }
+}
+
+/// Tracks an exponential moving average of model parameter tensors.
+///
+/// # Usage pattern
+///
+/// 1. Create [`ModelEma`] at start of training.
+/// 2. Call [`update`](ModelEma::update) after each optimizer step.
+/// 3. Before evaluation: call [`apply_shadow`](ModelEma::apply_shadow) to swap in EMA weights.
+/// 4. After evaluation: call [`restore`](ModelEma::restore) to swap back training weights.
+pub struct ModelEma {
+    config: ModelEmaConfig,
+    /// Shadow (EMA) parameters: name → flat `Vec<f32>`.
+    shadow: HashMap<String, Vec<f32>>,
+    /// Backup of original training parameters while shadow is applied.
+    backup: HashMap<String, Vec<f32>>,
+    /// Whether shadow is currently applied to the model.
+    shadow_applied: bool,
+    /// Training step counter (for warmup / bias correction).
+    step: usize,
+}
+
+impl ModelEma {
+    /// Create a new [`ModelEma`] with the given configuration.
+    pub fn new(config: ModelEmaConfig) -> Self {
+        Self {
+            config,
+            shadow: HashMap::new(),
+            backup: HashMap::new(),
+            shadow_applied: false,
+            step: 0,
+        }
+    }
+
+    /// Convenience constructor: standard high-quality EMA (decay=0.9999, warmup=100).
+    pub fn standard() -> Self {
+        Self::new(ModelEmaConfig {
+            decay: 0.9999,
+            warmup_steps: 100,
+            use_bias_correction: false,
+        })
+    }
+
+    /// Convenience constructor: fast EMA for shorter training runs (decay=0.999, warmup=50).
+    pub fn fast() -> Self {
+        Self::new(ModelEmaConfig {
+            decay: 0.999,
+            warmup_steps: 50,
+            use_bias_correction: false,
+        })
+    }
+
+    /// Update shadow parameters from the current named parameters.
+    ///
+    /// `params` is a list of `(name, flat_data)` pairs from the current model.
+    /// Initialises the shadow from `params` on the very first call (or during warmup
+    /// when step == 0). Increments the internal step counter after each call.
+    pub fn update(&mut self, params: &[(&str, &[f32])]) {
+        let decay = self.effective_decay();
+
+        for (name, data) in params {
+            match self.shadow.get_mut(*name) {
+                Some(shadow_data) => {
+                    // Blend: shadow = decay * shadow + (1 - decay) * param
+                    let one_minus_decay = 1.0 - decay;
+                    for (s, &p) in shadow_data.iter_mut().zip(data.iter()) {
+                        *s = decay * *s + one_minus_decay * p;
+                    }
+                }
+                None => {
+                    // First time we see this parameter — initialise directly from current value.
+                    self.shadow.insert(name.to_string(), data.to_vec());
+                }
+            }
+        }
+
+        self.step += 1;
+    }
+
+    /// Effective decay at the current step (accounts for warmup and bias correction).
+    ///
+    /// - During warmup (`step < warmup_steps`): linearly ramps from 0 to `config.decay`.
+    /// - After warmup, without bias correction: returns `config.decay`.
+    /// - After warmup, with bias correction: returns
+    ///   `decay * (1 - decay^step) / (1 - decay)`, capped at `config.decay`.
+    pub fn effective_decay(&self) -> f32 {
+        let d = self.config.decay;
+
+        if self.config.warmup_steps > 0 && self.step < self.config.warmup_steps {
+            // Linear warmup from 0 → decay.
+            return (self.step as f32 / self.config.warmup_steps as f32) * d;
+        }
+
+        if self.config.use_bias_correction {
+            let bias_corrected =
+                d * (1.0 - d.powi(self.step as i32)) / (1.0 - d).max(1e-8);
+            bias_corrected.min(d)
+        } else {
+            d
+        }
+    }
+
+    /// Apply shadow parameters to a model (stores a backup of current params).
+    ///
+    /// Returns the shadow parameters as a map of name → cloned `Vec<f32>`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the shadow is already applied. Call [`restore`](ModelEma::restore) first.
+    pub fn apply_shadow(&mut self) -> HashMap<String, Vec<f32>> {
+        assert!(
+            !self.shadow_applied,
+            "ModelEma::apply_shadow called while shadow is already applied; call restore() first"
+        );
+
+        self.shadow_applied = true;
+        // The backup slot is empty here — caller is expected to populate it externally
+        // (or it will remain empty until restore() is called).  We return a clone of
+        // the shadow so the caller can write the shadow values into the actual model.
+        self.shadow.clone()
+    }
+
+    /// Store the caller's current (live) parameters into the internal backup so that
+    /// [`restore`](ModelEma::restore) can return them.
+    ///
+    /// This is an optional companion to [`apply_shadow`](ModelEma::apply_shadow): callers
+    /// that manage their own weight storage can use it to record live weights before
+    /// overwriting them with shadow values.
+    pub fn store_backup(&mut self, params: &[(&str, &[f32])]) {
+        self.backup.clear();
+        for (name, data) in params {
+            self.backup.insert(name.to_string(), data.to_vec());
+        }
+    }
+
+    /// Restore backup parameters (undo [`apply_shadow`](ModelEma::apply_shadow)).
+    ///
+    /// Returns the backup parameter map (the live training weights saved before
+    /// [`apply_shadow`](ModelEma::apply_shadow) was called). Returns an empty map and is
+    /// a no-op if shadow is not currently applied.
+    pub fn restore(&mut self) -> HashMap<String, Vec<f32>> {
+        if !self.shadow_applied {
+            return HashMap::new();
+        }
+
+        self.shadow_applied = false;
+        let backup = std::mem::take(&mut self.backup);
+        backup
+    }
+
+    /// Whether the shadow is currently applied.
+    pub fn shadow_applied(&self) -> bool {
+        self.shadow_applied
+    }
+
+    /// Number of tracked parameter tensors.
+    pub fn num_params(&self) -> usize {
+        self.shadow.len()
+    }
+
+    /// Total number of shadow parameter elements across all tracked tensors.
+    pub fn total_elements(&self) -> usize {
+        self.shadow.values().map(|v| v.len()).sum()
+    }
+
+    /// Current step count (number of [`update`](ModelEma::update) calls completed).
+    pub fn steps(&self) -> usize {
+        self.step
+    }
+
+    /// Get the shadow value for a named parameter, or `None` if not tracked.
+    pub fn shadow_for(&self, name: &str) -> Option<&[f32]> {
+        self.shadow.get(name).map(Vec::as_slice)
+    }
+
+    /// Compute the L2 distance between shadow weights and current training weights.
+    ///
+    /// Useful for monitoring how far the EMA weights have diverged from the live model.
+    /// Returns 0.0 if no shadow parameters are tracked.
+    pub fn shadow_drift(&self, current_params: &[(&str, &[f32])]) -> f32 {
+        let mut sum_sq = 0.0_f32;
+
+        for (name, current) in current_params {
+            if let Some(shadow) = self.shadow.get(*name) {
+                for (&s, &c) in shadow.iter().zip(current.iter()) {
+                    let diff = s - c;
+                    sum_sq += diff * diff;
+                }
+            }
+        }
+
+        sum_sq.sqrt()
+    }
+
+    /// Reset — clears all shadow state, backup state, and the step counter.
+    pub fn reset(&mut self) {
+        self.shadow.clear();
+        self.backup.clear();
+        self.shadow_applied = false;
+        self.step = 0;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Unit tests (RED → GREEN following TDD discipline)
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Helper: construct a simple EMA with no warmup for deterministic arithmetic.
+    fn ema_no_warmup(decay: f32) -> ModelEma {
+        ModelEma::new(ModelEmaConfig {
+            decay,
+            warmup_steps: 0,
+            use_bias_correction: false,
+        })
+    }
+
+    // ------------------------------------------------------------------
+    // 1. After first update, shadow == params
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_shadow_initialized_from_first_update() {
+        let mut ema = ema_no_warmup(0.9);
+        let params = [("w", &[1.0_f32, 2.0, 3.0][..])];
+        ema.update(&params);
+
+        // At step 0, effective_decay() was called *before* step was incremented.
+        // With warmup_steps == 0 and step == 0 the shadow slot did not exist yet —
+        // it is initialised verbatim from the data.
+        let shadow = ema.shadow_for("w").unwrap();
+        assert_eq!(shadow, &[1.0_f32, 2.0, 3.0]);
+    }
+
+    // ------------------------------------------------------------------
+    // 2. Second update blends values toward a weighted average
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_ema_decay_applied() {
+        let decay = 0.9_f32;
+        let mut ema = ema_no_warmup(decay);
+
+        // First update: shadow = [1, 2, 3]
+        ema.update(&[("w", &[1.0_f32, 2.0, 3.0])]);
+        // Second update with different values
+        ema.update(&[("w", &[3.0_f32, 4.0, 5.0])]);
+
+        // Expected: decay * old + (1-decay) * new
+        let expected: Vec<f32> = vec![
+            decay * 1.0 + (1.0 - decay) * 3.0,
+            decay * 2.0 + (1.0 - decay) * 4.0,
+            decay * 3.0 + (1.0 - decay) * 5.0,
+        ];
+        let shadow = ema.shadow_for("w").unwrap();
+
+        for (s, e) in shadow.iter().zip(expected.iter()) {
+            assert!(
+                (s - e).abs() < 1e-6,
+                "shadow={s} expected={e}"
+            );
+        }
+
+        // Shadow should be strictly between old and new values
+        for (&s, (&old, &new)) in shadow.iter().zip([1.0_f32, 2.0, 3.0].iter().zip([3.0_f32, 4.0, 5.0].iter())) {
+            assert!(s > old && s < new, "EMA should lie between old and new: s={s} old={old} new={new}");
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // 3. During warmup, step 0 → effective_decay ≈ 0
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_effective_decay_zero_at_start_warmup() {
+        let mut ema = ModelEma::new(ModelEmaConfig {
+            decay: 0.9999,
+            warmup_steps: 100,
+            use_bias_correction: false,
+        });
+        // step == 0 before any update
+        let d = ema.effective_decay();
+        assert!(
+            d < 0.01,
+            "effective_decay at step 0 with warmup should be near 0, got {d}"
+        );
+
+        // step == 1 after first update
+        ema.update(&[("w", &[1.0_f32])]);
+        let d1 = ema.effective_decay();
+        assert!(
+            d1 < 0.02,
+            "effective_decay at step 1 with 100-step warmup should be very small, got {d1}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // 4. After warmup, effective_decay == config.decay
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_effective_decay_full_after_warmup() {
+        let config_decay = 0.9999_f32;
+        let mut ema = ModelEma::new(ModelEmaConfig {
+            decay: config_decay,
+            warmup_steps: 5,
+            use_bias_correction: false,
+        });
+
+        // Drive past the warmup window
+        for _ in 0..10 {
+            ema.update(&[("w", &[1.0_f32])]);
+        }
+
+        let d = ema.effective_decay();
+        assert!(
+            (d - config_decay).abs() < 1e-7,
+            "effective_decay after warmup should equal config.decay={config_decay}, got {d}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // 5. apply_shadow() returns EMA values
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_apply_shadow_returns_shadow_values() {
+        let mut ema = ema_no_warmup(0.0); // decay=0: shadow always == last update
+        ema.update(&[("layer.weight", &[10.0_f32, 20.0, 30.0])]);
+
+        let shadow_map = ema.apply_shadow();
+        let vals = shadow_map.get("layer.weight").expect("key missing");
+        assert_eq!(vals.as_slice(), &[10.0_f32, 20.0, 30.0]);
+    }
+
+    // ------------------------------------------------------------------
+    // 6. restore() returns the stored backup
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_restore_returns_backup() {
+        let mut ema = ema_no_warmup(0.9);
+        ema.update(&[("w", &[1.0_f32, 2.0])]);
+
+        // Store the "live" training params as backup before applying shadow.
+        let live_params = [("w", &[5.0_f32, 6.0][..])];
+        ema.store_backup(&live_params);
+        let _shadow = ema.apply_shadow();
+
+        let backup = ema.restore();
+        let restored = backup.get("w").expect("backup missing");
+        assert_eq!(restored.as_slice(), &[5.0_f32, 6.0]);
+    }
+
+    // ------------------------------------------------------------------
+    // 7. shadow_applied flag tracks apply/restore correctly
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_shadow_applied_flag() {
+        let mut ema = ema_no_warmup(0.9);
+        ema.update(&[("w", &[1.0_f32])]);
+
+        assert!(!ema.shadow_applied(), "shadow should not be applied initially");
+
+        let _ = ema.apply_shadow();
+        assert!(ema.shadow_applied(), "shadow should be applied after apply_shadow()");
+
+        let _ = ema.restore();
+        assert!(!ema.shadow_applied(), "shadow should not be applied after restore()");
+    }
+
+    // ------------------------------------------------------------------
+    // 8. shadow_drift == 0 when shadow and current are identical
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_shadow_drift_zero_when_identical() {
+        let mut ema = ema_no_warmup(0.0); // decay=0 → shadow == last param
+        let data = [1.0_f32, 2.0, 3.0];
+        ema.update(&[("w", &data)]);
+
+        // Shadow equals the data that was just pushed in.
+        let drift = ema.shadow_drift(&[("w", &data)]);
+        assert!(
+            drift.abs() < 1e-6,
+            "drift should be 0 when shadow == current, got {drift}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // 9. shadow_drift > 0 when shadow differs from current
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_shadow_drift_positive() {
+        let mut ema = ema_no_warmup(1.0); // decay=1 → shadow never changes after init
+        ema.update(&[("w", &[0.0_f32, 0.0, 0.0])]);
+        // Update again; shadow stays at [0,0,0] because decay=1
+        ema.update(&[("w", &[1.0_f32, 1.0, 1.0])]);
+
+        let current = [("w", &[1.0_f32, 1.0, 1.0][..])];
+        let drift = ema.shadow_drift(&current);
+        assert!(drift > 0.0, "drift should be positive when shadow != current, got {drift}");
+    }
+
+    // ------------------------------------------------------------------
+    // 10. num_params counts distinct parameter tensors
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_num_params_counts_tensors() {
+        let mut ema = ema_no_warmup(0.9);
+        assert_eq!(ema.num_params(), 0);
+
+        ema.update(&[
+            ("a", &[1.0_f32, 2.0]),
+            ("b", &[3.0_f32]),
+            ("c", &[4.0_f32, 5.0, 6.0]),
+        ]);
+
+        assert_eq!(ema.num_params(), 3);
+    }
+
+    // ------------------------------------------------------------------
+    // 11. total_elements sums element counts across all params
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_total_elements_sum() {
+        let mut ema = ema_no_warmup(0.9);
+        ema.update(&[
+            ("a", &[1.0_f32, 2.0]),        // 2 elements
+            ("b", &[3.0_f32]),              // 1 element
+            ("c", &[4.0_f32, 5.0, 6.0]),   // 3 elements
+        ]);
+
+        assert_eq!(ema.total_elements(), 6);
+    }
+
+    // ------------------------------------------------------------------
+    // 12. reset() clears all state including step counter
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_reset_clears_state() {
+        let mut ema = ema_no_warmup(0.9);
+        ema.update(&[("w", &[1.0_f32, 2.0, 3.0])]);
+        ema.update(&[("w", &[4.0_f32, 5.0, 6.0])]);
+
+        assert_eq!(ema.num_params(), 1);
+        assert_eq!(ema.steps(), 2);
+
+        ema.reset();
+
+        assert_eq!(ema.num_params(), 0, "reset should clear shadow params");
+        assert_eq!(ema.steps(), 0, "reset should clear step counter");
+        assert!(!ema.shadow_applied(), "reset should clear shadow_applied flag");
+        assert!(ema.shadow_for("w").is_none(), "reset should remove all shadow entries");
+    }
+
+    // ------------------------------------------------------------------
+    // 13. standard() convenience constructor uses expected defaults
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_standard_convenience_ctor() {
+        let ema = ModelEma::standard();
+        assert!(
+            (ema.config.decay - 0.9999_f32).abs() < 1e-7,
+            "standard() should have decay=0.9999, got {}",
+            ema.config.decay
+        );
+        assert_eq!(ema.config.warmup_steps, 100);
+        assert!(!ema.config.use_bias_correction);
+        assert_eq!(ema.num_params(), 0);
+        assert_eq!(ema.steps(), 0);
+    }
+
+    // ------------------------------------------------------------------
+    // 14. Bias correction warms decay up from near-0 to config.decay
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_bias_correction_warmup_behavior() {
+        let mut ema = ModelEma::new(ModelEmaConfig {
+            decay: 0.9_f32,
+            warmup_steps: 0,
+            use_bias_correction: true,
+        });
+
+        // At step 0, (1 - 0.9^0) = 0 → effective = 0
+        let d0 = ema.effective_decay();
+        assert!(d0.abs() < 1e-6, "bias-corrected decay at step 0 should be ~0, got {d0}");
+
+        // Drive forward many steps
+        for _ in 0..1000 {
+            ema.update(&[("w", &[1.0_f32])]);
+        }
+
+        // After many steps, bias-corrected decay should have saturated to config.decay
+        let d_late = ema.effective_decay();
+        assert!(
+            (d_late - 0.9_f32).abs() < 1e-4,
+            "bias-corrected decay should approach config.decay after many steps, got {d_late}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // 15. restore() is a no-op when shadow is not applied
+    // ------------------------------------------------------------------
+    #[test]
+    fn test_restore_noop_when_not_applied() {
+        let mut ema = ema_no_warmup(0.9);
+        ema.update(&[("w", &[1.0_f32])]);
+
+        // Shadow not applied — restore should return empty map
+        let result = ema.restore();
+        assert!(result.is_empty(), "restore when not applied should return empty map");
+        assert!(!ema.shadow_applied());
+    }
+}