improve from_pretrained for zero3 multi gpus mode
Author: 1ytic
Pull request: yes

# Decrease RAM consumption during Deepspeed Zero 3 model initialisation with multiple GPUs.

This simple PR will save a ton of RAM in case with multi GPUs and Zero 3 Deepspeed scenario.

The idea is simple. We do not need to load checkpoints for all instances, because `deepspeed.zero.GatheredParameters` will copy weights from 0 rank.

## Issues

Related issue #12273

## Who can review?

- @stas00
- @pacman100

Comments:
Thank you for the PR, @1ytic 

The idea is excellent, but we need to think about multiple use-cases here.

what happens if zero3 is used, but not `zero.Init` - won't this code fail if you try to access any of the meta weights before `deepspeed.initialize` gets called at which point deepspeed will partition weights from rank 0 to other ranks. I'm flagging the issue of timing of when they are sharded

Unfortunately, when I initially designed this I didn't think anybody would want to use zero-3 w/o `zero.Init` - so I lumped them together - @pacman100 improved upon it in Accelerate to have 2 separate possibilities - is zero3 and is zero.init enabled, so it gives a more refined control for such optimizations.

So let's wait for Sourab to weigh in.

Meanwhile please test 
1. that this code works with pytorch-1.9 (I don't remember when `meta` was made to work)
2. use `USE_SLOW=1 pytest tests/deepspeed` to do coverage testing, since the PR CI doesn't run deepspeed tests. 

Thank you for feedback, @stas00 

Just to clarify a little bit.

My changes effect only state_dict from checkpoints. The function `deepspeed.zero.Init()` doesn't care about state_dict. The real magic happened [here](https://github.com/huggingface/transformers/blob/9ef5256dfb28f5419649af8ea94e82573a6b2e77/src/transformers/modeling_utils.py#L539), when we exit from `deepspeed.zero.GatheredParameters()` context.

I know `modelling_utils.py` is 4k lines monster and maybe I missed something, but seems like I effect only one scenario when we load checkpoints for already partitioned zero3 model. At least, I tested this scenario with 10GB checkpoint and 4 GPUs. I was able to decrease RAM consumption from 45GB to 17GB on single node.

Exactly, the model instantiation is super complex, that's why I wrote the above.

The deepspeed integration test suite has a very high coverage so if you try running it and it succeeds then most likely it's all good. The size of the checkpoint doesn't matter for the purpose of accepting the PR, what matters is to ensure it doesn't break things.

and btw, you actually don't need to even load the checkpoint if you're resuming from a deepspeed zero checkpoint. In another project I hacked to have the model created w/o loading the model and then just used deepspeed checkpoint loading directly, which should already be doing that efficiently, since each gpu will only read its own shard of weights.

But, alas, making it generic enough so that it'd satisfy everybody is very difficult, that's why the general case is to ensure ease of use out of the box often at the cost of slow startup and more memory consumption.

Ideally the protocol should be like this:

1. create a model on meta (~0 secs)
2. load each shard into the gpu it belongs to (a few secs)

this should be extremely fast even for a huge model like BLOOM-176B

In the case of new training, there should be a way to pre-shard the model before loading it, so resume and new training will be identical model loading-wise. This is eventually will be done when universal checkpoint will be implemented for ZeRO (currently it's only available in Megatron-Deepspeed) https://github.com/microsoft/DeepSpeed/issues/2921

So lots and lots of things to improve there.

And more things to fix on the deepspeed side, e.g. this is very wasteful https://github.com/microsoft/DeepSpeed/issues/1971

so practically please run the integration tests I described in the first reply of mine and if possible with pytorch-1.9 (minimal supported pytorch version).

Just to quickly chime int, the minimal version is actually 1.10 now ;-) The meta device is in 1.9+ so that shouldn't be an issue.

Thank you for this insight, Sylvain.

So then any recent pt version should be ok to test with, @1ytic 

Hello,

The trainer's behaviour isn't changed at all because the DeepSpeed config is still created using `HfTrainerDeepSpeedConfig` which sets the weakref `_hf_deepspeed_config_weak_ref` which is used in `is_deepspeed_zero3_enabled` to check if it is Stage-3. So, from trainer's perspective, this should work fine.

From Accelerate's perspective, when user specifies `zero3_init_flag=False`, the weakref `_hf_deepspeed_config_weak_ref` isn't created and as such the `is_deepspeed_zero3_enabled` will return `False` even if it is using Stage-3 because the user doesn't want to use `deepspeed.zero.Init` context manager. So, in this case too, this PR should work fine as `map_location = "cpu"` due to absence of weakref.

So, the changes of this PR look good if all the slow tests pass.

_The documentation is not available anymore as the PR was closed or merged._

@stas00 you were right, I caught an uninitialized error while testing. After fixing the tests passed:

`RUN_SLOW=1 pytest -rs tests/deepspeed/`

```
================================================= short test summary info ==================================================
SKIPPED [1] tests/deepspeed/test_deepspeed.py:949: test requires bfloat16 hardware support
================================= 108 passed, 1 skipped, 98 warnings in 2273.67s (0:37:53) =================================
```

I also added one more import. If it's too much, I can revert it.

@tjruwase, please kindly have a look - do you see any problems to this approach of loading weights only on rank 0 and relying on partitioning to distribute the weights to the rest of the ranks under zero3? Could this somehow cause problem in the future?

The idea is to skip loading weights on all ranks but rank 0, since they will be discarded anyway.

Thank you!

@1ytic, this is pretty neat. LGTM. Thanks! 

> Thanks for your PR! Can we just leave `torch.distributed` as it was? `dist` is way less obvious as a name.

I thought `dist` quite common name for `torch.distributed`, but up to you. I will rename it back.

It is quite common and I would have no problem if this was in the Trainer file, but this file is not a distributed script and can be read by people less used to this. That's why it's better to spell it out IMO.

Thanks for bearing with me!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24964). All of your documentation changes will be reflected on that endpoint.