[INT8] BLOOM series model loading back issue
Author: lanking520
Pull request: no

### System Info

8x A100 GPUs with CUDA 11.3 driver

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Use the following script to save a INT8 quantized and try to load it back.
```
import os
import torch
import logging
import math
from transformers import AutoConfig, pipeline, AutoModelForCausalLM, AutoTokenizer

def get_max_memory_per_gpu_dict(dtype, model_name):
    """try to generate the memory map based on what we know about the model and the available hardware"""

    # figure out the memory map - the minimum per gpu required to load the model
    n_gpus = torch.cuda.device_count()

    try:
        # model_params calculation, as we don't have a model yet to do:
        # model_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())

        config = AutoConfig.from_pretrained(model_name)
        h = config.hidden_size
        l = config.n_layer
        v = config.vocab_size
        # from https://github.com/bigscience-workshop/bigscience/tree/6917a3b5fefcf439d3485ca184b4d9f6ab605150/math#model-sizing
        model_params = l * (12 * h ** 2 + 13 * h) + v * h + 4 * h
    except:
        logging.info(f"The model {model_name} has a broken config file. Please notify the owner")
        raise

    if dtype == torch.int8:
        bytes = 1
    else:
        bytes = torch.finfo(dtype).bits / 8
    param_memory_total_in_bytes = model_params * bytes
    # add 5% since weight sizes aren't the same and some GPU may need more memory
    param_memory_per_gpu_in_bytes = int(param_memory_total_in_bytes / n_gpus * 1.10)
    logging.info(f"Estimating {param_memory_per_gpu_in_bytes / 2 ** 30:0.2f}GB per gpu for weights")

    # check the real available memory
    # load cuda kernels first and only measure the real free memory after loading (shorter by ~2GB)
    torch.ones(1).cuda()
    max_memory_per_gpu_in_bytes = torch.cuda.mem_get_info(0)[0]
    if max_memory_per_gpu_in_bytes < param_memory_per_gpu_in_bytes:
        raise ValueError(
            f"Unable to generate the memory map automatically as the needed estimated memory per gpu ({param_memory_per_gpu_in_bytes / 2 ** 30:0.2f}GB) is bigger than the available per gpu memory ({max_memory_per_gpu_in_bytes / 2 ** 30:0.2f}GB)"
        )

    max_memory_per_gpu = {i: param_memory_per_gpu_in_bytes for i in range(torch.cuda.device_count())}
    print("Max memory per gpu:", max_memory_per_gpu)
    return max_memory_per_gpu


def load_model():
    world_size = torch.cuda.device_count()
    model_name = "bigscience/bloom"
    logging.info(f"Using {world_size} gpus")
    logging.info(f"Loading model {model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    dtype = torch.int8
    kwargs = dict(
        device_map="auto",
        max_memory=get_max_memory_per_gpu_dict(dtype, model_name),
    )
    logging.info("Using `load_in_8bit=True` to use quanitized model")
    kwargs["load_in_8bit"] = True
    model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
    return model, tokenizer

model, tokenizer = load_model()

model.save_pretrained("int8_model/", max_shard_size="8GB")
```

When loading from the directory, having the error on:

```
RuntimeError: Only Tensors of floating point dtype can require gradients
```

During the initialization of the model.

```
import torch
import torch.distributed as dist
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

model_name = 'int8_model/'
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.int8)
```

### Expected behavior

The loading should pass.  Looking for a workaround on it... 

Comments:
Maybe @sgugger has some insights

It's hard to know which part fails without the whole traceback. I suspect it's when we set the default dtype to `torch_dtype`, which only works for floating dtypes. If that's the case, there is a probably a workaround possible by only setting the default dtype when the `torch_dtype` passed is a floating type.

Also cc @younesbelkada since it's related to int8 format.

Hey @lanking520 ! 
Thanks for your issue 💪 
Let's try to debug this step by step. I suggest first to make your script run on the model `bigscience/bigscience-small-testing` - when running your script I got incorrect max_memory maps, so I had to overwrite the `max_memory` dict with the following `max_memory={0:"10GB", 1:"10GB"},` (I am running my tests on 2x NVIDIA T4).
After that, the loading script gives me the following error:
```

│ /home/younes_huggingface_co/debug_issues/code/transformers/src/transformers/modeling_utils.py:10 │
│ 49 in _set_default_torch_dtype                                                                   │
│                                                                                                  │
│   1046 │   │   `torch.int64` is passed. So if a non-float `dtype` is passed this functions will  │
│   1047 │   │   """                                                                               │
│   1048 │   │   if not dtype.is_floating_point:                                                   │
│ ❱ 1049 │   │   │   raise ValueError(                                                             │
│   1050 │   │   │   │   f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is  │
│   1051 │   │   │   )                                                                             │
│   1052                                                                                           │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: Can't instantiate BloomForCausalLM model under dtype=torch.int8 since it is not a floating point dtype
```
The "hack" as @sgugger suggested is to "force-load" the weights in a floating point format - if you run:
```
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto",  device_map="auto")
```
You can check that the weights are indeed `int8` weights but casted in half-precision
```
>>> model.transformer.h[0].mlp.dense_h_to_4h.weight
Parameter containing:
tensor([[  96.,  -90.,   -9.,  ...,   36.,  -25.,  -25.],
        [   0.,  -11.,  -51.,  ...,    9.,   20.,   38.],
        [  -8.,    8.,    2.,  ...,   36.,  -88.,  -12.],
        ...,
        [  -6.,   33.,  -41.,  ...,  -32.,  -18.,  -45.],
        [ -11.,  -43.,  -34.,  ...,  -14.,   -1.,  -50.],
        [ -42.,   44.,  108.,  ...,   80., -119.,   54.]], dtype=torch.float16,
       requires_grad=True)
```
However, please note that even if you manage to run an inference with these weights, you will not be able to retrieve the same accuracy / performance than the 8-bit model that is created by `load_in_8bit=True` from the `fp16` model. This is because of how the `Linear8bitLt` layer is constructed. 
The crucial components of this module are the quantization statistics that are stored in[ `self.state.SCB` ](https://github.com/TimDettmers/bitsandbytes/blob/b844e104b79ddc06161ff975aa93ffa9a7ec4801/bitsandbytes/nn/modules.py#L246). The problem when saving the `state_dict` from a `Linear8bitLt` is that it does not save these statistics that are needed at inference. So when you will load 8bit weights, the module will compute new quantization statistics based on the `int8` weights - which will lead to wrong results and computations. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi all - I'm also trying to save and re-load a BLOOM model in 8-bit format, see [https://github.com/TimDettmers/bitsandbytes/issues/80](https://github.com/TimDettmers/bitsandbytes/issues/80).

I'm quite new to the topic and not sure I'm able to follow everything @younesbelkada mentioned, but my understanding is that this is not possible yet, is that correct?