deepspeedai / deepspeedai/DeepSpeed
[BUG] zero_to_fp32 ordering files incorrectly for combining shards
@tjruwase is already working on this.
Since Jun 7, 2022.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Note
I found a bug and a fix. However, rather than directly submitting a fix and PR, I'm reporting the issue. If I have time I'll also submit a PR. However, I'm afraid patching the bug in just zero_to_fp32.py will leave it open in other places too (e.g., deepspeed.utils.zero_to_fp32. get_fp32_state_dict_from_zero_checkpoint).
Describe the bug
When combining more than 10 shards (ZeRO Stage 2) using zero_to_fp32.py (created/copied in the save directory), the combined shards produce a model which outputs garbage.
This same bug also appears present in deepspeed.utils.zero_to_fp32. get_fp32_state_dict_from_zero_checkpoint
Local Fix
The offending line in zero_to_fp32.py (not sure where else this same assumption is made, hence the bug report):
optim_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")))`
This only works for 10 or fewer shards. In my case, there were 48 files: labeled "0", "1", …, "9", "10", ... "47", which causes the sort to return them in non-numerical order.
I define a new sorting method to replace sorted:
def sorted_nicely( list_to_sort ):
""" Sort the given iterable in the way that humans expect."""
convert = lambda text: int(text) if text.isdigit() else text
alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ]
return sorted(list_to_sort, key = alphanum_key)
Then replace the offending line with
optim_files = sorted_nicely(glob.glob(os.path.join(checkpoint_dir, "*_optim_states.pt")))
and also
import re
To Reproduce
- Use Deepspeed Engine with ZeRO Stage 2
- Train enough so the output isn't garbage
- Save sharded model using: model.save_checkpoint(save_dir, tag)
- TO SEE THE BUG AFFECTING OTHER PIECES OF CODE USE deepspeed.utils.zero_to_fp32. get_fp32_state_dict_from_zero_checkpoint
- Test model --> baseline metrics
- Quit the python session
- navigate to save_dir
- python zero_to_fp32.py . combined.pt
- load up combined.pt
- Test model --> modified metrics
Expected behavior
I expected modified metrics to be significantly different from random chance and should be very close to baseline metrics
ds_report output
DeepSpeed C++/CUDA extension op report
NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.
JIT compiled ops requires ninja
ninja .................. [OKAY]
op name ................ installed .. compatible
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-dev package with apt
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
DeepSpeed general environment info:
torch install path ............... ['/home/mvalenzuela/miniconda3/lib/python3.9/site-packages/torch']
torch version .................... 1.9.1+cu102
torch cuda version ............... 10.2
nvcc version ..................... 10.0
deepspeed install path ........... ['/home/mvalenzuela/miniconda3/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.5.10, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.9, cuda 10.2
Screenshots
NA
System info (please complete the following information):
- OS: Ubuntu 18.04.2 LTS
- GPU x48 V100
- Interconnects (if applicable) six machines each with eight GPUs connected with 100 Gbps IB
- Python version: Python 3.9.5
- Any other relevant info about your setup
Launcher context
deepspeed
Docker context
Spans multiple docker images
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.