From 2265f1c18b0272c47a3c002afa7277362cc13ff9 Mon Sep 17 00:00:00 2001 From: sharmis-tp Date: Tue, 10 Nov 2026 07:40:17 +0000 Subject: [PATCH 1/2] batch tensor I/O into single read/write calls in checkpoint serialization --- src/serialize/checkpoint.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/serialize/checkpoint.rs b/src/serialize/checkpoint.rs index 2e4b336..c025ed3 100644 --- a/src/serialize/checkpoint.rs +++ b/src/serialize/checkpoint.rs @@ -87,9 +87,12 @@ pub fn save_tensors(tensors: &Checkpoint, path: impl AsRef) -> Result<()> for &dim in tensor.shape() { write_u32(&mut out, dim as u32)?; } - for value in tensor.to_vec() { - out.write_all(&value.to_le_bytes())?; + let values = tensor.to_vec(); + let mut buf = vec![0u8; values.len() * 4]; + for (i, value) in values.iter().enumerate() { + buf[i * 4..i * 4 + 4].copy_from_slice(&value.to_le_bytes()); } + out.write_all(&buf)?; } out.flush()?; @@ -125,11 +128,17 @@ pub fn load_tensors(path: impl AsRef) -> Result { .map(|_| read_u32(&mut input).map(|d| d as usize)) .collect::>()?; - let mut data = vec![0.0f32; shape.iter().product()]; - for value in &mut data { - let mut bytes = [0u8; 4]; - input.read_exact(&mut bytes)?; - *value = f32::from_le_bytes(bytes); + let numel: usize = shape.iter().product(); + let mut buf = vec![0u8; numel * 4]; + input.read_exact(&mut buf)?; + let mut data = vec![0.0f32; numel]; + for (i, value) in data.iter_mut().enumerate() { + *value = f32::from_le_bytes([ + buf[i * 4], + buf[i * 4 + 1], + buf[i * 4 + 2], + buf[i * 4 + 3], + ]); } tensors.insert(name, Tensor::from_vec(data, &shape)); From b271215ccadb9af243ba78231b023cbad2c6d4e0 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Thu, 12 Nov 2026 12:40:38 +0300 Subject: [PATCH 2/2] format --- src/serialize/checkpoint.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/serialize/checkpoint.rs b/src/serialize/checkpoint.rs index c025ed3..4d076c1 100644 --- a/src/serialize/checkpoint.rs +++ b/src/serialize/checkpoint.rs @@ -133,12 +133,8 @@ pub fn load_tensors(path: impl AsRef) -> Result { input.read_exact(&mut buf)?; let mut data = vec![0.0f32; numel]; for (i, value) in data.iter_mut().enumerate() { - *value = f32::from_le_bytes([ - buf[i * 4], - buf[i * 4 + 1], - buf[i * 4 + 2], - buf[i * 4 + 3], - ]); + *value = + f32::from_le_bytes([buf[i * 4], buf[i * 4 + 1], buf[i * 4 + 2], buf[i * 4 + 3]]); } tensors.insert(name, Tensor::from_vec(data, &shape));