Use BF16 on metal when possible. (#2378)

candle-core/src/device.rs

@@ -171,6 +171,22 @@ impl Device {
  matches!(self, Self::Metal(_))
  }
 
+ pub fn supports_bf16(&self) -> bool {
+ match self {
+ Self::Cuda(_) | Self::Metal(_) => true,
+ Self::Cpu => false,
+ }
+ }
+
+ /// Return `BF16` for devices that support it, otherwise default to `F32`.
+ pub fn bf16_default_to_f32(&self) -> DType {
+ if self.supports_bf16() {
+ DType::BF16
+ } else {
+ DType::F32
+ }
+ }
+
  pub fn cuda_if_available(ordinal: usize) -> Result<Self> {
  if crate::utils::cuda_is_available() {
  Self::new_cuda(ordinal)

candle-examples/examples/mixtral/main.rs

@@ -217,11 +217,7 @@ fn main() -> Result<()> {
  let start = std::time::Instant::now();
  let config = Config::v0_1_8x7b(args.use_flash_attn);
  let device = candle_examples::device(args.cpu)?;
- let dtype = if device.is_cuda() {
- DType::BF16
- } else {
- DType::F32
- };
+ let dtype = device.bf16_default_to_f32();
  let vb = unsafe { VarBuilder::from_mmaped_safetensors(&filenames, dtype, &device)? };
  let model = Model::new(&config, vb)?;
  println!("loaded the model in {:?}", start.elapsed());