Patch ISQ for Mixtral (#730)

* Correctly recast xs to original dtype * Reset * See if it works now * Apply fix to xlora model
EricLBuehler · Sep 1, 2024 · 3f5d0a9 · 3f5d0a9
1 parent db88c11
commit 3f5d0a9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/mistralrs-core/src/models/mixtral.rs b/mistralrs-core/src/models/mixtral.rs
@@ -449,7 +449,8 @@ impl DecoderLayer {
         let residual = &xs;
         let xs = xs
             .apply(&self.post_attention_layernorm)?
-            .apply(&self.block_sparse_moe)?;
+            .apply(&self.block_sparse_moe)?
+            .to_dtype(residual.dtype())?;
         residual + xs
     }
 }

diff --git a/mistralrs-core/src/xlora_models/mixtral.rs b/mistralrs-core/src/xlora_models/mixtral.rs
@@ -560,12 +560,15 @@ impl DecoderLayer {
         )?;
         let xs = (xs + residual)?;
         let residual = &xs;
-        let xs = self.block_sparse_moe.forward(
-            &xs.apply(&self.post_attention_layernorm)?,
-            scalings.clone(),
-            global_scaling_weight,
-            is_scaling_pass,
-        )?;
+        let xs = self
+            .block_sparse_moe
+            .forward(
+                &xs.apply(&self.post_attention_layernorm)?,
+                scalings.clone(),
+                global_scaling_weight,
+                is_scaling_pass,
+            )?
+            .to_dtype(residual.dtype())?;
         residual + xs
     }
 }