THUDM

/

cogvlm-chat-hf

like 193

Follow

Knowledge Engineering Group (KEG) & Data Mining at Tsinghua University 1.5k

I have been trying to increase throughput of CogVLM by using batches like other issue here advised and with a A100 I've been able to use a batch size of 16 and output on average 3.5 seconds per image. This still needs to be faster for my use case so I tried other methods and the only one that seemed faster to get to work is to load the model in 4 or 8 bit and increase the batch size, but I can't seem to get the inputs data types correct.I get the following message:

UserWarning: Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.
  warnings.warn(f'Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.')

Any other common datatype raises an error... Here's the code:

import torch
import requests
from PIL import Image
from transformers import AutoModelForCausalLM, LlamaTokenizer
import time
import os

tokenizer = LlamaTokenizer.from_pretrained('lmsys/vicuna-7b-v1.5')

model = AutoModelForCausalLM.from_pretrained(
    'THUDM/cogvlm-chat-hf',
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    load_in_4bit=True
).eval()

images = os.listdir("/tmp/textcaps/train_images/")

input_list = [model.build_conversation_input_ids(
    tokenizer,
    images=[Image.open("/tmp/textcaps/train_images/"+i).convert('RGB'),],
    query='Describe this image in detail',  # Q1
    history=[], 
    )
    for i in images[:16]]

def recur_move_to(item, tgt, criterion_func):
    if criterion_func(item):
        device_copy = item.to(tgt)
        return device_copy
    elif isinstance(item, list):
        return [recur_move_to(v, tgt, criterion_func) for v in item]
    elif isinstance(item, tuple):
        return tuple([recur_move_to(v, tgt, criterion_func) for v in item])
    elif isinstance(item, dict):
        return {k: recur_move_to(v, tgt, criterion_func) for k, v in item.items()}
    else:
        return item

def collate_fn(features, tokenizer) -> dict:
    images = [feature.pop('images') for feature in features]
    tokenizer.padding_side = 'left'
    padded_features = tokenizer.pad(features)
    inputs = {**padded_features, 'images': images}
    return inputs

input_batch = collate_fn(input_list, tokenizer)
input_batch = recur_move_to(input_batch, 'cuda', lambda x: isinstance(x, torch.Tensor))
input_batch = recur_move_to(input_batch, torch.float16, lambda x: isinstance(x, torch.Tensor) and torch.is_floating_point(x))

gen_kwargs = {"max_length": 2048, "do_sample": False}

t1 = time.time()
with torch.no_grad():
    outputs = model.generate(**input_batch, **gen_kwargs)
    outputs = outputs[:, input_batch['input_ids'].shape[1]:]
    print(tokenizer.batch_decode(outputs))

print(time.time() - t1)

If there are any other suggestions for increasing throughput I am also more than willing to accept. I have a node of 4 A100's and 6 GPUs with 50Gb of memory each.