From 08b3df31ec46549b188f56f5e9dcbe93a9dd68c5 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 23 May 2024 17:33:38 -0700 Subject: [PATCH 1/3] Remove test_ops.py warning spew --- test/test_ops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/test_ops.py b/test/test_ops.py index e260e86f0f..39ae32dd65 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,3 +1,4 @@ + import torch from torch.testing._internal.common_utils import TestCase, IS_FBCODE from torch.testing._internal.optests import opcheck @@ -5,10 +6,12 @@ from torchao.quantization.utils import TORCH_VERSION_AFTER_2_4 import unittest from parameterized import parameterized +import pytest # torch.testing._internal.optests.generate_tests.OpCheckError: opcheck(op, ...): # test_faketensor failed with module 'torch' has no attribute '_custom_ops' (scroll up for stack trace) +@pytest.mark.filterwarnings("ignore:create_unbacked_symint is deprecated, please use new_dynamic_size instead:UserWarning") @unittest.skipIf(IS_FBCODE, "Skipping the test in fbcode since we don't have TARGET file for kernels") class TestOps(TestCase): def _create_tensors_with_iou(self, N, iou_thresh): From c9b12692929bcacd705e5985006edeb3c8f766e9 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 23 May 2024 17:34:44 -0700 Subject: [PATCH 2/3] yolo --- test/test_ops.py | 1 - torchao/ops.py | 10 +++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 39ae32dd65..d73ae536ac 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,4 +1,3 @@ - import torch from torch.testing._internal.common_utils import TestCase, IS_FBCODE from torch.testing._internal.optests import opcheck diff --git a/torchao/ops.py b/torchao/ops.py index 3a25dbf6db..61d10d1a69 100644 --- a/torchao/ops.py +++ b/torchao/ops.py @@ -9,7 +9,7 @@ def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: # Defines the meta kernel / fake kernel / abstract impl -@torch.library.impl_abstract("torchao::nms") +@torch.library.register_fake("torchao::nms") def _(dets, scores, iou_threshold): torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D") torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}") @@ -36,7 +36,7 @@ def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor: return torch.ops.torchao.prepack_fp6_weight.default(fp6_weight) -@torch.library.impl_abstract("torchao::prepack_fp6_weight") +@torch.library.register_fake("torchao::prepack_fp6_weight") def _(fp6_weight): torch._check(fp6_weight.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_weight.dim()}D") return torch.empty_like(fp6_weight) @@ -49,7 +49,7 @@ def fp16_to_fp6(fp16_tensor: Tensor) -> Tensor: return torch.ops.torchao.fp16_to_fp6.default(fp16_tensor) -@torch.library.impl_abstract("torchao::fp16_to_fp6") +@torch.library.register_fake("torchao::fp16_to_fp6") def _(fp16_tensor): torch._check(fp16_tensor.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp16_tensor.dim()}D") torch._check(fp16_tensor.dtype is torch.float16, lambda: f"weight must be FP16, got {fp16_tensor.dtype}") @@ -74,7 +74,7 @@ def fp16act_fp6weight_linear(_in_feats: Tensor, _weights: Tensor, _scales: Tenso return torch.ops.torchao.fp16act_fp6weight_linear.default(_in_feats, _weights, _scales, splitK) -@torch.library.impl_abstract("torchao::fp16act_fp6weight_linear") +@torch.library.register_fake("torchao::fp16act_fp6weight_linear") def _(_in_feats, _weights, _scales, splitK = 1): torch._check(_in_feats.dim() == 2, lambda: f"input should be a 2d tensor, got {_in_feats.dim()}D") torch._check(_in_feats.dtype is torch.float16, lambda: f"weight must be FP16, got {_in_feats.dtype}") @@ -95,7 +95,7 @@ def fp6_weight_dequant(fp6_tensor: Tensor, fp16_scale: Tensor) -> Tensor: return torch.ops.torchao.fp6_weight_dequant.default(fp6_tensor, fp16_scale) -@torch.library.impl_abstract("torchao::fp6_weight_dequant") +@torch.library.register_fake("torchao::fp6_weight_dequant") def _(fp6_tensor, fp16_scale): torch._check(fp6_tensor.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_tensor.dim()}D") torch._check(fp6_tensor.dtype is torch.int32, lambda: f"weight must be INT32, got {fp6_tensor.dtype}") From fea73601e2c48fef61a09892c46d7d6c2d8c157b Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 23 May 2024 17:48:55 -0700 Subject: [PATCH 3/3] add custom op registration guards --- torchao/ops.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/torchao/ops.py b/torchao/ops.py index 61d10d1a69..fcc6ae9364 100644 --- a/torchao/ops.py +++ b/torchao/ops.py @@ -1,5 +1,14 @@ import torch from torch import Tensor +from torchao.quantization.utils import TORCH_VERSION_AFTER_2_4 + +def register_custom_op(name): + def decorator(func): + if TORCH_VERSION_AFTER_2_4: + return torch.library.register_fake(f"{name}")(func) + else: + return torch.library.impl_abstract(f"{name}")(func) + return decorator def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: """ @@ -9,7 +18,7 @@ def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor: # Defines the meta kernel / fake kernel / abstract impl -@torch.library.register_fake("torchao::nms") +@register_custom_op("torchao::nms") def _(dets, scores, iou_threshold): torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D") torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}") @@ -36,7 +45,7 @@ def prepack_fp6_weight(fp6_weight: Tensor) -> Tensor: return torch.ops.torchao.prepack_fp6_weight.default(fp6_weight) -@torch.library.register_fake("torchao::prepack_fp6_weight") +@register_custom_op("torchao::prepack_fp6_weight") def _(fp6_weight): torch._check(fp6_weight.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_weight.dim()}D") return torch.empty_like(fp6_weight) @@ -49,7 +58,7 @@ def fp16_to_fp6(fp16_tensor: Tensor) -> Tensor: return torch.ops.torchao.fp16_to_fp6.default(fp16_tensor) -@torch.library.register_fake("torchao::fp16_to_fp6") +@register_custom_op("torchao::fp16_to_fp6") def _(fp16_tensor): torch._check(fp16_tensor.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp16_tensor.dim()}D") torch._check(fp16_tensor.dtype is torch.float16, lambda: f"weight must be FP16, got {fp16_tensor.dtype}") @@ -74,7 +83,7 @@ def fp16act_fp6weight_linear(_in_feats: Tensor, _weights: Tensor, _scales: Tenso return torch.ops.torchao.fp16act_fp6weight_linear.default(_in_feats, _weights, _scales, splitK) -@torch.library.register_fake("torchao::fp16act_fp6weight_linear") +@register_custom_op("torchao::fp16act_fp6weight_linear") def _(_in_feats, _weights, _scales, splitK = 1): torch._check(_in_feats.dim() == 2, lambda: f"input should be a 2d tensor, got {_in_feats.dim()}D") torch._check(_in_feats.dtype is torch.float16, lambda: f"weight must be FP16, got {_in_feats.dtype}") @@ -95,7 +104,7 @@ def fp6_weight_dequant(fp6_tensor: Tensor, fp16_scale: Tensor) -> Tensor: return torch.ops.torchao.fp6_weight_dequant.default(fp6_tensor, fp16_scale) -@torch.library.register_fake("torchao::fp6_weight_dequant") +@register_custom_op("torchao::fp6_weight_dequant") def _(fp6_tensor, fp16_scale): torch._check(fp6_tensor.dim() == 2, lambda: f"weight should be a 2d tensor, got {fp6_tensor.dim()}D") torch._check(fp6_tensor.dtype is torch.int32, lambda: f"weight must be INT32, got {fp6_tensor.dtype}")