Ported quantize.cpp

.vscode/settings.json

@@ -1,3 +1,3 @@
 {
- "rust-analyzer.cargo.features": ["convert"]
+ "rust-analyzer.cargo.features": ["convert", "quantize"]
 }

Cargo.toml

@@ -2,7 +2,7 @@
 members = [
  "ggml-sys",
  "ggml",
- "ggml-loader",
+ "ggml-format",
  "llama-rs",
  "llama-cli",
  "generate-ggml-bindings"

ggml-loader/Cargo.toml → ggml-format/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "ggml-loader"
+name = "ggml-format"
 version = "0.1.0"
 edition = "2021"
 
@@ -8,3 +8,6 @@ edition = "2021"
 [dependencies]
 ggml = { path = "../ggml" }
 thiserror = "1.0"
+
+[dev-dependencies]
+rand = "0.8"

ggml-format/src/lib.rs

@@ -0,0 +1,45 @@
+#![deny(missing_docs)]
+//! A reader and writer for the `ggml` model format.
+//!
+//! The reader supports the GGML, GGMF and GGJT container formats, but
+//! only single-part models.
+//!
+//! The writer isn't implemented yet. It will support the GGJT container
+//! format only.
+
+/// Utilities for reading and writing.
+pub mod util;
+
+mod loader;
+mod saver;
+#[cfg(test)]
+mod tests;
+
+pub use loader::{
+ data_size, load_model, LoadError, LoadHandler, PartialHyperparameters, TensorInfo,
+};
+pub use saver::{save_model, SaveError, SaveHandler, TensorData};
+
+/// The type of a tensor element.
+pub type ElementType = ggml::Type;
+
+#[derive(Debug, PartialEq, Clone, Copy)]
+/// The format of the file containing the model.
+pub enum ContainerType {
+ /// `GGML`: legacy format, oldest ggml tensor file format
+ Ggml,
+ /// `GGMF`: also legacy format. Introduces versioning. Newer than GGML, older than GGJT.
+ Ggmf,
+ /// `GGJT`: mmap-able format.
+ Ggjt,
+}
+impl ContainerType {
+ /// Does this container type support mmap?
+ pub fn support_mmap(&self) -> bool {
+ match self {
+ ContainerType::Ggml => false,
+ ContainerType::Ggmf => false,
+ ContainerType::Ggjt => true,
+ }
+ }
+}

ggml-format/src/loader.rs

@@ -0,0 +1,240 @@
+use std::{
+ error::Error,
+ io::{BufRead, Seek, SeekFrom},
+};
+
+use crate::{
+ util::{has_data_left, read_bytes_with_len, read_f32, read_i32, read_u32},
+ ContainerType, ElementType,
+};
+
+#[derive(Debug, thiserror::Error)]
+/// Errors that can occur while loading a model.
+pub enum LoadError<E: Error> {
+ #[error("invalid file magic number: {0}")]
+ /// The file magic number is invalid.
+ InvalidMagic(u32),
+ #[error("invalid ggml format: format={0:?} version={1}")]
+ /// An unsupported format version was found.
+ InvalidFormatVersion(ContainerType, u32),
+ #[error("non-specific I/O error")]
+ /// A non-specific IO error.
+ Io(#[from] std::io::Error),
+ #[error("could not convert bytes to a UTF-8 string")]
+ /// One of the strings encountered was not valid UTF-8.
+ InvalidUtf8(#[from] std::string::FromUtf8Error),
+ #[error("invalid integer conversion")]
+ /// One of the integers encountered could not be converted to a more appropriate type.
+ InvalidIntegerConversion(#[from] std::num::TryFromIntError),
+ #[error("implementation error")]
+ /// An error `E` was returned by the implementation of the loader.
+ ImplementationError(#[source] E),
+ #[error("unsupported tensor type {ftype} for tensor {tensor_name}")]
+ /// One of the tensors encountered had an unsupported data type.
+ UnsupportedElementType {
+ /// The name of the tensor.
+ tensor_name: String,
+ /// The format type that was encountered.
+ ftype: i32,
+ },
+ #[error("invariant broken: {0}")]
+ /// An invariant was broken.
+ InvariantBroken(String),
+}
+
+#[derive(Debug, Clone)]
+/// Information about a tensor that is read.
+pub struct TensorInfo {
+ /// The name of the tensor.
+ pub name: String,
+ /// The number of dimensions in the tensor.
+ pub n_dims: usize,
+ /// The dimensions of the tensor.
+ pub dims: [usize; 2],
+ /// The number of elements in the tensor.
+ pub n_elements: usize,
+ /// The type of the elements in the tensor.
+ pub element_type: ElementType,
+ /// start of tensor - start of file
+ pub start_offset: u64,
+}
+impl TensorInfo {
+ /// Get the dimensions of the tensor.
+ pub fn dims(&self) -> &[usize] {
+ &self.dims[0..self.n_dims]
+ }
+
+ /// Calculate the size of the tensor's values in bytes.
+ pub fn calc_size(&self) -> usize {
+ data_size(self.element_type, self.dims().iter().product())
+ }
+
+ /// Reads the tensor's data from the given reader in an owned fashion.
+ ///
+ /// The behaviour is undefined if the reader does not correspond to this info.
+ ///
+ /// Do not use this if loading with `mmap`.
+ pub fn read_data<R: BufRead + Seek>(&self, reader: &mut R) -> std::io::Result<Vec<u8>> {
+ let n_bytes = self.n_elements * ggml::type_size(self.element_type);
+ let mut data = vec![0; n_bytes];
+ reader.seek(SeekFrom::Start(self.start_offset))?;
+ reader.read_exact(&mut data)?;
+ Ok(data)
+ }
+}
+
+/// Returns the size occupied by a tensor's data in bytes given the element type and number of elements.
+pub fn data_size(element_type: ElementType, n_elements: usize) -> usize {
+ (ggml::type_size(element_type) * n_elements) / ggml::blck_size(element_type)
+}
+
+#[derive(Debug, Clone)]
+/// Information present within the hyperparameters that is required to continue loading the model.
+pub struct PartialHyperparameters {
+ /// The number of vocabulary tokens.
+ pub n_vocab: usize,
+}
+
+/// A handler for loading a model.
+pub trait LoadHandler<E: Error> {
+ /// Called when the container type is read.
+ fn container_type(&mut self, container_type: ContainerType) -> Result<(), E>;
+ /// Called when a vocabulary token is read.
+ fn vocabulary_token(&mut self, i: usize, token: Vec<u8>, score: f32) -> Result<(), E>;
+ /// Called when the hyperparameters need to be read.
+ /// You must read the hyperparameters for your model here.
+ fn read_hyperparameters(
+ &mut self,
+ reader: &mut dyn BufRead,
+ ) -> Result<PartialHyperparameters, E>;
+ /// Called when a new tensor is found.
+ fn tensor_buffer(&mut self, info: TensorInfo) -> Result<(), E>;
+}
+
+/// Load a model from a `reader` with the `handler`, which will be called when certain events occur.
+pub fn load_model<E: Error, R: BufRead + Seek>(
+ reader: &mut R,
+ handler: &mut impl LoadHandler<E>,
+) -> Result<(), LoadError<E>> {
+ // Verify magic
+ let container_type: ContainerType = match read_u32(reader)? {
+ ggml::FILE_MAGIC_GGMF => ContainerType::Ggmf,
+ ggml::FILE_MAGIC_GGJT => ContainerType::Ggjt,
+ ggml::FILE_MAGIC_UNVERSIONED => ContainerType::Ggml,
+ magic => return Err(LoadError::InvalidMagic(magic)),
+ };
+ handler
+ .container_type(container_type)
+ .map_err(LoadError::ImplementationError)?;
+
+ // Load format version
+ match container_type {
+ ContainerType::Ggmf | ContainerType::Ggjt => {
+ let _version: u32 = match read_u32(reader)? {
+ ggml::FORMAT_VERSION => ggml::FORMAT_VERSION,
+ version => return Err(LoadError::InvalidFormatVersion(container_type, version)),
+ };
+ }
+ ContainerType::Ggml => {}
+ }
+
+ // Load hyper params
+ let hparams = handler
+ .read_hyperparameters(reader)
+ .map_err(LoadError::ImplementationError)?;
+ let n_vocab = hparams.n_vocab;
+
+ // Load vocabulary
+ for i in 0..n_vocab {
+ let len = read_u32(reader)?.try_into()?;
+ let token = read_bytes_with_len(reader, len)?;
+ let token_score = match container_type {
+ ContainerType::Ggmf | ContainerType::Ggjt => read_f32(reader)?,
+ ContainerType::Ggml => {
+ // Legacy model, set empty score
+ 0.
+ }
+ };
+ handler
+ .vocabulary_token(i, token, token_score)
+ .map_err(LoadError::ImplementationError)?;
+ }
+
+ // Load tensor data
+ match container_type {
+ ContainerType::Ggmf | ContainerType::Ggml => load_weights(reader, handler, false),
+ ContainerType::Ggjt => load_weights(reader, handler, true),
+ }
+}
+
+/// # Params
+///
+/// `align`
+/// align to 4 bytes before reading tensor weights
+fn load_weights<E: Error, R: BufRead + Seek>(
+ reader: &mut R,
+ handler: &mut impl LoadHandler<E>,
+ align: bool,
+) -> Result<(), LoadError<E>> {
+ while has_data_left(reader)? {
+ // load tensor header
+ let n_dims: usize = read_i32(reader)?.try_into()?;
+ let name_len = read_i32(reader)?;
+ let ftype = read_i32(reader)?;
+
+ let mut n_elements: usize = 1;
+ let mut dims = [1usize, 1];
+ let ne_len = dims.len();
+ if n_dims > ne_len {
+ return Err(LoadError::InvariantBroken(format!("{n_dims} <= {ne_len}")));
+ }
+
+ #[allow(clippy::needless_range_loop)]
+ for i in 0..n_dims {
+ let dim: usize = read_i32(reader)?.try_into()?;
+ dims[i] = dim;
+ n_elements *= dim;
+ }
+
+ // load tensor name
+ let name = String::from_utf8(read_bytes_with_len(reader, name_len.try_into()?)?)?;
+ let ftype = ggml::Type::try_from(ftype).map_err(|_| LoadError::UnsupportedElementType {
+ tensor_name: name.clone(),
+ ftype,
+ })?;
+
+ // sanity check
+ match ftype {
+ ElementType::Q4_0 | ElementType::Q4_1 => {
+ if dims[0] % 64 != 0 {
+ return Err(LoadError::InvariantBroken(format!("{dims:?}[0] % 64 == 0")));
+ }
+ }
+ _ => {}
+ }
+
+ // load tensor weights
+ let offset_curr = reader.stream_position()?;
+ let offset_aligned: u64 = if align {
+ (offset_curr + 31) & !31
+ } else {
+ offset_curr
+ };
+
+ let tensor_info = TensorInfo {
+ name,
+ dims,
+ n_dims,
+ n_elements,
+ element_type: ftype,
+ start_offset: offset_aligned,
+ };
+ let n_bytes = tensor_info.calc_size();
+ handler
+ .tensor_buffer(tensor_info)
+ .map_err(LoadError::ImplementationError)?;
+ reader.seek(SeekFrom::Start(offset_aligned + n_bytes as u64))?;
+ }
+
+ Ok(())
+}