macOS fixes

Author: dxqb

modules/module/quantized/LinearGGUFA8.py

@@ -1,8 +1,8 @@
+from modules.util.mm_8bit import mm_8bit as mm_8bit
 from modules.util.quantization_util import (
     quantize_fp8_axiswise,
     quantize_int8_axiswise,
 )
-from modules.util.triton_mm_8bit import mm_8bit as triton_mm_8bit
 
 import torch
 from torch import Tensor
@@ -39,14 +39,14 @@ def fp8_forward_axiswise(x: Tensor, weight: Tensor, bias: Tensor | None, compute
 def int8_backward_axiswise(output: Tensor, weight: Tensor) -> Tensor:
     output_8, output_scale = quantize_int8_axiswise(output, dim=-1)
     w_8, w_scale = quantize_int8_axiswise(weight, dim=0)
-    mm_res = triton_mm_8bit(output_8.contiguous(), w_8)
+    mm_res = mm_8bit(output_8.contiguous(), w_8)
     return mm_res.float().mul_(w_scale).mul_(output_scale).to(output.dtype)
 
 @torch.no_grad()
 def fp8_backward_axiswise(output: Tensor, weight: Tensor) -> Tensor:
     output_8, output_scale = quantize_fp8_axiswise(output, dim=-1)
     w_8, w_scale = quantize_fp8_axiswise(weight, dim=0)
-    mm_res = triton_mm_8bit(output_8.contiguous(), w_8)
+    mm_res = mm_8bit(output_8.contiguous(), w_8)
     return mm_res.float().mul_(w_scale).mul_(output_scale).to(output.dtype)
 
 class LinearGGUFIntA8RequantFunction(torch.autograd.Function):

modules/module/quantized/LinearW8A8.py

@@ -1,14 +1,14 @@
 
 from modules.module.quantized.mixin.QuantizedLinearMixin import QuantizedLinearMixin
 from modules.module.quantized.mixin.QuantizedModuleMixin import QuantizedModuleMixin
+from modules.util.mm_8bit import mm_8bit as mm_8bit
 from modules.util.quantization_util import (
     dequantize,
     quantize_fp8_axiswise,
     quantize_fp8_tensorwise,
     quantize_int8_axiswise,
     quantize_int8_tensorwise,
 )
-from modules.util.triton_mm_8bit import mm_8bit as triton_mm_8bit
 
 import torch
 from torch import Tensor, nn
@@ -37,13 +37,13 @@ def fp8_forward_tokenwise(x: Tensor, weight: Tensor, weight_scale: float, bias:
 def int8_backward_axiswise(output: Tensor, weight: Tensor, weight_scale: float) -> Tensor:
     output_8, output_scale = quantize_int8_axiswise(output, dim=-1)
     #almost always, grad outputs are already contiguous and this is a no-op. But there are some grad outputs from SDXL that are non-contiguous:
-    mm_res = triton_mm_8bit(output_8.contiguous(), weight)
+    mm_res = mm_8bit(output_8.contiguous(), weight)
     return mm_res.float().mul_(weight_scale * output_scale).to(output.dtype)
 
 @torch.no_grad()
 def fp8_backward_axiswise(output: Tensor, weight: Tensor, weight_scale: float) -> Tensor:
     output_8, output_scale = quantize_fp8_axiswise(output, dim=-1)
-    mm_res = triton_mm_8bit(output_8.contiguous(), weight)
+    mm_res = mm_8bit(output_8.contiguous(), weight)
     return mm_res.float().mul_(weight_scale * output_scale).to(output.dtype)
 
 
@@ -158,11 +158,11 @@ def benchmark_int8(m, k, n, device = 'cuda'):
 
 
     run_benchmark(lambda: torch._int_mm(x_8, w_8.T), "torch mm int")
-    run_benchmark(lambda: triton_mm_8bit(x_8, w_8.T), "triton mm int")
+    run_benchmark(lambda: mm_8bit(x_8, w_8.T), "triton mm int")
     def torch_backward(a, b):
         torch._int_mm(a, b.T.contiguous().T)
     run_benchmark(lambda: torch_backward(y_8, w_8), "torch mm backward int8")
-    run_benchmark(lambda: triton_mm_8bit(y_8, w_8), "triton mm backward int8")
+    run_benchmark(lambda: mm_8bit(y_8, w_8), "triton mm backward int8")
 
     run_benchmark(lambda: int8_forward_tokenwise(x, w_8, w_scale), "torch forward int", compile=True)
     run_benchmark(lambda: int8_backward_axiswise(y, w_8, w_scale), "triton backward int", compile=True)
@@ -179,11 +179,11 @@ def benchmark_fp8(m, k, n, device = 'cuda'):
     one_scale = torch.ones(1, device=device)
 
     run_benchmark(lambda: torch._scaled_mm(x_8, w_8.T, out_dtype=torch.bfloat16, scale_a=one_scale.float(), scale_b=w_scale.float()), "torch mm fp8")
-    run_benchmark(lambda: triton_mm_8bit(x_8, w_8.T), "triton mm fp8")
+    run_benchmark(lambda: mm_8bit(x_8, w_8.T), "triton mm fp8")
     def torch_backward(a, b):
         torch._scaled_mm(a, b.T.contiguous().T, out_dtype=torch.bfloat16, scale_a=one_scale.float(), scale_b=w_scale.float())
     run_benchmark(lambda: torch_backward(y_8, w_8), "torch mm backward fp8")
-    run_benchmark(lambda: triton_mm_8bit(y_8, w_8), "triton mm backward fp8")
+    run_benchmark(lambda: mm_8bit(y_8, w_8), "triton mm backward fp8")
     run_benchmark(lambda: fp8_forward_tokenwise(x, w_8, w_scale), "torch forward fp8", compile=True)
     run_benchmark(lambda: fp8_backward_axiswise(y, w_8, w_scale), "triton backward fp8", compile=True)
 

modules/util/dtype_util.py

@@ -33,9 +33,9 @@ def create_autocast_context(
 ) -> tuple[torch.autocast | nullcontext, DataType]:
     if torch.backends.mps.is_available():
         if any(train_dtype != dt for dt in weight_dtypes if dt is not None):
-            raise RuntimeError("macOS needs all dtypes to be the same.")
-
-        return nullcontext(), train_dtype
+            print("Warning: Mixed precision training is untested on macOS. Consider setting all dtypes to be the same.")
+        else:
+            return nullcontext(), train_dtype
 
     weight_dtypes = list(weight_dtypes)
     weight_dtypes = list(filter(lambda dtype: dtype != DataType.NONE and dtype is not None, weight_dtypes))

modules/util/mm_8bit.py

@@ -0,0 +1,15 @@
+try:
+    from modules.util.triton_mm_8bit import mm_8bit
+except ImportError as e:
+    print(str(e) + ", continueing without triton")
+    import torch
+    def mm_8bit(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        assert a.shape[1] == b.shape[0], "Incompatible dimensions"
+        assert a.is_contiguous(), "Matrix A must be contiguous"
+        assert a.dtype == b.dtype, "Incompatible dtypes"
+        assert a.dtype in [torch.int8, torch.float8_e4m3fn]
+        if a.dtype == torch.int8:
+            return torch._int_mm(a, b)
+        else:
+            one = torch.ones(1, device=a.device)
+            return torch._scaled_mm(a, b.T.contiguous().T, scale_a=one, scale_b=one)