diff --git a/crates/ml-alpha/Cargo.toml b/crates/ml-alpha/Cargo.toml index 22d7f7dab..55d63608c 100644 --- a/crates/ml-alpha/Cargo.toml +++ b/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"] } diff --git a/crates/ml-alpha/src/trainer/optim.rs b/crates/ml-alpha/src/trainer/optim.rs index d7401d95f..e31db2939 100644 --- a/crates/ml-alpha/src/trainer/optim.rs +++ b/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(()) + } }