Add a new script to reshard model parallel parts

[tangbinh]

Summary of Changes

The existing script for resharding model parallel parts (i.e. metaseq/scripts/reshard_model_parallel.py) loads all checkpoint parts at once and might result in OOM issues under RAM constraints, especially for very large models. Here, we rewrite the script and optimize for memory usage by first allocating an unsharded model state dict and iteratively merging model parallel parts into it.

Previously, peak memory usage was close to 2X model size as we needed to hold input and output state dicts, but theoretically it's closer to 1x model size now thanks to the iterative process.

The new script produces the same output as metaseq/scripts/reshard_model_parallel.py. We delete it to avoid duplication and note that the old script still remains accessible in the internal repo (see this script).

Test Plan

• Run the script with an OPT 2.7B checkpoint to reshard 4 MP parts into 8 MP parts and make sure the resulting checkpoint performs reasonably:

  seq 0 3 | parallel --line-buffer 'python metaseq/scripts/reshard_fsdp.py --input "/data/checkpoints/opt-2.7b/raw/checkpoint_last-model_part-{}-shard*.pt" --output "/data/checkpoints/opt-2.7b/reshard-no-os/reshard-model_part-{}.pt" --skip-optimizer-state True --unflatten-weights True --output-dtype fp16'
  python -m metaseq.scripts.reshard_mp --input "/data/checkpoints/opt-2.7b/reshard_no_os/reshard-model_part-*.pt" --output "/data/checkpoints/opt-2.7b/reshard_no_os_mp8/reshard-model_part-{i}.pt" --num-output-parts 8
  

  python metaseq/scripts/interactive.py --merges-filename /data/checkpoints/gpt2-merges.txt --vocab-filename /data/checkpoints/gpt2-vocab.json --path /data/checkpoints/opt-2.7b/reshard_no_os_mp8/reshard.pt --model-parallel-size 8 --distributed-world-size 8  --beam 3 --max-source-positions 4 --max-target-positions 128
  
  > Prompt: What is the meaning of life?
  Output: To be happy.
  

• We compare performance with metaseq/scripts/reshard_model_parallel.py while resharding an OPT-175B checkpoint from 8 MP parts into 16 MP parts. The old script takes 849.40 seconds and results in a peak RSS delta of 668,301 MB while the new script takes 891.65 seconds and has RSS delta of 458,185 MB (a 46% reduction in RAM usage).

  python metaseq/scripts/reshard_model_parallel.py --pth_prefix /data/checkpoints/opt-175b/reshard_no_os_unflat/reshard.pt --new-model-parts 16 --save-prefix  /data/checkpoints/opt-175b/reshard_no_os_unflat_mp16_ref/reshard.pt
  python -m metaseq.scripts.reshard_mp --input "/data/checkpoints/opt-175b/reshard_no_os_unflat/reshard-model_part-*.pt" --output "/data/checkpoints/opt-175b/reshard_no_os_unflat_mp16/reshard-model_part-{i}.pt" --num-output-parts 16
  

[ruanslv]

Is it possible to run $ git mv so that the new script becomes a second version of the original one? Makes it easier to review.

Also lint failed, maybe run black?

[tangbinh]

Is it possible to run $ git mv so that the new script becomes a second version of the original one? Makes it easier to review.

Also lint failed, maybe run black?

Thank you for the review. I did use git mv but it didn't trigger the comparison of file contents (perhaps they are too different). Please see updated PR for the linter fix.

[ruanslv]

Great job in getting this to work! The logic here is tricky.

I don't think we had fully validated the previous script either, so I think a good exercise for us to do once all the scripts are merged is to make sure the resharded model behind the API matches the validation ppl seen during training