Repository: pytorch/pytorch

Issue:
Sharded checkpointing fails on load for certain tensor sizes
### 🐛 Describe the bug

Sharded checkpointing (particularly with FSDP for the optimizer state) uses ChunkShardingSpec to save/load tensors. ChunkShardingSpec's behavior is similar to torch.chunk and will result in some chunks of size 0.

This can be reproduced by trying to save a tensor of size 6 with 4 gpus. This tensor is sharded across the first 3 gpus. The resulting size of the chunks will look like [2, 2, 2, 0]. On save, it seems like ChunkShardingSpec is aware of which gpus contain shards, so it saves the tensor with shard metadata showing the size of chunks to be [2, 2, 2].

The problem occurs when attempting to load the sharded checkpoint. ChunkShardingSpec attempts to rebuild the metadata, this time being unaware of how many gpus originally contained shards. It knows that there is a tensor of size 6 and 4 gpus though, so it generates shard metadata with chunk sizes [2, 2, 2], [skipping the last gpu since it has size 0](https://github.com/pytorch/pytorch/blob/ddbaad6d746d34b56f2d2cc0e02d7e0d0301e626/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py#L94). Then when attempting to shard the tensor, the 4th gpu has no shard metadata, so a [local_tensor is never created](https://github.com/pytorch/pytorch/blob/ddbaad6d746d34b56f2d2cc0e02d7e0d0301e626/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py#L165), resulting in an [assertion error on the 4th rank](https://github.com/pytorch/pytorch/blob/ddbaad6d746d34b56f2d2cc0e02d7e0d0301e626/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py#L172) and a [type error on all other ranks](https://github.com/pytorch/pytorch/blob/ddbaad6d746d34b56f2d2cc0e02d7e0d0301e626/torch/distributed/distributed_c10d.py#L808) because None is contained in the scatter_list.

There are a couple possible solutions to this.
1. Add shardmetadata for all gpus that allow for a tensor to be size 0
2. Change ChunkShardingSpec to distribute a tensor evenly across gpus (e.g. [2, 2, 1, 1] instead of [2, 2, 2, 0])

I've implemented and tested both solutions and both are backwards compatible with previously saved sharded checkpoints on versions 2.0.1, 2.1.0-rc3, and 8/27 nightly (e.g. save on 2.0.1 with the old ChunkShardingSpec and load on 2.0.1 with the new ChunkShardingSpec). Both solutions are also cross-version compatible for 2.0.1->2.1.0-rc3, 2.0.1->8/27 nightly, and 8/27 nightly->2.0.1 (e.g. save on 2.0.1 with the old ChunkShardingSpec and load on 2.1.0 with the new ChunkShardingSpec). The solutions might be version/cross-version compatible for other combinations, but I haven't tested those.

### Versions

This happens with pytorch 2.0.1, 2.1.0-rc3, and 8/27 nightly.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu


Discussion:
cc: @fegin @kumpera @wz337 
Also this is a breaking bug with sharded checkpoints where sharded checkpoints can be saved but never loaded if a gpu doesn't contain a shard. I'm very surprised no one has encountered this issue before.

Resolving patch:
diff --git a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
--- a/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
+++ b/torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py
@@ -91,20 +91,17 @@ def build_metadata(self,
         for idx, placement in enumerate(self.placements):
             # generate ShardMetadata for each placement device
             chunked_dim_size = get_chunked_dim_size(sharding_dim_size, split_size, idx)
-            if chunked_dim_size > 0:
-                shard_size = list(tensor_sizes)
-                current_offsets = [0] * tensor_num_dim
-                current_offsets[self.dim] = split_size * idx  # type: ignore[index]
-                shard_size[self.dim] = chunked_dim_size  # type: ignore[index]
-
-                shard_metadata = ShardMetadata(
-                    shard_offsets=current_offsets,
-                    shard_sizes=shard_size,
-                    placement=placement,
-                )
-                shards_metadata.append(shard_metadata)
-
-                # current_offsets[self.dim] += chunked_dim_size  # type: ignore[index]
+            shard_size = list(tensor_sizes)
+            current_offsets = [0] * tensor_num_dim
+            current_offsets[self.dim] = split_size * idx  # type: ignore[index]
+            shard_size[self.dim] = chunked_dim_size  # type: ignore[index]
+
+            shard_metadata = ShardMetadata(
+                shard_offsets=current_offsets,
+                shard_sizes=shard_size,
+                placement=placement,
+            )
+            shards_metadata.append(shard_metadata)
 
         return sharded_tensor_meta.ShardedTensorMetadata(
             shards_metadata,


Test patch: