[Float8] Fix serialization of dynamic activation fp8

stack-info: PR: #838, branch: drisspg/stack/12

Author: drisspg

test/dtypes/test_affine_quantized_float.py

@@ -24,6 +24,7 @@
 from functools import partial
 from typing import Tuple
 from contextlib import nullcontext
+import io
 
 
 random.seed(0)
@@ -142,6 +143,67 @@ def test_per_row_with_float32(self):
                 model, float8_dynamic_activation_float8_weight(granularity=PerRow())
             )
 
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    @unittest.skipIf(not is_cuda_8_9, "Requires GPU with compute capability >= 8.9")
+    @common_utils.parametrize("mode", ["dynamic", "weight-only"])
+    def test_serialization(self, mode: str):
+        # Create and quantize the model
+        model = ToyLinearModel(16, 32).to(device="cuda")
+        if mode == "dynamic":
+            factory = float8_dynamic_activation_float8_weight()
+        else:
+            factory = float8_weight_only()
+        quantize_(model, factory)
+
+        # Save the state dict to an in-memory buffer
+        buffer = io.BytesIO()
+        torch.save(model.state_dict(), buffer)
+
+        # Reset the buffer position
+        buffer.seek(0)
+
+        # Load the state dict from the buffer
+        loaded_state_dict = torch.load(buffer)
+
+        # Create a new model and load the state dict
+        with torch.device("meta"):
+            new_model = ToyLinearModel(16, 32)
+            new_model.load_state_dict(loaded_state_dict, assign=True)
+
+        # Compare the original and loaded models
+        if mode == "weight-only":
+            model_weight_1 = model.linear1.weight.layout_tensor.float8_data.to(
+                torch.float32
+            )
+            new_model_weight_1 = new_model.linear1.weight.layout_tensor.float8_data.to(
+                torch.float32
+            )
+
+            model_weight_2 = model.linear2.weight.layout_tensor.float8_data.to(
+                torch.float32
+            )
+            new_model_weight_2 = new_model.linear2.weight.layout_tensor.float8_data.to(
+                torch.float32
+            )
+
+        else:
+            model_weight_1 = model.linear1.weight.original_weight_tensor.layout_tensor.float8_data.to(
+                torch.float32
+            )
+            new_model_weight_1 = new_model.linear1.weight.original_weight_tensor.layout_tensor.float8_data.to(
+                torch.float32
+            )
+
+            model_weight_2 = model.linear2.weight.original_weight_tensor.layout_tensor.float8_data.to(
+                torch.float32
+            )
+            new_model_weight_2 = new_model.linear2.weight.original_weight_tensor.layout_tensor.float8_data.to(
+                torch.float32
+            )
+
+        assert torch.allclose(model_weight_1, new_model_weight_1)
+        assert torch.allclose(model_weight_2, new_model_weight_2)
+
 
 common_utils.instantiate_parametrized_tests(TestAffineQuantizedFloat8Compile)
 

torchao/quantization/quant_api.py

@@ -25,6 +25,7 @@
 from torchao.dtypes.uintx.Uintx import UintxLayoutType
 from torchao.dtypes import (
     to_affine_quantized_intx,
+    to_affine_quantized_floatx,
     TensorCoreTiledLayoutType,
     PlainLayoutType,
     AffineQuantizedTensor,
@@ -670,6 +671,35 @@ def _validate_granularity(
     else:
         raise ValueError(f"Invalid granularity specification: {granularity}, only PerTensor or PerRow are supported.")
 
+def _get_block_size(x: torch.Tensor, granularity: _fp8_granularities):
+    if isinstance(granularity, PerTensor):
+        return x.shape
+    elif isinstance(granularity, PerRow):
+        return (1,) * (x.dim() - 1) + (x.shape[-1],)
+    else:
+        raise ValueError(f"Unsupported granularity: {granularity}")
+
+
+def _input_quant_func_dyanmic_fp8(
+    x: torch.Tensor,
+    activation_granularity: _fp8_granularities,
+    activation_dtype: torch.dtype,
+):
+    if isinstance(activation_granularity, PerRow):
+        assert (
+            x.dtype == torch.bfloat16
+        ), "PerRow quantization only works for bfloat16 precision input activation"
+
+    block_size = _get_block_size(x, activation_granularity)
+    activation = to_affine_quantized_floatx(
+        input_float=x,
+        block_size=block_size,
+        target_dtype=activation_dtype,
+        scale_dtype=torch.float32,
+        layout_type=Float8LayoutType(mm_config=None),  # Config is stored on weight
+    )
+    return activation
+
 
 def float8_dynamic_activation_float8_weight(
     activation_dtype: torch.dtype = torch.float8_e4m3fn,
@@ -693,28 +723,18 @@ def float8_dynamic_activation_float8_weight(
         mm_config (Float8MMConfig): Configuration for the matrix multiplication. Default uses fast accumulation.
 
     """
-    from torchao.dtypes import to_affine_quantized_floatx
-
     if mm_config is None:
         mm_config = Float8MMConfig(use_fast_accum=True)
 
     activation_granularity, weight_granularity = _validate_granularity(granularity)
 
-    def get_block_size(x: torch.Tensor, granularity: _fp8_granularities):
-        if isinstance(granularity, PerTensor):
-            return x.shape
-        elif isinstance(granularity, PerRow):
-            return (1,) * (x.dim() - 1) + (x.shape[-1],)
-        else:
-            raise ValueError(f"Unsupported granularity: {granularity}")
-
     def apply_float8_dynamic_activation_quant(weight: torch.Tensor):
         if isinstance(weight_granularity, PerRow):
             assert (
                 weight.dtype == torch.bfloat16
             ), "PerRow quantization only works for bfloat16 precision input weight"
 
-        block_size = get_block_size(weight, weight_granularity)
+        block_size = _get_block_size(weight, weight_granularity)
         quantized_weight = to_affine_quantized_floatx(
             input_float=weight,
             block_size=block_size,
@@ -723,23 +743,11 @@ def apply_float8_dynamic_activation_quant(weight: torch.Tensor):
             layout_type=Float8LayoutType(mm_config=mm_config),
         )
 
-        def input_quant_func(x: torch.Tensor):
-            if isinstance(activation_granularity, PerRow):
-                assert (
-                    x.dtype == torch.bfloat16
-                ), "PerRow quantization only works for bfloat16 precision input activation"
-
-            block_size = get_block_size(x, activation_granularity)
-            activation = to_affine_quantized_floatx(
-                input_float=x,
-                block_size=block_size,
-                target_dtype=activation_dtype,
-                scale_dtype=torch.float32,
-                layout_type=Float8LayoutType(
-                    mm_config=None
-                ),  # Config is stored on weight
-            )
-            return activation
+        input_quant_func = partial(
+            _input_quant_func_dyanmic_fp8,
+            activation_granularity=activation_granularity,
+            activation_dtype=activation_dtype,
+        )
 
         quantized_weight = to_linear_activation_quantized(
             quantized_weight, input_quant_func