migrate prototype/quantized_training to configs

Summary:

As titled

Test Plan:

```
pytest test/prototype/test_quantized_training.py -s -x
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: e4b9e86
ghstack-comment-id: 2706898680
Pull Request resolved: #1855

Author: vkuzo

torchao/prototype/quantized_training/bitnet.py

@@ -12,7 +12,10 @@
 from torch.distributed._tensor import DTensor
 from torch.utils._triton import has_triton
 
-from torchao.quantization.quant_api import _get_linear_subclass_inserter
+from torchao.core.config import AOBaseConfig
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
 from torchao.utils import TorchAOBaseTensor
 
 from .int8 import quantize_int8_rowwise
@@ -232,10 +235,22 @@ def backward(ctx, grad_output):
         return grad_input, grad_weight, grad_bias
 
 
-def bitnet_training():
-    return _get_linear_subclass_inserter(
-        BitNetTrainingLinearWeight, allow_requires_grad=True
-    )
+class BitNetTrainingConfig(AOBaseConfig):
+    pass
+
+
+# for bc
+bitnet_training = BitNetTrainingConfig
+
+
+@register_quantize_module_handler(BitNetTrainingConfig)
+def _bitnet_training_transform(
+    module: torch.nn.Module,
+    config: BitNetTrainingConfig,
+) -> torch.nn.Module:
+    new_weight = BitNetTrainingLinearWeight(module.weight)
+    module.weight = torch.nn.Parameter(new_weight, requires_grad=True)
+    return module
 
 
 def _pack_i2_in_i8(x: Tensor):

torchao/prototype/quantized_training/int8.py

@@ -4,7 +4,10 @@
 from torch import Tensor
 from torch.utils._python_dispatch import return_and_correct_aliasing
 
-from torchao.quantization.quant_api import _get_linear_subclass_inserter
+from torchao.core.config import AOBaseConfig
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
 from torchao.utils import TorchAOBaseTensor
 
 aten = torch.ops.aten
@@ -293,7 +296,19 @@ def _(func, types, args, kwargs):
     return return_and_correct_aliasing(func, args, kwargs, out)
 
 
-def int8_weight_only_quantized_training():
-    return _get_linear_subclass_inserter(
-        Int8QuantizedTrainingLinearWeight.from_float, allow_requires_grad=True
-    )
+class Int8WeightOnlyQuantizedTrainingConfig(AOBaseConfig):
+    pass
+
+
+# for bc
+int8_weight_only_quantized_training = Int8WeightOnlyQuantizedTrainingConfig
+
+
+@register_quantize_module_handler(Int8WeightOnlyQuantizedTrainingConfig)
+def _int8_weight_only_quantized_training_transform(
+    module: torch.nn.Module,
+    config: Int8WeightOnlyQuantizedTrainingConfig,
+) -> torch.nn.Module:
+    new_weight = Int8QuantizedTrainingLinearWeight.from_float(module.weight)
+    module.weight = torch.nn.Parameter(new_weight, requires_grad=True)
+    return module

torchao/prototype/quantized_training/int8_mixed_precision.py

@@ -1,11 +1,15 @@
-from typing import Any, NamedTuple, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple, Union
 
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor, nn
 from torch.utils._triton import has_triton
 
-from torchao.quantization.quant_api import _get_linear_subclass_inserter
+from torchao.core.config import AOBaseConfig
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
 from torchao.utils import TorchAOBaseTensor
 
 from .int8 import quantize_int8_rowwise
@@ -23,10 +27,16 @@ def scaled_int8_mm(
         return torch._int_mm(A, B) * col_scale.view(-1) * row_scale.view(-1, 1)
 
 
-class Int8MixedPrecisionTrainingConfig(NamedTuple):
+@dataclass
+class Int8MixedPrecisionTrainingConfig(AOBaseConfig):
     output: bool = True
     grad_input: bool = True
     grad_weight: bool = True
+    module_swap: bool = False
+
+
+# for bc
+int8_mixed_precision_training = Int8MixedPrecisionTrainingConfig
 
 
 _DEFAULT_CONFIG = Int8MixedPrecisionTrainingConfig()
@@ -265,25 +275,23 @@ def backward(ctx, grad_output):
         return grad_input, grad_weight, grad_bias, None
 
 
-def int8_mixed_precision_training(
-    config: Int8MixedPrecisionTrainingConfig = _DEFAULT_CONFIG,
-    *,
-    module_swap: bool = False,
+@register_quantize_module_handler(Int8MixedPrecisionTrainingConfig)
+def _int8_mixed_precision_training_transform(
+    module: torch.nn.Module,
+    config: Int8MixedPrecisionTrainingConfig,
 ):
+    module_swap = config.module_swap
+
     # TODO: skip small layers that don't have perf gain.
     if module_swap:
         # module swap implementation
-        def convert_linear(linear: nn.Linear):
-            linear.__class__ = Int8MixedPrecisionTrainingLinear
-            linear.config = config
-            return linear
-
-        return convert_linear
+        module.__class__ = Int8MixedPrecisionTrainingLinear
+        module.config = config
+        return module
 
     else:
         # tensor subclass implementation
-        return _get_linear_subclass_inserter(
-            Int8MixedPrecisionTrainingLinearWeight,
-            config=config,
-            allow_requires_grad=True,
-        )
+
+        new_weight = Int8MixedPrecisionTrainingLinearWeight(module.weight, config)
+        module.weight = torch.nn.Parameter(new_weight, requires_grad=True)
+        return module