Summary
NamedShardingMetadata.to_serialized_string serializes the device mesh via dataclasses.asdict(self.device_mesh) (sharding.py#L261), and it runs once per array being saved. A checkpoint save typically has many arrays that all share the same mesh, so this re-walks and deep-copies the identical mesh once per array - O(num_arrays) redundant work that grows with mesh size and array count, and can become a significant fraction of save time at scale.
Proposal
Since all arrays in a save share one mesh, the conversion can be memoized so the deep copy happens once per distinct mesh instead of once per array:
# Build one shared DeviceMetadataMesh per distinct mesh
@classmethod
@functools.lru_cache(maxsize=128)
def from_jax_mesh(cls, mesh: jax.sharding.Mesh) -> Optional["DeviceMetadataMesh"]:
...
# Memoize the asdict deep-copy on that shared instance
@functools.cached_property
def as_serialized_dict(self) -> dict[str, Any]:
return dataclasses.asdict(self)
# to_serialized_string uses the cached dict instead of asdict per call
sharding_data[_DEVICE_MESH] = self.device_mesh.as_serialized_dictThe two work together: from_jax_mesh returns a single shared instance per distinct mesh, which lets as_serialized_dict memoize the deep copy exactly once — collapsing O(num_arrays) copies to O(num_unique_meshes).
Output is byte-identical to dataclasses.asdict, so existing checkpoints and the restore path are unaffected. The cache is stored on the instance (cached_property in __dict__), leaving dataclasses.asdict / __eq__ semantics unchanged.
Questions for maintainers
- Is memoizing the mesh serialization this way acceptable in general?
lru_cache on from_jax_mesh keyed by jax.sharding.Mesh relies on Mesh hashability/equality — is that a safe assumption across backends?- Any concern with the shared-instance semantics or cache lifetime (e.g. long-lived meshes retained by the
lru_cache) that we should account for?
Summary
NamedShardingMetadata.to_serialized_stringserializes the device mesh viadataclasses.asdict(self.device_mesh)(sharding.py#L261), and it runs once per array being saved. A checkpoint save typically has many arrays that all share the same mesh, so this re-walks and deep-copies the identical mesh once per array - O(num_arrays) redundant work that grows with mesh size and array count, and can become a significant fraction of save time at scale.Proposal
Since all arrays in a save share one mesh, the conversion can be memoized so the deep copy happens once per distinct mesh instead of once per array:
The two work together:
from_jax_meshreturns a single shared instance per distinct mesh, which letsas_serialized_dictmemoize the deep copy exactly once — collapsing O(num_arrays) copies to O(num_unique_meshes).Output is byte-identical to
dataclasses.asdict, so existing checkpoints and the restore path are unaffected. The cache is stored on the instance (cached_propertyin__dict__), leavingdataclasses.asdict/__eq__semantics unchanged.Questions for maintainers
lru_cacheonfrom_jax_meshkeyed byjax.sharding.Meshrelies on Mesh hashability/equality — is that a safe assumption across backends?lru_cache) that we should account for?