368 lines
13 KiB
Rust
368 lines
13 KiB
Rust
use super::*;
|
|
|
|
impl Executor {
|
|
pub(in crate::engine) fn save_checkpoint(
|
|
&mut self,
|
|
path: &Path,
|
|
tag: [u8; 32],
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
|
|
}
|
|
let temporary = path.with_extension("tmp");
|
|
let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
|
|
let mut reported = 0_u64;
|
|
let mut pending = 0_u64;
|
|
self.write_checkpoint(&mut file, tag, &mut |bytes| {
|
|
reported += bytes;
|
|
pending += bytes;
|
|
if pending >= CHECKPOINT_IO_CHUNK as u64 {
|
|
progress(pending);
|
|
pending = 0;
|
|
}
|
|
})?;
|
|
progress(pending);
|
|
let total = file.metadata().map_err(|error| error.to_string())?.len();
|
|
progress(total.saturating_sub(reported));
|
|
file.sync_all().map_err(|error| error.to_string())?;
|
|
fs::rename(&temporary, path).map_err(|error| error.to_string())?;
|
|
self.checkpoint_tag = tag;
|
|
Ok(())
|
|
}
|
|
|
|
pub(in crate::engine) fn load_checkpoint(
|
|
&mut self,
|
|
path: &Path,
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<bool, String> {
|
|
let mut file = match File::open(path) {
|
|
Ok(file) => file,
|
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
|
Err(error) => return Err(error.to_string()),
|
|
};
|
|
let mut reported = 0_u64;
|
|
let mut pending = 0_u64;
|
|
self.read_checkpoint(&mut file, &mut |bytes| {
|
|
reported += bytes;
|
|
pending += bytes;
|
|
if pending >= CHECKPOINT_IO_CHUNK as u64 {
|
|
progress(pending);
|
|
pending = 0;
|
|
}
|
|
})?;
|
|
progress(pending);
|
|
let total = file.metadata().map_err(|error| error.to_string())?.len();
|
|
progress(total.saturating_sub(reported));
|
|
Ok(true)
|
|
}
|
|
|
|
fn write_checkpoint(
|
|
&self,
|
|
file: &mut File,
|
|
tag: [u8; 32],
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
let shape = self.model.shape;
|
|
file.write_all(CHECKPOINT_MAGIC)
|
|
.map_err(|error| error.to_string())?;
|
|
for value in [
|
|
CHECKPOINT_VERSION,
|
|
self.session.context,
|
|
self.session.raw_cap,
|
|
shape.layers,
|
|
shape.head_dim as u32,
|
|
shape.indexer_head_dim as u32,
|
|
shape.vocab as u32,
|
|
u32::from(self.quality),
|
|
] {
|
|
write_u32(file, value)?;
|
|
}
|
|
write_u64(file, self.model.main.len())?;
|
|
write_u64(file, self.model_modified.0)?;
|
|
write_u32(file, self.model_modified.1)?;
|
|
for weight in [self.weights.token_embedding, self.weights.output] {
|
|
write_u64(file, weight.offset)?;
|
|
write_u64(file, weight.bytes)?;
|
|
write_u32(file, weight.kind)?;
|
|
}
|
|
file.write_all(&tag).map_err(|error| error.to_string())?;
|
|
let token_count = u32::try_from(self.tokens.len())
|
|
.map_err(|_| "KV checkpoint has too many tokens".to_owned())?;
|
|
let raw_live = token_count.min(self.session.raw_cap);
|
|
write_u32(file, token_count)?;
|
|
write_u32(file, raw_live)?;
|
|
for &token in &self.tokens {
|
|
write_u32(file, token as u32)?;
|
|
}
|
|
for &logit in &self.logits {
|
|
write_u32(file, logit.to_bits())?;
|
|
}
|
|
for layer in &self.session.layers {
|
|
write_u32(
|
|
file,
|
|
layer.compression.as_ref().map_or(0, |state| state.rows),
|
|
)?;
|
|
}
|
|
for layer in &self.session.layers {
|
|
write_u32(file, layer.indexer.as_ref().map_or(0, |state| state.rows))?;
|
|
}
|
|
|
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
|
let raw_first = token_count - raw_live;
|
|
for layer in &self.session.layers {
|
|
for position in raw_first..token_count {
|
|
let physical = position % self.session.raw_cap;
|
|
write_buffer(
|
|
file,
|
|
&layer.raw_cache,
|
|
u64::from(physical) * shape.head_dim * 4,
|
|
shape.head_dim * 4,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
}
|
|
if let Some(state) = &layer.compression {
|
|
write_buffer(
|
|
file,
|
|
&state.cache,
|
|
0,
|
|
u64::from(state.rows) * shape.head_dim * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
|
|
write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
|
write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
|
}
|
|
if let Some(state) = &layer.indexer {
|
|
write_buffer(
|
|
file,
|
|
&state.cache,
|
|
0,
|
|
u64::from(state.rows) * shape.indexer_head_dim * 4,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
|
|
write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
|
write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn read_checkpoint(
|
|
&mut self,
|
|
file: &mut File,
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
let mut magic = [0; 8];
|
|
file.read_exact(&mut magic)
|
|
.map_err(|error| error.to_string())?;
|
|
if &magic != CHECKPOINT_MAGIC {
|
|
return Err("KV checkpoint has an invalid signature".into());
|
|
}
|
|
let shape = self.model.shape;
|
|
let header = [
|
|
CHECKPOINT_VERSION,
|
|
self.session.context,
|
|
self.session.raw_cap,
|
|
shape.layers,
|
|
shape.head_dim as u32,
|
|
shape.indexer_head_dim as u32,
|
|
shape.vocab as u32,
|
|
u32::from(self.quality),
|
|
];
|
|
for expected in header {
|
|
if read_u32(file)? != expected {
|
|
return Err("KV checkpoint does not match the current executor".into());
|
|
}
|
|
}
|
|
if read_u64(file)? != self.model.main.len() {
|
|
return Err("KV checkpoint was written for a different model".into());
|
|
}
|
|
if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 {
|
|
return Err("KV checkpoint model file has changed".into());
|
|
}
|
|
for weight in [self.weights.token_embedding, self.weights.output] {
|
|
if read_u64(file)? != weight.offset
|
|
|| read_u64(file)? != weight.bytes
|
|
|| read_u32(file)? != weight.kind
|
|
{
|
|
return Err("KV checkpoint was written for a different model layout".into());
|
|
}
|
|
}
|
|
let mut checkpoint_tag = [0; 32];
|
|
file.read_exact(&mut checkpoint_tag)
|
|
.map_err(|error| error.to_string())?;
|
|
let token_count = read_u32(file)?;
|
|
let raw_live = read_u32(file)?;
|
|
if token_count > self.session.context || raw_live != token_count.min(self.session.raw_cap) {
|
|
return Err("KV checkpoint token count is invalid".into());
|
|
}
|
|
let mut tokens = Vec::with_capacity(token_count as usize);
|
|
for _ in 0..token_count {
|
|
let token = read_u32(file)?;
|
|
if u64::from(token) >= shape.vocab {
|
|
return Err("KV checkpoint contains an invalid token".into());
|
|
}
|
|
tokens.push(token as i32);
|
|
}
|
|
let mut logits = Vec::with_capacity(shape.vocab as usize);
|
|
for _ in 0..shape.vocab {
|
|
logits.push(f32::from_bits(read_u32(file)?));
|
|
}
|
|
let mut compressed_rows = Vec::with_capacity(shape.layers as usize);
|
|
let mut indexer_rows = Vec::with_capacity(shape.layers as usize);
|
|
for layer in 0..shape.layers {
|
|
let rows = read_u32(file)?;
|
|
let ratio = compression_ratio(layer);
|
|
if rows != token_count.checked_div(ratio).unwrap_or(0) {
|
|
return Err("KV checkpoint compressed row count is invalid".into());
|
|
}
|
|
compressed_rows.push(rows);
|
|
}
|
|
for layer in 0..shape.layers {
|
|
let rows = read_u32(file)?;
|
|
let expected = if compression_ratio(layer) == 4 {
|
|
token_count / 4
|
|
} else {
|
|
0
|
|
};
|
|
if rows != expected {
|
|
return Err("KV checkpoint indexer row count is invalid".into());
|
|
}
|
|
indexer_rows.push(rows);
|
|
}
|
|
|
|
self.reset()?;
|
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
|
let raw_first = token_count - raw_live;
|
|
for (index, layer) in self.session.layers.iter_mut().enumerate() {
|
|
for position in raw_first..token_count {
|
|
let physical = position % self.session.raw_cap;
|
|
read_buffer(
|
|
file,
|
|
&layer.raw_cache,
|
|
u64::from(physical) * shape.head_dim * 4,
|
|
shape.head_dim * 4,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
}
|
|
if let Some(state) = &mut layer.compression {
|
|
state.rows = compressed_rows[index];
|
|
read_buffer(
|
|
file,
|
|
&state.cache,
|
|
0,
|
|
u64::from(state.rows) * shape.head_dim * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
|
|
read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
|
read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
|
}
|
|
if let Some(state) = &mut layer.indexer {
|
|
state.rows = indexer_rows[index];
|
|
read_buffer(
|
|
file,
|
|
&state.cache,
|
|
0,
|
|
u64::from(state.rows) * shape.indexer_head_dim * 4,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
|
|
read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?;
|
|
read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?;
|
|
}
|
|
}
|
|
let mut trailing = [0];
|
|
if file
|
|
.read(&mut trailing)
|
|
.map_err(|error| error.to_string())?
|
|
!= 0
|
|
{
|
|
self.reset()?;
|
|
return Err("KV checkpoint has trailing data".into());
|
|
}
|
|
self.session.position = token_count;
|
|
self.tokens = tokens;
|
|
self.logits = logits;
|
|
self.checkpoint_tag = checkpoint_tag;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 {
|
|
let coefficient = if ratio == 4 { 2 } else { 1 };
|
|
coefficient * head_dim * coefficient * u64::from(ratio) * 4
|
|
}
|
|
|
|
fn write_buffer(
|
|
file: &mut File,
|
|
buffer: &Buffer,
|
|
mut offset: u64,
|
|
mut bytes: u64,
|
|
chunk: &mut [u8],
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
while bytes != 0 {
|
|
let length = bytes.min(chunk.len() as u64) as usize;
|
|
buffer.read(offset, &mut chunk[..length])?;
|
|
file.write_all(&chunk[..length])
|
|
.map_err(|error| error.to_string())?;
|
|
progress(length as u64);
|
|
offset += length as u64;
|
|
bytes -= length as u64;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn read_buffer(
|
|
file: &mut File,
|
|
buffer: &Buffer,
|
|
mut offset: u64,
|
|
mut bytes: u64,
|
|
chunk: &mut [u8],
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
while bytes != 0 {
|
|
let length = bytes.min(chunk.len() as u64) as usize;
|
|
file.read_exact(&mut chunk[..length])
|
|
.map_err(|error| error.to_string())?;
|
|
buffer.write(offset, &chunk[..length])?;
|
|
progress(length as u64);
|
|
offset += length as u64;
|
|
bytes -= length as u64;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
|
|
file.write_all(&value.to_le_bytes())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
|
|
file.write_all(&value.to_le_bytes())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn read_u32(file: &mut File) -> Result<u32, String> {
|
|
let mut bytes = [0; 4];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| error.to_string())?;
|
|
Ok(u32::from_le_bytes(bytes))
|
|
}
|
|
|
|
fn read_u64(file: &mut File) -> Result<u64, String> {
|
|
let mut bytes = [0; 8];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| error.to_string())?;
|
|
Ok(u64::from_le_bytes(bytes))
|
|
}
|