mxtensor: add serialization support

Summary:

This is needed to save HF models with MX weights.

Test Plan:

```
pytest test/prototype/mx_formats/test_inference_workflow.py -s -k test_inference_workflow_mx
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 60a78e5
ghstack-comment-id: 3338428571
Pull Request resolved: #3078

Author: vkuzo

test/prototype/mx_formats/test_inference_workflow.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import copy
+import tempfile
 
 import pytest
 import torch
@@ -100,6 +101,16 @@ def test_inference_workflow_mx(elem_dtype, bias: bool, compile: bool, emulate: b
         f"Got a sqnr of {sqnr} for {elem_dtype} and bias={bias}"
     )
 
+    # serialization
+    with tempfile.NamedTemporaryFile() as f:
+        torch.save(m_mx.state_dict(), f)
+        f.seek(0)
+
+        # temporary workaround for https://github.com/pytorch/ao/issues/3077
+        torch.serialization.add_safe_globals([getattr])
+
+        _ = torch.load(f, weights_only=True)
+
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 @pytest.mark.skipif(

torchao/prototype/mx_formats/inference_workflow.py

@@ -17,7 +17,11 @@
     _validate_elem_dtype,
     _validate_gemm_kernel_choice,
 )
-from torchao.prototype.mx_formats.mx_tensor import MXTensor, QuantizeTensorToMXKwargs
+from torchao.prototype.mx_formats.mx_tensor import (
+    MXTensor,
+    QuantizeTensorToMXKwargs,
+    ScaleCalculationMode,
+)
 from torchao.prototype.mx_formats.nvfp4_tensor import (
     NVFP4MMConfig,
     NVFP4Tensor,
@@ -206,6 +210,8 @@ def _nvfp4_inference_linear_transform(
         NVFP4Tensor,
         NVFP4MMConfig,
         MXGemmKernelChoice,
+        QuantizeTensorToMXKwargs,
+        ScaleCalculationMode,
     ]
 )