Repository: pytorch
Issue: RuntimeError in run_streaming_llama.py When Using Accelerate with Streaming LLMa Model on A4500 GPU
Labels: needs reproduction,module: error checking,triaged

### 🐛 Describe the bug

Referring to the issue https://github.com/mit-han-lab/streaming-llm/issues/37#issue-1940692615

Description:
When running the run_streaming_llama.py script with the --enable_streaming flag, I encountered a RuntimeError related to CUDA and the PyTorch Accelerate library.

Steps to Reproduce:
Set the environment variable: CUDA_VISIBLE_DEVICES=0
Run the following command:
```
python examples/run_streaming_llama.py --enable_streaming
```

Expected Behavior:
The script should run successfully and provide streaming inference results.

Actual Behavior:
The script crashes with the following error:
```
RuntimeError: handle_0 INTERNAL ASSERT FAILED at "../c10/cuda/driver_api.cpp":15, please report a bug to PyTorch.
```

GPU information
```
CUDA version: 11.7
GPU: NVIDIA RTX A4500
GPU driver information: 
NVRM version: NVIDIA UNIX x86_64 Kernel Module  525.78.01  Mon Dec 26 05:58:42 UTC 2022
GCC version:  gcc version 9.4.0 (Ubuntu 9.4.0-1ubuntu1~20.04.1)
```

OS information:
```
Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.5 LTS
Release:	20.04
Codename:	focal
```

Full error log:
```
(streaming) zli@aes:~/streaming-llm$ CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py  --enable_streaming
Loading model from lmsys/vicuna-13b-v1.3 ...
You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. If you see this, DO NOT PANIC! This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thouroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:14<00:00,  4.74s/it]
Loading data from data/mt_bench.jsonl ...
StartRecentKVCache: 4, 2000

USER: Compose an engaging travel blog post about a recent trip to Hawaii, highlighting cultural experiences and must-see attractions.

ASSISTANT: 🌴🌺 Aloha from the beautiful islands of Hawaii! I recently had the opportunity to explore the diverse culture and stunning landscapes of this tropical paradise, and I can't wait to share my adventures with you.

🏝️ First things first, let's talk about the beaches. Hawaii is known for its pristine white sand and crystal clear waters, and I was lucky enough to visit some of the most breathtaking beaches on the islands. My favorite was Waimea Bay on the North Shore of Oahu. The waves were perfect for surfing, and the surrounding cliffs and lush greenery made for a postcard-worthy backdrop.

🌊 But Hawaii isn't just about the beaches. The culture of the islands is rich and vibrant, and I made it a point to immerse myself in it during my trip. I attended a traditional Hawaiian luau on the Big Island, where I learned about the history and customs of the native people. I also tried my hand at hula dancing and even got to sample some delicious Hawaiian cuisine, including poke bowls and shaved ice.

🏯 Another must-see attraction in Hawaii is the Hawaii Volcanoes National Park on the Big Island. The park is home to two active volcanoes, Kilauea and Mauna Loa, and offers a unique opportunity to see the power of nature up close. I hiked through the park and was amazed by the steam vents, lava tubes, and volcanic craters. It was truly a once-in-a-lifetime experience.

🌺 In addition to the natural wonders of Hawaii, the islands are also home to some of the most beautiful gardens and parks in the world. The Maui Tropical Plantation offered a glimpse into the island's agricultural history, while the Hilo Tropical Botanical Garden showcased the stunning flora and fauna of the area.

🏯 In conclusion, Hawaii is a destination that offers something for everyone. From the stunning beaches to the rich culture and unique attractions, there's no shortage of things to see and do on these beautiful islands. If you're planning a trip to Hawaii, be sure to add these must-see attractions to your itinerary. Mahalo for reading, and aloha!

USER: Rewrite your previous response. Start every sentence with the letter A.

ASSISTANT: 🌴🌺 Aloha from the beautiful islands of Hawaii! I recently had the opportunity to explore the diverse culture and stunning landscapes of this tropical paradise, and I can't wait to share my adventures with you.

🏝️ First things first, let's talk about the beaches. Hawaii is known for its pristine white sand and crystal clear waters, and I was lucky enough to visit some of the most breathtaking beaches on the islands. My favorite was Waimea Bay on the North Shore of Oahu. The waves were perfect for surfing, and the surrounding cliffs and lush greenery made for a postcard-worthy backdrop.

🌊 But Hawaii isn't just about the beaches. The culture of the islands is rich and vibrant, and I made it a point to immerse myself in it during my trip. I attended a traditional Hawaiian luau on the Big Island, where I learned about the history and customs of the native people. I also tried my hand at hula dancing and even got to sample some delicious Hawaiian cuisine, including poke bowls and shaved ice.

🌺 In addition to the natural wonders of Hawaii, the islands are also home to some of the most stunning gardens and parks in Traceback (most recent call last):
  File "examples/run_streaming_llama.py", line 122, in <module>
    main(args)
  File "examples/run_streaming_llama.py", line 103, in main
    streaming_inference(
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "examples/run_streaming_llama.py", line 73, in streaming_inference
    past_key_values = greedy_generate(
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "examples/run_streaming_llama.py", line 30, in greedy_generate
    outputs = model(
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 838, in forward
    logits = self.lm_head(hidden_states)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/accelerate/hooks.py", line 160, in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/accelerate/hooks.py", line 286, in pre_forward
    set_module_tensor_to_device(
  File "/home/zli/anaconda3/envs/streaming/lib/python3.8/site-packages/accelerate/utils/modeling.py", line 317, in set_module_tensor_to_device
    new_value = value.to(device)
RuntimeError: handle_0 INTERNAL ASSERT FAILED at "../c10/cuda/driver_api.cpp":15, please report a bug to PyTorch.
```

### Versions

torch version:                   2.1.0

conda create -yn streaming python=3.8
conda activate streaming
pip install torch torchvision torchaudio
pip install transformers==4.33.0 accelerate datasets evaluate wandb scikit-learn scipy sentencepiece
python setup.py develop
git clone git@github.com:mit-han-lab/streaming-llm.git
cd streaming-llm/
python setup.py develop

cc @malfet