Enable training Llama with model or pipeline parallelism
Author: kooshi
Pull request: yes

# What does this PR do?

This PR enables model and pipeline parallelism for Llama models.

The change moves the target tensor to the output device, if needed, for loss calculation.


## 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._

Great jobs!

Could you share a snippet of code that is failing prior to this PR?

Certainly, here's a relatively minimal example:
```python
import transformers, datasets
from peft import (
    LoraConfig,
    get_peft_model,
)

CHECKPOINT = "decapoda-research/llama-7b-hf"
model = transformers.LlamaForCausalLM.from_pretrained(
    CHECKPOINT,
    device_map="auto",
    max_memory={0:"15GB", 1:"15GB"}
)
tokenizer = transformers.LlamaTokenizer.from_pretrained(CHECKPOINT, add_eos_token=True)
tokenizer.pad_token_id = 0

config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj","v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)

data = datasets.load_dataset("laion/OIG", data_files="unified_chip2.jsonl", split="train")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=256,
        padding="max_length",
    )

data = data.map(tokenize)

#Tell Trainer not to attempt DataParallel
model.is_parallelizable = True
model.model_parallel = True
trainer = transformers.Trainer(
    model=model,
    train_dataset=data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        learning_rate=3e-4,
        logging_steps=10,
        evaluation_strategy="no",
        save_strategy="no",
        output_dir="/tmp/"
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
```

Before the change, this fails with
`RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument target in method wrapper_CUDA_nll_loss_forward)`

I found the solution here: [Pytorch Pipeline Tutorial](https://pytorch.org/tutorials/intermediate/pipeline_tutorial.html#run-the-model)
> \# Need to move targets to the device where the output of the pipeline resides.

And after the change, the code example above runs as expected.