From b916570ebd3ced05685901e40b289a489f1a3521 Mon Sep 17 00:00:00 2001
From: Aarav <211792885+AaravMehta-07@users.noreply.github.com>
Date: Sun, 11 Oct 2026 00:11:59 +0530
Subject: [PATCH] Support root parameters in safetensors state dictionaries

Share prefix construction between flattening and reconstruction while preserving nested names. Cover real CPU plain/Int8 round trips, forward equivalence and deferred shard loading.

Co-authored-by: OpenAI Codex
Signed-off-by: Aarav <211792885+AaravMehta-07@users.noreply.github.com>
---
 .../safetensors/test_safetensors_support.py   | 58 +++++++++++++++++++
 .../safetensors/safetensors_support.py        | 12 ++--
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/test/prototype/safetensors/test_safetensors_support.py b/test/prototype/safetensors/test_safetensors_support.py
index e40f8cb8622..5ccb1f83c2e 100644
--- a/test/prototype/safetensors/test_safetensors_support.py
+++ b/test/prototype/safetensors/test_safetensors_support.py
@@ -163,6 +163,64 @@ def test_safetensors_sharded(self, config, device, act_pre_scale=False):
                     assert key in leftover_tensor_data_dict
 
 
+class TestSafeTensorsRootParameters(TestCase):
+    @parametrize("depth", [0, 1, 2])
+    @parametrize("quantized", [False, True])
+    def test_root_and_nested_parameters_round_trip(self, depth, quantized):
+        model = torch.nn.Linear(32, 16)
+        for _ in range(depth):
+            model = torch.nn.Sequential(model)
+        if quantized:
+            quantize_(model, Int8WeightOnlyConfig(version=2))
+        inputs = torch.randn(2, 32)
+        expected = model(inputs)
+        state_dict = model.state_dict()
+        tensors_data, metadata = flatten_tensor_state_dict(state_dict)
+        prefix = "0." * depth
+        if quantized:
+            self.assertIn(f"{prefix}_weight_qdata", tensors_data)
+        self.assertIn(f"{prefix}bias", tensors_data)
+
+        with tempfile.TemporaryDirectory() as directory:
+            path = f"{directory}/model.safetensors"
+            save_file(tensors_data, path, metadata=metadata)
+            loaded_tensors, loaded_metadata = load_data(path, "cpu")
+            restored, leftover = unflatten_tensor_state_dict(
+                loaded_tensors, loaded_metadata
+            )
+
+        self.assertFalse(leftover)
+        self.assertEqual(set(restored), set(state_dict))
+        restored_model = torch.nn.Linear(32, 16)
+        for _ in range(depth):
+            restored_model = torch.nn.Sequential(restored_model)
+        restored_model.load_state_dict(restored, assign=True)
+        self.assertEqual(restored_model(inputs), expected)
+        self.assertEqual(
+            type(restored[f"{prefix}weight"]), type(state_dict[f"{prefix}weight"])
+        )
+
+    @parametrize("depth", [0, 1])
+    def test_partial_quantized_state_waits_for_remaining_data(self, depth):
+        model = torch.nn.Linear(32, 16)
+        for _ in range(depth):
+            model = torch.nn.Sequential(model)
+        quantize_(model, Int8WeightOnlyConfig(version=2))
+        tensors_data, metadata = flatten_tensor_state_dict(model.state_dict())
+        prefix = "0." * depth
+        qdata = tensors_data.pop(f"{prefix}_weight_qdata")
+
+        restored, leftover = unflatten_tensor_state_dict(tensors_data, metadata)
+        self.assertEqual(set(restored), {f"{prefix}bias"})
+        self.assertEqual(set(leftover), set(tensors_data) - {f"{prefix}bias"})
+        remaining = {**leftover, f"{prefix}_weight_qdata": qdata}
+        restored_weight, leftover = unflatten_tensor_state_dict(remaining, metadata)
+        self.assertEqual(set(restored_weight), {f"{prefix}weight"})
+        self.assertFalse(leftover)
+        self.assertEqual(restored_weight[f"{prefix}weight"].qdata, qdata)
+
+
+instantiate_parametrized_tests(TestSafeTensorsRootParameters)
 instantiate_parametrized_tests(TestSafeTensors)
 
 if __name__ == "__main__":
diff --git a/torchao/prototype/safetensors/safetensors_support.py b/torchao/prototype/safetensors/safetensors_support.py
index 08040f17b68..a8c25d7b45a 100644
--- a/torchao/prototype/safetensors/safetensors_support.py
+++ b/torchao/prototype/safetensors/safetensors_support.py
@@ -13,6 +13,11 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+def _tensor_data_prefix(tensor_name: str) -> str:
+    module_fqn, separator, weight_name = tensor_name.rpartition(".")
+    return f"{module_fqn}{separator}_{weight_name}_"
+
+
 def unflatten_tensor_state_dict(
     tensors_data_dict: Dict[str, Any],
     metadata: Dict[str, Any],
@@ -71,9 +76,7 @@ def unflatten_tensor_state_dict(
     for tensor_name in tensor_names:
         processed_tensors = []
 
-        module_fqn, weight_name = tensor_name.rsplit(".", 1)
-
-        prefix = f"{module_fqn}._{weight_name}_"
+        prefix = _tensor_data_prefix(tensor_name)
         tensor_tensors = {}
 
         for key, value in combined_data.items():
@@ -181,8 +184,9 @@ def flatten_tensor_state_dict(
             tensor_metadata = json.dumps(tensor, cls=TensorSubclassAttributeJSONEncoder)
 
             # Clone tensors to avoid memory sharing issues
+            prefix = _tensor_data_prefix(tensor_name)
             tensors_dict_to_save = {
-                f"{tensor_name.rsplit('.', 1)[0]}._{tensor_name.rsplit('.', 1)[1]}_{key}": (
+                f"{prefix}{key}": (
                     value.detach().clone() if isinstance(value, torch.Tensor) else value
                 )
                 for key, value in tensor_dict.items()