Multi-thread inference failed when load_in_8bit with chatglm2
Author: MickeyJ1002
Pull request: no

              > till persists

I have tried almost all version combinations and this problem still persists.
Only when transformers==4.31.0, [Multi-thread issue ](https://github.com/huggingface/transformers/issues/25197)was resolved but the GPU memory was not reduced [# 25228](https://github.com/huggingface/transformers/issues/25228)

_Originally posted by @MickeyJ1002 in https://github.com/huggingface/transformers/issues/25197#issuecomment-1813702796_
            

Comments:
Hi @MickeyJ1002, thanks for raising this issue! 

Could you make sure to follow the issue template, and provide a minimal code snippet we can reproduce the error and the running environment it occurs in? Are my statements about the respective bugs [here](https://github.com/huggingface/transformers/issues/25197#issuecomment-1665352607) still correct?

The linked issues are useful. However to be able to debug we really need to know precisely the error your encountering, with a detailed description of the bug or unexpected behaviour, any observations and error messages, as well as what you've tried so far. 

system info:

Python version: 3.8.10
gpus_num: 4

(transformers==4.33.1 accelerate==0.23.0 bitsandbytes==0.42.1)
(transformers==4.33.1 accelerate==0.21.0 bitsandbytes==0.37.1)
(transformers==4.30.2 accelerate==0.21.0 bitsandbytes==0.37.1)
(transformers==4.32.0 accelerate==0.21.0 bitsandbytes==0.37.1)
..... I've tried all of the above combinations, and the same error info arised (Multi-thread inference failed).

My test code:
```
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
from transformers import AutoTokenizer, AutoModel, TextIteratorStreamer
from threading import Thread, currentThread

model_path = "THUDM/chatglm2-6b"

device_map = {'transformer.embedding': 0, 'transformer.encoder.final_layernorm': 0, 
              'transformer.output_layer': 0, 'transformer.rotary_pos_emb': 0, 
              'lm_head': 0, 
              'transformer.encoder.layers.0': 0, 'transformer.encoder.layers.1': 0, 
              'transformer.encoder.layers.2': 0, 'transformer.encoder.layers.3': 1, 
              'transformer.encoder.layers.4': 1, 'transformer.encoder.layers.5': 1, 
              'transformer.encoder.layers.6': 1, 'transformer.encoder.layers.7': 1, 
              'transformer.encoder.layers.8': 1, 'transformer.encoder.layers.9': 1, 
              'transformer.encoder.layers.10': 1, 'transformer.encoder.layers.11': 1, 
              'transformer.encoder.layers.12': 2, 'transformer.encoder.layers.13': 2, 
              'transformer.encoder.layers.14': 2, 'transformer.encoder.layers.15': 2, 
              'transformer.encoder.layers.16': 2, 'transformer.encoder.layers.17': 2, 
              'transformer.encoder.layers.18': 2, 'transformer.encoder.layers.19': 2, 
              'transformer.encoder.layers.20': 2, 'transformer.encoder.layers.21': 3, 
              'transformer.encoder.layers.22': 3, 'transformer.encoder.layers.23': 3, 
              'transformer.encoder.layers.24': 3, 'transformer.encoder.layers.25': 3, 
              'transformer.encoder.layers.26': 3, 'transformer.encoder.layers.27': 3}

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, device_map='auto', trust_remote_code=True, load_in_8bit=True)

def infer(prompt):
    inputs = tokenizer(prompt, return_tensors='pt')
    inputs = inputs.to(model.device)
    print('before generate')
    t = currentThread()
    streamer = TextIteratorStreamer(tokenizer)
    generation_kwargs = dict(inputs, streamer=streamer, max_length=2048)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    print("------******-------")
    for new_text in streamer:
        print("thread id:", t.ident ,"new text:",new_text)
    print("------******-------")

if __name__ == '__main__':
    prompt1 = '写一篇关于黄鹤楼的800字作文'
    prompt2 = 'Describe each state in the United States in detail'
    t1 = Thread(target=infer, args=(prompt1,))
    t2 = Thread(target=infer, args=(prompt2,))
    t1.start()
    # time.sleep(5)
    t2.start()
    t1.join()
    t2.join()
```

error info:
```
Traceback (most recent call last):
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/transformers/generation/utils.py", line 1602, in generate
    return self.greedy_search(
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/transformers/generation/utils.py", line 2450, in greedy_search
    outputs = self(
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/wujianjian/.cache/huggingface/modules/transformers_modules/base/modeling_chatglm.py", line 928, in forward
    transformer_outputs = self.transformer(
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wujianjian/.cache/huggingface/modules/transformers_modules/base/modeling_chatglm.py", line 824, in forward
    hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wujianjian/.cache/huggingface/modules/transformers_modules/base/modeling_chatglm.py", line 637, in forward
    layer_ret = layer(
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/wujianjian/.cache/huggingface/modules/transformers_modules/base/modeling_chatglm.py", line 562, in forward
    mlp_output = self.mlp(layernorm_output)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/wujianjian/.cache/huggingface/modules/transformers_modules/base/modeling_chatglm.py", line 498, in forward
    output = self.dense_4h_to_h(intermediate_parallel)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/bitsandbytes/nn/modules.py", line 242, in forward
    out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/bitsandbytes/autograd/_functions.py", line 488, in matmul
    return MatMul8bitLt.apply(A, B, out, bias, state)
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/data/qinbaoshuai/software/anaconda3/envs/qlora_infer/lib/python3.8/site-packages/bitsandbytes/autograd/_functions.py", line 397, in forward
    output += torch.matmul(subA, state.subB)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x15 and 6x4096)
```

@amyeroberts 

> Hi @MickeyJ1002, thanks for raising this issue!
> 
> Could you make sure to follow the issue template, and provide a minimal code snippet we can reproduce the error and the running environment it occurs in? Are my statements about the respective bugs [here](https://github.com/huggingface/transformers/issues/25197#issuecomment-1665352607) still correct?
> 
> The linked issues are useful. However to be able to debug we really need to know precisely the error your encountering, with a detailed description of the bug or unexpected behaviour, any observations and error messages, as well as what you've tried so far.

Exactly, the statements about the respective bugs [here](https://github.com/huggingface/transformers/issues/25197#issuecomment-1665352607) are still correct

In this case - I'm going to hand over to @younesbelkada who was on the case last time 

Hi @MickeyJ1002 
Apologies for the delay !
To load 8-bit / 4-bit models with multi-thread you need to load entirely the model on each GPU. Can you try out the snippet below with latest bitsandbytes / accelerate / transformers (`pip install -U accelerate transformers bitsandbytes`)

```python
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
from transformers import AutoTokenizer, AutoModel, TextIteratorStreamer
from accelerate import PartialState
from threading import Thread, currentThread

model_path = "THUDM/chatglm2-6b"

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, device_map={'': PartialState().process_index}, trust_remote_code=True, load_in_4bit=True)

def infer(prompt):
    inputs = tokenizer(prompt, return_tensors='pt')
    inputs = inputs.to(model.device)
    print('before generate')
    t = currentThread()
    streamer = TextIteratorStreamer(tokenizer)
    generation_kwargs = dict(inputs, streamer=streamer, max_length=2048)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    print("------******-------")
    for new_text in streamer:
        print("thread id:", t.ident ,"new text:",new_text)
    print("------******-------")

if __name__ == '__main__':
    prompt1 = '写一篇关于黄鹤楼的800字作文'
    prompt2 = 'Describe each state in the United States in detail'
    t1 = Thread(target=infer, args=(prompt1,))
    t2 = Thread(target=infer, args=(prompt2,))
    t1.start()
    # time.sleep(5)
    t2.start()
    t1.join()
    t2.join()
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.