NVIDIA / NVIDIA/Megatron-LM

[BUG] Can not load _extra_state with TorchDistLoadShardedStrategy

Open
#1,497 4 comments 0 reactions 0 assignees View on GitHub
bug community-request waiting-on-customer
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.

![Image](https://github.com/user-attachments/assets/d0081cac-f858-4058-be55-f9e1d05bde5d)

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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.