candding 0.1.0

//! Weight loading that casts on the CPU before a tensor reaches the device: candle casts on the device, and its
//! Metal allocator keeps both copies resident, which past about 15 GB silently turns the outputs to zeros.

pub mod npy;
pub(crate) mod size;

use std::path::PathBuf;

use candle_core::DType;
use candle_core::Device;
use candle_core::Result;
use candle_core::Shape;
use candle_core::Tensor;
use candle_core::pickle::PthTensors;
use candle_core::safetensors::MmapedSafetensors;
use candle_nn::Init;
use candle_nn::VarBuilder;
use candle_nn::var_builder::SimpleBackend;

struct CastOnCpu(MmapedSafetensors);

impl SimpleBackend for CastOnCpu {
  fn get(&self, s: Shape, name: &str, _: Init, dtype: DType, dev: &Device) -> Result<Tensor> {
    let tensor = self.get_unchecked(name, dtype, dev)?;
    if tensor.shape() != &s {
      candle_core::bail!(
        "shape mismatch for {name}: got {:?}, expected {:?}",
        tensor.shape(),
        s
      );
    }
    Ok(tensor)
  }

  fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
    self
      .0
      .load(name, &Device::Cpu)?
      .to_dtype(dtype)?
      .to_device(dev)
  }

  fn contains_tensor(&self, name: &str) -> bool {
    self.0.get(name).is_ok()
  }
}

struct CastOnCpuPickle(PthTensors);

impl SimpleBackend for CastOnCpuPickle {
  fn get(&self, s: Shape, name: &str, _: Init, dtype: DType, dev: &Device) -> Result<Tensor> {
    let tensor = self.get_unchecked(name, dtype, dev)?;
    if tensor.shape() != &s {
      candle_core::bail!(
        "shape mismatch for {name}: got {:?}, expected {:?}",
        tensor.shape(),
        s
      );
    }
    Ok(tensor)
  }

  /// `PthTensors::get` returns a CPU tensor, so the cast happens before the move. A fortran-ordered storage comes
  /// back permuted, so the tensor is made contiguous first.
  fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
    let tensor = self.0.get(name)?.ok_or_else(|| {
      candle_core::Error::CannotFindTensor {
        path: name.to_string(),
      }
      .bt()
    })?;
    tensor.contiguous()?.to_dtype(dtype)?.to_device(dev)
  }

  /// The header the reader parsed at load already lists every name, so a probe costs no zip read.
  fn contains_tensor(&self, name: &str) -> bool {
    self.0.tensor_infos().contains_key(name)
  }
}

/// # Safety
/// The memory-mapped safetensors files must not change while the builder or any tensor loaded from it lives.
pub unsafe fn var_builder(
  files: &[PathBuf],
  dtype: DType,
  device: &Device,
) -> Result<VarBuilder<'static>> {
  if let [single] = files
    && single.extension().is_some_and(|ext| ext == "bin")
  {
    let pth = PthTensors::new(single, None)?;
    return Ok(VarBuilder::from_backend(
      Box::new(CastOnCpuPickle(pth)),
      dtype,
      device.clone(),
    ));
  }
  if let Some(pickle) = files.iter().find(|f| {
    f.extension()
      .is_some_and(|ext| ext == "bin" || ext == "pt" || ext == "pth")
  }) {
    candle_core::bail!(
      "{} is a torch pickle, which candding reads only as a repository's single weight file",
      pickle.display()
    );
  }
  let mmaped = unsafe { MmapedSafetensors::multi(files)? };
  Ok(VarBuilder::from_backend(
    Box::new(CastOnCpu(mmaped)),
    dtype,
    device.clone(),
  ))
}

#[cfg(test)]
mod tests {
  use std::collections::HashMap;

  use candle_core::Device;
  use candle_core::Tensor;
  use tempfile::tempdir;

  use super::*;

  fn pickle_path(name: &str) -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
      .join("tests")
      .join("fixtures")
      .join("pickle")
      .join(name)
  }

  #[test]
  fn casts_on_cpu_and_checks_shape() {
    let dir = tempdir().unwrap();
    let path = dir.path().join("model.safetensors");
    let a = Tensor::new(&[1.5f32, -2.0], &Device::Cpu).unwrap();
    let w = Tensor::new(&[[1f32, 2.0, 3.0], [4f32, 5.0, 6.0]], &Device::Cpu).unwrap();
    candle_core::safetensors::save(
      &HashMap::from([("a".to_string(), a), ("w".to_string(), w)]),
      &path,
    )
    .unwrap();

    let vb = unsafe { var_builder(&[path], DType::F16, &Device::Cpu).unwrap() };
    assert!(vb.contains_tensor("a"));
    assert!(!vb.contains_tensor("nope"));

    let got = vb.get(2, "a").unwrap();
    assert_eq!(got.dtype(), DType::F16);
    assert_eq!(
      got.to_dtype(DType::F32).unwrap().to_vec1::<f32>().unwrap(),
      vec![1.5, -2.0]
    );

    assert!(vb.get((2, 3), "w").is_ok());
    assert!(vb.get((3, 2), "w").is_err());
  }

  /// On the CPU the move is a no-op, so this cannot prove the cast came first; only the dtype and the shape are checked.
  #[test]
  fn pickle_casts_on_cpu_and_checks_shape() {
    let path = pickle_path("two_tensors.bin");
    let vb = unsafe { var_builder(&[path], DType::F16, &Device::Cpu).unwrap() };
    assert!(vb.contains_tensor("a"));
    assert!(!vb.contains_tensor("nope"));

    let got = vb.get(2, "a").unwrap();
    assert_eq!(got.dtype(), DType::F16);
    assert_eq!(
      got.to_dtype(DType::F32).unwrap().to_vec1::<f32>().unwrap(),
      vec![1.5, -2.0]
    );

    assert!(vb.get((2, 3), "w").is_ok());
    assert!(vb.get((3, 2), "w").is_err());
  }

  /// `transposed.bin` holds `arange(6).reshape(2, 3).t()` at stride `(1, 3)`, which a same-dtype cast would pass through,
  /// so the backend makes it contiguous itself.
  #[test]
  fn pickle_hands_back_a_transposed_tensor_contiguous() {
    let vb =
      unsafe { var_builder(&[pickle_path("transposed.bin")], DType::F32, &Device::Cpu).unwrap() };
    let got = vb.get((3, 2), "t").unwrap();
    assert!(got.is_contiguous(), "layout {:?}", got.layout());
    assert_eq!(
      got.to_vec2::<f32>().unwrap(),
      vec![vec![0.0, 3.0], vec![1.0, 4.0], vec![2.0, 5.0]]
    );
  }

  /// A pickle in a multi-file list would otherwise reach the safetensors reader, whose error names neither file nor reason.
  #[test]
  fn a_pickle_among_several_weight_files_is_refused_by_name() {
    let dir = tempdir().unwrap();
    let shard = dir.path().join("model-00001-of-00002.safetensors");
    candle_core::safetensors::save(
      &HashMap::from([(
        "a".to_string(),
        Tensor::new(&[1.5f32, -2.0], &Device::Cpu).unwrap(),
      )]),
      &shard,
    )
    .unwrap();
    for name in ["pytorch_model.bin", "model.pt", "model.pth"] {
      let pickle = dir.path().join(name);
      std::fs::copy(pickle_path("two_tensors.bin"), &pickle).unwrap();
      let files = vec![shard.clone(), pickle];
      let Err(err) = (unsafe { var_builder(&files, DType::F32, &Device::Cpu) }) else {
        panic!("{name} beside a shard must be refused");
      };
      let err = err.to_string();
      assert!(err.contains(name) && err.contains("torch pickle"), "{err}");
    }
  }

  #[test]
  fn pickle_reports_a_missing_tensor() {
    let path = pickle_path("two_tensors.bin");
    let vb = unsafe { var_builder(&[path], DType::F32, &Device::Cpu).unwrap() };
    let err = vb.get(2, "nope").unwrap_err().to_string();
    assert!(err.contains("nope"), "{err}");
  }
}