[modeling utils] revamp `from_pretrained(..., low_cpu_mem_usage=True)` + tests
Author: stas00
Pull request: yes

The initial `from_pretrained(..., low_cpu_mem_usage=True)` implementation was a quick hack to enable loading gptj models on low CPU memory setups. It didn't work with all models.

This PR takes it one step further. It revamps the implementation to support many features it wasn't supporting by revamping the implementation and delegating all the work to the normal `from_pretrained` code path except the final step of `state_dict` => model param overwrite.

This PR:
1. revamps `low_cpu_mem_usage=True`
3. adds a functional test that checks `from_pretrained(mname, low_cpu_mem_usage=True)` works with sharded and non-sharded checkpoint
4. adds a quality test that measures CPU memory and checks that indeed `low_cpu_mem_usage=True` uses less memory.
5. adds various testing utils helper functions to support the new tests

The low cpu memory usage code path is still not 100% complete feature-wise, but it's getting there. Though I'm contemplating a different approach to solving the issue of low cpu memory. That is by introducing several new `from_pretrained` args that should allow loading the model and/or `state_dict` directly on GPU for single GPU or DDP. But that's for another PR.

@sgugger, @LysandreJik, @patrickvonplaten 


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

Hmm, so trying to write a test that shows the memory saving proved to be a puzzle. Getting inconsistent results between my desktop and the CI. That was using `process.memory_info().rss` then I also tried `tracemalloc`, but I think that one is problematic if pytorch uses some kernels that don't go through python memory allocation.

I think I may try `/usr/bin/time -f %M` via an external process. As it gives me an RSS max peak for the whole independent process.

The results are very peculiar:

- for a 1GB model there is no saving
- for a 10GB model, low_mem saves 1/4 of memory, 
- for a 40GB model, low_mem saves 1/2 of memory, 

but at least that explains why my memory tracking wasn't showing the saving consistently since I was using a 0.5GB model for the test.

So what I'm doing is:
```

# 420 MB https://huggingface.co/bert-base-uncased

/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bert-base-uncased", low_cpu_mem_usage=True)'
1139516
/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bert-base-uncased", low_cpu_mem_usage=False)'
1140324

# 1.25GB https://huggingface.co/bert-large-uncased/

/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bert-large-uncased", low_cpu_mem_usage=True)'
2906584
/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bert-large-uncased", low_cpu_mem_usage=False)'
2908236

# 10.6 GB https://huggingface.co/bigscience/T0_3B

/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bigscience/T0_3B", low_cpu_mem_usage=True)'
16122900
/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bigscience/T0_3B", low_cpu_mem_usage=False)'
22299560

# 41.5 GB https://huggingface.co/bigscience/T0

/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bigscience/T0", low_cpu_mem_usage=True)'
43788452
/usr/bin/time -f %M python -c 'from transformers import AutoModel; AutoModel.from_pretrained("bigscience/T0", low_cpu_mem_usage=False)'
86765944
```

**update**: The culprit proved be that my original low_cpu_mem code was not able to handle models with a custom prefix in its keys like `bert.` - this PR fixes it. 

The quality test will not work as the original implementation doesn't work with bert or any other model with its custom `bert.` prefix. It doesn't put the model to the meta device and thus doesn't save any memory.

I totally hear you about the complexity and that the PR is difficult to review in several places.

So I propose this plan:

1. a new PR that just refactors `_find_mismatched_keys`
2. merge (1) then re-base this PR and revisit?

does that sound OK?

That sounds right, thanks for understanding!

Step 1 is ready: https://github.com/huggingface/transformers/pull/16706

@sgugger, as planned I rebased on https://github.com/huggingface/transformers/pull/16706, so the diff should be much easier to read now.