Torchao version check changes/BC import of TensorCoreTiledLayout

torchtune/training/quantization.py

@@ -6,7 +6,13 @@
 
 from typing import Callable, Optional
 
-from torchao.dtypes import TensorCoreTiledLayoutType
+from torchtune.utils._import_guard import _NEW_TENSOR_CORE_TILED_LAYOUT_API
+
+if _NEW_TENSOR_CORE_TILED_LAYOUT_API:
+    from torchao.dtypes import TensorCoreTiledLayout
+else:
+    from torchao.dtypes import TensorCoreTiledLayoutType as TensorCoreTiledLayout
+
 from torchao.quantization import (
     int4_weight_only,
     int8_dynamic_activation_int4_weight,
@@ -88,7 +94,7 @@ def __init__(self, groupsize: int = 128, inner_k_tiles: int = 8):
         self.inner_k_tiles = inner_k_tiles
 
     def quantize(self, model):
-        layout_type = TensorCoreTiledLayoutType(self.inner_k_tiles)
+        layout_type = TensorCoreTiledLayout(self.inner_k_tiles)
         quantize_fn = int4_weight_only(self.groupsize, layout_type)
         quantize_(model, quantize_fn)
         return model

torchtune/utils/_import_guard.py

@@ -5,11 +5,23 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torchtune.utils._version import torch_version_ge
+from torchtune.utils._version import (
+    _get_torchao_version,
+    _is_fbcode,
+    _nightly_version_ge,
+    torch_version_ge,
+)
 
 # We can only use flex attention / BlockMask if torch version >= 2.5.0 and GPU is Turing / SM75 and above
 _SUPPORTS_FLEX_ATTENTION = (
     torch_version_ge("2.5.0")
     and torch.cuda.is_available()
     and torch.cuda.get_device_capability() >= (7, 5)
 )
+
+torchao_version = _get_torchao_version()
+
+_NEW_TENSOR_CORE_TILED_LAYOUT_API = not _is_fbcode() and (
+    ("dev" not in torchao_version and torchao_version >= "0.6.0")
+    or ("dev" in torchao_version and _nightly_version_ge(torchao_version, "2024-10-10"))
+)

torchtune/utils/_version.py

@@ -3,8 +3,14 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+from datetime import datetime
+from typing import Optional
+
 import torch
 
+import torchao
+
 
 def torch_version_ge(version: str) -> bool:
     """
@@ -23,3 +29,36 @@ def torch_version_ge(version: str) -> bool:
         True
     """
     return version in torch.__version__ or torch.__version__ >= version
+
+
+def _is_fbcode():
+    return not hasattr(torch.version, "git_version")
+
+
+def _nightly_version_ge(ao_version_str: str, date: str) -> bool:
+    """
+    Compare a torchao nightly version to a date of the form
+    %Y-%m-%d.
+
+    Returns True if the nightly version is greater than or equal to
+        the date, False otherwise
+    """
+    ao_datetime = datetime.strptime(
+        ao_version_str.split("+")[0].split("dev")[1], "%Y%m%d"
+    )
+    return ao_datetime >= datetime.strptime(date, "%Y-%m-%d")
+
+
+def _get_torchao_version() -> Optional[str]:
+    """
+    Get torchao version.
+
+    Checks:
+        1) is_fbcode, then
+        2) torchao.__version__ (only defined for torchao >= 0.3.0), then
+
+    """
+    if _is_fbcode():
+        return None
+    else:
+        return torchao.__version__