Multi-gpu inference llama-3.2 vision with QLoRA

shon-ot
Tue, 15 Apr 2025 18:23:56 +0000

Hello

After fine-tuning meta-llama/Llama-3.2-11B-Vision-Instruct I run into a weird error while running inference with multi-gpu.

This is how I loads the model:

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype="bfloat16",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_storage='bfloat16'
)

model = MllamaForConditionalGeneration.from_pretrained(
    model_path_or_name,
     quantization_config=bnb_config,
     device_map='auto',
     low_cpu_mem_usage=True
 )
model = PeftModel.from_pretrained(model, adapter_path)

The model is loaded into all available gpu - in my case 4.

Code for generation:

for i in trange(0, len(dataset), batch_size, desc="Inference"):
    batch = dataset[i: i + batch_size]
    batch = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
    device_type = next(iter(model.parameters())).device.type

    with torch.no_grad():
        tokenized_batch = collator(batch)
        tokenized_batch = tokenized_batch.to(device_type)

        # text generation part
        raw_texts = generate(model, processor, tokenized_batch, terminators)

and I running on device mismatch error :
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!

for other model like Qwen for example all works perfectly! this is only with MllamaForConditionalGeneration I was able to work around this by remove device_map from from_pretrained but that not a good one beacuse this is running only with single gpu.

BTW I running the scripts with python inference.py no accelerate or something else
any idea/suggestions?