feat(rl): AdamW save/load for checkpoint persistence

Serializes m, v moments + step_count + hyperparams to a binary writer.
Uses mapped-pinned DtoD (not raw memcpy_htod/dtoh) per
feedback_no_htod_htoh_only_mapped_pinned.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/Cargo.toml

@@ -33,6 +33,7 @@ tracing-subscriber.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
 bincode.workspace = true
+bytemuck = { workspace = true }
 thiserror.workspace = true
 clap = { workspace = true, features = ["derive"] }
 

crates/ml-alpha/src/trainer/optim.rs

@@ -10,6 +10,7 @@
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
+use bytemuck;
 use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
 use cudarc::driver::sys::CUstream;
 use ml_core::device::MlDevice;
@@ -126,4 +127,98 @@ impl AdamW {
     pub fn step_count(&self) -> i32 {
         self.step_count_host
     }
+
+    /// Serialize Adam state (m, v, step_count, hyperparams) to writer.
+    ///
+    /// Wire format (little-endian):
+    ///   u64  n          — number of parameters
+    ///   i32  step_count
+    ///   f32  lr, beta1, beta2, eps, wd
+    ///   [f32; n]  m     — first moment
+    ///   [f32; n]  v     — second moment
+    pub fn save(&self, w: &mut impl std::io::Write) -> Result<()> {
+        use crate::pinned_mem::MappedF32Buffer;
+        use crate::trainer::raw_launch::raw_memcpy_dtod_async;
+        let n = self.m.len();
+        let staging = unsafe { MappedF32Buffer::new(n) }
+            .map_err(|e| anyhow::anyhow!("AdamW save staging: {e}"))?;
+        let raw_s = self.raw_stream;
+        unsafe {
+            raw_memcpy_dtod_async(staging.dev_ptr, self.m.raw_ptr(), n * 4, raw_s)
+                .map_err(|e| anyhow::anyhow!("AdamW save m dtod: {e:?}"))?;
+        }
+        self._stream.synchronize()
+            .map_err(|e| anyhow::anyhow!("AdamW save sync m: {e}"))?;
+        let m_host = staging.read_all();
+        unsafe {
+            raw_memcpy_dtod_async(staging.dev_ptr, self.v.raw_ptr(), n * 4, raw_s)
+                .map_err(|e| anyhow::anyhow!("AdamW save v dtod: {e:?}"))?;
+        }
+        self._stream.synchronize()
+            .map_err(|e| anyhow::anyhow!("AdamW save sync v: {e}"))?;
+        let v_host = staging.read_all();
+        w.write_all(&(n as u64).to_le_bytes())?;
+        w.write_all(&self.step_count_host.to_le_bytes())?;
+        w.write_all(&self.lr.to_le_bytes())?;
+        w.write_all(&self.beta1.to_le_bytes())?;
+        w.write_all(&self.beta2.to_le_bytes())?;
+        w.write_all(&self.eps.to_le_bytes())?;
+        w.write_all(&self.wd.to_le_bytes())?;
+        w.write_all(bytemuck::cast_slice(&m_host))?;
+        w.write_all(bytemuck::cast_slice(&v_host))?;
+        Ok(())
+    }
+
+    /// Restore Adam state from reader. Parameter count must match allocation.
+    pub fn load(&mut self, r: &mut impl std::io::Read) -> Result<()> {
+        let mut buf8 = [0u8; 8];
+        let mut buf4 = [0u8; 4];
+        r.read_exact(&mut buf8)?;
+        let n = u64::from_le_bytes(buf8) as usize;
+        anyhow::ensure!(
+            n == self.m.len(),
+            "AdamW load: m size mismatch {n} vs {}",
+            self.m.len()
+        );
+        r.read_exact(&mut buf4)?;
+        self.step_count_host = i32::from_le_bytes(buf4);
+        r.read_exact(&mut buf4)?;
+        self.lr = f32::from_le_bytes(buf4);
+        r.read_exact(&mut buf4)?;
+        self.beta1 = f32::from_le_bytes(buf4);
+        r.read_exact(&mut buf4)?;
+        self.beta2 = f32::from_le_bytes(buf4);
+        r.read_exact(&mut buf4)?;
+        self.eps = f32::from_le_bytes(buf4);
+        r.read_exact(&mut buf4)?;
+        self.wd = f32::from_le_bytes(buf4);
+        use crate::pinned_mem::MappedF32Buffer;
+        use crate::trainer::raw_launch::raw_memcpy_dtod_async;
+        let staging = unsafe { MappedF32Buffer::new(n) }
+            .map_err(|e| anyhow::anyhow!("AdamW load staging: {e}"))?;
+        let raw_s = self.raw_stream;
+        // m: read file → mapped-pinned host_ptr → DtoD → device
+        {
+            let host_slice = unsafe { std::slice::from_raw_parts_mut(staging.host_ptr, n) };
+            r.read_exact(bytemuck::cast_slice_mut(host_slice))?;
+            unsafe {
+                raw_memcpy_dtod_async(self.m.raw_ptr(), staging.dev_ptr, n * 4, raw_s)
+                    .map_err(|e| anyhow::anyhow!("AdamW load m dtod: {e:?}"))?;
+            }
+            self._stream.synchronize()
+                .map_err(|e| anyhow::anyhow!("AdamW load sync m: {e}"))?;
+        }
+        // v: same pattern
+        {
+            let host_slice = unsafe { std::slice::from_raw_parts_mut(staging.host_ptr, n) };
+            r.read_exact(bytemuck::cast_slice_mut(host_slice))?;
+            unsafe {
+                raw_memcpy_dtod_async(self.v.raw_ptr(), staging.dev_ptr, n * 4, raw_s)
+                    .map_err(|e| anyhow::anyhow!("AdamW load v dtod: {e:?}"))?;
+            }
+            self._stream.synchronize()
+                .map_err(|e| anyhow::anyhow!("AdamW load sync v: {e}"))?;
+        }
+        Ok(())
+    }
 }