import transformers
import torch
from composer.utils import dist
def main():
dist.initialize_dist('gpu')
name = 'meta-llama/Meta-Llama-3-8B-Instruct'
tokenizer = transformers.AutoTokenizer.from_pretrained(name)
pad_token_id = tokenizer.eos_token_id
model = transformers.AutoModelForCausalLM.from_pretrained(name)
rank = dist.get_global_rank()
model.to(f'cuda:{rank}')
if dist.get_global_rank() == 0:
content = 'Write one short sentence.'
else:
content = 'Write one long paragraph.'
messages = [
{
'role': 'user',
'content': content,
}
]
tokenized_messages = tokenizer.apply_chat_template(messages, return_tensors='pt')
padded_messages = torch.cat(
[
torch.LongTensor((4096 - 20) * [pad_token_id]),
tokenized_messages[0], # [seq]
],
dim=0,
)
padded_messages = padded_messages.unsqueeze(0)
padded_messages = padded_messages.to(f'cuda:{rank}')
attention_mask = ~(padded_messages == pad_token_id)
attention_mask = attention_mask.to(f'cuda:{rank}')
output = model.generate(input_ids=padded_messages, attention_mask=attention_mask, synced_gpus=True, max_new_tokens=200)
print(tokenizer.decode(output[0]))
if __name__ == '__main__':
main()Traceback (most recent call last):
File "/mnt/workdisk/danielking/github/multi-gpu.py", line 47, in <module>
main()
File "/mnt/workdisk/danielking/github/multi-gpu.py", line 42, in main
output = model.generate(input_ids=padded_messages, attention_mask=attention_mask, synced_gpus=True, max_new_tokens=200)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
result = self._sample(
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
outputs = self(**model_inputs, return_dict=True)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1189, in forward
outputs = self.model(
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1001, in forward
layer_outputs = decoder_layer(
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 734, in forward
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 660, in forward
attn_output = torch.nn.functional.scaled_dot_product_attention(
RuntimeError: The expanded size of the tensor (4105) must match the existing size (4104) at non-singleton dimension 3. Target sizes: [1, 32, 1, 4105]. Tensor sizes: [1, 1, 1, 4104]
Multi GPU generate does not error..
System Info
transformersversion: 4.44.0composeras the distributed launcherWho can help?
@gante @ArthurZucker
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
Run llama with
synced_gpus=Trueand an attention mask. This worked fine on transformers 4.40.2 (and 4.41.x), but no longer works. The use of Composer for the dist stuff is just convenience, shouldn't affect anything to swap in a different distributed launcher, etc.This results in
Expected behavior
Multi GPU generate does not error..