FIX / Quantizaiton: Revert back previous replacement logic
Author: younesbelkada
Pull request: yes

# What does this PR do?

Fixes: https://github.com/huggingface/peft/issues/1720

https://github.com/huggingface/transformers/pull/29958 introduced a new logic for linear layer replacement for bnb which led to a corner case.

Before that PR, for a BertForSequenceClassification model: 

```python
from transformers import AutoModelForSequenceClassification, BitsAndBytesConfig

config = BitsAndBytesConfig(
    load_in_4bit=True, # quantize the model to 4-bits when you load it
    bnb_4bit_quant_type="nf4", # use a special 4-bit data type for weights initialized from a normal distribution
    bnb_4bit_use_double_quant=True, # nested quantization scheme to quantize the already quantized weights
    bnb_4bit_compute_dtype=torch.bfloat16, # use bfloat16 for faster computation,
)

model = AutoModelForSequenceClassification.from_pretrained(
    model_checkpoint,
    quantization_config=config
)

print(model)
```

Gives:

```bash
DistilBertForSequenceClassification(
  (distilbert): DistilBertModel(
    (embeddings): Embeddings(
      (word_embeddings): Embedding(30522, 768, padding_idx=0)
      (position_embeddings): Embedding(512, 768)
      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
      (dropout): Dropout(p=0.1, inplace=False)
    )
    (transformer): Transformer(
      (layer): ModuleList(
        (0-5): 6 x TransformerBlock(
          (attention): MultiHeadSelfAttention(
            (dropout): Dropout(p=0.1, inplace=False)
            (q_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (k_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (v_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (out_lin): Linear4bit(in_features=768, out_features=768, bias=True)
          )
          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          (ffn): FFN(
            (dropout): Dropout(p=0.1, inplace=False)
            (lin1): Linear4bit(in_features=768, out_features=3072, bias=True)
            (lin2): Linear4bit(in_features=3072, out_features=768, bias=True)
            (activation): GELUActivation()
          )
          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
        )
      )
    )
  )
  (pre_classifier): Linear(in_features=768, out_features=768, bias=True)
  (classifier): Linear(in_features=768, out_features=2, bias=True)
  (dropout): Dropout(p=0.2, inplace=False)
)
```

But after #29958 : 

```bash
DistilBertForSequenceClassification(
  (distilbert): DistilBertModel(
    (embeddings): Embeddings(
      (word_embeddings): Embedding(30522, 768, padding_idx=0)
      (position_embeddings): Embedding(512, 768)
      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
      (dropout): Dropout(p=0.1, inplace=False)
    )
    (transformer): Transformer(
      (layer): ModuleList(
        (0-5): 6 x TransformerBlock(
          (attention): MultiHeadSelfAttention(
            (dropout): Dropout(p=0.1, inplace=False)
            (q_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (k_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (v_lin): Linear4bit(in_features=768, out_features=768, bias=True)
            (out_lin): Linear4bit(in_features=768, out_features=768, bias=True)
          )
          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          (ffn): FFN(
            (dropout): Dropout(p=0.1, inplace=False)
            (lin1): Linear4bit(in_features=768, out_features=3072, bias=True)
            (lin2): Linear4bit(in_features=3072, out_features=768, bias=True)
            (activation): GELUActivation()
          )
          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
        )
      )
    )
  )
  (pre_classifier): Linear4bit(in_features=768, out_features=768, bias=True)
  (classifier): Linear(in_features=768, out_features=2, bias=True)
  (dropout): Dropout(p=0.2, inplace=False)
)
```

The pre_classifier gets also quantized. In order to preserve BC (even though it might not be 100% what is intended) I propose to add that extra condition and make sure previous behaviour is maintained

cc @SunMarc @BenjaminBossan 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30779). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Before the PR, `pre_classifier` was not quantized because we have `modules_to_not_convert = ['classifier']` and `classifier` is in `pre_classifier`. I'm not sure about the right fix. One solution would be to fix `get_keys_to_not_convert` in such way that we get `modules_to_not_convert = ['classifier', 'pre_classifier']` or we can just pass `llm_int8_skip_modules=["classifier", "pre_classifier"]`. However, going back to the previous logic is not optimal. 

I agree, ideally we don't want to do that .. I also mentioned here: https://github.com/huggingface/peft/issues/1720#issuecomment-2107878793 that the user should use `modules_to_not_convert`, IMO `pre_classifier` shouldn't be un-quantized at first place .

Closing as I agree there is no point reverting a bug!