Fix broken circular dep error

Differential Revision: D75980282

Pull Request resolved: #2320

Author: jerryzh168

torchao/float8/__init__.py

@@ -16,6 +16,7 @@
 )
 from torchao.float8.fsdp_utils import precompute_float8_dynamic_scale_for_fsdp
 from torchao.float8.inference import Float8MMConfig
+from torchao.float8.types import FP8Granularity
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 if TORCH_VERSION_AT_LEAST_2_5:
@@ -41,5 +42,7 @@
     # top level UX
     "convert_to_float8_training",
     "precompute_float8_dynamic_scale_for_fsdp",
+    # types
+    "FP8Granularity",
     # note: Float8Tensor and Float8Linear are not public APIs
 ]

torchao/float8/inference.py

@@ -12,6 +12,7 @@
 import torch
 
 from torchao.float8.float8_utils import is_row_major, pad_tensor_for_matmul
+from torchao.float8.types import FP8Granularity
 from torchao.quantization.granularity import (
     PerRow,
     PerTensor,
@@ -116,9 +117,6 @@ def _is_rowwise_scaled(x) -> bool:
     return x.block_size == (1,) * (x.dim() - 1) + (x.shape[-1],)
 
 
-FP8Granularity = Union[PerTensor, PerRow]
-
-
 def _normalize_granularity(
     granularity: Optional[
         Union[

torchao/float8/types.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Common types for float8 quantization
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from torchao.quantization.granularity import PerRow, PerTensor
+
+
+# Define FP8Granularity type alias to break circular import dependencies
+FP8Granularity = Union["PerTensor", "PerRow"]