deepspeedai / deepspeedai/DeepSpeed

[BUG] String sort based on ``rank_repr`` may lead to problems?

Open
#1,381 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Describe the bug
This is more of question than a known bug, and if there is a bug, it would probably only occur for very large models or extreme cases. There may be an issue based on how checkpoint file names are sorted.

The file names include a rank_repr encoding:

https://github.com/microsoft/DeepSpeed/blob/4ad8019cdf5b48ff169bacc2ca3dce5f80671709/deepspeed/runtime/pipe/module.py#L537

which uses two-digit formatting as shown here:

https://github.com/microsoft/DeepSpeed/blob/4ad8019cdf5b48ff169bacc2ca3dce5f80671709/deepspeed/runtime/pipe/topology.py#L99

So one gets a set of file names like:

layer_01-model_00-model_states.pt
layer_01-model_01-model_states.pt
layer_01-model_02-model_states.pt
layer_01-model_03-model_states.pt

where the model_00 part is generated by the rank_repr encoding.

When reading those checkpoint files back, they are returned in an ordered list, where a normal string sort is used to order the files, like here:

https://github.com/microsoft/DeepSpeed/blob/4ad8019cdf5b48ff169bacc2ca3dce5f80671709/deepspeed/runtime/pipe/module.py#L549

I think that could lead to problems if one were to use really large models, e.g., if the tensor parallelism reached 100+, since files would be ordered like so (where 100 comes before 11):

layer_01-model_09-model_states.pt
layer_01-model_10-model_states.pt
layer_01-model_100-model_states.pt
layer_01-model_11-model_states.pt

I believe that would mix weights, perhaps in a silent way that is hard for the user to detect. For example, this line might cause rank 11 to read in the weights from rank 100:

https://github.com/microsoft/DeepSpeed/blob/4ad8019cdf5b48ff169bacc2ca3dce5f80671709/deepspeed/runtime/state_dict_factory.py#L80

In the short term, perhaps an assert could be added to at least alert the user there is a problem if they exceed two digits on any rank_repr encoding?

To Reproduce
Steps to reproduce the behavior:

  1. Go to '...'
  2. Click on '....'
  3. Scroll down to '....'
  4. See error

Expected behavior
A clear and concise description of what you expected to happen.

ds_report output
Please run ds_report to give us details about your setup.

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

  • OS: [e.g. Ubuntu 18.04]
  • GPU count and types [e.g. two machines with x8 A100s each]
  • Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
  • Python version
  • Any other relevant info about your setup

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with rank_repr in deepspeed/runtime/pipe/module.py and its formatting in deepspeed/runtime/pipe/topology.py, then trace checkpoint ordering in module.py and deepspeed/runtime/state_dict_factory.py. Reproduce the load path with rank values above 99 and verify that sorted checkpoint files map each rank to the correct weights; done means the issue is detected or ordering is shown to be safe.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.