LLAMA 2 Distributed Training Support
Author: BiEchi
Pull request: no

### Feature request

LLAMA 2 support for `device_map=True`

### Motivation

The current LLAMA 2 does not include support for `device_map=True`.

```
Traceback (most recent call last):
  File "/u/haob2/saliency4alce/salience_llama_ecco.py", line 38, in <module>
    output = lm.generate(text, generate=3, beam_size=1, do_sample=True, attribution=['ig'])
  File "/u/haob2/saliency4alce/ecco/src/ecco/lm.py", line 221, in generate
    output = self.model.generate(
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/generation/utils.py", line 1588, in generate
    return self.sample(
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/generation/utils.py", line 2642, in sample
    outputs = self(
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 806, in forward
    outputs = self.model(
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 693, in forward
    layer_outputs = decoder_layer(
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 405, in forward
    hidden_states = self.input_layernorm(hidden_states)
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/u/haob2/miniconda3/envs/salience/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/u/haob2/.local/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 89, in forward
    return self.weight * hidden_states.to(input_dtype)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!
```

### Your contribution

I'm looking for suggestions and possible help from distributed training.

Comments:
Without a code reproducer of the error you encounter, there is little we will be able to do to help.

Just figured it out. So after we set `device_map=True`, we can't move the model to a specific device using `model = model.to(device)`, because the model is already dynamically allocated on all available devices.
Sry for bothering @sgugger, and thanks a lot for your prompt reply! I'll give detailed reproduction steps next time (it also helps me identify the bug).