Skip to content
Docs
Working With Models
Sharded checkpoints

Sharded Checkpoints

🚧
This page is written from general distributed-tensor domain knowledge (the placement vocabulary is standard across PyTorch DTensor/FSDP-style training), not adapted from a source document - the local PyTorch reference checkout in this workspace covers pytorch.org's marketing/get-started site, not the distributed-tensor API reference, so there is no local source to draw from for framework-specific claims.

A model trained across many GPUs is usually checkpointed the same way it was trained: split. A "sharded checkpoint" is a set of files where a single logical tensor is divided across shards rather than stored whole in any one file - sometimes purely for file-size reasons (an index file lists which shard holds which tensor, with no tensor actually divided), sometimes because a tensor was genuinely partitioned across devices during training (each shard holds a slice of one logical tensor, not the whole thing).

Two different things called "sharding"

It matters which kind you're looking at, because Tensormorph handles them differently:

  • File-level sharding - a large checkpoint split into multiple safetensors files purely so no single file is unmanageably large, with an index (commonly model.safetensors.index.json) mapping each tensor name to the shard file that holds it whole. Every tensor still exists complete in exactly one file. Tensormorph reads the index once and treats the whole set as a single logical checkpoint - the split is invisible everywhere except a "this tensor lives in shard 3 of 12" note in the Inspector panel.
  • Placement sharding - a tensor that was genuinely partitioned during training, so that no single file (or even any single device) ever holds the complete tensor. Reconstructing it means combining slices from multiple shards according to a placement description - commonly a device mesh (the logical grid of ranks a tensor was distributed across) plus a placement per mesh dimension (a shard dimension, where each rank holds a distinct slice; a replicate dimension, where every rank holds the full tensor; or a partial dimension, awaiting a reduction across ranks).

Viewing a placement-sharded tensor

Tensormorph reassembles a placement-sharded tensor's logical view - the Architecture and Matrix/Heatmap views render it as one seamless tensor - without necessarily materializing the whole thing in memory on any single device. Which shard actually holds the data you're looking at is visible in the Inspector panel as a local() vs. full() distinction: local() shows only the data resident on one shard, full() requests the logical whole (potentially the expensive path for a very large tensor, and gated the same way any exact-vs-approximate tradeoff is - see Level of Detail and Streaming).

Sharding strategy affects what a diff means

Two checkpoints trained with different sharding strategies (for example, one fully replicating a tensor across ranks, another fully sharding it) can still be logically identical once reassembled - the strategy is a training-time communication/memory tradeoff, not a property of the model itself. Tensormorph's weight diff always compares reassembled logical tensors, never raw per-shard bytes, specifically so a sharding-strategy change alone never shows up as a false difference.

Related resources

  • Shards and Storage - the underlying storage model this page's placement vocabulary builds on.
  • GPU Backend Compatibility - reassembling a placement-sharded tensor across a heterogeneous set of backends.
  • Runtime Traces - inspecting cross-rank communication introduced by a sharding strategy during an actual run.