`run_benchmark.py` runs error when using `--distribute-model`

When running run_benchmark.py without --distribute-model, all works well.
While after adding --distribute-model and keeping others same as before, it occured the following error.
@nailimixaM So how to support tensor parallel in run_benchmark.py? Thanks in advance.

Running benchmarking of a sliced model.
PyTorch device: cuda
Number of available cuda devices: 2
Loading sliced meta-llama/Llama-2 model from /sliced_models/llama2-7b/Llama-2_0.2.pt with sparsity 0.2
...
Preparing dataloader done
Benchmarking:   0%|                                                                                                                                                                                        | 1/4096 [00:01<1:26:06,  1.26s/it]
Traceback (most recent call last):
  File "/TransformerCompression/experiments/run_benchmark.py", line 151, in <module>
    benchmarking_main(benchmarking_args)
  File "/TransformerCompression/experiments/run_benchmark.py", line 138, in benchmarking_main
    results = gpu_utils.benchmark(model_adapter, next(iter(train_loader)))
  File "/TransformerCompression/src/slicegpt/gpu_utils.py", line 130, in benchmark
    output = model_adapter.model(input_ids_i, past_key_values=cache["past"], attention_mask=attention_mask_i)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 165, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py", line 1183, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py", line 1070, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1561, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 165, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/bigdata/zhaoyang/github/TransformerCompression/src/slicegpt/adapters/llama_adapter.py", line 55, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 165, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py", line 710, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/cache_utils.py", line 127, in update
    self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_CUDA_cat)

Start at experiments/run_benchmark.py in benchmarking_main, then trace gpu_utils.benchmark into the model call and src/slicegpt/adapters/llama_adapter.py. Reproduce with --distribute-model and inspect the cache update reported in transformers/cache_utils.py for the cuda:0/cuda:1 mismatch. Done means the distributed benchmark runs without this device error.