Merge branch 'main' into fix_llama

Author: yiliu30

benchmarks/benchmark_fp6_llm.py renamed to benchmarks/benchmark_fp6.py

@@ -1,16 +1,15 @@
 import torch
 import pandas as pd
 import torch.nn.functional as F
-from torchao.prototype.quant_llm import QuantLlmLinearWeight
+from torchao.dtypes import to_affine_quantized_fpx
+from torchao.dtypes.fpx import FpxTensorCoreAQTLayout, FpxTensorCoreLayoutType
 from torchao.utils import benchmark_torch_function_in_microseconds
 from tqdm import tqdm
 
 
 def benchmark(m: int, k: int, n: int):
-    fp6_data = torch.randint(256, size=(n, k * 3 // 4), dtype=torch.uint8, device="cuda")
-    scale = torch.rand(n, dtype=torch.half, device="cuda") + 0.5
-    fp6_weight = QuantLlmLinearWeight(fp6_data, scale, 3, 2)
-
+    float_data = torch.randn(n, k, dtype=torch.half, device="cuda")
+    fp6_weight = to_affine_quantized_fpx(float_data, FpxTensorCoreLayoutType(3, 2))
     fp16_weight = fp6_weight.dequantize(torch.half)
 
     fp16_act = torch.randn(m, k, dtype=torch.half, device="cuda")

benchmarks/float8/bench_linear_float8.py

@@ -91,6 +91,8 @@ def float8_pct_top_peak(self):
         return self.float8_tops_sec / dtype_to_peak_tops[torch.float8_e4m3fn]
 
 
+# TODO(future PR): add option to measure GPU kernel time, as in other
+# scripts in this folder
 def main(
     sweep_path: Optional[Path] = None,
     compile: bool = True,
@@ -112,10 +114,33 @@ def main(
     scaling_type_input = ScalingType(scaling_type_input)
     scaling_type_weight = ScalingType(scaling_type_weight)
     scaling_type_grad_output = ScalingType(scaling_type_grad_output)
+
+    if scaling_type_input is ScalingType.STATIC:
+        cast_config_input=CastConfig(
+            scaling_type=scaling_type_input,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_input=CastConfig(scaling_type=scaling_type_input)
+    if scaling_type_weight is ScalingType.STATIC:
+        cast_config_weight=CastConfig(
+            scaling_type=scaling_type_weight,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_weight=CastConfig(scaling_type=scaling_type_weight)
+    if scaling_type_grad_output is ScalingType.STATIC:
+        cast_config_grad_output=CastConfig(
+            scaling_type=scaling_type_grad_output,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output)
+
     config = Float8LinearConfig(
-        cast_config_input=CastConfig(scaling_type=scaling_type_input),
-        cast_config_weight=CastConfig(scaling_type=scaling_type_weight),
-        cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output),
+        cast_config_input=cast_config_input,
+        cast_config_weight=cast_config_weight,
+        cast_config_grad_output=cast_config_grad_output,
     )
 
     name_to_shapes = get_name_to_shapes_iter(shape_gen_name, M, K, N)

benchmarks/float8/profile_linear_float8.py

@@ -263,13 +263,35 @@ def main(
     scaling_type_input = ScalingType(scaling_type_input)
     scaling_type_weight = ScalingType(scaling_type_weight)
     scaling_type_grad_output = ScalingType(scaling_type_grad_output)
+
+    if scaling_type_input is ScalingType.STATIC:
+        cast_config_input=CastConfig(
+            scaling_type=scaling_type_input,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_input=CastConfig(scaling_type=scaling_type_input)
+    if scaling_type_weight is ScalingType.STATIC:
+        cast_config_weight=CastConfig(
+            scaling_type=scaling_type_weight,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_weight=CastConfig(scaling_type=scaling_type_weight)
+    if scaling_type_grad_output is ScalingType.STATIC:
+        cast_config_grad_output=CastConfig(
+            scaling_type=scaling_type_grad_output,
+            static_scale=torch.tensor([1.0], device="cuda"),
+        )
+    else:
+        cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output)
+
     config = Float8LinearConfig(
-        cast_config_input=CastConfig(scaling_type=scaling_type_input),
-        cast_config_weight=CastConfig(scaling_type=scaling_type_weight),
-        cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output),
-        enable_amax_init=False,
-        enable_pre_and_post_forward=False,
+        cast_config_input=cast_config_input,
+        cast_config_weight=cast_config_weight,
+        cast_config_grad_output=cast_config_grad_output,
     )
+
     scaling_repr = "_".join(
         [
             s.short_str()

docs/source/api_ref_dtypes.rst

@@ -11,7 +11,9 @@ torchao.dtypes
     :nosignatures:
 
     to_nf4
-    to_affine_quantized
+    to_affine_quantized_intx
+    to_affine_quantized_floatx
+    to_affine_quantized_intx_static
     AffineQuantizedTensor
 
 ..

scripts/hf_eval.py

@@ -20,6 +20,7 @@
     int8_dynamic_activation_int8_weight,
     quantize_,
     autoquant,
+    fpx_weight_only,
 )
 from torchao.sparsity import (
     sparsify_,
@@ -59,6 +60,8 @@ def run_evaluation(repo_id, tasks, limit, device, precision, quantization, spars
     elif quantization == "int4wo":
         # note cannot quantize this model on cpu and run it on cuda at this time
         quantize_(model.to(device=device), int4_weight_only())
+    elif quantization == "fp6":
+        quantize_(model, fpx_weight_only(3, 2))
     elif quantization == "autoquant":
         model = autoquant(model.to(device=device))
 
@@ -79,7 +82,7 @@ def all_linear(mod, name):
             return False
         torch.sparse.semi_structured._FORCE_CUTLASS = False
         sparsify_(model, semi_sparse_weight(), filter_fn=all_linear)
-    
+
     if sparsity and compile:
         model = torch.compile(model, mode="max-autotune", fullgraph=True)
 
@@ -111,7 +114,7 @@ def all_linear(mod, name):
     parser.add_argument('--limit', type=int, default=None, help='Number of eval samples to evaluate')
     parser.add_argument('--precision', type=lambda x: getattr(torch, x.split(".")[-1]), default=torch.bfloat16, help='dtype precision to use')
     parser.add_argument('--device', type=str, default="cuda", help='Device to use for evaluation')
-    parser.add_argument('-q', '--quantization', default = "None", choices=["int8dq", "int8wo", "int4wo","autoquant", "None"], help='Which quantization technique to apply')
+    parser.add_argument('-q', '--quantization', default = "None", choices=["int8dq", "int8wo", "int4wo","autoquant", "fp6", "None"], help='Which quantization technique to apply')
     parser.add_argument('-s', '--sparsity', default = "None", choices=["semi_sparse", "semi_sparse_mlp_only", "None"], help='Which sparsity technique to apply')
     parser.add_argument('--compile', action='store_true', help='Whether to compile the model.')
     parser.add_argument('--save', action='store_true', help='Whether to save the model.')

test/dtypes/test_affine_quantized.py

@@ -8,16 +8,15 @@
     int8_dynamic_activation_int4_weight,
     int8_dynamic_activation_int8_weight,
     int8_dynamic_activation_int8_semi_sparse_weight,
-)
-from torchao.dtypes import (
-    to_affine_quantized,
+    float8_weight_only,
 )
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 import torch
 import unittest
 import tempfile
 
+
 class TestAffineQuantized(TestCase):
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_tensor_core_layout_transpose(self):
@@ -40,7 +39,8 @@ def test_tensor_core_layout_transpose(self):
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_weights_only(self):
-        for apply_quant in [int4_weight_only(group_size=32), int8_weight_only(), int8_dynamic_activation_int4_weight(), int8_dynamic_activation_int8_weight(), int8_dynamic_activation_int8_semi_sparse_weight()]:
+        for apply_quant in [int4_weight_only(group_size=32), int8_weight_only(), int8_dynamic_activation_int4_weight(),
+                            int8_dynamic_activation_int8_weight(), int8_dynamic_activation_int8_semi_sparse_weight(), float8_weight_only()]:
             l = torch.nn.Linear(128, 256, dtype=torch.bfloat16, device="cuda")
             ql = apply_quant(l)
             with tempfile.NamedTemporaryFile() as f:
@@ -69,6 +69,5 @@ def test_to_device(self):
             ql.cuda()
 
 
-
 if __name__ == "__main__":
     run_tests()

test/prototype/test_quant_llm.py renamed to test/dtypes/test_fpx.py

@@ -8,23 +8,27 @@
     parametrize,
     run_tests,
 )
-from torchao.prototype.quant_llm import (
-    QuantLlmLinearWeight,
-    quant_llm_fpx_weight_only,
-    fp6_llm_weight_only,
+from torchao.dtypes.fpx import (
+    FpxTensorCoreAQTLayout,
+    FpxTensorCoreLayoutType,
     to_scaled_tc_fpx,
     from_scaled_tc_fpx,
 )
-from torchao.prototype.quant_llm.quant_llm import _pack_tc_fpx, _pack_tc_fp6
+from torchao.dtypes.fpx.fpx import _pack_tc_fpx, _pack_tc_fp6
 from torchao.prototype.custom_fp_utils import _f32_to_fpx_unpacked, _fpx_unpacked_to_f32
-from torchao.quantization.quant_api import quantize_
+from torchao.quantization import (
+    quantize_,
+    fpx_weight_only,
+)
+
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 
 _DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
 _FPx_DTYPES = [(3, 2), (2, 2)]
 
 
-class TestQuantLlmLinearWeight(TestCase):
+class TestFpxTensorCoreAQTLayout(TestCase):
     @parametrize("device", _DEVICES)
     def test_pack_tc_fp6_correctness(self, device):
         x = torch.randint(256, size=(256, 64), dtype=torch.uint8, device=device)
@@ -69,61 +73,40 @@ def test_from_scaled_tc_fpx_compile(self, ebits, mbits, device):
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @parametrize("ebits,mbits", _FPx_DTYPES)
     def test_to_copy_device(self, ebits, mbits):
-        x = torch.randn(256, 64)
-        fpx = QuantLlmLinearWeight.from_float(x, ebits, mbits).cuda()
-        assert fpx.device.type == "cuda"
-        fpx = fpx.cpu()
-        assert fpx.device.type == "cpu"
+        from torchao.quantization.quant_primitives import (
+            choose_qparams_affine_fpx,
+            quantize_affine_fpx,
+        )
 
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    @parametrize("ebits,mbits", _FPx_DTYPES)
-    @parametrize("leading_dims", [(4,), (2, 4)])
-    @parametrize("bias", [False, True])
-    def test_quant_llm_linear_weight(self, ebits, mbits, bias, leading_dims):
-        OC, IC = 256, 64
-        device = "cuda"
-
-        fp16_weight = torch.randn(OC, IC, device=device, dtype=torch.half)
-        fp16_bias = torch.randn(OC, device=device, dtype=torch.half) if bias else None
-
-        fpx_weight = QuantLlmLinearWeight.from_float(fp16_weight, ebits, mbits)
-
-        x = torch.randn(*leading_dims, IC, device=device, dtype=torch.half)
-        out = torch.nn.functional.linear(x, fpx_weight, fp16_bias)
-        assert out.shape == leading_dims + (OC,)
+        x = torch.randn(256, 64)
+        scale = choose_qparams_affine_fpx(x, ebits, mbits)
+        x = quantize_affine_fpx(x, scale, ebits, mbits)
+        layout_type = FpxTensorCoreLayoutType(ebits, mbits)
+        fpx_layout_tensor = FpxTensorCoreAQTLayout.from_plain(x, scale, None, layout_type).cuda()
+        assert fpx_layout_tensor.device.type == "cuda"
+        fpx_layout_tensor = fpx_layout_tensor.cpu()
+        assert fpx_layout_tensor.device.type == "cpu"
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_5, reason="quantization only works with torch.compile for 2.5+")
     @parametrize("ebits,mbits", _FPx_DTYPES)
     @parametrize("bias", [False, True])
-    def test_quant_llm_quantize(self, ebits, mbits, bias):
-        N, OC, IC = 4, 256, 64
-        device = "cuda"
-
-        linear = torch.nn.Linear(IC, OC, bias=bias, device=device)
-        fpx_linear = copy.deepcopy(linear)
-        quantize_(fpx_linear, quant_llm_fpx_weight_only(ebits, mbits))
-
-        x = torch.randn(N, IC, device=device, dtype=torch.half)
-        expected = fpx_linear(x)
-        actual = torch.compile(fpx_linear, fullgraph=True)(x)
-        torch.testing.assert_close(actual, expected)
-
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    def test_fp6_llm_quantize(self):
+    def test_fpx_weight_only(self, ebits, mbits, bias):
         N, OC, IC = 4, 256, 64
         device = "cuda"
 
-        linear = torch.nn.Linear(IC, OC, device=device)
+        linear = torch.nn.Linear(IC, OC, bias=bias, device=device, dtype=torch.half)
         fpx_linear = copy.deepcopy(linear)
-        quantize_(fpx_linear, fp6_llm_weight_only())
+        quantize_(fpx_linear, fpx_weight_only(ebits, mbits))
 
         x = torch.randn(N, IC, device=device, dtype=torch.half)
         expected = fpx_linear(x)
         actual = torch.compile(fpx_linear, fullgraph=True)(x)
+        # somehow compile now changes the result a bit
         torch.testing.assert_close(actual, expected)
 
 
-instantiate_parametrized_tests(TestQuantLlmLinearWeight)
+instantiate_parametrized_tests(TestFpxTensorCoreAQTLayout)
 
 
 if __name__ == "__main__":

test/float8/test_base.py

@@ -134,6 +134,7 @@ def test_copy_(self):
         fp8_b.copy_(fp8_a)
         torch.testing.assert_close(fp8_a._data, fp8_b._data)
 
+    @pytest.mark.skip("broken")
     def test_weights_only_load(self):
         module = nn.Linear(16, 16)
         # Save model state dict
@@ -226,14 +227,16 @@ def _test_linear_impl(
     @pytest.mark.parametrize("emulate", [True, False] if is_cuda_8_9 else [True])
     @pytest.mark.parametrize("x_shape", [(16, 16), (2, 16, 16), (3, 2, 16, 16)])
     @pytest.mark.parametrize(
-        "scaling_type_input", [ScalingType.DELAYED, ScalingType.DYNAMIC]
+        "scaling_type_input", 
+        [ScalingType.DELAYED, ScalingType.DYNAMIC, ScalingType.STATIC]
     )
     @pytest.mark.parametrize(
-        "scaling_type_weight", [ScalingType.DELAYED, ScalingType.DYNAMIC]
+        "scaling_type_weight", 
+        [ScalingType.DELAYED, ScalingType.DYNAMIC, ScalingType.STATIC]
     )
     @pytest.mark.parametrize(
         "scaling_type_grad_output",
-        [ScalingType.DELAYED, ScalingType.DYNAMIC],
+        [ScalingType.DELAYED, ScalingType.DYNAMIC, ScalingType.STATIC],
     )
     @pytest.mark.parametrize("linear_dtype", [torch.bfloat16, torch.float32])
     @pytest.mark.parametrize("linear_bias", [False, True])
@@ -259,10 +262,33 @@ def test_linear(
                 pytest.skip()
         x = torch.randn(*x_shape, device="cuda", dtype=linear_dtype)
         m_ref = nn.Linear(16, 32, bias=linear_bias, device="cuda", dtype=linear_dtype)
+
+        if scaling_type_input is ScalingType.STATIC:
+            cast_config_input = CastConfig(
+                scaling_type=scaling_type_input,
+                static_scale=torch.tensor([1.0], device="cuda"),
+            )
+        else:
+            cast_config_input = CastConfig(scaling_type=scaling_type_input)
+        if scaling_type_weight is ScalingType.STATIC:
+            cast_config_weight = CastConfig(
+                scaling_type=scaling_type_weight,
+                static_scale=torch.tensor([1.0], device="cuda"),
+            )
+        else:
+            cast_config_weight = CastConfig(scaling_type=scaling_type_weight)
+        if scaling_type_grad_output is ScalingType.STATIC:
+            cast_config_grad_output = CastConfig(
+                scaling_type=scaling_type_grad_output,
+                static_scale=torch.tensor([1.0], device="cuda"),
+            )
+        else:
+            cast_config_grad_output = CastConfig(scaling_type=scaling_type_grad_output)
+
         config = Float8LinearConfig(
-            cast_config_input=CastConfig(scaling_type=scaling_type_input),
-            cast_config_weight=CastConfig(scaling_type=scaling_type_weight),
-            cast_config_grad_output=CastConfig(scaling_type=scaling_type_grad_output),
+            cast_config_input=cast_config_input,
+            cast_config_weight=cast_config_weight,
+            cast_config_grad_output=cast_config_grad_output,
             emulate=emulate,
         )
         self._test_linear_impl(