int8 tests passing

Author: matthewdouglas

bitsandbytes/backends/default/ops.py

@@ -88,13 +88,17 @@ def _(A: torch.Tensor, threshold=0.0):
     rows = prod(A.shape[:-1])
     outlier_cols = None
 
+    outlier_restore = None
+
     if threshold > 0.0:
         outliers = A.abs() >= threshold
 
         if outliers.any():
             # Determine which columns contain outliers, and zero out the
-            # outliers ahead of quantization.
+            # outliers ahead of quantization. We need to keep a backup of these
+            # outliers to restore them after quantization.
             outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1)
+            outlier_restore = A[outliers].clone()
             A[outliers] = 0
         else:
             # Needed for torch.compile support.
@@ -110,4 +114,8 @@ def _(A: torch.Tensor, threshold=0.0):
     if rows > 1 and outlier_cols is not None:
         out_row[:, outlier_cols] = 0
 
+    # Restore outliers.
+    if outlier_restore is not None:
+        A[outliers] = outlier_restore
+
     return out_row, row_stats, outlier_cols