Remove FbgemmConfig and remaining Fbgemm tensors (#3032)

Summary:
This is used for prototype previously, not used now, we now expose fbgemm
kernels through Int4WeightOnlyConfig (for int4) and Float8DynamicActivationFloat8WeightConfig (for FP8)

Not considering this BC breaking since we haven't publicized the API yet

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

docs/source/torchao_vllm_integration.md

@@ -171,7 +171,7 @@ class MyNewQuantConfig(AOBaseConfig):
     VERSION: ClassVar[int] = 1
 
 class MyQuantizedTensor(TorchAOBaseTensor):
-    """Example based on FbgemmFp8Tensor - stores quantized data + scale"""
+    """Example based on Float8Tensor - stores quantized data + scale"""
 
     tensor_data_attrs = ["quantized_data", "scale"]
     tensor_attributes = ["dtype"]

test/core/test_config.py

@@ -24,7 +24,6 @@
     AWQStep,
 )
 from torchao.quantization.quant_api import (
-    FbgemmConfig,
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
     Float8WeightOnlyConfig,
@@ -92,7 +91,6 @@
     ),
     AWQConfig(Int4WeightOnlyConfig(group_size=128), step=AWQStep.PREPARE_FOR_LOADING),
     AWQConfig(Int4WeightOnlyConfig(group_size=128), step="prepare_for_loading"),
-    FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16, [1, 1, 256]),
 ]
 
 

test/dtypes/test_affine_quantized.py

@@ -24,9 +24,7 @@
     to_affine_quantized_intx,
     to_affine_quantized_intx_static,
 )
-from torchao.float8.config import e4m3_dtype
 from torchao.quantization import (
-    FbgemmConfig,
     Float8WeightOnlyConfig,
     GemliteUIntXWeightOnlyConfig,
     Int4DynamicActivationInt4WeightConfig,
@@ -44,7 +42,6 @@
     is_fbcode,
     is_ROCM,
     is_sm_at_least_89,
-    is_sm_at_least_90,
 )
 
 is_cusparselt_available = (
@@ -100,10 +97,6 @@ def get_quantization_functions(
     if is_sm_at_least_89():
         base_functions.append(Float8WeightOnlyConfig())
 
-    if is_sm_at_least_90():
-        base_functions.append(FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16))
-        base_functions.append(FbgemmConfig(e4m3_dtype, e4m3_dtype, torch.bfloat16))
-
     return base_functions
 
 

torchao/_models/llama/generate.py

@@ -434,25 +434,6 @@ def ffn_or_attn_only(mod, fqn):
                 model,
                 Int4WeightOnlyConfig(group_size=group_size, use_hqq=use_hqq, version=1),
             )
-        elif "fbgemm" in quantization and "int4" in quantization:
-            from torchao.quantization import FbgemmConfig
-
-            _, precision, group_size = quantization.split("-")
-            group_size = int(group_size)
-            block_size = [1, group_size]
-            assert precision == "int4", f"FbegemmConfig({precision=}) not supported yet"
-            quantize_(
-                model,
-                FbgemmConfig(torch.bfloat16, torch.int4, torch.bfloat16, block_size),
-            )
-        elif "fbgemm" in quantization and "fp8" in quantization:
-            from torchao.float8.config import e4m3_dtype
-            from torchao.quantization import FbgemmConfig
-
-            quantize_(
-                model,
-                FbgemmConfig(e4m3_dtype, e4m3_dtype, torch.bfloat16),
-            )
         elif "int4dq-" in quantization:
             from torchao.dtypes import CutlassInt4PackedLayout
 

torchao/dtypes/__init__.py

@@ -8,7 +8,6 @@
     to_affine_quantized_intx,
     to_affine_quantized_intx_static,
 )
-from .fbgemm_fp8_tensor import FbgemmFp8Tensor, to_fbgemm_fp8
 from .floatx import (
     CutlassSemiSparseLayout,
     Float8Layout,

torchao/dtypes/fbgemm_fp8_tensor.py

@@ -43,7 +43,6 @@
 )
 from .quant_api import (
     CutlassInt4PackedLayout,
-    FbgemmConfig,
     Float8DynamicActivationFloat8SemiSparseWeightConfig,
     Float8DynamicActivationFloat8WeightConfig,
     Float8DynamicActivationInt4WeightConfig,
@@ -161,7 +160,6 @@
     "GemliteUIntXWeightOnlyConfig",
     "AOPerModuleConfig",
     "ModuleFqnToConfig",
-    "FbgemmConfig",
     # tensor subclasses
     "Int4Tensor",
     "Int4PlainInt32Tensor",