How can we use CPU offloading when using AutoModelForCausalLM and THUDM/cogvlm2-llama3-chat-19B
Author: FurkanGozukara
Pull request: no

It is working great with below however still not sufficient. Uses around 16 GB VRAM

I want to further lower the requirement if possible

How can I achieve that?

model path is : `THUDM/cogvlm2-llama3-chat-19B`

```
                    model = AutoModelForCausalLM.from_pretrained(
                        MODEL_PATH,
                        torch_dtype=TORCH_TYPE,
                        trust_remote_code=True,
                        quantization_config=BitsAndBytesConfig(load_in_4bit=True),
                        low_cpu_mem_usage=True
                    ).eval()
```

### Who can help?

text models: @ArthurZucker
vision models: @amyeroberts, @qubvel
pipelines: @Rocketknight1
quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber 


Comments:
Hey @FurkanGozukara,
To offload the model to the CPU, we can create a `device_map` dictionary that specifies which model weights should be allocated to CUDA or the CPU. Additionally, we need to enable a few arguments in the `BitsAndBytesConfig` which can help reduce VRAM usage further.

For the given model, the `device_map` can be defined as follows:
```py
device_map = {
    "transformer.word_embeddings": 0,
    "transformer.word_embeddings_layernorm": 0,
    "lm_head": 0,
    "model.embed_tokens": 0,
    "model.norm": 0,
    "model.vision": "cpu",
    "model.layers": 0,
    "transformer.h": 0,
    "transformer.ln_f": 0,
}
```
For `THUDM/cogvlm2-llama3-chat-19B` model, with this mapping, the vision model weights are offloaded to the CPU, while the rest of the model weights are allocated on the GPU. This setup results in 9.5 GB of VRAM usage on the GPU and 14.2 GB of RAM usage on the CPU, reducing the GPU's memory load from 16GB of VRAM usage. This can be seen here:

![Image](https://github.com/user-attachments/assets/28102240-0eaa-4893-b88e-d46915c738ce)

To replicate this setup, Here is the full code:
```py
import torch
from transformers import BitsAndBytesConfig
from transformers import AutoTokenizer, AutoModelForCausalLM
from PIL import Image

repo = "THUDM/cogvlm2-llama3-chat-19B"
device = 'cuda'
torch_type = torch.float16

device_map = {
    "transformer.word_embeddings": 0,
    "transformer.word_embeddings_layernorm": 0,
    "lm_head": 0,
    "model.embed_tokens":0,
    "model.norm":0,
    "model.vision":"cpu",
    "model.layers":0,
    "transformer.h": 0,
    "transformer.ln_f": 0,
}

quantization_config = BitsAndBytesConfig(load_in_4bit=True, 
                                         bnb_4bit_compute_dtype=torch.float16,
                                         bnb_4bit_use_double_quant=True,
                                         llm_int8_enable_fp32_cpu_offload=True,
                                        )

model = AutoModelForCausalLM.from_pretrained(repo,
                                             torch_dtype=torch.float16,
                                             trust_remote_code=True,
                                             device_map=device_map,
                                             quantization_config=quantization_config,
                                            ).eval()

tokenizer = AutoTokenizer.from_pretrained(repo)

text_only_template = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {} ASSISTANT:"

while True:
    image_path = input("image path >>>>> ")
    if image_path == '':
        print('You did not enter image path, the following will be a plain text conversation.')
        image = None
        text_only_first_query = True
    else:
        image = Image.open(image_path).convert('RGB')

    history = []

    while True:
        query = input("Human:")
        if query == "clear":
            break

        if image is None:
            if text_only_first_query:
                query = text_only_template.format(query)
                text_only_first_query = False
            else:
                old_prompt = ''
                for _, (old_query, response) in enumerate(history):
                    old_prompt += old_query + " " + response + "\n"
                query = old_prompt + "USER: {} ASSISTANT:".format(query)
        if image is None:
            input_by_model = model.build_conversation_input_ids(
                tokenizer,
                query=query,
                history=history,
                template_version='chat'
            )
        else:
            input_by_model = model.build_conversation_input_ids(
                tokenizer,
                query=query,
                history=history,
                images=[image],
                template_version='chat'
            )
        inputs = {
            'input_ids': input_by_model['input_ids'].unsqueeze(0).to(device),
            'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(device),
            'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(device),
            'images': [[input_by_model['images'][0].to(device).to(torch_type)]] if image is not None else None,
        }
        gen_kwargs = {
            "max_new_tokens": 2048,
            "pad_token_id": 128002,  
        }
        with torch.no_grad():
            outputs = model.generate(**inputs, **gen_kwargs)
            outputs = outputs[:, inputs['input_ids'].shape[1]:]
            response = tokenizer.decode(outputs[0])
            response = response.split("
")[0]
            print("\nCogVLM2:", response)
        history.append((query, response))
```
This should offload model weights to cpu and help lower your VRAM usage. 
Hope this helps 🤗!

Wow this is a legendary answer @nnilayy . Hopefully I will try and implement asap and let you know

Closing for now because this is a support question rather than a bug or feature request