AssertionError when trying to use tensor parallelism

[1protagoras1]

Hello, I'm having this issue where I can't make tensor parelism work on my pc. My system is running Ubuntu 22.04, two rtx 3090 and CUDA 12.04. I tryed several models but I had no luck so far. The one in this example is Qwen2.5 8.0 bpw, this is the code Im running:

import flash_attn
from exllamav2 import (
    ExLlamaV2,
    ExLlamaV2Config,
    ExLlamaV2Cache_Q4,
    ExLlamaV2Tokenizer,
)
from exllamav2.generator import (
    ExLlamaV2BaseGenerator,
    ExLlamaV2Sampler
)

model_directory = '/home/Models/decoder/Qwen2.5-Coder'

config = ExLlamaV2Config()
config.model_dir = model_directory
config.prepare()
config.max_seq_len = 16384

model = ExLlamaV2(config)
model.load_tp()


tokenizer = ExLlamaV2Tokenizer(config)

cache = ExLlamaV2Cache_Q4(model)

settings = ExLlamaV2Sampler.Settings()
generator = ExLlamaV2BaseGenerator(model, cache, tokenizer)

max_new_tokens = 2000

prompt = 'Write a Python function that calculates the sum of all even numbers between 5 and 100'
output = generator.generate_simple(prompt, settings, max_new_tokens, seed = 1234, stop_token=151643, completion_only=True)
print(output)

The error message I get is:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[2], line 2
      1 prompt = 'Hey'
----> 2 output = generator.generate_simple(prompt, settings, max_new_tokens, seed = 1234, stop_token=151643, completion_only=True)
      3 print(output)

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/exllamav2/generator/base.py:260, in ExLlamaV2BaseGenerator.generate_simple(self, prompt, gen_settings, num_tokens, seed, token_healing, encode_special_tokens, decode_special_tokens, loras, stop_token, add_bos, abort_event, input_embeddings, completion_only, filters, filter_prefer_eos)
    257 if self.abort_event and self.abort_event.is_set():
    258     break
--> 260 logits = self.model.forward(self.sequence_ids[:, -1:],
    261                             self.cache,
    262                             input_mask = mask,
    263                             loras = loras,
    264                             position_offsets = position_offsets,
    265                             indexed_embeddings = input_embeddings).float().cpu()
    267 token, ptokens, pprobs, prob, eos = \
    268 ExLlamaV2Sampler.sample(
    269     logits,
   (...)
    276     filter_prefer_eos = filter_prefer_eos
    277 )
    279 if unhealed_token is not None:

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/exllamav2/model.py:871, in ExLlamaV2.forward(self, input_ids, cache, input_mask, preprocess_only, last_id_only, loras, return_last_state, position_offsets, abort_event, cpu_logits, **kwargs)
    868 _last_id_only = last_id_only
    869 _preprocess_only = preprocess_only or (chunk_end < q_len and last_id_only)
--> 871 r = self.forward_chunk(
    872     input_ids = input_ids[:, chunk_begin : chunk_end],
    873     cache = cache,
    874     input_mask = input_mask,
    875     preprocess_only = _preprocess_only,
    876     last_id_only = _last_id_only,
    877     loras = loras,
    878     return_last_state = return_last_state and remaining_q_len <= chunk_size,
    879     position_offsets = position_offsets,
    880     abort_event = abort_event,
    881     **kwargs
    882 )
    884 if abort_event and abort_event.is_set(): return
    886 if not _preprocess_only:

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/exllamav2/model.py:977, in ExLlamaV2.forward_chunk(self, input_ids, cache, input_mask, preprocess_only, last_id_only, loras, return_last_state, position_offsets, abort_event, attn_params, extract_state_indices, **kwargs)
    974 if n_device is not None and n_device != device and n_device >= 0:
    975     x = safe_move_tensor(x, n_device, non_blocking = True)
--> 977 x = module.forward(x, cache = cache, attn_params = attn_params, past_len = past_len, loras = loras, **kwargs)
    979 if preprocess_only and idx == self.last_kv_layer_idx:
    980     x = None

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/exllamav2/attn.py:980, in ExLlamaV2Attention.forward(self, hidden_states, cache, attn_params, past_len, intermediates, loras, **kwargs)
    978 if self.is_tp:
    979     if cache is not None and use_flash_attn:
--> 980         return self.forward_tp(
    981             hidden_states,
    982             cache,
    983             attn_params,
    984             past_len,
    985             intermediates,
    986             loras,
    987             **kwargs,
    988         )
    989     else:
    990         # TODO: Can't use the optimized forward function because it writes directly to a fixed output
    991         #   tensor, and flash-attn currently has a bug that prevents that from working when q_len == 1
    992         return self.forward_tp_old(
    993             hidden_states,
    994             cache,
   (...)
    999             **kwargs,
   1000         )

File ~/miniconda3/envs/NLP/lib/python3.10/site-packages/exllamav2/attn.py:1135, in ExLlamaV2Attention.forward_tp(self, hidden_states, cache, attn_params, past_len, intermediates, loras, **kwargs)
   1132 cfg = self.model.config
   1133 ctx = self.model.tp_context
-> 1135 assert not cache or cache.q_block != 1, \
   1136     "Models with odd key/value dims not supported in TP mode with quantized cache"
   1137 assert not self.sliding_window, \
   1138     "Sliding window not supported in TP mode"
   1140 attn_params.prep_tp(self.model)

AssertionError: Models with odd key/value dims not supported in TP mode with quantized cache

I apologize beforehand if if I'm missing something trivial, I tryed to look for similar errors but couldn't find any. I'm not a coder so I have a lot of trouble figuring out what I'm doing wrong.

[1protagoras1]

Actually I found the solution, apparently I wasn't able to find the example in exllamav2 directory. It turns out I had the wrong cache type all along, when I should have used "ExLlamaV2Cache_TP".