migrate nf4 to configs

test/dtypes/test_nf4.py

@@ -26,11 +26,11 @@
 
 import torchao
 from packaging import version
+from torchao.dtypes._nf4tensor_api import nf4_weight_only
 from torchao.dtypes.nf4tensor import (
     _INNER_TENSOR_NAMES_FOR_SHARDING,
     NF4Tensor,
     linear_nf4,
-    nf4_weight_only,
     to_nf4,
 )
 from torchao.testing.utils import skip_if_rocm

torchao/dtypes/_nf4tensor_api.py

@@ -0,0 +1,34 @@
+import torch
+
+from torchao.core.config import AOBaseConfig
+from torchao.dtypes.nf4tensor import NF4Tensor
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
+
+
+class NF4WeightOnlyConfig(AOBaseConfig):
+    """
+    Note: the file location of this workflow is temporary.
+    TODO(future PR): integrate this properly into torchao's directory structure
+    """
+
+    block_size: int = 64
+    scaler_block_size: int = 256
+
+
+# for bc
+nf4_weight_only = NF4WeightOnlyConfig
+
+
+@register_quantize_module_handler(NF4WeightOnlyConfig)
+def _nf4_weight_only_transform(
+    module: torch.nn.Module,
+    config: NF4WeightOnlyConfig,
+) -> torch.nn.Module:
+    block_size = config.block_size
+    scaler_block_size = config.scaler_block_size
+
+    new_weight = NF4Tensor.from_tensor(module.weight, block_size, scaler_block_size)
+    module.weight = torch.nn.Parameter(new_weight, requires_grad=False)
+    return module

torchao/dtypes/nf4tensor.py

@@ -985,15 +985,6 @@ def to_nf4(tensor, block_size: int = 64, scaler_block_size: int = 256):
     return NF4Tensor.from_tensor(tensor, block_size, scaler_block_size)
 
 
-def nf4_weight_only(block_size: int = 64, scaler_block_size: int = 256):
-    from torchao.quantization.quant_api import _get_linear_subclass_inserter
-
-    def _to_nf4(tensor: torch.Tensor):
-        return NF4Tensor.from_tensor(tensor, block_size, scaler_block_size)
-
-    return _get_linear_subclass_inserter(_to_nf4)
-
-
 NF4_TORCH_FUNCTIONS = {}