[`bnb`] Fine-tuning HF 8-bit models
Author: younesbelkada
Pull request: yes

# What does this PR do?

This PR attempts to add the official support of fine-tuning 8-bit models using `transformers`, `bitsandbytes` and adaptors (such as LoRA), supported by `peft`.

With this PR, it will be possible to fine-tune large models with no cost, for e.g. it will be possible to fine-tune `opt-6.7b` in a single Google Colab instance. This would also enable fine-tuning Whisper and large flan-t5 in 8bit.

In order to perform this fine-tuning, a user needs to load the model with the flag `enable_memory_efficient_backward=True`, freeze the parameters of the model and use `peft` to inject adaptators inside the model.

The PR comes also with `Trainer` integration of this feature, that is supported at least in a single GPU setup.
Here is a script (based on an old notebook from @justheuristic) a user can try to run with this PR: 
```
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-6.7b", 
    load_in_8bit=True,
    device_map='auto', 
    torch_dtype=torch.float16, 
    enable_memory_efficient_backward=True
)


tokenizer = AutoTokenizer.from_pretrained("facebook/opt-6.7b")


for param in model.parameters():
  param.requires_grad = False  # freeze the model - train adapters later
  if param.ndim == 1:
    # cast the small parameters (e.g. layernorm) to fp32 for stability
    param.data = param.data.to(torch.float32)

model.gradient_checkpointing_enable()  # reduce number of stored activations
model.model.decoder.project_in = lambda x: x.requires_grad_(True)

class CastOutputToFloat(nn.Sequential):
  def forward(self, x): return super().forward(x).to(torch.float32)
model.lm_head = CastOutputToFloat(model.lm_head)

class LoRALayer(nn.Module):
    """Wraps a linear layer with LoRA-like adapter"""
    def __init__(self, module: nn.Module, rank: int):
        super().__init__()
        self.module = module
        self.adapter = nn.Sequential(nn.Linear(module.in_features, rank, bias=False),
                                     nn.Linear(rank, module.out_features, bias=False))
        small_std = (2. / (5 * min(module.in_features, module.out_features))) ** 0.5
        nn.init.normal_(self.adapter[0].weight, std=small_std)
        nn.init.zeros_(self.adapter[1].weight)
        self.adapter.to(module.weight.device)

    def forward(self, input, *args, **kwargs):
        return self.module(input, *args, **kwargs) + self.adapter(input)

for name, module in model.named_modules():
  if 'OPTAttention' in repr(type(module)):
    module.q_proj = LoRALayer(module.q_proj, rank=16)
    module.k_proj = LoRALayer(module.k_proj, rank=16)
    module.v_proj = LoRALayer(module.v_proj, rank=16)

assert sum(isinstance(module, LoRALayer) for module in model.modules()) == 96

import transformers
from datasets import load_dataset
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: tokenizer(samples['quote']), batched=True)

trainer = transformers.Trainer(
    model=model, train_dataset=data['train'],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=4, gradient_accumulation_steps=4,
        warmup_steps=250, max_steps=1000, learning_rate=2e-4, fp16=True,
        logging_steps=1, output_dir='outputs'),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train()
```

## TODOs:

- [x] clear notebook
- [x] how to share adapters weights using `peft`
- [x] tests

cc @pacman100 @TimDettmers @sgugger 

 


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

As the training support has just been added in `bitsandbytes==0.37.0` I proposed new changes, I also added new tests
This PR is now ready for review