Implement DS4 model intake
This commit is contained in:
Georg Bauer
509
src/engine/gguf.rs
Normal file
509
src/engine/gguf.rs
Normal file
@@ -0,0 +1,509 @@
|
||||
use memmap2::{Mmap, MmapOptions};
|
||||
use std::collections::HashMap;
|
||||
use std::fs::File;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const MAGIC: u32 = 0x4655_4747;
|
||||
const MAX_COUNT: u64 = 1_000_000;
|
||||
const MAX_DIMS: usize = 8;
|
||||
const MAX_DEPTH: u8 = 8;
|
||||
|
||||
pub(super) const F32: u32 = 0;
|
||||
pub(super) const F16: u32 = 1;
|
||||
pub(super) const Q4_0: u32 = 2;
|
||||
pub(super) const Q8_0: u32 = 8;
|
||||
pub(super) const Q2_K: u32 = 10;
|
||||
pub(super) const Q4_K: u32 = 12;
|
||||
pub(super) const Q5_K: u32 = 13;
|
||||
pub(super) const Q6_K: u32 = 14;
|
||||
pub(super) const IQ2_XXS: u32 = 16;
|
||||
pub(super) const I32: u32 = 26;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(super) enum Value {
|
||||
U32(u32),
|
||||
I32(i32),
|
||||
U64(u64),
|
||||
I64(i64),
|
||||
F32(f32),
|
||||
F64(f64),
|
||||
Bool(bool),
|
||||
Bytes(Vec<u8>),
|
||||
U32s(Vec<u32>),
|
||||
F32s(Vec<f32>),
|
||||
Strings(Vec<Vec<u8>>),
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(super) struct Tensor {
|
||||
pub(super) kind: u32,
|
||||
pub(super) dims: Vec<u64>,
|
||||
pub(super) offset: u64,
|
||||
pub(super) bytes: u64,
|
||||
}
|
||||
|
||||
pub(super) struct Gguf {
|
||||
path: PathBuf,
|
||||
map: Mmap,
|
||||
pub(super) metadata: HashMap<String, Value>,
|
||||
pub(super) tensors: HashMap<String, Tensor>,
|
||||
}
|
||||
|
||||
impl Gguf {
|
||||
pub(super) fn open(path: &Path) -> Result<Self, String> {
|
||||
let file =
|
||||
File::open(path).map_err(|error| format!("Cannot open {}: {error}", path.display()))?;
|
||||
if file.metadata().map_err(|error| error.to_string())?.len() < 32 {
|
||||
return Err(format!("{} is too small to be GGUF", path.display()));
|
||||
}
|
||||
// SAFETY: managed model artifacts are opened read-only and are never mutated
|
||||
// while a Model owns this mapping.
|
||||
let map = unsafe { MmapOptions::new().map(&file) }
|
||||
.map_err(|error| format!("Cannot map {}: {error}", path.display()))?;
|
||||
let mut cursor = Cursor::new(&map);
|
||||
if cursor.u32()? != MAGIC {
|
||||
return Err(format!("{} is not a GGUF file", path.display()));
|
||||
}
|
||||
if cursor.u32()? != 3 {
|
||||
return Err(format!("{} is not GGUF v3", path.display()));
|
||||
}
|
||||
let tensor_count = cursor.u64()?;
|
||||
let metadata_count = cursor.u64()?;
|
||||
if tensor_count > MAX_COUNT || metadata_count > MAX_COUNT {
|
||||
return Err("GGUF directory count is too large".into());
|
||||
}
|
||||
|
||||
let mut metadata = HashMap::with_capacity(metadata_count as usize);
|
||||
let mut alignment = 32_u64;
|
||||
for _ in 0..metadata_count {
|
||||
let key = String::from_utf8(cursor.string()?.to_vec())
|
||||
.map_err(|_| "GGUF metadata key is not UTF-8".to_owned())?;
|
||||
let kind = cursor.u32()?;
|
||||
let value = cursor.value(kind, &key, 0)?;
|
||||
if key == "general.alignment"
|
||||
&& let Value::U32(value) = value
|
||||
{
|
||||
alignment = u64::from(value);
|
||||
}
|
||||
if metadata.insert(key.clone(), value).is_some() {
|
||||
return Err(format!("duplicate GGUF metadata key: {key}"));
|
||||
}
|
||||
}
|
||||
if alignment == 0 {
|
||||
return Err("GGUF alignment is zero".into());
|
||||
}
|
||||
|
||||
let mut tensors = HashMap::with_capacity(tensor_count as usize);
|
||||
for _ in 0..tensor_count {
|
||||
let name = String::from_utf8(cursor.string()?.to_vec())
|
||||
.map_err(|_| "GGUF tensor name is not UTF-8".to_owned())?;
|
||||
let ndim = cursor.u32()? as usize;
|
||||
if ndim == 0 || ndim > MAX_DIMS {
|
||||
return Err(format!(
|
||||
"tensor {name} has unsupported dimension count {ndim}"
|
||||
));
|
||||
}
|
||||
let mut dims = Vec::with_capacity(ndim);
|
||||
let mut elements = 1_u64;
|
||||
for _ in 0..ndim {
|
||||
let dim = cursor.u64()?;
|
||||
elements = elements
|
||||
.checked_mul(dim)
|
||||
.ok_or_else(|| format!("tensor {name} element count overflows"))?;
|
||||
dims.push(dim);
|
||||
}
|
||||
let kind = cursor.u32()?;
|
||||
let relative_offset = cursor.u64()?;
|
||||
let (block_elements, block_bytes) = tensor_type(kind)
|
||||
.ok_or_else(|| format!("tensor {name} uses unsupported GGUF type {kind}"))?;
|
||||
let blocks = elements
|
||||
.checked_add(block_elements - 1)
|
||||
.ok_or_else(|| format!("tensor {name} size overflows"))?
|
||||
/ block_elements;
|
||||
let bytes = blocks
|
||||
.checked_mul(block_bytes)
|
||||
.ok_or_else(|| format!("tensor {name} size overflows"))?;
|
||||
if tensors
|
||||
.insert(
|
||||
name.clone(),
|
||||
Tensor {
|
||||
kind,
|
||||
dims,
|
||||
offset: relative_offset,
|
||||
bytes,
|
||||
},
|
||||
)
|
||||
.is_some()
|
||||
{
|
||||
return Err(format!("duplicate GGUF tensor: {name}"));
|
||||
}
|
||||
}
|
||||
|
||||
let data_start = align(cursor.position() as u64, alignment)?;
|
||||
for (name, tensor) in &mut tensors {
|
||||
tensor.offset = data_start
|
||||
.checked_add(tensor.offset)
|
||||
.ok_or_else(|| format!("tensor {name} offset overflows"))?;
|
||||
let end = tensor
|
||||
.offset
|
||||
.checked_add(tensor.bytes)
|
||||
.ok_or_else(|| format!("tensor {name} end overflows"))?;
|
||||
if end > map.len() as u64 {
|
||||
return Err(format!("tensor {name} points outside the GGUF file"));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
path: path.to_owned(),
|
||||
map,
|
||||
metadata,
|
||||
tensors,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub(super) fn len(&self) -> u64 {
|
||||
self.map.len() as u64
|
||||
}
|
||||
|
||||
pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
|
||||
self.tensors
|
||||
.get(name)
|
||||
.ok_or_else(|| format!("required tensor is missing: {name}"))
|
||||
}
|
||||
|
||||
pub(super) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
|
||||
let tensor = self.tensor(name)?;
|
||||
let start = tensor.offset as usize;
|
||||
let end = start + tensor.bytes as usize;
|
||||
Ok(&self.map[start..end])
|
||||
}
|
||||
|
||||
pub(super) fn u32(&self, key: &str) -> Result<u32, String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::U32(value)) => Ok(*value),
|
||||
_ => Err(format!("required uint32 metadata key is missing: {key}")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn u64(&self, key: &str) -> Result<u64, String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::U64(value)) => Ok(*value),
|
||||
Some(Value::U32(value)) => Ok(u64::from(*value)),
|
||||
_ => Err(format!("required integer metadata key is missing: {key}")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn f32(&self, key: &str) -> Result<f32, String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::F32(value)) => Ok(*value),
|
||||
Some(Value::F64(value)) => Ok(*value as f32),
|
||||
Some(Value::U32(value)) => Ok(*value as f32),
|
||||
Some(Value::I32(value)) => Ok(*value as f32),
|
||||
_ => Err(format!("required numeric metadata key is missing: {key}")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn token_id(&self, key: &str) -> Result<i32, String> {
|
||||
let value = match self.metadata.get(key) {
|
||||
Some(Value::U32(value)) => i64::from(*value),
|
||||
Some(Value::I32(value)) => i64::from(*value),
|
||||
Some(Value::U64(value)) => i64::try_from(*value).unwrap_or(-1),
|
||||
Some(Value::I64(value)) => *value,
|
||||
_ => -1,
|
||||
};
|
||||
i32::try_from(value)
|
||||
.ok()
|
||||
.filter(|value| *value >= 0)
|
||||
.ok_or_else(|| format!("required tokenizer token ID is missing: {key}"))
|
||||
}
|
||||
|
||||
pub(super) fn boolean(&self, key: &str) -> Result<bool, String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::Bool(value)) => Ok(*value),
|
||||
_ => Err(format!("required boolean metadata key is missing: {key}")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn bytes(&self, key: &str) -> Result<&[u8], String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::Bytes(value)) => Ok(value),
|
||||
_ => Err(format!("required string metadata key is missing: {key}")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn u32s(&self, key: &str) -> Result<&[u32], String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::U32s(value)) => Ok(value),
|
||||
_ => Err(format!(
|
||||
"required integer-array metadata key is missing: {key}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn f32s(&self, key: &str) -> Result<&[f32], String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::F32s(value)) => Ok(value),
|
||||
_ => Err(format!(
|
||||
"required float-array metadata key is missing: {key}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn strings(&self, key: &str) -> Result<&[Vec<u8>], String> {
|
||||
match self.metadata.get(key) {
|
||||
Some(Value::Strings(value)) => Ok(value),
|
||||
_ => Err(format!(
|
||||
"required string-array metadata key is missing: {key}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tensor_type(kind: u32) -> Option<(u64, u64)> {
|
||||
Some(match kind {
|
||||
0 => (1, 4),
|
||||
1 => (1, 2),
|
||||
2 => (32, 18),
|
||||
3 => (32, 20),
|
||||
6 => (32, 22),
|
||||
7 => (32, 24),
|
||||
8 => (32, 34),
|
||||
9 => (32, 40),
|
||||
10 => (256, 84),
|
||||
11 => (256, 110),
|
||||
12 => (256, 144),
|
||||
13 => (256, 176),
|
||||
14 => (256, 210),
|
||||
15 => (256, 292),
|
||||
16 => (256, 66),
|
||||
17 => (256, 74),
|
||||
18 => (256, 98),
|
||||
19 => (256, 110),
|
||||
20 => (256, 50),
|
||||
21 => (256, 110),
|
||||
22 => (256, 82),
|
||||
23 => (256, 136),
|
||||
24 => (1, 1),
|
||||
25 => (1, 2),
|
||||
26 => (1, 4),
|
||||
27 => (1, 8),
|
||||
28 => (1, 8),
|
||||
29 => (256, 56),
|
||||
30 => (1, 2),
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
fn align(value: u64, alignment: u64) -> Result<u64, String> {
|
||||
let remainder = value % alignment;
|
||||
if remainder == 0 {
|
||||
Ok(value)
|
||||
} else {
|
||||
value
|
||||
.checked_add(alignment - remainder)
|
||||
.ok_or_else(|| "GGUF alignment overflows".into())
|
||||
}
|
||||
}
|
||||
|
||||
struct Cursor<'a> {
|
||||
bytes: &'a [u8],
|
||||
position: usize,
|
||||
}
|
||||
|
||||
impl<'a> Cursor<'a> {
|
||||
fn new(bytes: &'a [u8]) -> Self {
|
||||
Self { bytes, position: 0 }
|
||||
}
|
||||
|
||||
fn position(&self) -> usize {
|
||||
self.position
|
||||
}
|
||||
|
||||
fn take(&mut self, count: u64) -> Result<&'a [u8], String> {
|
||||
let count = usize::try_from(count).map_err(|_| "GGUF value is too large")?;
|
||||
let end = self
|
||||
.position
|
||||
.checked_add(count)
|
||||
.filter(|end| *end <= self.bytes.len())
|
||||
.ok_or_else(|| format!("truncated GGUF at byte {}", self.position))?;
|
||||
let value = &self.bytes[self.position..end];
|
||||
self.position = end;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn u8(&mut self) -> Result<u8, String> {
|
||||
Ok(self.take(1)?[0])
|
||||
}
|
||||
|
||||
fn u16(&mut self) -> Result<u16, String> {
|
||||
Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn u32(&mut self) -> Result<u32, String> {
|
||||
Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn i32(&mut self) -> Result<i32, String> {
|
||||
Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn u64(&mut self) -> Result<u64, String> {
|
||||
Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn i64(&mut self) -> Result<i64, String> {
|
||||
Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn f32(&mut self) -> Result<f32, String> {
|
||||
Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn f64(&mut self) -> Result<f64, String> {
|
||||
Ok(f64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn string(&mut self) -> Result<&'a [u8], String> {
|
||||
let len = self.u64()?;
|
||||
self.take(len)
|
||||
}
|
||||
|
||||
fn value(&mut self, kind: u32, key: &str, depth: u8) -> Result<Value, String> {
|
||||
if depth > MAX_DEPTH {
|
||||
return Err("GGUF metadata arrays are nested too deeply".into());
|
||||
}
|
||||
Ok(match kind {
|
||||
0 => {
|
||||
self.u8()?;
|
||||
Value::Other
|
||||
}
|
||||
1 => {
|
||||
self.u8()?;
|
||||
Value::Other
|
||||
}
|
||||
2 | 3 => {
|
||||
self.u16()?;
|
||||
Value::Other
|
||||
}
|
||||
4 => Value::U32(self.u32()?),
|
||||
5 => Value::I32(self.i32()?),
|
||||
6 => Value::F32(self.f32()?),
|
||||
7 => Value::Bool(self.u8()? != 0),
|
||||
8 => Value::Bytes(self.string()?.to_vec()),
|
||||
9 => self.array(key, depth + 1)?,
|
||||
10 => Value::U64(self.u64()?),
|
||||
11 => Value::I64(self.i64()?),
|
||||
12 => Value::F64(self.f64()?),
|
||||
_ => return Err(format!("unknown GGUF metadata type {kind}")),
|
||||
})
|
||||
}
|
||||
|
||||
fn array(&mut self, key: &str, depth: u8) -> Result<Value, String> {
|
||||
let item = self.u32()?;
|
||||
let len = self.u64()?;
|
||||
if len > MAX_COUNT {
|
||||
return Err(format!("GGUF metadata array is too large: {key}"));
|
||||
}
|
||||
let keep = matches!(
|
||||
key,
|
||||
"tokenizer.ggml.tokens"
|
||||
| "tokenizer.ggml.merges"
|
||||
| "deepseek4.attention.compress_ratios"
|
||||
| "deepseek4.swiglu_clamp_exp"
|
||||
| "deepseek4.dspark.target_layer_ids"
|
||||
| "deepseek4.dspark_target_layer_ids"
|
||||
| "dspark.target_layer_ids"
|
||||
);
|
||||
if !keep {
|
||||
for _ in 0..len {
|
||||
self.value(item, key, depth)?;
|
||||
}
|
||||
return Ok(Value::Other);
|
||||
}
|
||||
match item {
|
||||
4 | 5 => {
|
||||
let mut values = Vec::with_capacity(len as usize);
|
||||
for _ in 0..len {
|
||||
let value = if item == 4 {
|
||||
self.u32()?
|
||||
} else {
|
||||
self.i32()?
|
||||
.try_into()
|
||||
.map_err(|_| format!("negative value in {key}"))?
|
||||
};
|
||||
values.push(value);
|
||||
}
|
||||
Ok(Value::U32s(values))
|
||||
}
|
||||
6 | 12 => {
|
||||
let mut values = Vec::with_capacity(len as usize);
|
||||
for _ in 0..len {
|
||||
values.push(if item == 6 {
|
||||
self.f32()?
|
||||
} else {
|
||||
self.f64()? as f32
|
||||
});
|
||||
}
|
||||
Ok(Value::F32s(values))
|
||||
}
|
||||
8 => {
|
||||
let mut values = Vec::with_capacity(len as usize);
|
||||
for _ in 0..len {
|
||||
values.push(self.string()?.to_vec());
|
||||
}
|
||||
Ok(Value::Strings(values))
|
||||
}
|
||||
_ => Err(format!("unsupported array type {item} for {key}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
#[test]
|
||||
fn maps_metadata_and_tensor_payload_without_copying_weights() {
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"ds4-server-gguf-{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
));
|
||||
let mut bytes = Vec::new();
|
||||
bytes.extend(MAGIC.to_le_bytes());
|
||||
bytes.extend(3_u32.to_le_bytes());
|
||||
bytes.extend(1_u64.to_le_bytes());
|
||||
bytes.extend(1_u64.to_le_bytes());
|
||||
push_string(&mut bytes, b"general.architecture");
|
||||
bytes.extend(8_u32.to_le_bytes());
|
||||
push_string(&mut bytes, b"deepseek4");
|
||||
push_string(&mut bytes, b"weight");
|
||||
bytes.extend(1_u32.to_le_bytes());
|
||||
bytes.extend(1_u64.to_le_bytes());
|
||||
bytes.extend(F32.to_le_bytes());
|
||||
bytes.extend(0_u64.to_le_bytes());
|
||||
bytes.resize(bytes.len().div_ceil(32) * 32, 0);
|
||||
bytes.extend(1_f32.to_le_bytes());
|
||||
fs::write(&path, bytes).unwrap();
|
||||
|
||||
let model = Gguf::open(&path).unwrap();
|
||||
assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
|
||||
assert_eq!(model.tensor("weight").unwrap().dims, [1]);
|
||||
assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
|
||||
fs::remove_file(path).unwrap();
|
||||
}
|
||||
|
||||
fn push_string(bytes: &mut Vec<u8>, value: &[u8]) {
|
||||
bytes.extend((value.len() as u64).to_le_bytes());
|
||||
bytes.extend(value);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user