load_in_8bit now respects 'balanced' device maps in multi-gpu environments
Author: kooshi
Pull request: yes

# What does this PR do?

Fixes `max_memory` generation for 'auto' 'balanced' and 'balanced_low_0' `device_map`s for models being loaded in 8bit


Fixes # (N/A)
no issues found, but one guy made a [comment](https://github.com/TimDettmers/bitsandbytes/issues/177#issuecomment-1481609654) about it in the bnb issues, and it caused confusion and workarounds elsewhere.
The problem was easily worked around by manually passing a device map or max memory config.

Before this change, the following code would attempt to load the whole model on the first GPU in a two gpu setup, potentially causing OOM errors. After the change, it loads it evenly across GPUs, as intended.
```python
model = AutoModelForCausalLM.from_pretrained(
    checkpoint,
    load_in_8bit=True,
    device_map="auto",
)
```


Additionally, I removed 'pipeline' from my earlier comment in a previous PR, as true Pipeline Parallelism would require some more non-trivial changes to the model.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: HF Trainer: @stas00, Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger, @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: @sgugger
- TensorFlow: @Rocketknight1

 -->


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

I'm not sure if this is caused by your work but I met this problem:

I'm training LLaMA-13B with PEFT, using lora + modules_to_save=['model.embed_tokens', 'lm_head']
And it can run training normally
But the finally model file doesn't contain the lm_head part, only lora + embed_tokens.
And When I use DDP or single card. It can save all the thing normally.

(I'm using your fork + your modified version alpaca lora code)

@KohakuBlueleaf that's a bit disingenuous, as you've changed quite a few other things 😉

I think I was able to reproduce what you were talking about on [your repo](https://github.com/KohakuBlueleaf/guanaco-lora) though. Do you mean that when you run export_hf_checkpoint.py, `head_changed` shows `False`?

That's what I am seeing. What's strange, is that the weights *are* in the lora, they're just named "base_model.model.lm_head.0.weight" instead of "base_model.model.lm_head.weight"

If you add `adapters_weights["base_model.model.lm_head.weight"] = adapters_weights["base_model.model.lm_head.0.weight"]` to peft_model.py right after the lora is loaded, but before it is merged with the base_model, then you can get `head_changed: True` 

This might be from my change... as the lm_head is a special case when loading in 8bit, but I'm not sure. I see the same result for *both* single or multi-gpu.  I've run out of time today to investigate though, so I'll have to dig in more later, possibly tonight.

> @KohakuBlueleaf that's a bit disingenuous, as you've changed quite a few other things 😉
> 
> I think I was able to reproduce what you were talking about on [your repo](https://github.com/KohakuBlueleaf/guanaco-lora) though. Do you mean that when you run export_hf_checkpoint.py, `head_changed` shows `False`?
> 
> That's what I am seeing. What's strange, is that the weights _are_ in the lora, they're just named "base_model.model.lm_head.0.weight" instead of "base_model.model.lm_head.weight"
> 
> If you add `adapters_weights["base_model.model.lm_head.weight"] = adapters_weights["base_model.model.lm_head.0.weight"]` to peft_model.py right after the lora is loaded, but before it is merged with the base_model, then you can get `head_changed: True`
> 
> This might be from my change... as the lm_head is a special case when loading in 8bit, but I'm not sure. I see the same result for _both_ single or multi-gpu. I've run out of time today to investigate though, so I'll have to dig in more later, possibly tonight.

Yeah I also figured it out
But thx for your reply!