update 250807

atorch/auto/opt_lib/module_replace_optimization.py

@@ -67,13 +67,15 @@ def decorator(cls):
 
 # decorator mode register. Not doing this in cls definition
 # because importing `register_replace_pair` there incurs circular import
-register_replace_pair("HF_BertAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(BertAttentionFA)
-register_replace_pair("HF_CLIPAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(CLIPAttentionFA)
-register_replace_pair("MultiheadAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(MultiheadAttentionFA)
 if package_version_smaller_than("transformers", "4.38.0"):
-    # transformers 4.38.0 changed LlamaAttention interface, so check version first.
+    # Not support new transformer version, thus only apply to older version.
     register_replace_pair("HF_LlamaAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(LlamaAttentionFA)
-register_replace_pair("HF_GPT2Attention_FA", supported_dtypes={torch.float16, torch.bfloat16})(GPT2AttentionFA)
+    register_replace_pair("HF_BertAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(BertAttentionFA)
+    register_replace_pair("HF_CLIPAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(CLIPAttentionFA)
+    register_replace_pair("MultiheadAttention_FA", supported_dtypes={torch.float16, torch.bfloat16})(
+        MultiheadAttentionFA
+    )
+    register_replace_pair("HF_GPT2Attention_FA", supported_dtypes={torch.float16, torch.bfloat16})(GPT2AttentionFA)
 
 
 def _check_model_params_device(model):

atorch/kernels/triton_jit/cross_entropy.py

@@ -64,7 +64,7 @@ def cross_entropy_fwd_kernel(
     else:
         label_idx -= class_start_idx
         if label_idx >= col_block_idx * BLOCK_SIZE and label_idx < min(n_cols, (col_block_idx + 1) * BLOCK_SIZE):
-            logits_label = tl.load(logits_ptr + label_idx)
+            logits_label = tl.load(logits_ptr + label_idx).to(tl.float32)  # pragma: no cover
             if HAS_SMOOTHING:
                 loss = (
                     (lse if not SPLIT else 0.0)

atorch/modules/fp8/cuda_kernel.py

@@ -20,11 +20,11 @@ def tile_quant(
     dtype=torch.float8_e4m3fn,
     block_size: int = 128,
     pow_2_scale: bool = False,
+    eps: float = 0.0,
     return_transpose: bool = False,
     use_cublas=False,
 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
     # return qx, sx, qx_t, sx_t
-    eps = 1e-10
     # only 128 is supported for block_size for now
     assert block_size == 128
     return ops.quantize_vector_blockwise(
@@ -43,11 +43,11 @@ def block_quant(
     dtype=torch.float8_e4m3fn,
     block_size: int = 128,
     pow_2_scale: bool = False,
+    eps: float = 0.0,
     return_transpose: bool = False,
     use_cublas=False,
 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
     # return qx, sx, qx_t, sx_t
-    eps = 1e-10
     # only 128 is supported for block_size for now
     assert block_size == 128
     return ops.quantize_square_blockwise(

atorch/modules/fp8/quantize.py

@@ -130,33 +130,45 @@ def get_linear_tileblock_quantize_params(block_size):
 class Fp8Quantization:
     E4M3_MAX_POS = torch.finfo(torch.float8_e4m3fn).max if hasattr(torch, "float8_e4m3fn") else 448.0
     E5M2_MAX_POS = torch.finfo(torch.float8_e5m2).max if hasattr(torch, "float8_e5m2") else 57344.0
-    EPS: float = 1e-12
 
     @staticmethod
-    def _amax_to_scale(amax: torch.Tensor, float8_dtype: torch.dtype) -> torch.Tensor:
+    def _amax_to_scale(amax: torch.Tensor, float8_dtype: torch.dtype, eps: float = 0.0) -> torch.Tensor:
         with torch.no_grad():
-            amax = amax.float()
+            amax = torch.clamp(amax.float(), min=eps)
+            if amax.numel() == 1:
+                if amax == 0.0:
+                    amax.fill_(1.0)
+            else:
+                amax[amax == 0.0] = 1.0
+
             if float8_dtype == torch.float8_e4m3fn:
-                res = Fp8Quantization.E4M3_MAX_POS / torch.clamp(amax, min=Fp8Quantization.EPS)
+                res = Fp8Quantization.E4M3_MAX_POS / amax
             else:  # e5m2
-                res = Fp8Quantization.E5M2_MAX_POS / torch.clamp(amax, min=Fp8Quantization.EPS)
+                res = Fp8Quantization.E5M2_MAX_POS / amax
             return res
 
     @staticmethod
     def quantize_tensorwise(
-        x: torch.Tensor, float8_dtype: torch.dtype, method: ScaleComputMethod = ScaleComputMethod.DEFAULT
+        x: torch.Tensor,
+        float8_dtype: torch.dtype,
+        method: ScaleComputMethod = ScaleComputMethod.DEFAULT,
+        eps: float = 0.0,
     ):
         if method == ScaleComputMethod.DEFAULT or method == ScaleComputMethod.PYTORCH:
-            return Fp8Quantization.quantize_tensorwise_pt(x, float8_dtype)
+            return Fp8Quantization.quantize_tensorwise_pt(x, float8_dtype, eps)
         else:
             assert 0, f"{ScaleComputMethod} not implemented"
 
     @staticmethod
     def quantize_axiswise(
-        x: torch.Tensor, float8_dtype: torch.dtype, dim=1, method: ScaleComputMethod = ScaleComputMethod.DEFAULT
+        x: torch.Tensor,
+        float8_dtype: torch.dtype,
+        dim=1,
+        method: ScaleComputMethod = ScaleComputMethod.DEFAULT,
+        eps: float = 0.0,
     ):
         if method == ScaleComputMethod.DEFAULT or method == ScaleComputMethod.PYTORCH or x.dtype == torch.float16:
-            return Fp8Quantization.quantize_axiswise_pt(x, float8_dtype, dim)
+            return Fp8Quantization.quantize_axiswise_pt(x, float8_dtype, dim, eps)
         else:
             assert 0, f"{ScaleComputMethod} not implemented"
 
@@ -167,10 +179,11 @@ def quantize_tilewise(
         block_size,
         method: ScaleComputMethod = ScaleComputMethod.DEFAULT,
         return_transpose=False,
+        eps=0.0,
     ):
         # Only CUTLASS kernel supports return_transpose
         if method == ScaleComputMethod.DEFAULT or method == ScaleComputMethod.TRITON:
-            return Fp8Quantization.quantize_tilewise_triton(x, float8_dtype, block_size=block_size)
+            return Fp8Quantization.quantize_tilewise_triton(x, float8_dtype, block_size=block_size, eps=eps)
         elif (
             method == ScaleComputMethod.CUTLASS
             or method == ScaleComputMethod.CUBLAS
@@ -180,6 +193,7 @@ def quantize_tilewise(
                 x,
                 float8_dtype,
                 block_size=block_size,
+                eps=eps,
                 return_transpose=return_transpose,
                 use_cublas=(method != ScaleComputMethod.CUTLASS),
             )
@@ -193,6 +207,7 @@ def quantize_blockwise(
         block_size,
         method: ScaleComputMethod = ScaleComputMethod.DEFAULT,
         return_transpose=False,
+        eps=0.0,
     ):
         # Only CUTLASS kernel supports return_transpose
         if (
@@ -203,65 +218,76 @@ def quantize_blockwise(
             # Only support square block shape
             assert isinstance(block_size, int) or block_size[0] == block_size[1]
             bsize = block_size if isinstance(block_size, int) else block_size[0]
-            return Fp8Quantization.quantize_blockwise_triton(x, float8_dtype, block_size=bsize)
+            return Fp8Quantization.quantize_blockwise_triton(x, float8_dtype, block_size=bsize, eps=eps)
         elif method == ScaleComputMethod.CUTLASS or method == ScaleComputMethod.CUBLAS:
             return Fp8Quantization.quantize_blockwise_cuda(
                 x,
                 float8_dtype,
                 block_size=block_size,
+                eps=eps,
                 return_transpose=return_transpose,
                 use_cublas=method == ScaleComputMethod.CUBLAS,
             )
         else:
             assert 0, f"{ScaleComputMethod} not implemented"
 
     @staticmethod
-    def quantize_tensorwise_pt(x: torch.Tensor, float8_dtype: torch.dtype):
+    def quantize_tensorwise_pt(x: torch.Tensor, float8_dtype: torch.dtype, eps=0.0):
         amax = torch.max(torch.abs(x))
-        scale = Fp8Quantization._amax_to_scale(amax, float8_dtype)
+        scale = Fp8Quantization._amax_to_scale(amax, float8_dtype, eps)
         x_fp8 = (x * scale).to(float8_dtype)
         inverse_scale = scale.reciprocal()
         return x_fp8, inverse_scale
 
     @staticmethod
-    def quantize_axiswise_pt(x: torch.Tensor, float8_dtype: torch.dtype, dim=1):
+    def quantize_axiswise_pt(x: torch.Tensor, float8_dtype: torch.dtype, dim=1, eps=0.0):
         # set dim=1 for rowwise, dim=0 for colwise.
         amax = torch.max(torch.abs(x), dim=dim, keepdim=True).values
-        scale = Fp8Quantization._amax_to_scale(amax, float8_dtype)
+        scale = Fp8Quantization._amax_to_scale(amax, float8_dtype, eps)
         x_fp8 = (x * scale).to(float8_dtype)
         inverse_scale = scale.reciprocal()
         return x_fp8, inverse_scale
 
     @staticmethod
-    def quantize_tilewise_triton(x: torch.Tensor, float8_dtype: torch.dtype, block_size=128):
+    def quantize_tilewise_triton(x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, eps=0.0):
         from .triton_kernel import tile_quant
 
-        return tile_quant(x, dtype=float8_dtype, block_size=block_size)
+        return tile_quant(x, dtype=float8_dtype, block_size=block_size, eps=eps)
 
     @staticmethod
     def quantize_tilewise_cuda(
-        x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, return_transpose=False, use_cublas=False
+        x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, eps=0.0, return_transpose=False, use_cublas=False
     ):
         from .cuda_kernel import tile_quant
 
         return tile_quant(
-            x, dtype=float8_dtype, block_size=block_size, return_transpose=return_transpose, use_cublas=use_cublas
+            x,
+            dtype=float8_dtype,
+            block_size=block_size,
+            eps=eps,
+            return_transpose=return_transpose,
+            use_cublas=use_cublas,
         )
 
     @staticmethod
-    def quantize_blockwise_triton(x: torch.Tensor, float8_dtype: torch.dtype, block_size=128):
+    def quantize_blockwise_triton(x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, eps=0.0):
         from .triton_kernel import block_quant
 
-        return block_quant(x, dtype=float8_dtype, block_size=block_size)
+        return block_quant(x, dtype=float8_dtype, block_size=block_size, eps=eps)
 
     @staticmethod
     def quantize_blockwise_cuda(
-        x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, return_transpose=False, use_cublas=False
+        x: torch.Tensor, float8_dtype: torch.dtype, block_size=128, eps=0.0, return_transpose=False, use_cublas=False
     ):
         from .cuda_kernel import block_quant
 
         return block_quant(
-            x, dtype=float8_dtype, block_size=block_size, return_transpose=return_transpose, use_cublas=use_cublas
+            x,
+            dtype=float8_dtype,
+            block_size=block_size,
+            eps=eps,
+            return_transpose=return_transpose,
+            use_cublas=use_cublas,
         )
 
 
@@ -276,32 +302,40 @@ def get_fp8_quantize_underflows(
     blockwise_required=True,
     block_size=128,
     quantize_method="DEFAULT",
+    eps=0.0,
 ):
     # Return a dict of underflow percentages for different quantize methods.
     if fp8_dtype is None:
         fp8_dtype = torch.float8_e4m3fn
     data = data.contiguous().view(-1, data.shape[-1])
     total_zero = (data == 0).sum().item()
     total_nonzero = data.numel() - total_zero
+
     results = {}
     if tensorwise_required:
-        tensor_fp8, _ = Fp8Quantization.quantize_tensorwise(data, fp8_dtype)
+        tensor_fp8, _ = Fp8Quantization.quantize_tensorwise(data, fp8_dtype, eps=eps)
         tensor_zero = (tensor_fp8 == 0).sum().item()
-        tenor_underflow = (tensor_zero - total_zero) / total_nonzero * 100.0
+        tenor_underflow = 0
+        if total_nonzero > 0:
+            tenor_underflow = (tensor_zero - total_zero) / total_nonzero * 100.0
         results["tensorwise"] = tenor_underflow
     if rowwise_required:
         row_fp8, _ = Fp8Quantization.quantize_axiswise(
             data, fp8_dtype, dim=1, method=ScaleComputMethod(quantize_method)
         )
         row_zero = (row_fp8 == 0).sum().item()
-        row_underflow = (row_zero - total_zero) / total_nonzero * 100.0
+        row_underflow = 0
+        if total_nonzero > 0:
+            row_underflow = (row_zero - total_zero) / total_nonzero * 100.0
         results["rowwise"] = row_underflow
     if colwise_required:
         col_fp8, _ = Fp8Quantization.quantize_axiswise(
             data, fp8_dtype, dim=0, method=ScaleComputMethod(quantize_method)
         )
         col_zero = (col_fp8 == 0).sum().item()
-        col_underflow = (col_zero - total_zero) / total_nonzero * 100.0
+        col_underflow = 0
+        if total_nonzero > 0:
+            col_underflow = (col_zero - total_zero) / total_nonzero * 100.0
         results["colwise"] = col_underflow
 
     padded_tensor = data
@@ -312,25 +346,31 @@ def get_fp8_quantize_underflows(
 
     if tilewise_required:
         tile_fp8, _ = Fp8Quantization.quantize_tilewise(
-            padded_tensor, fp8_dtype, block_size, method=ScaleComputMethod(quantize_method)
+            padded_tensor, fp8_dtype, block_size, eps=eps, method=ScaleComputMethod(quantize_method)
         )
         tile_zero = (tile_fp8 == 0).sum().item()
-        tile_underflow = (tile_zero - total_zero) / total_nonzero * 100.0
+        tile_underflow = 0
+        if total_nonzero > 0:
+            tile_underflow = (tile_zero - total_zero) / total_nonzero * 100.0
         results["tilewise"] = tile_underflow
 
     if v_tilewise_required:
         tile_fp8, _ = Fp8Quantization.quantize_tilewise(
-            padded_tensor.t().contiguous(), fp8_dtype, block_size, method=ScaleComputMethod(quantize_method)
+            padded_tensor.t().contiguous(), fp8_dtype, block_size, eps=eps, method=ScaleComputMethod(quantize_method)
         )
         tile_zero = (tile_fp8 == 0).sum().item()
-        tile_underflow = (tile_zero - total_zero) / total_nonzero * 100.0
+        tile_underflow = 0
+        if total_nonzero > 0:
+            tile_underflow = (tile_zero - total_zero) / total_nonzero * 100.0
         results["v_tilewise"] = tile_underflow
 
     if blockwise_required:
         block_fp8, _ = Fp8Quantization.quantize_blockwise(
-            padded_tensor, fp8_dtype, block_size, method=ScaleComputMethod(quantize_method)
+            padded_tensor, fp8_dtype, block_size, eps=eps, method=ScaleComputMethod(quantize_method)
         )
         block_zero = (block_fp8 == 0).sum().item()
-        block_underflow = (block_zero - total_zero) / total_nonzero * 100.0
+        block_underflow = 0
+        if total_nonzero > 0:
+            block_underflow = (block_zero - total_zero) / total_nonzero * 100.0
         results["blockwise"] = block_underflow
     return results

atorch/modules/fp8/triton_kernel.py

@@ -19,7 +19,7 @@ class RoundingMode(IntEnum):
 
 
 @triton.jit
-def block_quant_kernel(x_ptr, y_ptr, s_ptr, M, N, BLOCK_SIZE: tl.constexpr):  # pragma: no cover
+def block_quant_kernel(x_ptr, y_ptr, s_ptr, M, N, BLOCK_SIZE: tl.constexpr, EPS: tl.constexpr):  # pragma: no cover
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
     n = tl.cdiv(N, BLOCK_SIZE)
@@ -28,20 +28,21 @@ def block_quant_kernel(x_ptr, y_ptr, s_ptr, M, N, BLOCK_SIZE: tl.constexpr):  #
     offs = offs_m[:, None] * N + offs_n[None, :]
     mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
     x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
-    eps = 1e-10
-    s = tl.maximum(tl.max(tl.abs(x)), eps) / 448.0
+    s = tl.maximum(tl.max(tl.abs(x)), EPS) / 448.0
+    if s == 0.0:
+        s = 1.0
     y = x / s
     y = y.to(y_ptr.dtype.element_ty)
     tl.store(y_ptr + offs, y, mask=mask)
     tl.store(s_ptr + pid_m * n + pid_n, s)
 
 
-def block_quant(x: torch.Tensor, dtype=torch.float8_e4m3fn, block_size: int = 128) -> torch.Tensor:
+def block_quant(x: torch.Tensor, dtype=torch.float8_e4m3fn, block_size: int = 128, eps: float = 0.0) -> torch.Tensor:
     M, N = x.size()
     y = torch.empty_like(x, dtype=dtype)
     s = x.new_empty(x.size(-2) // block_size, x.size(-1) // block_size, dtype=torch.float32)
     grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]))  # noqa: E731
-    block_quant_kernel[grid](x, y, s, M, N, BLOCK_SIZE=block_size)
+    block_quant_kernel[grid](x, y, s, M, N, BLOCK_SIZE=block_size, EPS=eps)
     return y, s
 
 
@@ -92,13 +93,19 @@ def dequant(x: torch.Tensor, s: torch.Tensor, dtype=torch.float32, block_size: i
 
 @triton.jit
 def tile_quant_kernel(
-    x_ptr, y_ptr, s_ptr, scale_rounding_mode: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    x_ptr,
+    y_ptr,
+    s_ptr,
+    scale_rounding_mode: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    EPS: tl.constexpr,
 ):  # pragma: no cover
     pid = tl.program_id(axis=0)
     offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     x = tl.load(x_ptr + offs).to(tl.float32)
-    eps = 1e-10
-    s = tl.maximum(tl.max(tl.abs(x)), eps) / 448.0
+    s = tl.maximum(tl.max(tl.abs(x)), EPS) / 448.0
+    if s == 0.0:
+        s = 1.0
     if scale_rounding_mode == 1:
         # ceil rounding to power of 2
         s = tl.ceil(tl.log2(s))
@@ -118,14 +125,18 @@ def tile_quant_kernel(
 
 
 def tile_quant(
-    x: torch.Tensor, dtype=torch.float8_e4m3fn, block_size: int = 128, scale_rounding_mode=RoundingMode.none
+    x: torch.Tensor,
+    dtype=torch.float8_e4m3fn,
+    block_size: int = 128,
+    scale_rounding_mode=RoundingMode.none,
+    eps: float = 0.0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     assert x.is_contiguous()
     assert x.size(-1) % block_size == 0
     y = torch.empty_like(x, dtype=dtype)
     s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
     grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)  # noqa: E731
-    tile_quant_kernel[grid](x, y, s, scale_rounding_mode, BLOCK_SIZE=block_size)
+    tile_quant_kernel[grid](x, y, s, scale_rounding_mode, BLOCK_SIZE=block_size, EPS=eps)
     return y, s
 
 

atorch/ops/git_version_info_installed.py

@@ -0,0 +1,6 @@
+version = "1.6.4dev8+45eb14116"
+git_hash = "45eb14116"
+git_branch = "1.6.4"
+installed_ops = {"quantization_optimizer": False, "quantizer": False}
+compatible_ops = {"quantization_optimizer": True, "quantizer": True, "atorch_not_implemented": False}
+torch_info = {"version": "1.13", "bf16_support": False, "cuda_version": "11.6", "nccl_version": "2.14"}