Closed kosstbarz closed 3 weeks ago
cc @gante
The issue arises because tensors are being allocated on different devices (CUDA and CPU), causing a RuntimeError
. This often happens when using multi-GPU setups and certain layers or operations are not properly configured to use the same device.
Possible Solution : Ensure all tensors are on the same device (CUDA) before running the
generate
function.
import transformers
# Load the model and tokenizer
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
# Define the prompt and tokenize
prompt = "Hey, are you conscious? Can you talk to me?\n\n"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# Generate the output
generate_ids = model.generate(inputs.input_ids, max_length=200, dola_layers='low')
# Decode and print the result
result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
print(result)
By moving the input tensors to CUDA using
.to("cuda")
, we ensure all operations are performed on the same device, thus resolving the error.
Thank you @ashishpatel26 for your response.
I tried .to("cuda")
and get
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:3!
The same error I get with .to("cuda:0")
and .to("cuda:3")
The same result when I use model.generate(inputs.input_ids.to("cuda:3")
Frankly I expected that model with device_map == "auto" would put the input tensor on the correct device. Because base Mixtral model (with no dola_layers) works fine with input tensor on cpu.
Gentle ping @gante
hehe a .to
was missing in the DoLa body, opening a PR to fix it
(TL;DR the logits were on the same device as the model inputs, but the auxiliary DoLa tensors were on the same device as the lm head. On multi-gpu settings, that's not the same device.)
DoLa decoding on Mixtral model with multi-GPU setup returns error:
System Info
transformers
version: 4.43.0.dev0Who can help?
No response
Information
Tasks
examples
folder (such as GLUE/SQuAD, ...)Reproduction
This script uses DoLa decoding new feature with multi GPU setup. Run the file test.py:
with this command
CUDA_VISIBLE_DEVICES=0,1,2,3 python test.py
Expected behavior
No errors.