Multi GPU Inference Failure: Rotary Embedding Bug

Nobody has claimed this yet.

> **Machine:**

8x NVIDIA RTX A6000 - 49GiB

> **Command:**

CUDA_LAUNCH_BLOCKING=1 srun --nodelist=worker_name --gres=gpu:2 -c 16 python3 -m fastchat.serve.cli --model-path lmsys/longchat-7b-32k-v1.5 --num-gpus 2 --max-gpu-memory 45GiB

> **Error Log:**

/opt/conda/conda-bld/pytorch_1659484810403/work/aten/src/ATen/native/cuda/IndexKernel.cu:91: operator(): block: [6,0,0], thread: [127,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
Traceback (most recent call last):
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/runpy.py", line 194, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/runpy.py", line 87, in _run_code
exec(code, run_globals)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/fastchat/serve/cli.py", line 304, in
main(args)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/fastchat/serve/cli.py", line 227, in main
chat_loop(
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/fastchat/serve/inference.py", line 532, in chat_loop
outputs = chatio.stream_output(output_stream)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/fastchat/serve/cli.py", line 63, in stream_output
for outputs in output_stream:
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 43, in generator_context
response = gen.send(None)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/fastchat/serve/inference.py", line 132, in generate_stream
out = model(input_ids=start_ids, use_cache=True)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 1183, in forward
outputs = self.model(
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 1070, in forward
layer_outputs = decoder_layer(
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 798, in forward
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
output = module._old_forward(*args, **kwargs)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 404, in forward
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
File "/home/stud/hannan/anaconda3/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 232, in apply_rotary_pos_emb
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
RuntimeError: CUDA error: device-side assert triggered

Reproduce the reported command, then trace the inference path through fastchat/serve/cli.py and fastchat/serve/inference.py into transformers/models/llama/modeling_llama.py, where the rotary embedding indexing error is shown. Done means the reported multi-GPU inference completes without the CUDA device-side assertion.

Written by the indexing model from the issue text.