[Float8Configs] Make named tuples have better docs + public

stack-info: PR: #808, branch: drisspg/stack/9
pytorch · Sep 4, 2024 · 07b71e4 · 07b71e4
1 parent f5703b0
commit 07b71e4
Show file tree

Hide file tree

Showing 9 changed files with 108 additions and 614 deletions.
diff --git a/ruff.toml b/ruff.toml
@@ -7,4 +7,5 @@ include = [
     "torchao/float8/float8_utils.py",
     "torchao/dtypes/nf4tensor.py",
     "test/dtypes/test_nf4.py",
+    "torchao/float8/float8_tensor.py",
 ]
diff --git a/test/float8/test_base.py b/test/float8/test_base.py
@@ -45,11 +45,6 @@
     FP8_TYPES,
     tensor_to_scale,
 )
-from torchao.float8.inference import (
-    ActivationCasting,
-    QuantConfig,
-    quantize_to_float8,
-)
 
 random.seed(0)
 torch.manual_seed(0)
@@ -134,21 +129,6 @@ def test_copy_(self):
         fp8_b.copy_(fp8_a)
         torch.testing.assert_close(fp8_a._data, fp8_b._data)
 
-    @pytest.mark.skip("broken")
-    def test_weights_only_load(self):
-        module = nn.Linear(16, 16)
-        # Save model state dict
-        buffer = io.BytesIO()
-        fp8_module = quantize_to_float8(
-            module,
-            QuantConfig(
-                ActivationCasting.DYNAMIC,
-            ),
-        )
-
-        torch.save(fp8_module.state_dict(), buffer)
-        buffer.seek(0)
-        _ = torch.load(buffer, weights_only=True)
 
 
 class TestFloat8Linear:

diff --git a/test/float8/test_everything.sh b/test/float8/test_everything.sh
@@ -6,7 +6,6 @@ IS_ROCM=$(rocm-smi --version || true)
 
 pytest test/float8/test_base.py
 pytest test/float8/test_compile.py
-pytest test/float8/test_inference_flows.py
 pytest test/float8/test_numerics_integration.py
 
 # These tests do not work on ROCm yet

diff --git a/test/float8/test_inference_flows.py b/test/float8/test_inference_flows.py