From 29617b32a44d3aa30cbc52f66ffdfeed2646a909 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 30 Aug 2024 15:10:12 -0700 Subject: [PATCH 1/6] Move more utils to TorchAOBaseTensor Summary: This moves over _implements, _dispatch__torch_dispatch__, _dispatch__torch_function__, _register_layout_cls and _get_layout_tensor_constructor to `TorchAOBaseTensor` so when people inherit from this, they can get these utils directly Test Plan: python test/quantization/test_quant_api.py python test/integration/test_integration.py rely on CI for other tests Reviewers: Subscribers: Tasks: Tags: --- torchao/dtypes/affine_quantized_tensor.py | 18 +-- torchao/dtypes/uintx/Uintx.py | 10 +- torchao/dtypes/utils.py | 112 ------------- .../prototype/low_bit_optim/subclass_4bit.py | 7 +- .../prototype/low_bit_optim/subclass_8bit.py | 7 +- .../prototype/low_bit_optim/subclass_fp8.py | 7 +- torchao/prototype/quantized_training/int8.py | 8 +- .../linear_activation_quantized_tensor.py | 9 -- .../prototype/superblock/blocksparse.py | 16 +- torchao/utils.py | 148 ++++++++++++++++++ .../my_dtype_tensor_subclass.py | 16 +- 11 files changed, 172 insertions(+), 186 deletions(-) diff --git a/torchao/dtypes/affine_quantized_tensor.py b/torchao/dtypes/affine_quantized_tensor.py index 11b9356adf..751622e299 100644 --- a/torchao/dtypes/affine_quantized_tensor.py +++ b/torchao/dtypes/affine_quantized_tensor.py @@ -21,11 +21,6 @@ ) from torch.utils._python_dispatch import return_and_correct_aliasing from torchao.dtypes.utils import ( - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, - _register_layout_cls, - _get_layout_tensor_constructor, LayoutType, PlainLayoutType, is_device, @@ -405,7 +400,8 @@ def _apply_fn_to_data(self, fn): strides=self.stride(), ) - implements = classmethod(_implements) + # following are the comments for __torch_function__/__torch_dispatch__, we can clean this up + # a bit later # Note: we only added cpu path here for 8da4w, this is for executorch, in the future # 1. we'll add cpu/cuda version (int4mm etc.) # 2. we'll need to hide the 8da4w executorch version under things like layouts (we also have multiple impl for cpu kernel as Michael mentioned), so it will be something like @@ -417,19 +413,13 @@ def _apply_fn_to_data(self, fn): # 1 - when tensor is on CUDA: we'll add this later, we'll also enable dispatching to optimized # kernels in CPU as well, see the note above # 2 - we're given non-floats - quantizing long to int8 is crazy - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - __torch_function__ = classmethod(_dispatch__torch_function__) ###################################################### # LayoutType and Layout Tensor Subclass Registration # ###################################################### - -def register_layout_cls(layout_type_class: type(LayoutType)): - return _register_layout_cls(AffineQuantizedTensor, layout_type_class) - -def get_layout_tensor_constructor(layout_type_class: type(LayoutType)): - return _get_layout_tensor_constructor(AffineQuantizedTensor, layout_type_class) +register_layout_cls = AffineQuantizedTensor.register_layout_cls +get_layout_tensor_constructor = AffineQuantizedTensor.get_layout_tensor_constructor @dataclass(frozen=True) class SemiSparseLayoutType(LayoutType): diff --git a/torchao/dtypes/uintx/Uintx.py b/torchao/dtypes/uintx/Uintx.py index 067f4924fd..cfe75f4dc7 100644 --- a/torchao/dtypes/uintx/Uintx.py +++ b/torchao/dtypes/uintx/Uintx.py @@ -6,10 +6,8 @@ from .bitpacking import pack, unpack from torchao.dtypes.utils import ( LayoutType, - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, ) +from torchao.utils import TorchAOBaseTensor from torchao.dtypes.affine_quantized_tensor import PlainAQTLayout, register_layout_cls from torchao.utils import TORCH_VERSION_AT_LEAST_2_3 @@ -35,7 +33,7 @@ print("uintx feature need torch 2.3+, please upgrade pytorch") -class UintxTensor(torch.Tensor): +class UintxTensor(TorchAOBaseTensor): """ Splits int data into packed shards based on bit size fields: @@ -99,10 +97,6 @@ def __tensor_unflatten__( packed_shape, bit_width, pack_dim = tensor_attributes return cls(shards, packed_shape, bit_width, pack_dim) - implements = classmethod(_implements) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - __torch_function__ = classmethod(_dispatch__torch_function__) - def get_plain(self): return unpack(self.get_shards(), self.bit_width, dim = self.pack_dim) diff --git a/torchao/dtypes/utils.py b/torchao/dtypes/utils.py index 3a197da05e..1cf034296a 100644 --- a/torchao/dtypes/utils.py +++ b/torchao/dtypes/utils.py @@ -5,72 +5,6 @@ from dataclasses import dataclass from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 -""" -Helper function for implementing aten op or torch function dispatch -and dispatching to these implementations. -""" -def _implements(cls, aten_ops_or_torch_fns): - """Use this decorator to implement a function for an aten ops in __torch_dispatch__ - (if user passed in a list of ops) - or torch function in __torch_function__ (if user passed in a single object) - - class MyTensor(torch.Tensor): - ... - implements = classmethod(_implements) - - implements = MyTensor.implements - - @implements(torch.nn.functional.linear): - def _(func, types, args, kwargs): - ... - - """ - if not hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE"): - cls._ATEN_OP_OR_TORCH_FN_TABLE = {} - - if not isinstance(aten_ops_or_torch_fns, (list, tuple)): - aten_ops_or_torch_fns = [aten_ops_or_torch_fns] - def decorator(func): - for op in aten_ops_or_torch_fns: - @functools.wraps(op) - def wrapper(f, types, args, kwargs): - return func(f, types, args, kwargs) - - cls._ATEN_OP_OR_TORCH_FN_TABLE[op] = wrapper - return func - return decorator - -def _dispatch__torch_function__(cls, func, types, args=(), kwargs=None): - """Use this util function for a common `__torch_function__` implementation - that dispatches to ops/functions registered with `_implements` - - class MyTensor(torch.Tensor): - ... - __torch_function__ = classmethod(_dispatch__torch_function__) - """ - kwargs = {} if kwargs is None else kwargs - if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ - func in cls._ATEN_OP_OR_TORCH_FN_TABLE: - return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) - - with torch._C.DisableTorchFunctionSubclass(): - return func(*args, **kwargs) - -def _dispatch__torch_dispatch__(cls, func, types, args, kwargs): - """Use this util function for a common `__torch_dispatch__` implementation - that dispatches to ops/functions registered with `_implements` - - class MyTensor(torch.Tensor): - ... - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - """ - if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ - func in cls._ATEN_OP_OR_TORCH_FN_TABLE: - return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) - - raise NotImplementedError(f"{cls.__name__} dispatch: attempting to run unimplemented operator/function: {func}") - - """ Base class for different LayoutType, should not be instantiated directly used to allow users to pass around configurations for the layout tensor, e.g. inner_k_tiles @@ -101,52 +35,6 @@ def extra_repr(self) -> str: class PlainLayoutType(LayoutType): pass -""" -layout tensor constructor registration for different tensor subclassesa - -first key is a tensor subclass type like AffineQuantizedTensor -second key is an extended layout string, like tensor_core_tiled -value is a constructor for the LayoutTensor class, e.g. TensorCoreTiledAQTLayout.from_plain -""" -_LAYOUT_CONSTRUCTOR_TABLE: Dict[Callable, Dict[type(LayoutType), Callable]] = defaultdict(dict) - -def _register_layout_cls(cls: Callable, layout_type_class: type(LayoutType)): - """Helper function for layout registrations, this is used to implement - register_layout_cls decorator for each tensor subclass, see aqt.py for example usage - - Args: - cls: Tensor subclass type - layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` - - Returns: - a decorator that registers the layout tensor constructor in the table - """ - def decorator(layout_cls): - _LAYOUT_CONSTRUCTOR_TABLE[cls][layout_type_class] = layout_cls.from_plain - if TORCH_VERSION_AT_LEAST_2_5: - # Allow serialization to work for models uses this layout tensor subclass - torch.serialization.add_safe_globals([layout_type_class, layout_cls]) - return layout_cls - return decorator - -def _get_layout_tensor_constructor(cls: Callable, layout_type_class: type(LayoutType)) -> Callable: - """Get Layout class constructor (LayoutClass.from_plain) for `cls` based on `layout_type_class` - `layout_type_class` means the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` - - Args: - cls: Tensor subclass type - layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` - - Returns: - layout tensor subclass constructor for the layout_type_class - """ - if cls not in _LAYOUT_CONSTRUCTOR_TABLE: - raise ValueError(f"no registered layout class constructor for: {cls}") - if layout_type_class not in _LAYOUT_CONSTRUCTOR_TABLE[cls]: - raise ValueError(f"layout_name: {layout_type_class} is not supported yet for {cls}") - - return _LAYOUT_CONSTRUCTOR_TABLE[cls][layout_type_class] - def is_device(target_device_str: str, device: Union[str, torch.device]): return torch.device(device).type == target_device_str diff --git a/torchao/prototype/low_bit_optim/subclass_4bit.py b/torchao/prototype/low_bit_optim/subclass_4bit.py index 5c83d83773..b87d1d4aba 100644 --- a/torchao/prototype/low_bit_optim/subclass_4bit.py +++ b/torchao/prototype/low_bit_optim/subclass_4bit.py @@ -2,7 +2,7 @@ import torch from torch import Tensor -from torchao.dtypes.utils import _implements, _dispatch__torch_dispatch__ +from torchao.utils import TorchAOBaseTensor from .quant_utils import create_dynamic_map, scale_tensor, quantize_4bit_with_qmap, dequant_with_qmap @@ -18,8 +18,7 @@ QMAP_UNSIGNED = torch.linspace(0, 1, 17)[1:].tolist() # no zero -class OptimState4bit(Tensor): - implements = classmethod(_implements) +class OptimState4bit(TorchAOBaseTensor): tensor_attrs = ["codes", "scale", "qmap"] @staticmethod @@ -80,8 +79,6 @@ def __repr__(self): f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" ) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - @OptimState4bit.implements(aten.copy_.default) def _(func, types, args, kwargs): diff --git a/torchao/prototype/low_bit_optim/subclass_8bit.py b/torchao/prototype/low_bit_optim/subclass_8bit.py index 9a4f54f71b..865498a57e 100644 --- a/torchao/prototype/low_bit_optim/subclass_8bit.py +++ b/torchao/prototype/low_bit_optim/subclass_8bit.py @@ -1,6 +1,6 @@ import torch from torch import Tensor -from torchao.dtypes.utils import _implements, _dispatch__torch_dispatch__ +from torchao.utils import TorchAOBaseTensor from .quant_utils import create_dynamic_map, scale_tensor, quantize_8bit_with_qmap, dequant_with_qmap @@ -13,8 +13,7 @@ QMAP_UNSIGNED = create_dynamic_map(signed=False) -class OptimState8bit(Tensor): - implements = classmethod(_implements) +class OptimState8bit(TorchAOBaseTensor): tensor_attrs = ["codes", "scale", "qmap"] @staticmethod @@ -66,8 +65,6 @@ def __repr__(self): f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" ) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - @OptimState8bit.implements(aten.copy_.default) def _(func, types, args, kwargs): diff --git a/torchao/prototype/low_bit_optim/subclass_fp8.py b/torchao/prototype/low_bit_optim/subclass_fp8.py index 883d331187..805c516f4e 100644 --- a/torchao/prototype/low_bit_optim/subclass_fp8.py +++ b/torchao/prototype/low_bit_optim/subclass_fp8.py @@ -1,6 +1,6 @@ import torch from torch import Tensor -from torchao.dtypes.utils import _implements, _dispatch__torch_dispatch__ +from torchao.utils import TorchAOBaseTensor aten = torch.ops.aten @@ -21,8 +21,7 @@ def quantize_fp8(input: Tensor, block_size: int): # NOTE: FP8 sign bit is redundant for unsigned optim state. # we may investigate how to use it to increase range/precision for unsigned optim state. -class OptimStateFp8(Tensor): - implements = classmethod(_implements) +class OptimStateFp8(TorchAOBaseTensor): tensor_attrs = ["codes", "scale"] @staticmethod @@ -72,8 +71,6 @@ def __repr__(self): f"shape={tuple(self.shape)}, device={self.device}, requires_grad={self.requires_grad})" ) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - @OptimStateFp8.implements(aten.copy_.default) def _(func, types, args, kwargs): diff --git a/torchao/prototype/quantized_training/int8.py b/torchao/prototype/quantized_training/int8.py index c301f011c2..fa8805fce7 100644 --- a/torchao/prototype/quantized_training/int8.py +++ b/torchao/prototype/quantized_training/int8.py @@ -4,7 +4,7 @@ from torch import Tensor, nn from torch.utils._python_dispatch import return_and_correct_aliasing -from torchao.dtypes.utils import _dispatch__torch_dispatch__, _dispatch__torch_function__, _implements +from torchao.utils import TorchAOBaseTensor aten = torch.ops.aten @@ -12,7 +12,7 @@ _c10d_functional = torch.ops._c10d_functional -class Int8QTLinearWeight(Tensor): +class Int8QTLinearWeight(TorchAOBaseTensor): """INT8 symmetric quantization weight, with absmax scaling [-127, 127]. The main difference of this tensor subclass from AffineQuantizedTensor: 1. `F.linear` is differentiable i.e. backward is defined. @@ -22,10 +22,6 @@ class Int8QTLinearWeight(Tensor): for more details. """ - implements = classmethod(_implements) - __torch_function__ = classmethod(_dispatch__torch_function__) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - @staticmethod @torch._dynamo.disable def __new__(cls, int_data: Tensor, scale: Tensor): diff --git a/torchao/quantization/linear_activation_quantized_tensor.py b/torchao/quantization/linear_activation_quantized_tensor.py index b88286b219..7fe76b20fa 100644 --- a/torchao/quantization/linear_activation_quantized_tensor.py +++ b/torchao/quantization/linear_activation_quantized_tensor.py @@ -1,9 +1,4 @@ import torch -from torchao.dtypes.utils import ( - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, -) from typing import Callable from torch.utils._python_dispatch import return_and_correct_aliasing from torchao.utils import ( @@ -94,10 +89,6 @@ def to(self, *args, **kwargs): self.input_quant_func, ) - implements = classmethod(_implements) - __torch_function__ = classmethod(_dispatch__torch_function__) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - implements = LinearActivationQuantizedTensor.implements @implements(torch.nn.functional.linear) diff --git a/torchao/sparsity/prototype/superblock/blocksparse.py b/torchao/sparsity/prototype/superblock/blocksparse.py index 0ed526f7b8..06b3548c55 100644 --- a/torchao/sparsity/prototype/superblock/blocksparse.py +++ b/torchao/sparsity/prototype/superblock/blocksparse.py @@ -3,11 +3,7 @@ import torch from typing import Optional, Tuple, List, Dict, Any, Callable from torch.utils._python_dispatch import return_and_correct_aliasing -from torchao.dtypes.utils import ( - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, -) +from torchao.utils import TorchAOBaseTensor from torchao.quantization.quant_api import _get_linear_subclass_inserter aten = torch.ops.aten @@ -24,16 +20,12 @@ def blocksparse_linear_abstract(A: torch.Tensor, crow_indices: torch.Tensor, col return torch.empty(new_shape, dtype=A.dtype, device=A.device) # Subclass definition -class BlockSparseTensor(torch.Tensor): +class BlockSparseTensor(TorchAOBaseTensor): bsr_crow_indices: Optional[torch.Tensor] bsr_col_indices: Optional[torch.Tensor] bsr_values: Optional[torch.Tensor] - __slots__ = ["bsr_crow_indices", "bsr_col_indices", "bsr_values"] - - implements = classmethod(_implements) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - __torch_function__ = classmethod(_dispatch__torch_function__) + __slots__ = ["bsr_crow_indices", "bsr_col_indices", "bsr_values"] @staticmethod def __new__( # noqa: PYI034 @@ -109,7 +101,7 @@ def apply_fn_to_shard(self, func): requires_grad=self.requires_grad, ) -# Subclass op dispatch registration +# Subclass op dispatch registration implements = BlockSparseTensor.implements @implements(aten.detach.default) diff --git a/torchao/utils.py b/torchao/utils.py index e55dc3dc00..53d76d4be2 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -285,9 +285,157 @@ def unwrap_tensor_subclass(model, filter_fn=None): unwrap_tensor_subclass(child) return model + +""" +Helper function for implementing aten op or torch function dispatch +and dispatching to these implementations. +""" +def _implements(cls, aten_ops_or_torch_fns): + """Use this decorator to implement a function for an aten ops in __torch_dispatch__ + (if user passed in a list of ops) + or torch function in __torch_function__ (if user passed in a single object) + + class MyTensor(torch.Tensor): + ... + implements = classmethod(_implements) + + implements = MyTensor.implements + + @implements(torch.nn.functional.linear): + def _(func, types, args, kwargs): + ... + + """ + if not hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE"): + cls._ATEN_OP_OR_TORCH_FN_TABLE = {} + + if not isinstance(aten_ops_or_torch_fns, (list, tuple)): + aten_ops_or_torch_fns = [aten_ops_or_torch_fns] + def decorator(func): + for op in aten_ops_or_torch_fns: + @functools.wraps(op) + def wrapper(f, types, args, kwargs): + return func(f, types, args, kwargs) + + cls._ATEN_OP_OR_TORCH_FN_TABLE[op] = wrapper + return func + return decorator + +def _dispatch__torch_function__(cls, func, types, args=(), kwargs=None): + """Use this util function for a common `__torch_function__` implementation + that dispatches to ops/functions registered with `_implements` + + class MyTensor(torch.Tensor): + ... + __torch_function__ = classmethod(_dispatch__torch_function__) + """ + kwargs = {} if kwargs is None else kwargs + if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ + func in cls._ATEN_OP_OR_TORCH_FN_TABLE: + return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) + + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **kwargs) + +def _dispatch__torch_dispatch__(cls, func, types, args, kwargs): + """Use this util function for a common `__torch_dispatch__` implementation + that dispatches to ops/functions registered with `_implements` + + class MyTensor(torch.Tensor): + ... + __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) + """ + if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ + func in cls._ATEN_OP_OR_TORCH_FN_TABLE: + return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) + + raise NotImplementedError(f"{cls.__name__} dispatch: attempting to run unimplemented operator/function: {func}") + +def _register_layout_cls(cls: Callable, layout_type_class: type(LayoutType)): + """Helper function for layout registrations, this is used to implement + register_layout_cls decorator for each tensor subclass, see aqt.py for example usage + + Args: + cls: Tensor subclass type + layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` + + Returns: + a decorator that registers the layout tensor constructor in the table + """ + + # cls._LAYOUT_CONSTRUCTOR_TABLE is a map from layout_type_class like TensorCoreTiledLayout + # to layout class constructor like TensorCoreTiledAQTLayout.from_plain that can construct a layout_tensor + # from plain data like (quantized, unpacked) `data`, `scale`, `zero_point` + if not hasattr(cls, "_LAYOUT_CONSTRUCTOR_TABLE"): + cls._LAYOUT_CONSTRUCTOR_TABLE = {} + + def decorator(layout_class): + cls._LAYOUT_CONSTRUCTOR_TABLE[layout_type_class] = layout_class.from_plain + if TORCH_VERSION_AT_LEAST_2_5: + # Allow serialization to work for models uses this layout tensor subclass + torch.serialization.add_safe_globals([layout_type_class, layout_class]) + return layout_class + return decorator + +def _get_layout_tensor_constructor(cls: Callable, layout_type_class: type(LayoutType)) -> Callable: + """Get Layout class constructor (LayoutClass.from_plain) for `cls` based on `layout_type_class` + `layout_type_class` means the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` + + Args: + cls: Tensor subclass type + layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` + + Returns: + layout tensor subclass constructor for the layout_type_class + """ + if not hasattr(cls, "_LAYOUT_CONSTRUCTOR_TABLE"): + raise ValueError(f"no registered layout class constructor for: {cls}") + if layout_type_class not in cls._LAYOUT_CONSTRUCTOR_TABLE: + raise ValueError(f"layout_name: {layout_type_class} is not supported yet for {cls}") + + return cls._LAYOUT_CONSTRUCTOR_TABLE[layout_type_class] + + class TorchAOBaseTensor(torch.Tensor): """A util tensor subclass that provides commonly used functions + new tensor subclass can inherit it to get all the utility functions + + class MyTensor(TorchAOBaseTensor): + pass + + This includes: + `_get_to_kwargs` that can get the kwargs for `to` + class MyTensor(TorchAOBaseTensor): + def to(self, *args, **kwargs): + kwargs = _get_to_kwargs(*args, **kwargs) + ... + `implements`: + implements = MyTensor.implements + + @implements(torch.nn.functional.linear): + def _(func, types, args, kwargs): + ... + + `register_layout_cls`: + register_layout_cls = MyTensor.register_layout_cls + + @register_layout_cls(PlainLayoutType) + class PlainAQTLayout(...): + ... + + `get_layout_tensor_constructor`: + get_layout_tensor_constructor = MyTensor.get_layout_tensor_constructor + # in constructor of MyTensor: + layout_tensor_ctr = get_layout_tensor_constructor(type(layout_type)) + layout_tensor = layout_tensor_ctr(data, scale, zero_point, layout_type) + """ + implements = classmethod(_implements) + __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) + __torch_function__ = classmethod(_dispatch__torch_function__) + register_layout_cls = classmethod(_register_layout_cls) + get_layout_tensor_constructor = classmethod(_get_layout_tensor_constructor) + def _get_to_kwargs(self, *args, **kwargs): # `torch._C._nn._parse_to` can't handle `layout` argument for arg in args: diff --git a/tutorials/developer_api_guide/my_dtype_tensor_subclass.py b/tutorials/developer_api_guide/my_dtype_tensor_subclass.py index f2ed169289..c599ca4f4d 100644 --- a/tutorials/developer_api_guide/my_dtype_tensor_subclass.py +++ b/tutorials/developer_api_guide/my_dtype_tensor_subclass.py @@ -17,14 +17,10 @@ from torch.utils._python_dispatch import return_and_correct_aliasing from torchao.quantization.quant_primitives import choose_qparams_affine, MappingType from torchao.dtypes.utils import ( - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, - _register_layout_cls, - _get_layout_tensor_constructor, LayoutType, PlainLayoutType, ) +from torchao.utils import TorchAOBaseTensor aten = torch.ops.aten @@ -65,7 +61,11 @@ def __repr__(self): # Tensor Subclass Definition # ############################## -class MyDTypeTensor(torch.Tensor): +class MyDTypeTensor(TorchAOBaseTensor): + """Inheriting from `TorchAOBaseTensor` gives us some helper functions, please see docs + for :class:`~torchao.utils.TorchAOBaseTensor` for more details + """ + """We need to define __new__ for constructing a new tensor subclass instance and __init__ for initialize the instance. There is no requirement on what the argument list should look like here, only requirement is that `__new__` must return a Tensor instance with `torch.Tensor._make_wrapper_subclass(cls, shape, ...)` call @@ -176,8 +176,6 @@ def _apply_fn_to_data(self, fn): self.dtype, ) - implements = classmethod(_implements) - """There are two entry points that we can modify the behavior of a pytorch op: torch_function and torch_dispatch: __torch_function__: will be called whenever a torch level function is called on the Tensor object, for example: torch.nn.functional.linear, @@ -188,8 +186,6 @@ def _apply_fn_to_data(self, fn): We have some helper functions that can dispatch to the functions registered with MyDTypeTensor.implements, but if the default implementation does not work for your use case, please feel free to customize it """ - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - __torch_function__ = classmethod(_dispatch__torch_function__) ###################################################### # LayoutType and Layout Tensor Subclass Registration # From b0768abb37ad9652098329eb993256103139c4e4 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 30 Aug 2024 15:27:20 -0700 Subject: [PATCH 2/6] fix import --- torchao/dtypes/utils.py | 5 +-- torchao/utils.py | 75 +++++++++++++++++++++-------------------- 2 files changed, 39 insertions(+), 41 deletions(-) diff --git a/torchao/dtypes/utils.py b/torchao/dtypes/utils.py index 1cf034296a..7771bc34c5 100644 --- a/torchao/dtypes/utils.py +++ b/torchao/dtypes/utils.py @@ -1,9 +1,6 @@ import torch -from typing import Dict, Callable, Union, Tuple, Optional -from collections import defaultdict -import functools +from typing import Union, Tuple from dataclasses import dataclass -from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 """ Base class for different LayoutType, should not be instantiated directly diff --git a/torchao/utils.py b/torchao/utils.py index 53d76d4be2..8cee5fc3df 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -1,6 +1,7 @@ import torch -from typing import Tuple, Any +from typing import Tuple, Any, Callable from functools import reduce +import functools from importlib.metadata import version from math import gcd import torch.nn.utils.parametrize as parametrize @@ -286,6 +287,42 @@ def unwrap_tensor_subclass(model, filter_fn=None): return model +def _is_float8_type(dtype: torch.dtype) -> bool: + fp8_types = { + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, + torch.float8_e5m2fnuz, + } + return dtype in fp8_types + + +def parse_version(version_string): + # Extract just the X.Y.Z part from the version string + match = re.match(r'(\d+\.\d+\.\d+)', version_string) + if match: + version = match.group(1) + return [int(x) for x in version.split('.')] + else: + raise ValueError(f"Invalid version string format: {version_string}") + +def compare_versions(v1, v2): + v1_parts = parse_version(v1) + v2_parts = parse_version(v2) + return (v1_parts > v2_parts) - (v1_parts < v2_parts) + +def is_fbcode(): + return not hasattr(torch.version, "git_version") + +def torch_version_at_least(min_version): + return is_fbcode() or compare_versions(torch.__version__, min_version) >= 0 + +TORCH_VERSION_AT_LEAST_2_5 = torch_version_at_least("2.5.0") +TORCH_VERSION_AT_LEAST_2_4 = torch_version_at_least("2.4.0") +TORCH_VERSION_AT_LEAST_2_3 = torch_version_at_least("2.3.0") +TORCH_VERSION_AT_LEAST_2_2 = torch_version_at_least("2.2.0") + + """ Helper function for implementing aten op or torch function dispatch and dispatching to these implementations. @@ -457,42 +494,6 @@ def _get_to_kwargs(self, *args, **kwargs): return kwargs -def _is_float8_type(dtype: torch.dtype) -> bool: - fp8_types = { - torch.float8_e4m3fn, - torch.float8_e4m3fnuz, - torch.float8_e5m2, - torch.float8_e5m2fnuz, - } - return dtype in fp8_types - - -def parse_version(version_string): - # Extract just the X.Y.Z part from the version string - match = re.match(r'(\d+\.\d+\.\d+)', version_string) - if match: - version = match.group(1) - return [int(x) for x in version.split('.')] - else: - raise ValueError(f"Invalid version string format: {version_string}") - -def compare_versions(v1, v2): - v1_parts = parse_version(v1) - v2_parts = parse_version(v2) - return (v1_parts > v2_parts) - (v1_parts < v2_parts) - -def is_fbcode(): - return not hasattr(torch.version, "git_version") - -def torch_version_at_least(min_version): - return is_fbcode() or compare_versions(torch.__version__, min_version) >= 0 - -TORCH_VERSION_AT_LEAST_2_5 = torch_version_at_least("2.5.0") -TORCH_VERSION_AT_LEAST_2_4 = torch_version_at_least("2.4.0") -TORCH_VERSION_AT_LEAST_2_3 = torch_version_at_least("2.3.0") -TORCH_VERSION_AT_LEAST_2_2 = torch_version_at_least("2.2.0") - - ## Deprecated, will be deleted in the future def _torch_version_at_least(min_version): return is_fbcode() or version("torch") >= min_version From 9aabaf770bbeac236b694072edd1d1451a6755b4 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 30 Aug 2024 15:47:44 -0700 Subject: [PATCH 3/6] fix doc build --- torchao/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchao/utils.py b/torchao/utils.py index 8cee5fc3df..1f4f66e1f4 100644 --- a/torchao/utils.py +++ b/torchao/utils.py @@ -388,7 +388,7 @@ class MyTensor(torch.Tensor): raise NotImplementedError(f"{cls.__name__} dispatch: attempting to run unimplemented operator/function: {func}") -def _register_layout_cls(cls: Callable, layout_type_class: type(LayoutType)): +def _register_layout_cls(cls: Callable, layout_type_class: Callable): """Helper function for layout registrations, this is used to implement register_layout_cls decorator for each tensor subclass, see aqt.py for example usage @@ -414,7 +414,7 @@ def decorator(layout_class): return layout_class return decorator -def _get_layout_tensor_constructor(cls: Callable, layout_type_class: type(LayoutType)) -> Callable: +def _get_layout_tensor_constructor(cls: Callable, layout_type_class: Callable) -> Callable: """Get Layout class constructor (LayoutClass.from_plain) for `cls` based on `layout_type_class` `layout_type_class` means the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` From a23586adf50a42bc36d208e09c0243f5a9d6a860 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 30 Aug 2024 15:57:59 -0700 Subject: [PATCH 4/6] refactor affine fake quantized tensor --- .../qat/affine_fake_quantized_tensor.py | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py b/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py index c5f8204771..9d7cdccfc3 100644 --- a/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py +++ b/torchao/quantization/prototype/qat/affine_fake_quantized_tensor.py @@ -9,11 +9,7 @@ MappingType, ) from torch.utils._python_dispatch import return_and_correct_aliasing -from torchao.dtypes.utils import ( - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, -) +from torchao.utils import TorchAOBaseTensor from .utils import ( _GenericFakeQuantize, _UnwrapAffineFakeQuantizedTensor, @@ -80,7 +76,7 @@ def backward(ctx, gy): return gy, None, None, None, None, None, None, None, None, None, None -class AffineFakeQuantizedTensor(torch.Tensor): +class AffineFakeQuantizedTensor(TorchAOBaseTensor): """ Affine fake quantized tensor subclass. Affine quantization means we quantize the floating point tensor with an affine transformation: @@ -179,15 +175,15 @@ def _get_to_kwargs(self, *args, **kwargs): device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs) device = self.device if device is None else device dtype = self.dtype if dtype is None else dtype - memory_format = ( + memory_format = ( memory_format if memory_format is not None else torch.preserve_format - ) - kwargs = { + ) + kwargs = { "device": device, "dtype": dtype, "memory_format": memory_format, "requires_grad": self.requires_grad, - } + } return kwargs def to(self, *args, **kwargs): @@ -227,10 +223,6 @@ def _create_new(self, new_value: torch.Tensor): requires_grad=False, ) - implements = classmethod(_implements) - __torch_function__ = classmethod(_dispatch__torch_function__) - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) - implements = AffineFakeQuantizedTensor.implements From 60444e79e275dff80af8a68a97a183658443e9d1 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 30 Aug 2024 16:11:45 -0700 Subject: [PATCH 5/6] remove import --- torchao/dtypes/fpx/fpx.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/torchao/dtypes/fpx/fpx.py b/torchao/dtypes/fpx/fpx.py index 00ee84b658..c4f818ea12 100644 --- a/torchao/dtypes/fpx/fpx.py +++ b/torchao/dtypes/fpx/fpx.py @@ -8,14 +8,10 @@ from torchao.ops import quant_llm_linear from torchao.dtypes.utils import ( LayoutType, - _implements, - _dispatch__torch_function__, - _dispatch__torch_dispatch__, ) from torchao.quantization.quant_api import _get_linear_subclass_inserter from dataclasses import dataclass from torchao.dtypes.affine_quantized_tensor import AQTLayout, register_layout_cls -from torchao.utils import TorchAOBaseTensor aten = torch.ops.aten From 8586f6a1e236febf65191cd9408f2faf7f40636f Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 3 Sep 2024 15:52:40 -0700 Subject: [PATCH 6/6] pt version --- test/prototype/test_low_bit_optim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py index 701d90e22c..0979dd5cf2 100644 --- a/test/prototype/test_low_bit_optim.py +++ b/test/prototype/test_low_bit_optim.py @@ -75,7 +75,7 @@ def test_quantize_4bit_with_qmap_compile(self, device): class TestOptim(TestCase): - @pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_3, reason="requires PyTorch >= 2.3") + @pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_4, reason="requires PyTorch >= 2.3") @parametrize("optim_name", ["Adam8bit", "AdamW8bit", "Adam4bit", "AdamW4bit", "AdamFp8", "AdamWFp8"]) @parametrize("dtype", [torch.float32, torch.bfloat16]) @parametrize("device", _DEVICES)