[dicp][ascend] bug fix for llama finetune (#626)

DeepLink-org · Jan 14, 2024 · 6b76f2e · 6b76f2e
1 parent b77713f
commit 6b76f2e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/dicp/dicp/vendor/AscendGraph/conversion.py b/dicp/dicp/vendor/AscendGraph/conversion.py
@@ -168,6 +168,11 @@ def binary_cmp_cast_input(self, x, y):
             y = self.get_param_proxy(y, const_dtype, y_shape)
             if x_dtype == torch.float16:
                 y = self.get_proxy(ascend_op.Cast, (y, "FLOAT16"))
+        else:
+            x_dtype = x.node.meta["val"].dtype
+            y_dtype = y.node.meta["val"].dtype
+            if x_dtype != y_dtype:
+                y = self.get_proxy(ascend_op.Cast, (y, get_ascend_dtype(x_dtype)))
         return x, y
 
     def shape_prod(self, shape):