[go: nahoru, domu]
Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve Auto dtype determination #438

Merged
merged 2 commits into from
Jun 15, 2024
Merged

Improve Auto dtype determination #438

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0fccdaa
Improve auto dtype determination
EricLBuehler Jun 15, 2024
434afd4
Clippy
EricLBuehler Jun 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Improve auto dtype determination
  • Loading branch information
@EricLBuehler
EricLBuehler committed Jun 15, 2024
commit 0fccdaa153206d3e5f6fe6c46e8fa76e48ef6b17
30 changes: 19 additions & 11 deletions mistralrs-core/src/utils/normal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::{fmt::Display, str::FromStr};

use anyhow::Result;
use candle_core::{DType, Device};
use candle_core::{DType, Device, Tensor};
use serde::Deserialize;
use tracing::info;

Expand All @@ -10,9 +10,7 @@
///
/// If the model is quantized, this is ignored so it is reasonable to use the [`Default`] impl.
///
/// ## `Auto` rules
/// - If CUDA device or CPU, use BF16
/// - Fallback to F16
/// Note: When using `Auto`, fallback pattern is: BF16 -> F16 -> 32
pub enum ModelDType {
#[default]
#[serde(rename = "auto")]
Expand Down Expand Up @@ -64,16 +62,26 @@
}
}

fn determine_auto_dtype(device: &Device) -> candle_core::Result<DType> {
for dtype in [DType::BF16, DType::F16] {
// Try a matmul
let x = Tensor::zeros((2, 2), dtype, device)?;
let y = x.matmul(&x);
match y {
Ok(_) => return Ok(dtype),
Err(e) => match e {
candle_core::Error::UnsupportedDTypeForOp(_, _) => continue,
other => return Err(other),
},
}
}
return Ok(DType::F32);

Check failure on line 78 in mistralrs-core/src/utils/normal.rs

View workflow job for this annotation

GitHub Actions / Clippy 

unneeded `return` statement
}

impl TryIntoDType for ModelDType {
fn try_into_dtype(&self, device: &Device) -> Result<DType> {
let dtype = match self {
Self::Auto => {
if device.is_cuda() || device.is_cpu() {
Ok(DType::BF16)
} else {
Ok(DType::F32)
}
}
Self::Auto => Ok(determine_auto_dtype(device).map_err(anyhow::Error::msg)?),
Self::BF16 => Ok(DType::BF16),
Self::F16 => Ok(DType::F16),
Self::F32 => Ok(DType::F32),
Expand Down
Loading