Avoid f64 to f32 cast in phirope for dtype on metal (#512)

EricLBuehler · Jul 1, 2024 · af18834 · af18834
1 parent 2899e10
commit af18834
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs
@@ -18,7 +18,6 @@ use candle_nn::{Linear, Module, VarBuilder};
 
 pub use crate::layers_masker::CausalMasker;
 pub use crate::layers_utils::{flash_attn, repeat_kv};
-
 use crate::{cublaslt::CUBLASLT_HANDLE, pipeline::Phi3RopeScaling, INHIBIT_GEMM_F16};
 
 #[derive(Debug, Clone)]
@@ -140,16 +139,20 @@ impl PhiRotaryEmbedding {
                 .step_by(2)
                 .enumerate()
                 .map(|(k, i)| {
-                    1f64 / (scaled_params.long_factor[k]
-                        * cfg.rope_theta.powf(i as f64 / dim as f64))
+                    (1f64
+                        / (scaled_params.long_factor[k]
+                            * cfg.rope_theta.powf(i as f64 / dim as f64)))
+                        as f32
                 })
                 .collect::<Vec<_>>();
             let inv_freq_short = (0..dim)
                 .step_by(2)
                 .enumerate()
                 .map(|(k, i)| {
-                    1f64 / (scaled_params.short_factor[k]
-                        * cfg.rope_theta.powf(i as f64 / dim as f64))
+                    (1f64
+                        / (scaled_params.short_factor[k]
+                            * cfg.rope_theta.powf(i as f64 / dim as f64)))
+                        as f32
                 })
                 .collect::<Vec<_>>();
             let inv_freq_len = inv_freq_long.len();
@@ -159,8 +162,7 @@ impl PhiRotaryEmbedding {
                 .reshape((max_seq_len, 1))?;
 
             // Calculate sin,cos for long
-            let inv_freq_long =
-                Tensor::from_vec(inv_freq_long, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?;
+            let inv_freq_long = Tensor::from_vec(inv_freq_long, (1, inv_freq_len), dev)?;
             let freqs_long = t.matmul(&inv_freq_long)?;
             let long_sin = freqs_long.sin()?.mul(scaling_factor)?.to_dtype(dtype)?;
             let long_cos = freqs_long.cos()?.mul(scaling_factor)?.to_dtype(dtype)?;