# Multi-gpu inference llama-3.2 vision with QLoRA

**URL:** <https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685>\
**Category:** 🤗Accelerate\
**Created:** [April 15, 2025, 6:23pm UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685 "2025-04-15T18:23:55Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![shon-ot](https://avatars.discourse-cdn.com/v4/letter/s/b2d939/32.png) [@shon-ot](https://discuss.huggingface.co/u/shon-ot)\
**Post date:** [April 15, 2025, 6:23pm UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685/1 "2025-04-15T18:23:56Z")

</div>

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:

```python
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:

```python
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?

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [April 16, 2025, 2:21am UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685/2 "2025-04-16T02:21:31Z")

</div>

Well, since Accelerate is called with device\_map=, it’s probably a problem with Accelerate and bitsandbytes…

And Llama 3.x sometimes has bugs when it’s not bfloat16. Also, device\_map=“auto” doesn’t seem to get along well with bitsandbytes.

Well, in your case, it seems like you have enough VRAM even without device\_map=, so I think that’s fine…

```py
model = MllamaForConditionalGeneration.from_pretrained(
     model_path_or_name,
     torch_dtype=torch.bfloat16, # added
     quantization_config=bnb_config,
     #device_map='auto',
     device_map='sequential',
     low_cpu_mem_usage=True
 )

```

> **[Loading big models into memory](https://huggingface.co/docs/accelerate/concept_guides/big_model_inference)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

---

<div class="post-metadata">

**Author:** ![shon-ot](https://avatars.discourse-cdn.com/v4/letter/s/b2d939/32.png) [@shon-ot](https://discuss.huggingface.co/u/shon-ot)\
**Post date:** [April 16, 2025, 7:41am UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685/4 "2025-04-16T07:41:50Z")

</div>

Hi @John6666 thanks for your inputs.

I don’t think the issue lies in loading the model itself, as I can clearly see the model is being properly sharded across devices when using `device_map='auto'`. I should also mention that everything works perfectly when running inference with just the base model — it gets correctly sharded and the batch is distributed across the GPUs. The problem only arises after adding the PEFT adapter.

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [April 16, 2025, 7:52am UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685/5 "2025-04-16T07:52:29Z")

</div>

I see, that certainly doesn’t seem to be the problem with **from\_pretrained**. It might be the PEFT or the **get\_peft\_model** function that’s suspicious.

For now, I’ve found something that looks like bad know-how, but it’s not a solution…😅

> <https://github.com/huggingface/peft/issues/1831>
>
> \### System Info
> 
> \`\`\`
> Transformers 4.41.2
> peft 0.11.1
> \`\`\`
> 
> Single \`T4\` GP…U. 
> 
> I am implementing QLoRA for fine-tuning a \`mistral-7b\` on a \`T4\` gpu. I loaded the model with the quantized configuration, however, when I attempt to test the model I keep getting an runtime error related to device. I ensured that the model and inputs are on \`cuda\`.
> 
> \`\`\`
> peft\_model.device
> device(type='cuda', index=0)
> 
> \# Test
> inputs = tokenizer("Do you have time", return\_tensors="pt").input\_ids.to(device)
> print("Inputs:", inputs)
> Inputs: tensor(\[\[1, 2378, 368, 506, 727\]\], device='cuda:0')
> \`\`\`
> 
> \### Who can help?
> 
> @sayakpaul @younesbelkada @BenjaminBossan 
> 
> \### Information
> 
> \- \[\] The official example scripts
> \- \[\] My own modified scripts
> 
> \### Tasks
> 
> \- \[\] An officially supported task in the \`examples\` folder
> \- \[\] My own task or dataset (give details below)
> 
> \### Reproduction
> 
> \`\`\`
> import bitsandbytes
> from peft import LoraConfig, get\_peft\_model
> 
> tokenizer = AutoTokenizer.from\_pretrained("mistralai/Mistral-7B-v0.1")
> 
> 
> bnb\_config = BitsAndBytesConfig(
> load\_in\_4bit=True, =
> bnb\_4bit\_quant\_type="nf4", 
> bnb\_4bit\_compute\_dtype=torch.bfloat16, 
> bnb\_4bit\_use\_double\_quant=True, 
> bnb\_4bit\_quant\_storage=torch.bfloat16 
> )
> 
> 
> model = AutoModelForCausalLM.from\_pretrained(
> "mistralai/Mistral-7B-v0.1"
> quantization\_config=bnb\_config,
> )
> 
> \# Lora
> config = LoraConfig(
> r=32, 
> lora\_alpha=64, 
> lora\_dropout=0.01, 
> target\_modules=\["q\_proj", "k\_proj", "v\_proj"\], 
> bias="none", 
> task\_type="CAUSAL\_LM"
> )
>                    
> peft\_model = get\_peft\_model(model, config)
> 
> inputs = tokenizer("Do you have time", return\_tensors="pt")
> print("Inputs:", inputs)
> 
> {'input\_ids': tensor(\[\[1, 2378, 368, 506, 727\]\], device='cuda:0'),
> 'attention\_mask': tensor(\[\[1, 1, 1, 1, 1\]\], device='cuda:0')} 
>  
> with torch.no\_grad():
> outputs = peft\_model(\*\*inputs)
> 
> print("Outputs:\\n", outputs.logits)
> \`\`\`
> 
> Error: 
> 
> \`\`\`
> \---------------------------------------------------------------------------
> RuntimeError Traceback (most recent call last)
> Cell In\[44\], line 2
> 1 with torch.no\_grad():
> \----\> 2 outputs = peft\_model(\*\*inputs)
> 4 print("Outputs:\\n", outputs.logits)
> 5 print("Outputs dimensions:", outputs.logits.shape) # shape: (batch\_size, num\_tokens, num\_classes)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1530 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1531 else:
> \-\> 1532 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1536 # If we don't have any hooks, we want to skip the rest of the logic in
> 1537 # this function, and just call forward.
> 1538 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1539 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1540 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1541 return forward\_call(\*args, \*\*kwargs)
> 1543 try:
> 1544 result = None
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/peft/peft\_model.py:1430, in PeftModelForCausalLM.forward(self, input\_ids, attention\_mask, inputs\_embeds, labels, output\_attentions, output\_hidden\_states, return\_dict, task\_ids, \*\*kwargs)
> 1428 with self.\_enable\_peft\_forward\_hooks(\*\*kwargs):
> 1429 kwargs = {k: v for k, v in kwargs.items() if k not in self.special\_peft\_forward\_args}
> \-\> 1430 return self.base\_model(
> 1431 input\_ids=input\_ids,
> 1432 attention\_mask=attention\_mask,
> 1433 inputs\_embeds=inputs\_embeds,
> 1434 labels=labels,
> 1435 output\_attentions=output\_attentions,
> 1436 output\_hidden\_states=output\_hidden\_states,
> 1437 return\_dict=return\_dict,
> 1438 \*\*kwargs,
> 1439 )
> 1441 batch\_size = \_get\_batch\_size(input\_ids, inputs\_embeds)
> 1442 if attention\_mask is not None:
> 1443 # concat prompt attention mask
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1530 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1531 else:
> \-\> 1532 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1536 # If we don't have any hooks, we want to skip the rest of the logic in
> 1537 # this function, and just call forward.
> 1538 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1539 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1540 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1541 return forward\_call(\*args, \*\*kwargs)
> 1543 try:
> 1544 result = None
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/peft/tuners/tuners\_utils.py:179, in BaseTuner.forward(self, \*args, \*\*kwargs)
> 178 def forward(self, \*args: Any, \*\*kwargs: Any):
> \--\> 179 return self.model.forward(\*args, \*\*kwargs)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/accelerate/hooks.py:166, in add\_hook\_to\_module.\<locals\>.new\_forward(module, \*args, \*\*kwargs)
> 164 output = module.\_old\_forward(\*args, \*\*kwargs)
> 165 else:
> \--\> 166 output = module.\_old\_forward(\*args, \*\*kwargs)
> 167 return module.\_hf\_hook.post\_forward(module, output)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/transformers/models/mistral/modeling\_mistral.py:1152, in MistralForCausalLM.forward(self, input\_ids, attention\_mask, position\_ids, past\_key\_values, inputs\_embeds, labels, use\_cache, output\_attentions, output\_hidden\_states, return\_dict)
> 1139 outputs = self.model(
> 1140 input\_ids=input\_ids,
> 1141 attention\_mask=attention\_mask,
> (...)
> 1148 return\_dict=return\_dict,
> 1149 )
> 1151 hidden\_states = outputs\[0\]
> \-\> 1152 logits = self.lm\_head(hidden\_states)
> 1153 logits = logits.float()
> 1155 loss = None
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1530 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1531 else:
> \-\> 1532 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1536 # If we don't have any hooks, we want to skip the rest of the logic in
> 1537 # this function, and just call forward.
> 1538 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1539 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1540 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1541 return forward\_call(\*args, \*\*kwargs)
> 1543 try:
> 1544 result = None
> 
> File ~/anaconda3/envs/python3/lib/python3.10/site-packages/torch/nn/modules/linear.py:116, in Linear.forward(self, input)
> 115 def forward(self, input: Tensor) -\> Tensor:
> \--\> 116 return F.linear(input, self.weight, self.bias)
> 
> RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper\_CUDA\_mm)
> \`\`\`
> 
> \### Expected behavior
> 
> Output logits.

> <https://github.com/huggingface/peft/issues/2129>
>
> \### System Info
> 
> System Info
> 
> transformers version: 4.45.1
> peft versio…n: 0.13.0
> Platform: Linux-5.10.225-213.878.amzn2.x86\_64-x86\_64-with-glibc2.35
> Python version: 3.10.14
> Huggingface\_hub version: 0.24.5
> Safetensors version: 0.4.4
> Accelerate version: 0.34.2
> Accelerate config: not found
> PyTorch version (GPU?): 2.1.0 (True)
> Tensorflow version (GPU?): 2.15.0 (True)
> Flax version (CPU?/GPU?/TPU?): not installed (NA)
> Jax version: not installed
> JaxLib version: not installed
> Using distributed or parallel set-up in script?: SFTTrainer, FSDP
> 
> 
> \### Who can help?
> 
> @BenjaminBossan @sayakpaul
> 
> \### Information
> 
> \- \[\] The official example scripts
> \- \[X\] My own modified scripts
> 
> \### Tasks
> 
> \- \[\] An officially supported task in the \`examples\` folder
> \- \[X\] My own task or dataset (give details below)
> 
> \### Reproduction
> 
> Script executed with instance type g5.xlarge (1 GPUs):
> 
> \`\`\`
> from accelerate import Accelerator
> from huggingface\_hub import login
> from peft import AutoPeftModelForCausalLM, LoraConfig, get\_peft\_model, prepare\_model\_for\_kbit\_training
> from sagemaker.remote\_function import remote
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set\_seed
> from trl import SFTConfig, SFTTrainer
> import transformers
> 
> def train\_fn(
> model\_name,
> train\_ds,
> test\_ds=None,
> lora\_r=8,
> lora\_alpha=16,
> lora\_dropout=0.1,
> per\_device\_train\_batch\_size=8,
> per\_device\_eval\_batch\_size=8,
> gradient\_accumulation\_steps=1,
> learning\_rate=2e-4,
> num\_train\_epochs=1,
> fsdp="",
> fsdp\_config=None,
> max\_seq\_length=2048,
> gradient\_checkpointing=False,
> merge\_weights=False,
> seed=42,
> token=None
> ):
> 
> set\_seed(seed)
> 
> accelerator = Accelerator()
> 
> if token is not None:
> login(token=token)
> 
> tokenizer = AutoTokenizer.from\_pretrained(model\_name)
> 
> # Set Tokenizer pad Token
> tokenizer.pad\_token = tokenizer.eos\_token
> 
> with accelerator.main\_process\_first():
> # tokenize and chunk dataset
> lm\_train\_dataset = train\_ds.map(
> lambda sample: tokenizer(sample\["text"\]), remove\_columns=list(train\_ds.features)
> )
> 
> print(f"Total number of train samples: {len(lm\_train\_dataset)}")
> 
> if test\_ds is not None:
> lm\_test\_dataset = test\_ds.map(
> lambda sample: tokenizer(sample\["text"\]), remove\_columns=list(test\_ds.features)
> )
> 
> print(f"Total number of test samples: {len(lm\_test\_dataset)}")
> else:
> lm\_test\_dataset = None
> 
> bnb\_config = BitsAndBytesConfig(
> load\_in\_4bit=True,
> bnb\_4bit\_use\_double\_quant=True,
> bnb\_4bit\_quant\_type="nf4",
> bnb\_4bit\_compute\_dtype=torch.bfloat16
> )
>     
> model = AutoModelForCausalLM.from\_pretrained(
> model\_name,
> trust\_remote\_code=True,
> quantization\_config=bnb\_config,
> attn\_implementation="flash\_attention\_2",
> use\_cache=False if gradient\_checkpointing else True,
> cache\_dir="/tmp/.cache"
> )
> 
> model = prepare\_model\_for\_kbit\_training(model, use\_gradient\_checkpointing=gradient\_checkpointing)
> 
> if gradient\_checkpointing:
> model.gradient\_checkpointing\_enable()
> 
> # get lora target modules
> modules = find\_all\_linear\_names(model)
> print(f"Found {len(modules)} modules to quantize: {modules}")
>     
> config = LoraConfig(
> r=lora\_r,
> lora\_alpha=lora\_alpha,
> target\_modules=modules,
> lora\_dropout=lora\_dropout,
> bias="none",
> task\_type="CAUSAL\_LM"
> )
> 
> model = get\_peft\_model(model, config)
> print\_trainable\_parameters(model)
>     
> trainer = transformers.Trainer(
> model=model,
> train\_dataset=lm\_train\_dataset,
> eval\_dataset=lm\_test\_dataset if lm\_test\_dataset is not None else None,
> args=transformers.TrainingArguments(
> per\_device\_train\_batch\_size=per\_device\_train\_batch\_size,
> per\_device\_eval\_batch\_size=per\_device\_eval\_batch\_size,
> gradient\_accumulation\_steps=gradient\_accumulation\_steps,
> gradient\_checkpointing=gradient\_checkpointing,
> logging\_strategy="steps",
> logging\_steps=1,
> log\_on\_each\_node=False,
> num\_train\_epochs=num\_train\_epochs,
> learning\_rate=learning\_rate,
> bf16=True,
> ddp\_find\_unused\_parameters=False,
> fsdp=fsdp,
> fsdp\_config=fsdp\_config,
> save\_strategy="no",
> output\_dir="outputs"
> ),
> data\_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
> )
> 
> trainer.train()
> 
> if trainer.is\_fsdp\_enabled:
> trainer.accelerator.state.fsdp\_plugin.set\_state\_dict\_type("FULL\_STATE\_DICT")
> 
> if merge\_weights:
> if accelerator.is\_main\_process:
> output\_dir = "/tmp/model"
>     
> # merge adapter weights with base model and save
> # save int 4 model
> trainer.model.save\_pretrained(output\_dir, safe\_serialization=False)
> # clear memory
> del model
> del trainer
>     
> torch.cuda.empty\_cache()
>     
> # load PEFT model in fp16
> model = AutoPeftModelForCausalLM.from\_pretrained(
> output\_dir,
> low\_cpu\_mem\_usage=True,
> torch\_dtype=torch.float16,
> cache\_dir="/tmp/.cache"
> )
>     
> # Merge LoRA and base model and save
> model = model.merge\_and\_unload()
> model.save\_pretrained(
> "/opt/ml/model", safe\_serialization=True, max\_shard\_size="2GB"
> )
> else:
> trainer.model.save\_pretrained("/opt/ml/model", safe\_serialization=True)
> 
> if accelerator.is\_main\_process:
> tokenizer.save\_pretrained("/opt/ml/model")
> 
> model\_id = "meta-llama/Meta-Llama-3.2-1B-Instruct"
> 
> train\_fn(
> model\_id,
> train\_ds=train\_dataset,
> test\_ds=test\_dataset,
> per\_device\_train\_batch\_size=4,
> per\_device\_eval\_batch\_size=1,
> gradient\_accumulation\_steps=4,
> gradient\_checkpointing=True,
> num\_train\_epochs=3,
> merge\_weights=True,
> token="\<HF\_TOKEN\>"
> )
> \`\`\`
> 
> train\_dataset:
> 
> \`\`\`
> Dataset({
> features: \['text'\],
> num\_rows: 140
> })
> \`\`\`
> 
> train\_dataset\[0\]\["text"\] (mock):
> 
> \`\`\`
> \<|begin\_of\_text|\>\<|start\_header\_id|\>user\<|end\_header\_id|\>What dimensions of the Llama 3.2 model are available?\<|eot\_id|\>\<|start\_header\_id|\>assistant\<|end\_header\_id|\>Llama 3.2 offers several model dimensions to cater to different use cases and computational requirements: Text-Only Models 1\\ 1B parameter model: This is the smallest Llama 3.2 model, designed for lightweight text processing tasks and on-device applications 2\\ 3B parameter model: A slightly larger text-only model that still maintains efficiency for edge devices and mobile applications. Multimodal Vision Models: 1\\ 11B parameter model: This is the smaller of the two vision-capable models, suitable for efficient deployment and development on consumer-grade GPUs 2\\ 90B parameter model: The largest Llama 3.2 model, designed for large-scale applications and advanced image reasoning tasks\<|end\_of\_text|\>\<|eot\_id|\>
> \`\`\`
> 
> Error returned:
> 
> \`\`\`
> \--\> 156 model = model.merge\_and\_unload()
> 157 model.save\_pretrained(
> 158 "\[/opt/ml/model\](https://ahwj8sa5y2ar84p.studio.us-east-1.sagemaker.aws/opt/ml/model)", safe\_serialization=True, max\_shard\_size="2GB"
> 159 )
> 160 else:
> 
> File \[/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/model.py:868\](https://ahwj8sa5y2ar84p.studio.us-east-1.sagemaker.aws/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/model.py#line=867), in merge\_and\_unload()
> 840 def merge\_and\_unload(
> 841 self, progressbar: bool = False, safe\_merge: bool = False, adapter\_names: Optional\[list\[str\]\] = None
> 842 ) -\> torch.nn.Module:
> 843 r"""
> 844 This method merges the LoRa layers into the base model. This is needed if someone wants to use the base model
> 845 as a standalone model.
> (...)
> 866 \`\`\`
> 867 """
> \--\> 868 return self.\_unload\_and\_optionally\_merge(
> 869 progressbar=progressbar, safe\_merge=safe\_merge, adapter\_names=adapter\_names
> 870 )
> 
> File \[/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/model.py:490\](https://ahwj8sa5y2ar84p.studio.us-east-1.sagemaker.aws/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/model.py#line=489), in \_unload\_and\_optionally\_merge()
> 488 if hasattr(target, "base\_layer"):
> 489 if merge:
> \--\> 490 target.merge(safe\_merge=safe\_merge, adapter\_names=adapter\_names)
> 491 self.\_replace\_module(parent, target\_name, target.get\_base\_layer(), target)
> 492 elif isinstance(target, ModulesToSaveWrapper):
> 493 # save any additional trainable modules part of \`modules\_to\_save\`
> 
> File \[/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/layer.py:483\](https://ahwj8sa5y2ar84p.studio.us-east-1.sagemaker.aws/opt/conda/lib/python3.10/site-packages/peft/tuners/lora/layer.py#line=482), in merge()
> 481 delta\_weight = self.get\_delta\_weight(active\_adapter)
> 482 if not self.use\_dora\[active\_adapter\]:
> \--\> 483 base\_layer.weight.data += delta\_weight
> 484 else:
> 485 # handle dora
> 486 # since delta\_weight already includes scaling, set it to 1 here
> 487 weight\_norm = (
> 488 self.lora\_magnitude\_vector\[active\_adapter\]
> 489 .get\_weight\_norm(
> (...)
> 492 .detach()
> 493 )
> 
> RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
> \`\`\`
> 
> requirements.txt:
> 
> \`\`\`
> transformers==4.45.1
> peft==0.13.0
> accelerate==0.34.2
> bitsandbytes==0.44.0
> evaluate==0.4.1
> safetensors\>=0.4.3
> sagemaker==2.232.1
> trl==0.11.1
> tokenizers\>=0.19.1
> py7zr
> \`\`\`
> 
> \### Expected behavior
> 
> With the older version of peft==0.10.0, the training was running perfectly. Now it is not able to merge the adapter without moving the model to cuda.
> How can I solve it? What am I doing wrong?

> <https://github.com/huggingface/transformers/issues/34695>
>
> \*\*Reproduction\*\*
> 
> I am trying to finetune Qwen2-0.5B model on some training da…ta using a multi-GPU setup. The same code (given further below) seems to work in a single-GPU setting (when i set CUDA\_VISIBLE\_DEVICES=0):
> 
> \`\`\`python
> \---------------------------------------------------------------------------
> RuntimeError Traceback (most recent call last)
> Cell In\[18\], line 4
> 2 import torch
> 3 torch.autograd.set\_detect\_anomaly(True)
> \----\> 4 main()
> 
> Cell In\[14\], line 15, in main()
> 8 trainer = Trainer(env\_params=env\_params,
> 9 model\_params=model\_params,
> 10 optimizer\_params=optimizer\_params,
> 11 trainer\_params=trainer\_params)
> 13 copy\_all\_src(trainer.result\_folder)
> \---\> 15 trainer.run()
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTrainerTransformer.py:92, in TSPTrainer.run(self)
> 89 self.scheduler.step()
> 91 # Train
> \---\> 92 train\_score, train\_loss = self.\_train\_one\_epoch(epoch)
> 93 self.result\_log.append('train\_score', epoch, train\_score)
> 94 self.result\_log.append('train\_loss', epoch, train\_loss)
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTrainerTransformer.py:151, in TSPTrainer.\_train\_one\_epoch(self, epoch)
> 148 remaining = train\_num\_episode - episode
> 149 batch\_size = min(self.trainer\_params\['train\_batch\_size'\], remaining)
> \--\> 151 avg\_score, avg\_loss = self.\_train\_one\_batch(batch\_size)
> 152 score\_AM.update(avg\_score, batch\_size)
> 153 loss\_AM.update(avg\_loss, batch\_size)
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTrainerTransformer.py:193, in TSPTrainer.\_train\_one\_batch(self, batch\_size)
> 191 state, reward, done = self.env.pre\_step()
> 192 while not done:
> \--\> 193 selected, prob = self.model.module(state)
> 194 # shape: (batch, pomo)
> 195 state, reward, done = self.env.step(selected)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTransformerModelQuant\_b.py:39, in TSPTransformer.forward(self, state)
> 37 return self.\_init\_sequence(batch\_size, pomo\_size)
> 38 else:
> \---\> 39 return self.\_continue\_sequence(state, batch\_size, pomo\_size)
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTransformerModelQuant\_b.py:84, in TSPTransformer.\_continue\_sequence(self, state, batch\_size, pomo\_size)
> 81 state.ninf\_mask = state.ninf\_mask.to(self.device)
> 83 # Get probabilities from decoder
> \---\> 84 probs = self.decoder(self.seq\_so\_far, self.input\_mask, state.ninf\_mask)
> 86 # Select next node
> 87 if self.training or self.model\_params\['eval\_type'\] == 'softmax':
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/neuralcombinatorialoptimization/NCO-master/NEW\_py\_ver/TSP/POMO/TSPTransformerModelQuant\_b.py:185, in Decoder.forward(self, seq\_so\_far, inp\_mask, ninf\_mask)
> 182 flat\_mask = inp\_mask.reshape(batch\_size \* pomo\_size, problem\_size)
> 184 # Get model outputs
> \--\> 185 outputs = self.model(inputs\_embeds=flat\_seq, attention\_mask=flat\_mask)
> 186 logits = outputs.logits
> 188 # Get last valid position
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/peft/peft\_model.py:1644, in PeftModelForCausalLM.forward(self, input\_ids, attention\_mask, inputs\_embeds, labels, output\_attentions, output\_hidden\_states, return\_dict, task\_ids, \*\*kwargs)
> 1642 with self.\_enable\_peft\_forward\_hooks(\*\*kwargs):
> 1643 kwargs = {k: v for k, v in kwargs.items() if k not in self.special\_peft\_forward\_args}
> \-\> 1644 return self.base\_model(
> 1645 input\_ids=input\_ids,
> 1646 attention\_mask=attention\_mask,
> 1647 inputs\_embeds=inputs\_embeds,
> 1648 labels=labels,
> 1649 output\_attentions=output\_attentions,
> 1650 output\_hidden\_states=output\_hidden\_states,
> 1651 return\_dict=return\_dict,
> 1652 \*\*kwargs,
> 1653 )
> 1655 batch\_size = \_get\_batch\_size(input\_ids, inputs\_embeds)
> 1656 if attention\_mask is not None:
> 1657 # concat prompt attention mask
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/peft/tuners/tuners\_utils.py:197, in BaseTuner.forward(self, \*args, \*\*kwargs)
> 196 def forward(self, \*args: Any, \*\*kwargs: Any):
> \--\> 197 return self.model.forward(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/accelerate/hooks.py:170, in add\_hook\_to\_module.\<locals\>.new\_forward(module, \*args, \*\*kwargs)
> 168 output = module.\_old\_forward(\*args, \*\*kwargs)
> 169 else:
> \--\> 170 output = module.\_old\_forward(\*args, \*\*kwargs)
> 171 return module.\_hf\_hook.post\_forward(module, output)
> 
> File ~/second/lib/python3.10/site-packages/transformers/models/qwen2/modeling\_qwen2.py:1170, in Qwen2ForCausalLM.forward(self, input\_ids, attention\_mask, position\_ids, past\_key\_values, inputs\_embeds, labels, use\_cache, output\_attentions, output\_hidden\_states, return\_dict, cache\_position, num\_logits\_to\_keep, \*\*loss\_kwargs)
> 1167 return\_dict = return\_dict if return\_dict is not None else self.config.use\_return\_dict
> 1169 # decoder outputs consists of (dec\_features, layer\_state, dec\_hidden, dec\_attn)
> \-\> 1170 outputs = self.model(
> 1171 input\_ids=input\_ids,
> 1172 attention\_mask=attention\_mask,
> 1173 position\_ids=position\_ids,
> 1174 past\_key\_values=past\_key\_values,
> 1175 inputs\_embeds=inputs\_embeds,
> 1176 use\_cache=use\_cache,
> 1177 output\_attentions=output\_attentions,
> 1178 output\_hidden\_states=output\_hidden\_states,
> 1179 return\_dict=return\_dict,
> 1180 cache\_position=cache\_position,
> 1181 )
> 1183 hidden\_states = outputs\[0\]
> 1184 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/transformers/models/qwen2/modeling\_qwen2.py:901, in Qwen2Model.forward(self, input\_ids, attention\_mask, position\_ids, past\_key\_values, inputs\_embeds, use\_cache, output\_attentions, output\_hidden\_states, return\_dict, cache\_position)
> 889 layer\_outputs = self.\_gradient\_checkpointing\_func(
> 890 decoder\_layer.\_\_call\_\_,
> 891 hidden\_states,
> (...)
> 898 position\_embeddings,
> 899 )
> 900 else:
> \--\> 901 layer\_outputs = decoder\_layer(
> 902 hidden\_states,
> 903 attention\_mask=causal\_mask,
> 904 position\_ids=position\_ids,
> 905 past\_key\_value=past\_key\_values,
> 906 output\_attentions=output\_attentions,
> 907 use\_cache=use\_cache,
> 908 cache\_position=cache\_position,
> 909 position\_embeddings=position\_embeddings,
> 910 )
> 912 hidden\_states = layer\_outputs\[0\]
> 914 if use\_cache:
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/accelerate/hooks.py:170, in add\_hook\_to\_module.\<locals\>.new\_forward(module, \*args, \*\*kwargs)
> 168 output = module.\_old\_forward(\*args, \*\*kwargs)
> 169 else:
> \--\> 170 output = module.\_old\_forward(\*args, \*\*kwargs)
> 171 return module.\_hf\_hook.post\_forward(module, output)
> 
> File ~/second/lib/python3.10/site-packages/transformers/models/qwen2/modeling\_qwen2.py:629, in Qwen2DecoderLayer.forward(self, hidden\_states, attention\_mask, position\_ids, past\_key\_value, output\_attentions, use\_cache, cache\_position, position\_embeddings, \*\*kwargs)
> 626 hidden\_states = self.input\_layernorm(hidden\_states)
> 628 # Self Attention
> \--\> 629 hidden\_states, self\_attn\_weights, present\_key\_value = self.self\_attn(
> 630 hidden\_states=hidden\_states,
> 631 attention\_mask=attention\_mask,
> 632 position\_ids=position\_ids,
> 633 past\_key\_value=past\_key\_value,
> 634 output\_attentions=output\_attentions,
> 635 use\_cache=use\_cache,
> 636 cache\_position=cache\_position,
> 637 position\_embeddings=position\_embeddings,
> 638 )
> 639 hidden\_states = residual + hidden\_states
> 641 # Fully Connected
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/accelerate/hooks.py:170, in add\_hook\_to\_module.\<locals\>.new\_forward(module, \*args, \*\*kwargs)
> 168 output = module.\_old\_forward(\*args, \*\*kwargs)
> 169 else:
> \--\> 170 output = module.\_old\_forward(\*args, \*\*kwargs)
> 171 return module.\_hf\_hook.post\_forward(module, output)
> 
> File ~/second/lib/python3.10/site-packages/transformers/models/qwen2/modeling\_qwen2.py:506, in Qwen2SdpaAttention.forward(self, hidden\_states, attention\_mask, position\_ids, past\_key\_value, output\_attentions, use\_cache, cache\_position, position\_embeddings)
> 495 return super().forward(
> 496 hidden\_states=hidden\_states,
> 497 attention\_mask=attention\_mask,
> (...)
> 501 use\_cache=use\_cache,
> 502 )
> 504 bsz, q\_len, \_ = hidden\_states.size()
> \--\> 506 query\_states = self.q\_proj(hidden\_states)
> 507 key\_states = self.k\_proj(hidden\_states)
> 508 value\_states = self.v\_proj(hidden\_states)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/peft/tuners/lora/layer.py:572, in Linear.forward(self, x, \*args, \*\*kwargs)
> 570 result = self.base\_layer(x, \*args, \*\*kwargs)
> 571 else:
> \--\> 572 result = self.base\_layer(x, \*args, \*\*kwargs)
> 573 torch\_result\_dtype = result.dtype
> 574 for active\_adapter in self.active\_adapters:
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs)
> 1734 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore\[misc\]
> 1735 else:
> \-\> 1736 return self.\_call\_impl(\*args, \*\*kwargs)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module.\_call\_impl(self, \*args, \*\*kwargs)
> 1742 # If we don't have any hooks, we want to skip the rest of the logic in
> 1743 # this function, and just call forward.
> 1744 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks
> 1745 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks
> 1746 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks):
> \-\> 1747 return forward\_call(\*args, \*\*kwargs)
> 1749 result = None
> 1750 called\_always\_called\_hooks = set()
> 
> File ~/second/lib/python3.10/site-packages/accelerate/hooks.py:170, in add\_hook\_to\_module.\<locals\>.new\_forward(module, \*args, \*\*kwargs)
> 168 output = module.\_old\_forward(\*args, \*\*kwargs)
> 169 else:
> \--\> 170 output = module.\_old\_forward(\*args, \*\*kwargs)
> 171 return module.\_hf\_hook.post\_forward(module, output)
> 
> File ~/second/lib/python3.10/site-packages/torch/nn/modules/linear.py:125, in Linear.forward(self, input)
> 124 def forward(self, input: Tensor) -\> Tensor:
> \--\> 125 return F.linear(input, self.weight, self.bias)
> 
> RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper\_CUDA\_addmm)
> 
> \`\`\`
> Code for the above error is given below:
> 
> \`\`\`python
> Trainer.py
> 
> import torch
> from logging import getLogger
> from torch.nn.parallel import DataParallel
> from TSPEnvQuant import TSPEnv as Env
> from TSPTransformerModelQuant\_b import TSPTransformer as Model
> 
> from torch.optim import Adam as Optimizer
> from torch.optim.lr\_scheduler import MultiStepLR as Scheduler
> 
> from utils.utils import \*
> 
> 
> class TSPTrainer:
> def \_\_init\_\_(self,
> env\_params,
> model\_params,
> optimizer\_params,
> trainer\_params):
> 
> # save arguments
> self.env\_params = env\_params
> self.model\_params = model\_params
> self.optimizer\_params = optimizer\_params
> self.trainer\_params = trainer\_params
> 
> # result folder, logger
> self.logger = getLogger(name='trainer')
> self.result\_folder = get\_result\_folder()
> self.result\_log = LogData()
> 
> # cuda
> USE\_CUDA = self.trainer\_params\['use\_cuda'\]
> if USE\_CUDA:
> cuda\_device\_num = self.trainer\_params\['cuda\_device\_num'\]
> torch.cuda.set\_device(cuda\_device\_num)
> device = torch.device('cuda', cuda\_device\_num)
> torch.set\_default\_tensor\_type('torch.cuda.FloatTensor')
> else:
> device = torch.device('cpu')
> torch.set\_default\_tensor\_type('torch.FloatTensor')
> 
> # Main Components
> self.model = Model(\*\*self.model\_params)
> if USE\_CUDA and torch.cuda.device\_count() \> 1:
> self.logger.info(f"Using {torch.cuda.device\_count()} GPUs!")
> self.model = DataParallel(self.model)
> self.model = self.model.to(device)
> 
> self.env = Env(\*\*self.env\_params)
> self.optimizer = Optimizer(self.model.parameters(), \*\*self.optimizer\_params\['optimizer'\])
> self.scheduler = Scheduler(self.optimizer, \*\*self.optimizer\_params\['scheduler'\])
> 
> # Restore
> self.start\_epoch = 1
> model\_load = trainer\_params\['model\_load'\]
> if model\_load\['enable'\]:
> checkpoint\_fullname = '{path}/checkpoint-{epoch}.pt'.format(\*\*model\_load)
> checkpoint = torch.load(checkpoint\_fullname, map\_location=device)
> # Handle loading state dict for DataParallel
> if isinstance(self.model, DataParallel):
> # If saved model wasn't using DataParallel but current model is
> if not any(key.startswith('module.') for key in checkpoint\['model\_state\_dict'\].keys()):
> new\_state\_dict = {'module.' + k: v for k, v in checkpoint\['model\_state\_dict'\].items()}
> self.model.load\_state\_dict(new\_state\_dict)
> else:
> self.model.load\_state\_dict(checkpoint\['model\_state\_dict'\])
> else:
> # If saved model was using DataParallel but current model isn't
> if any(key.startswith('module.') for key in checkpoint\['model\_state\_dict'\].keys()):
> new\_state\_dict = {k.replace('module.', ''): v for k, v in checkpoint\['model\_state\_dict'\].items()}
> self.model.load\_state\_dict(new\_state\_dict)
> else:
> self.model.load\_state\_dict(checkpoint\['model\_state\_dict'\])
> self.start\_epoch = 1 + model\_load\['epoch'\]
> self.result\_log.set\_raw\_data(checkpoint\['result\_log'\])
> self.optimizer.load\_state\_dict(checkpoint\['optimizer\_state\_dict'\])
> self.scheduler.last\_epoch = model\_load\['epoch'\]-1
> self.logger.info('Saved Model Loaded !!')
> 
> # utility
> self.time\_estimator = TimeEstimator()
> 
> def run(self):
> self.time\_estimator.reset(self.start\_epoch)
> for epoch in range(self.start\_epoch, self.trainer\_params\['epochs'\]+1):
> self.logger.info('=================================================================')
> 
> # LR Decay
> self.scheduler.step()
> 
> # Train
> train\_score, train\_loss = self.\_train\_one\_epoch(epoch)
> self.result\_log.append('train\_score', epoch, train\_score)
> self.result\_log.append('train\_loss', epoch, train\_loss)
> 
> ############################
> # Logs & Checkpoint
> ############################
> elapsed\_time\_str, remain\_time\_str = self.time\_estimator.get\_est\_string(epoch, self.trainer\_params\['epochs'\])
> self.logger.info("Epoch {:3d}/{:3d}: Time Est.: Elapsed\[{}\], Remain\[{}\]".format(
> epoch, self.trainer\_params\['epochs'\], elapsed\_time\_str, remain\_time\_str))
> 
> all\_done = (epoch == self.trainer\_params\['epochs'\])
> model\_save\_interval = self.trainer\_params\['logging'\]\['model\_save\_interval'\]
> img\_save\_interval = self.trainer\_params\['logging'\]\['img\_save\_interval'\]
> 
> if epoch \> 1: # save latest images, every epoch
> self.logger.info("Saving log\_image")
> image\_prefix = '{}/latest'.format(self.result\_folder)
> util\_save\_log\_image\_with\_label(image\_prefix, self.trainer\_params\['logging'\]\['log\_image\_params\_1'\],
> self.result\_log, labels=\['train\_score'\])
> util\_save\_log\_image\_with\_label(image\_prefix, self.trainer\_params\['logging'\]\['log\_image\_params\_2'\],
> self.result\_log, labels=\['train\_loss'\])
> 
> if all\_done or (epoch % model\_save\_interval) == 0:
> self.logger.info("Saving trained\_model")
> checkpoint\_dict = {
> 'epoch': epoch,
> 'model\_state\_dict': self.model.state\_dict(),
> 'optimizer\_state\_dict': self.optimizer.state\_dict(),
> 'scheduler\_state\_dict': self.scheduler.state\_dict(),
> 'result\_log': self.result\_log.get\_raw\_data()
> }
> torch.save(checkpoint\_dict, '{}/checkpoint-{}.pt'.format(self.result\_folder, epoch))
> 
> if all\_done or (epoch % img\_save\_interval) == 0:
> image\_prefix = '{}/img/checkpoint-{}'.format(self.result\_folder, epoch)
> util\_save\_log\_image\_with\_label(image\_prefix, self.trainer\_params\['logging'\]\['log\_image\_params\_1'\],
> self.result\_log, labels=\['train\_score'\])
> util\_save\_log\_image\_with\_label(image\_prefix, self.trainer\_params\['logging'\]\['log\_image\_params\_2'\],
> self.result\_log, labels=\['train\_loss'\])
> 
> if all\_done:
> self.logger.info(" \*\*\* Training Done \*\*\* ")
> self.logger.info("Now, printing log array...")
> util\_print\_log\_array(self.logger, self.result\_log)
> 
> def \_train\_one\_epoch(self, epoch):
> 
> score\_AM = AverageMeter()
> loss\_AM = AverageMeter()
> 
> train\_num\_episode = self.trainer\_params\['train\_episodes'\]
> episode = 0
> loop\_cnt = 0
> while episode \< train\_num\_episode:
> 
> remaining = train\_num\_episode - episode
> batch\_size = min(self.trainer\_params\['train\_batch\_size'\], remaining)
> 
> avg\_score, avg\_loss = self.\_train\_one\_batch(batch\_size)
> score\_AM.update(avg\_score, batch\_size)
> loss\_AM.update(avg\_loss, batch\_size)
> 
> episode += batch\_size
> 
> # Log First 10 Batch, only at the first epoch
> if epoch == self.start\_epoch:
> loop\_cnt += 1
> if loop\_cnt \<= 10:
> self.logger.info('Epoch {:3d}: Train {:3d}/{:3d}({:1.1f}%) Score: {:.4f}, Loss: {:.4f}'
> .format(epoch, episode, train\_num\_episode, 100. \* episode / train\_num\_episode,
> score\_AM.avg, loss\_AM.avg))
> 
> # Log Once, for each epoch
> self.logger.info('Epoch {:3d}: Train ({:3.0f}%) Score: {:.4f}, Loss: {:.4f}'
> .format(epoch, 100. \* episode / train\_num\_episode,
> score\_AM.avg, loss\_AM.avg))
> 
> return score\_AM.avg, loss\_AM.avg
> 
> def \_train\_one\_batch(self, batch\_size):
> 
> # Prep
> ###############################################
> self.model.train()
> self.env.load\_problems(batch\_size)
> reset\_state, \_, \_ = self.env.reset()
> # Handle pre\_forward for DataParallel
> if isinstance(self.model, DataParallel):
> print("Is DataParallel")
> self.model.module.pre\_forward(reset\_state)
> else:
> self.model.pre\_forward(reset\_state)
> 
> prob\_list = torch.zeros(size=(batch\_size, self.env.pomo\_size, 0))
> # shape: (batch, pomo, 0~problem)
> 
> # POMO Rollout
> ###############################################
> state, reward, done = self.env.pre\_step()
> while not done:
> selected, prob = self.model.module(state)
> # shape: (batch, pomo)
> state, reward, done = self.env.step(selected)
> prob\_list = torch.cat((prob\_list, prob\[:, :, None\]), dim=2)
> 
> # Loss
> ###############################################
> advantage = reward - reward.float().mean(dim=1, keepdims=True)
> # shape: (batch, pomo)
> log\_prob = prob\_list.log().sum(dim=2)
> # size = (batch, pomo)
> loss = -advantage \* log\_prob # Minus Sign: To Increase REWARD
> # shape: (batch, pomo)
> loss\_mean = loss.mean()
> 
> # Score
> ###############################################
> max\_pomo\_reward, \_ = reward.max(dim=1) # get best results from pomo
> score\_mean = -max\_pomo\_reward.float().mean() # negative sign to make positive value
> 
> # Step & Return
> ###############################################
> self.model.zero\_grad()
> loss\_mean.backward()
> self.optimizer.step()
> return score\_mean.item(), loss\_mean.item()
> 
> \`\`\`
> \`\`\`python
> Model.py
> 
> import torch
> import torch.nn as nn
> from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
> from transformers import AutoConfig
> from peft import LoraConfig, get\_peft\_model, prepare\_model\_for\_kbit\_training, PeftModel, TaskType
> from typing import Optional, Dict, Any, Tuple
> 
> class TSPTransformer(nn.Module):
> def \_\_init\_\_(self, \*\*kwargs):
> super().\_\_init\_\_()
> self.model\_params = kwargs
> self.encoder = Encoder(\*\*kwargs)
> self.embedding\_size = kwargs.get('embedding\_dim', 896)
>         
> # Load the model with LoRA and 4-bit quantization if needed
> self.model = load\_model(kwargs)
> self.decoder = Decoder(self.model, \*\*kwargs)
>         
> # Initialize state storage
> self.encoded\_nodes = None
> self.seq\_so\_far = None
> self.input\_mask = None
> self.t = None
> self.device = kwargs.get('device', torch.device("cuda" if torch.cuda.is\_available() else "cpu"))
> 
> def pre\_forward(self, reset\_state):
> """Initialize model state for new sequence"""
> self.encoded\_nodes = self.encoder(reset\_state.problems)
> self.problem\_size = reset\_state.problems.size(1)
> self.batch\_size = reset\_state.problems.size(0)
> 
> def forward(self, state) -\> Tuple\[torch.Tensor, Optional\[torch.Tensor\]\]:
> batch\_size = state.BATCH\_IDX.size(0)
> pomo\_size = state.BATCH\_IDX.size(1)
> 
> if state.current\_node is None:
> return self.\_init\_sequence(batch\_size, pomo\_size)
> else:
> return self.\_continue\_sequence(state, batch\_size, pomo\_size)
> 
> def \_init\_sequence(self, batch\_size: int, pomo\_size: int) -\> Tuple\[torch.Tensor, torch.Tensor\]:
> """Initialize sequence state"""
> self.t = 0 # Start at 0 instead of -1
>         
> # Create new tensors instead of modifying in place
> selected = torch.arange(pomo\_size, device=self.device).expand(batch\_size, pomo\_size)
> prob = torch.ones(size=(batch\_size, pomo\_size), device=self.device)
>         
> # Initialize sequence storage with proper dimensions
> self.seq\_so\_far = torch.zeros(
> (batch\_size, pomo\_size, self.problem\_size, self.embedding\_size),
> device=self.device
> )
>         
> self.input\_mask = torch.zeros(
> (batch\_size, pomo\_size, self.problem\_size),
> dtype=torch.bool,
> device=self.device
> )
>         
> return selected, prob
> 
> def \_continue\_sequence(self, state, batch\_size: int, pomo\_size: int) -\> Tuple\[torch.Tensor, Optional\[torch.Tensor\]\]:
> """Continue sequence generation"""
> # Get encoded representation of current node
> encoded\_current = self.\_get\_encoded\_node(state.current\_node)
>         
> # Create new tensor for updated sequence
> new\_seq = self.seq\_so\_far.clone()
> new\_seq\[:, :, self.t, :\] = encoded\_current
> self.seq\_so\_far = new\_seq
>         
> # Create new tensor for updated mask
> new\_mask = self.input\_mask.clone()
> new\_mask\[:, :, self.t\] = True
> self.input\_mask = new\_mask
>         
> # Move tensors to correct device
> self.seq\_so\_far = self.seq\_so\_far.to(self.device)
> self.input\_mask = self.input\_mask.to(self.device)
> state.ninf\_mask = state.ninf\_mask.to(self.device)
>         
> # Get probabilities from decoder
> probs = self.decoder(self.seq\_so\_far, self.input\_mask, state.ninf\_mask)
>         
> # Select next node
> if self.training or self.model\_params\['eval\_type'\] == 'softmax':
> selected, prob = self.\_sample\_node(probs, state, batch\_size, pomo\_size)
> else:
> selected = probs.argmax(dim=2)
> prob = None
>         
> self.t += 1
> return selected, prob
> 
> def \_get\_encoded\_node(self, node\_indices: torch.Tensor) -\> torch.Tensor:
> """Get encoded representation of nodes safely"""
> batch\_size, pomo\_size = node\_indices.shape
> embedding\_dim = self.encoded\_nodes.size(2)
>         
> # Create gathering indices
> gather\_idx = node\_indices\[:, :, None\].expand(batch\_size, pomo\_size, embedding\_dim)
> gather\_idx = gather\_idx.to(self.encoded\_nodes.device)
>         
> # Gather encoded representations
> return self.encoded\_nodes.gather(dim=1, index=gather\_idx)
> 
> def \_sample\_node(self, probs: torch.Tensor, state, batch\_size: int, pomo\_size: int) -\> Tuple\[torch.Tensor, torch.Tensor\]:
> """Sample next node with retry logic"""
> max\_attempts = 100
> for \_ in range(max\_attempts):
> # Reshape for sampling
> flat\_probs = probs.reshape(batch\_size \* pomo\_size, -1)
>             
> # Sample indices
> selected = flat\_probs.multinomial(1, replacement=True)
> selected = selected.reshape(batch\_size, pomo\_size)
>             
> # Calculate probabilities
> prob = probs\[state.BATCH\_IDX, state.POMO\_IDX, selected\]
> prob = prob.reshape(batch\_size, pomo\_size)
>             
> if (prob \> 0).all():
> return selected, prob
>         
> raise RuntimeError(f"Failed to sample valid nodes after {max\_attempts} attempts")
> 
> class Encoder(nn.Module):
> def \_\_init\_\_(self, \*\*kwargs):
> super().\_\_init\_\_()
> self.embedding\_dim = kwargs.get('embedding\_dim', 896) - 1
> self.embed\_layer = nn.Linear(2, self.embedding\_dim)
> self.device = kwargs.get('device', torch.device("cuda" if torch.cuda.is\_available() else "cpu"))
>     
> def forward(self, problems: torch.Tensor) -\> torch.Tensor:
> batch\_size, problem\_size = problems.shape\[:2\]
>         
> # Create position encodings
> ids = torch.arange(problem\_size, device=self.device).expand(batch\_size, problem\_size)
>         
> # Embed coordinates
> embedded = self.embed\_layer(problems.reshape(-1, 2))
> embedded = embedded.reshape(batch\_size, problem\_size, self.embedding\_dim)
>         
> # Concatenate position encodings
> return torch.cat(\[ids.unsqueeze(-1).float(), embedded\], dim=-1)
> 
> class Decoder(nn.Module):
> def \_\_init\_\_(self, model: nn.Module, \*\*kwargs):
> super().\_\_init\_\_()
> self.model = model
> self.problem\_size = kwargs.get('problem\_size', 20)
> self.use\_lora = kwargs.get('use\_lora', True)
>         
> self.\_setup\_model()
>     
> def \_setup\_model(self):
> """Configure model architecture"""
> # Modify output size
> self.model.lm\_head = nn.Linear(
> self.model.config.hidden\_size,
> self.problem\_size
> ).to(self.model.device)
>         
> # Apply LoRA if requested
> if self.use\_lora:
> lora\_config = LoraConfig(
> r=4,
> lora\_alpha=32,
> target\_modules=\["q\_proj", "v\_proj"\],
> lora\_dropout=0.1,
> bias="none",
> task\_type=TaskType.CAUSAL\_LM
> )
> self.model = get\_peft\_model(self.model, lora\_config)
>     
> def forward(self, seq\_so\_far: torch.Tensor, inp\_mask: torch.Tensor, ninf\_mask: torch.Tensor) -\> torch.Tensor:
> batch\_size, pomo\_size, problem\_size, embedding\_dim = seq\_so\_far.shape
>         
> # Reshape inputs
> flat\_seq = seq\_so\_far.reshape(batch\_size \* pomo\_size, problem\_size, embedding\_dim)
> flat\_mask = inp\_mask.reshape(batch\_size \* pomo\_size, problem\_size)
>         
> # Get model outputs
> outputs = self.model(inputs\_embeds=flat\_seq, attention\_mask=flat\_mask)
> logits = outputs.logits
>         
> # Get last valid position
> last\_positions = flat\_mask.sum(dim=1).long() - 1
>         
> # Gather logits for last positions
> batch\_indices = torch.arange(batch\_size \* pomo\_size, device=logits.device)
> gathered\_logits = logits\[batch\_indices, last\_positions\]
>         
> # Reshape and apply mask
> logits = gathered\_logits.reshape(batch\_size, pomo\_size, problem\_size)
> masked\_logits = logits + ninf\_mask.float()
>         
> # Return probabilities
> return torch.softmax(masked\_logits, dim=2)
> 
> def load\_model(config: Dict\[str, Any\]) -\> nn.Module:
> """Load model with proper configuration"""
> # print(config)
> device = config.get('device', torch.device("cuda" if torch.cuda.is\_available() else "cpu"))
>     
> if config.get('checkpoint\_path'):
> # print('checkpoint\_path')
> try:
> return PeftModel.from\_pretrained(
> config\['model\_name'\],
> config\['checkpoint\_path'\],
> is\_trainable=True
> ).to(device)
> except Exception as e:
> print(f"Error loading checkpoint: {e}")
> print("Falling back to base model...")
>     
> print(config)
> # print(config\['use\_4bit'\])
> if config.get('use\_4bit', True):
> print('use\_4bit')
> bnb\_config = BitsAndBytesConfig(
> load\_in\_4bit=True,
> bnb\_4bit\_compute\_dtype=torch.bfloat16,
> llm\_int8\_has\_fp16\_weight=False,
> bnb\_4bit\_quant\_type="nf4",
> llm\_int8\_threshold=6.0,
> bnb\_4bit\_use\_double\_quant=True,
> )
> # print(config\['model\_name'\])
> # print(type(config\['model\_name'\]))
> model = AutoModelForCausalLM.from\_pretrained(
> config\['model\_name'\],
> trust\_remote\_code=True,
> device\_map="auto",
> torch\_dtype=torch.bfloat16,
> quantization\_config=bnb\_config
> )
> model = prepare\_model\_for\_kbit\_training(model)
> model.config.use\_cache = False
> else:
> # print('else')
> model = AutoModelForCausalLM.from\_pretrained(
> config\['model\_name'\],
> torch\_dtype=torch.float32,
> trust\_remote\_code=True,
> device\_map="auto",
> ).to(device)
>     
> return model
> \`\`\`
> \*\*Expected behavior\*\*
> Expected behavior is that the model should train in a multi-GPU setting without throwing any errors. The same script works in single-GPU setting but throws the above error in a multi-GPU setting

---

<div class="post-metadata">

**Author:** ![BenjaminB](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/benjaminb/32/30898_2.png) [@BenjaminB](https://discuss.huggingface.co/u/BenjaminB)\
**Post date:** [April 25, 2025, 4:12pm UTC](https://discuss.huggingface.co/t/multi-gpu-inference-llama-3-2-vision-with-qlora/150685/6 "2025-04-25T16:12:54Z")

</div>

If could open an issue on PEFT ([GitHub · Where software is built](https://github.com/huggingface/peft/issues)) and provide a full reproducer, we will be happy to look into this. Please also attach the full error message.