_layout -> layout for public API

Author: jainapurva

test/dtypes/test_affine_quantized.py

@@ -31,7 +31,7 @@ def get_quantization_functions(do_sparse: bool, do_int4: bool):
         base_functions.append(int4_weight_only(group_size=32))
 
     if do_sparse:
-        base_functions.append(int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()))
+        base_functions.append(int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()))
 
     if is_cuda_8_9:
         base_functions.append(float8_weight_only())

test/integration/test_integration.py

@@ -876,7 +876,7 @@ def test_int4_weight_only_quant_subclass_api_grouped(self, device, dtype):
         for test_shape in ([(256, 256, 16)] + ([(256, 256, 8)] if device=='cuda' else [])):
             for groupsize in [64, 32]:
                 for inner_k_tiles in [4, 2]:
-                    kwargs = {"groupsize": groupsize, "_layout": TensorCoreTiledLayout(inner_k_tiles=inner_k_tiles)}
+                    kwargs = {"groupsize": groupsize, "layout": TensorCoreTiledLayout(inner_k_tiles=inner_k_tiles)}
 
                     def api(mod):
                         kwargs_copy = kwargs.copy()
@@ -888,7 +888,7 @@ def api(mod):
                                 unwrap_tensor_subclass(mod)
                         else:
                             kwargs_copy["inner_k_tiles"] = inner_k_tiles
-                            del kwargs_copy["_layout"]
+                            del kwargs_copy["layout"]
                             change_linear_weights_to_int4_woqtensors(mod, **kwargs_copy)
 
                     self._test_lin_weight_subclass_api_impl(

test/sparsity/test_marlin.py

@@ -50,7 +50,7 @@ def test_quant_sparse_marlin_layout_eager(self):
         dense_result = model_copy(self.input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(self.model, int4_weight_only(_layout=MarlinSparseLayout()))
+        quantize_(self.model, int4_weight_only(layout=MarlinSparseLayout()))
         sparse_result = self.model(self.input)
 
         assert torch.allclose(dense_result, sparse_result, atol=3e-1), "Results are not close"
@@ -67,7 +67,7 @@ def test_quant_sparse_marlin_layout_compile(self):
         dense_result = model_copy(self.input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(self.model, int4_weight_only(_layout=MarlinSparseLayout()))
+        quantize_(self.model, int4_weight_only(layout=MarlinSparseLayout()))
         self.model.forward = torch.compile(self.model.forward, fullgraph=True)
         sparse_result = self.model(self.input)
 

test/sparsity/test_sparse_api.py

@@ -74,7 +74,7 @@ def test_quant_semi_sparse(self, compile):
 
         quantize_(
             model,
-            int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()),
+            int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()),
         )
         if compile:
             model = torch.compile(model)
@@ -108,7 +108,7 @@ def test_sparse_marlin(self, compile):
         dense_result = model_copy(input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
+        quantize_(model, int4_weight_only(layout=MarlinSparseLayout()))
         if compile:
             model = torch.compile(model)
         sparse_result = model(input)
@@ -190,7 +190,7 @@ def test_sparse(self, compile):
         quantize_(
             model,
             int8_dynamic_activation_int8_weight(
-                _layout=BlockSparseLayout(blocksize=64)
+                layout=BlockSparseLayout(blocksize=64)
             ),
         )
         if compile:

torchao/_models/llama/eval.py

@@ -98,7 +98,7 @@ def run_evaluation(
             quantize_(model, uintx_weight_only(dtype, group_size, use_hqq=use_hqq))
         if "marlin" in quantization:
             from torchao.dtypes import MarlinSparseLayout
-            quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
+            quantize_(model, int4_weight_only(layout=MarlinSparseLayout()))
         if "int4wo" in quantization and "gptq" in quantization:
             # avoid circular imports
             from torchao._models._eval import InputRecorder

torchao/_models/llama/generate.py

@@ -231,7 +231,7 @@ def main(
             quantize_(model, int4_weight_only(group_size=groupsize))
         if "marlin" in quantization:
             from torchao.dtypes import MarlinSparseLayout
-            quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
+            quantize_(model, int4_weight_only(layout=MarlinSparseLayout()))
         if "fp6" in quantization:
             quantize_(model, fpx_weight_only(3, 2))
         if quantization.startswith("awq"):

torchao/_models/sam/eval_combo.py

@@ -315,7 +315,7 @@ def mlp_only(mod, name):
                   int8_dynamic_activation_int8_weight(),
                   attn_only)
         quantize_(predictor.model.image_encoder,
-                  int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()),
+                  int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()),
                   mlp_lin1_only)
         sparsify_(predictor.model.image_encoder,
                   semi_sparse_weight(),
@@ -330,7 +330,7 @@ def mlp_only(mod, name):
         quantize_(predictor.model.image_encoder,
                   int8_dynamic_activation_int8_weight(),
                   attn_only)
-        quantize_(predictor.model.image_encoder, int4_weight_only(_layout=MarlinSparseLayout()), mlp_lin1_only)
+        quantize_(predictor.model.image_encoder, int4_weight_only(layout=MarlinSparseLayout()), mlp_lin1_only)
         sparsify_(predictor.model.image_encoder,
                   semi_sparse_weight(),
                   mlp_lin2_only)

torchao/quantization/quant_api.py

@@ -511,7 +511,7 @@ def int8_dynamic_activation_int4_weight(group_size=32, mapping_type=MappingType.
     return _get_linear_subclass_inserter(apply_int8_dynamic_activation_int4_weight_quant, group_size=group_size, mapping_type=mapping_type)
 
 
-def int4_weight_only(group_size=128, _layout=TensorCoreTiledLayout(inner_k_tiles=8), use_hqq=False):
+def int4_weight_only(group_size=128, layout=TensorCoreTiledLayout(inner_k_tiles=8), use_hqq=False):
     """
     Applies uint4 weight-only asymmetric per-group quantization to linear layers, using
     "tensor_core_tiled" layout for speedup with tinygemm kernel
@@ -527,7 +527,7 @@ def int4_weight_only(group_size=128, _layout=TensorCoreTiledLayout(inner_k_tiles
     Args:
         `group_size`: parameter for quantization, controls the granularity of quantization, smaller
          size is more fine grained, choices are [256, 128, 64, 32]
-        `_layout`: layout type for quantized tensor, default is `TensorCoreTiledLayout(inner_k_tiles=8)`
+        `layout`: layout type for quantized tensor, default is `TensorCoreTiledLayout(inner_k_tiles=8)`
         `use_hqq`: whether to use hqq or default quantization mode, default is False
     """
     def apply_int4_weight_only_quant(weight):
@@ -550,12 +550,12 @@ def apply_int4_weight_only_quant(weight):
         # Sparse Marlin only supports symmetric quantization.
         # NOTE: If we start having lots of layouts that require different configurations,
         # we should consider moving this logic somewhere else.
-        if isinstance(_layout, MarlinSparseLayout):
+        if isinstance(layout, MarlinSparseLayout):
             mapping_type = MappingType.SYMMETRIC
             preserve_zero = True
             zero_point_domain = ZeroPointDomain.INT
 
-        return to_affine_quantized_intx(weight, mapping_type, block_size, target_dtype, quant_min, quant_max, eps, zero_point_dtype=zero_point_dtype, preserve_zero=preserve_zero, zero_point_domain=zero_point_domain, _layout=_layout, use_hqq=use_hqq)
+        return to_affine_quantized_intx(weight, mapping_type, block_size, target_dtype, quant_min, quant_max, eps, zero_point_dtype=zero_point_dtype, preserve_zero=preserve_zero, zero_point_domain=zero_point_domain, _layout=layout, use_hqq=use_hqq)
 
     return _get_linear_subclass_inserter(apply_int4_weight_only_quant)
 
@@ -583,7 +583,7 @@ def _int8_symm_per_token_reduced_range_quant(x: torch.Tensor) -> torch.Tensor:
     return to_affine_quantized_intx(x, mapping_type, _get_per_token_block_size(x), target_dtype, eps=eps, quant_min=quant_min, quant_max=quant_max, scale_dtype=torch.float32 if x.dtype == torch.float16 else None)
 
 
-def int8_dynamic_activation_int8_weight(_layout=PlainLayout()):
+def int8_dynamic_activation_int8_weight(layout=PlainLayout()):
     """
     Applies int8 dynamic symmetric per-token activation and int8 per-channel weight
     quantization to linear layers
@@ -609,7 +609,7 @@ def get_weight_block_size(x):
         input_quant_func = _int8_symm_per_token_reduced_range_quant
 
         block_size = get_weight_block_size(weight)
-        weight = to_affine_quantized_intx(weight, mapping_type, block_size, target_dtype, eps=eps, zero_point_dtype=zero_point_dtype, _layout=_layout)
+        weight = to_affine_quantized_intx(weight, mapping_type, block_size, target_dtype, eps=eps, zero_point_dtype=zero_point_dtype, _layout=layout)
         weight = to_linear_activation_quantized(weight, input_quant_func)
         return weight
 
@@ -621,12 +621,12 @@ def int8_dynamic_activation_int8_semi_sparse_weight():
     Applies int8 dnynamic symmetric per-token activation and int8 per-channel weight
     quantization + 2:4 sparsity to linear layers.
     """
-    warnings.warn("""int8_dyanmic_activation_int8_semi_sparse_weight() will be deprecated at a later release. Please use the _layout kwarg in int8_dynamic_activation_int8_weight instead.
+    warnings.warn("""int8_dyanmic_activation_int8_semi_sparse_weight() will be deprecated at a later release. Please use the layout kwarg in int8_dynamic_activation_int8_weight instead.
 
     from torchao.dtypes import SemiSparseLayout
-    int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()""")
+    int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()""")
 
-    return int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout())
+    return int8_dynamic_activation_int8_weight(layout=SemiSparseLayout())
 
 
 def float8_weight_only(weight_dtype: torch.dtype = torch.float8_e4m3fn):

torchao/sparsity/README.md

@@ -57,7 +57,7 @@ from torchao.dtypes import MarlinSparseLayout
 
 # Your FP16 model
 model = model.cuda().half()
-quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
+quantize_(model, int4_weight_only(layout=MarlinSparseLayout()))
 ```
 
 Note the existing API results in an extremely high accuracy degredation and is intended to be used in concert with an already sparsified+finetuned checkpoint where possible until we develop
@@ -72,7 +72,7 @@ from torchao.quantization.quant_api import quantize_, int8_dynamic_activation_in
 from torchao.dtypes import SemiSparseLayout
 
 model = model.cuda()
-quantize_(model, int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()))
+quantize_(model, int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()))
 ```
 
 ### 2:4 sparsity

torchao/sparsity/prototype/superblock/utils.py

@@ -164,7 +164,7 @@ def accelerate_with_sparsity(model, args):
 
             quantize_(
                 model,
-                int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()),
+                int8_dynamic_activation_int8_weight(layout=SemiSparseLayout()),
                 mlp_0_only,
             )
             sparsify_(model, semi_sparse_weight(), mlp_3_only)