allowlist WeightWithDynamicFloat8CastTensor for deserialization for checkpointing (#2573)

danielvegamyhre · web-flow · commit 11ce634705af · 2025-07-18T16:53:53.000-07:00
diff --git a/torchao/float8/fsdp_utils.py b/torchao/float8/fsdp_utils.py
@@ -266,3 +266,7 @@ def fsdp_post_all_gather(
             self._linear_mm_config,
             gemm_input_role=GemmInputRole.WEIGHT,
         ), (data,)
+
+
+# Needed to allowlist this subclass for deserialization used for restoring checkpoints.
+torch.serialization.add_safe_globals([WeightWithDynamicFloat8CastTensor])
