About reshape deepspeed checkpoint

[henan991201]

I want to ask why we can not make tp and pp of a checkpoint bigger? For example, make tp=4 when its original tp is 2. I tried to do this thing but got an assertion error. Thanks for any reply.

[misska1]

I need to finetune 7B model on V100(32GB)*4 can only use zero_stage 0 and convert transformers into megatron models for TP&PP 4. Looking forward to your work on this script to increase TP &PP size for Deepspeed checkpoint @stas00

[stas00]

Please see:
https://github.com/bigscience-workshop/bigscience/tree/master/train/tr11-176B-ml#checkpoint-reshaping
I know it's not generic at the moment, please let me know if you run into any difficulties following those instructions while adapting to your situation.

@tjruwase, it's probably a good idea to document the topology reshaping process as part of the functionality. Do you have plans to do that? You can totally re-use the steps I shared above. I know you were planning to move all of it into deepspeed.

[tjruwase]

@stas00, thanks for helping with this issue. Yes, documentation is on my TODO. I think linking to your recipe is a great solution for now.

[thies1006]

I tried to create an universal checkpoint of the 3B model (following the instructions in the link) and I get KeyError: 'param_slice_mappings'. It seems that this key is created by the BP16 optimizer and (if I'm not mistaken) the 3B model was trained in FP16 format. Is it possible to create the universal checkpoint for those models as well? Thank you! @stas00

[tjruwase]

@thies1006, there is ongoing work to support other models. Can you please try the following combination?

Megatron-DeepSpeed: #349
DeepSpeed: deepspeedai/DeepSpeed#2284

[thies1006]

I tried with those PRs and it seems that the problem stays the same.

Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 48, in mapstar
    return list(map(*args))
  File "ds_to_universal.py", line 108, in extract_zero_shards
    param_slice_mappings = optim_sd[PARAM_SLICE_MAPPINGS]
KeyError: 'param_slice_mappings'
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "ds_to_universal.py", line 290, in <module>
    main()
  File "ds_to_universal.py", line 269, in main
    _extract_zero_shard_files(args, ds_checkpoint, slice_shapes, temp_dir)
  File "ds_to_universal.py", line 231, in _extract_zero_shard_files
    _do_parallel_work(do_work, work_chunks, args.num_extract_workers)
  File "ds_to_universal.py", line 220, in _do_parallel_work
    pool.map(do_work, batch)
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 364, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 771, in get
    raise self._value
KeyError: 'param_slice_mappings'

[tjruwase]

Can you confirm that the model was originally trained using zero stage 1?

[thies1006]

I'm referring to this model:
https://huggingface.co/bigscience/bloom-3b-optimizer-states/tree/main/global_step337250
I cannot say for sure if this model was trained with Zero 1, maybe @stas00 can confirm?

[stas00]

only 176B was trained on A100 and thus bf16 (z?), everything else was trained on V100s, thus fp16, thus z1

[tjruwase]

Got it. In that case, the model was trained before we started adding universal checkpoint support to z1. Resolving this could require regenerating the checkpoints with the new z1 code. Is this feasible, @stas00?

[stas00]

I wasn't part of this training, @TevenLeScao do you by chance know who did the bloom-3b training? and if it's possible to update to the deepspeed@master in the conda env, run one more iteration, save the new checkpoint and update https://huggingface.co/bigscience/bloom-3b-optimizer-states/tree/main/global_step337250?

[heraclex12]

Hi, any updates? Does anyone overcome this issue?

[TevenLeScao]

Hey, sorry for missing this thread - @stas00 I'm trying that fix.

[ShinoharaHare]

Hi, any updates?

[wptoux]

I tried to generate a universal checkpoint for bloomz-7b on 8 x a100 40G using the following method.

1. using deepspeed_to_deepspeed.py, convert the data parallel of the checkpoint to 8.
2. in the training code, before starting training, insert the save_checkpoint_and_time(iteration, model, optimizer, lr_scheduler) function so that it can generate a new version of the checkpoint. then load the checkpoint generated in step 1. The reason for saving the checkpoint before training is that the program will OOM after starting training.
   

   Megatron-DeepSpeed/megatron/training.py

   Line 889 in e52bdab
3. use ds_to_universal.py to convert the checkpoint generated in step 2 into a universal checkpoint.