[BUG] Can not load _extra_state with TorchDistLoadShardedStrategy
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
I used the `TorchDistLoadShardedStrategy` loading strategy to load model weights in the `distcp` format. There are two different formats involved: `ShardedTensor` and `ShardedObject`. The former stores sliced weights, while the latter contains `_extra_state`, which holds some FP8-related information about the weights. The issue is that when the `TorchDistLoadShardedStrategy` attempts to read a `ShardedObject`, it throws an error.

I notice the code here:
```
checkpoint.load(
pyt_state_dict,
FileSystemReader(checkpoint_dir),
planner=MCoreLoadPlanner(
shapes_validation_sharded_tensors=flexible_shape_sharded_tensors
),
)
```
Megatron uses a binary method to read model weights, loading the results from `checkpoint_dir` and overwriting them into `sharded_state_dict`. However, it seems that it can only read content of the `ShardedTensor` type and is unable to read the contents stored in `ShardedObject`.
This results in the `extra_state_dict` being loaded as some `io.BytesIO` objects instead of the expected structured data.
Contributor guide
Assessment
This issue has not been assessed yet.