We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
exllama gibberish output
### Describe the bug

I'm running into issues using exllama with the textgen UI. I have exllama running well using its webui, but when I load the same model into textgen UI, I only get gibberish as output. Llama.cpp works fine as a backend. I looked at the output in the console with --verbose and the input prompt looks fine. I tried multiple different sampler presets with no change. I'm using an AMD Vega 56 and I talked with another person on Discord with the same problem on an AMD card so it is probably an AMD-specific issue. I followed the oobabooga ROCm rentry guide in the README to set up textgen. One thing I noticed is the bitsandbytes package in the rocm fork is a little outdated, so could that be the issue?

### Is there an existing issue for this?

- [X] I have searched the existing issues

### Reproduction

Set up textgen webUI using instructions in the README and the ROCm installation guide in Arch Linux on a system with an AMD card. Clone exllama into the repositories folder and install the requirements. Test that exllama works with its on WebGUI. Load a model using exllama in textgen webUI, then generate output from any prompt. 

### Screenshot

_No response_

### Logs

```shell
No log available
```


### System Info

```shell
Arch Linux
AMD 5600X
AMD Vega 56
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `modules/exllama_hf.py`
Content:
```
1 import os
2 from pathlib import Path
3 from typing import Any, Dict, Optional, Union
4 
5 import torch
6 from torch.nn import CrossEntropyLoss
7 from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel
8 from transformers.modeling_outputs import CausalLMOutputWithPast
9 
10 from modules import shared
11 from modules.logging_colors import logger
12 
13 try:
14     from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
15 except:
16     logger.warning('Exllama module failed to load. Will attempt to load from repositories.')
17     try:
18         from modules.relative_imports import RelativeImport
19 
20         with RelativeImport("repositories/exllama"):
21             from model import ExLlama, ExLlamaCache, ExLlamaConfig
22     except:
23         logger.error("Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.")
24         raise
25 
26 
27 class ExllamaHF(PreTrainedModel):
28     def __init__(self, config: ExLlamaConfig):
29         super().__init__(PretrainedConfig())
30         self.ex_config = config
31         self.ex_model = ExLlama(self.ex_config)
32         self.generation_config = GenerationConfig()
33         self.lora = None
34 
35     def _validate_model_class(self):
36         pass
37 
38     def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
39         pass
40 
41     def prepare_inputs_for_generation(self, input_ids, **kwargs):
42         return {'input_ids': input_ids, **kwargs}
43 
44     @property
45     def device(self) -> torch.device:
46         return torch.device(0)
47 
48     def __call__(self, *args, **kwargs):
49         # TODO: Some decoding methods (such as Contrastive Search) may not work at this time
50         assert len(args) == 0, 'no *args should be passed to forward'
51         use_cache = kwargs.get('use_cache', True)
52         labels = kwargs.get('labels', None)
53         seq = kwargs['input_ids'][0].tolist()
54         cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None
55         if cache is None:
56             cache = ExLlamaCache(self.ex_model)
57             self.ex_model.forward(torch.tensor([seq[:-1]], dtype=torch.long), cache, preprocess_only=True, lora=self.lora)
58 
59         logits = self.ex_model.forward(torch.tensor([seq[-1:]], dtype=torch.long), cache, lora=self.lora).to(kwargs['input_ids'].device)
60 
61         loss = None
62         if labels is not None:
63             # Shift so that tokens < n predict n
64             shift_logits = logits[..., :-1, :].contiguous()
65             shift_labels = labels[..., 1:].contiguous()
66             # Flatten the tokens
67             loss_fct = CrossEntropyLoss()
68             shift_logits = shift_logits.view(-1, logits.shape[-1])
69             shift_labels = shift_labels.view(-1)
70             # Enable model parallelism
71             shift_labels = shift_labels.to(shift_logits.device)
72             loss = loss_fct(shift_logits, shift_labels)
73 
74         return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None)
75 
76     @classmethod
77     def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
78         assert len(model_args) == 0 and len(kwargs) == 0, "extra args is currently not supported"
79         if isinstance(pretrained_model_name_or_path, str):
80             pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
81 
82         pretrained_model_name_or_path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)
83         config = ExLlamaConfig(pretrained_model_name_or_path / 'config.json')
84 
85         # from 'oobabooga/text-generation-webui/modules/exllama.py'
86         weight_path = None
87         for ext in ['.safetensors', '.pt', '.bin']:
88             found = list(pretrained_model_name_or_path.glob(f"*{ext}"))
89             if len(found) > 0:
90                 weight_path = found[-1]
91                 break
92         assert weight_path is not None, f'could not find weight in "{pretrained_model_name_or_path}"'
93 
94         config.model_path = str(weight_path)
95         config.max_seq_len = shared.args.max_seq_len
96         config.compress_pos_emb = shared.args.compress_pos_emb
97         if shared.args.gpu_split:
98             config.set_auto_map(shared.args.gpu_split)
99             config.gpu_peer_fix = True
100 
101         # This slowes down a bit but align better with autogptq generation.
102         # TODO: Should give user choice to tune the exllama config
103         # config.fused_attn = False
104         # config.fused_mlp_thd = 0
105 
106         return ExllamaHF(config)
107
```

Path: `modules/exllama.py`
Content:
```
1 import sys
2 from pathlib import Path
3 
4 from modules import shared
5 from modules.logging_colors import logger
6 
7 try:
8     from exllama.generator import ExLlamaGenerator
9     from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
10     from exllama.tokenizer import ExLlamaTokenizer
11 except:
12     logger.warning('Exllama module failed to load. Will attempt to load from repositories.')
13     try:
14         from modules.relative_imports import RelativeImport
15 
16         with RelativeImport("repositories/exllama"):
17             from generator import ExLlamaGenerator
18             from model import ExLlama, ExLlamaCache, ExLlamaConfig
19             from tokenizer import ExLlamaTokenizer
20     except:
21         logger.error("Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.")
22         raise
23 
24 
25 class ExllamaModel:
26     def __init__(self):
27         pass
28 
29     @classmethod
30     def from_pretrained(self, path_to_model):
31 
32         path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)
33         tokenizer_model_path = path_to_model / "tokenizer.model"
34         model_config_path = path_to_model / "config.json"
35 
36         # Find the model checkpoint
37         model_path = None
38         for ext in ['.safetensors', '.pt', '.bin']:
39             found = list(path_to_model.glob(f"*{ext}"))
40             if len(found) > 0:
41                 if len(found) > 1:
42                     logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')
43 
44                 model_path = found[-1]
45                 break
46 
47         config = ExLlamaConfig(str(model_config_path))
48         config.model_path = str(model_path)
49         config.max_seq_len = shared.args.max_seq_len
50         config.compress_pos_emb = shared.args.compress_pos_emb
51         if shared.args.gpu_split:
52             config.set_auto_map(shared.args.gpu_split)
53             config.gpu_peer_fix = True
54 
55         model = ExLlama(config)
56         tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))
57         cache = ExLlamaCache(model)
58         generator = ExLlamaGenerator(model, tokenizer, cache)
59 
60         result = self()
61         result.config = config
62         result.model = model
63         result.cache = cache
64         result.tokenizer = tokenizer
65         result.generator = generator
66         return result, result
67 
68     def generate_with_streaming(self, prompt, state):
69         self.generator.settings.temperature = state['temperature']
70         self.generator.settings.top_p = state['top_p']
71         self.generator.settings.top_k = state['top_k']
72         self.generator.settings.typical = state['typical_p']
73         self.generator.settings.token_repetition_penalty_max = state['repetition_penalty']
74         self.generator.settings.token_repetition_penalty_sustain = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']
75         if state['ban_eos_token']:
76             self.generator.disallow_tokens([self.tokenizer.eos_token_id])
77         else:
78             self.generator.disallow_tokens(None)
79 
80         self.generator.end_beam_search()
81         ids = self.generator.tokenizer.encode(prompt)
82         self.generator.gen_begin_reuse(ids)
83         initial_len = self.generator.sequence[0].shape[0]
84         has_leading_space = False
85         for i in range(state['max_new_tokens']):
86             token = self.generator.gen_single_token()
87             if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):
88                 has_leading_space = True
89 
90             decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])
91             if has_leading_space:
92                 decoded_text = ' ' + decoded_text
93 
94             yield decoded_text
95             if token.item() == self.generator.tokenizer.eos_token_id or shared.stop_everything:
96                 break
97 
98     def generate(self, prompt, state):
99         output = ''
100         for output in self.generate_with_streaming(prompt, state):
101             pass
102 
103         return output
104 
105     def encode(self, string, **kwargs):
106         return self.tokenizer.encode(string)
107
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/modules/exllama.py b/modules/exllama.py
--- a/modules/exllama.py
+++ b/modules/exllama.py
@@ -1,6 +1,8 @@
 import sys
 from pathlib import Path
 
+from torch import version as torch_version
+
 from modules import shared
 from modules.logging_colors import logger
 
@@ -51,6 +53,12 @@
         if shared.args.gpu_split:
             config.set_auto_map(shared.args.gpu_split)
             config.gpu_peer_fix = True
+        if torch_version.hip:
+            config.rmsnorm_no_half2 = True
+            config.rope_no_half2 = True
+            config.matmul_no_half2 = True
+            config.silu_no_half2 = True
+
 
         model = ExLlama(config)
         tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))
diff --git a/modules/exllama_hf.py b/modules/exllama_hf.py
--- a/modules/exllama_hf.py
+++ b/modules/exllama_hf.py
@@ -97,6 +97,11 @@
         if shared.args.gpu_split:
             config.set_auto_map(shared.args.gpu_split)
             config.gpu_peer_fix = True
+        if torch.version.hip:
+            config.rmsnorm_no_half2 = True
+            config.rope_no_half2 = True
+            config.matmul_no_half2 = True
+            config.silu_no_half2 = True
 
         # This slowes down a bit but align better with autogptq generation.
         # TODO: Should give user choice to tune the exllama config

{"golden_diff": "diff --git a/modules/exllama.py b/modules/exllama.py\n--- a/modules/exllama.py\n+++ b/modules/exllama.py\n@@ -1,6 +1,8 @@\n import sys\n from pathlib import Path\n \n+from torch import version as torch_version\n+\n from modules import shared\n from modules.logging_colors import logger\n \n@@ -51,6 +53,12 @@\n         if shared.args.gpu_split:\n             config.set_auto_map(shared.args.gpu_split)\n             config.gpu_peer_fix = True\n+        if torch_version.hip:\n+            config.rmsnorm_no_half2 = True\n+            config.rope_no_half2 = True\n+            config.matmul_no_half2 = True\n+            config.silu_no_half2 = True\n+\n \n         model = ExLlama(config)\n         tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))\ndiff --git a/modules/exllama_hf.py b/modules/exllama_hf.py\n--- a/modules/exllama_hf.py\n+++ b/modules/exllama_hf.py\n@@ -97,6 +97,11 @@\n         if shared.args.gpu_split:\n             config.set_auto_map(shared.args.gpu_split)\n             config.gpu_peer_fix = True\n+        if torch.version.hip:\n+            config.rmsnorm_no_half2 = True\n+            config.rope_no_half2 = True\n+            config.matmul_no_half2 = True\n+            config.silu_no_half2 = True\n \n         # This slowes down a bit but align better with autogptq generation.\n         # TODO: Should give user choice to tune the exllama config\n", "issue": "exllama gibberish output\n### Describe the bug\n\nI'm running into issues using exllama with the textgen UI. I have exllama running well using its webui, but when I load the same model into textgen UI, I only get gibberish as output. Llama.cpp works fine as a backend. I looked at the output in the console with --verbose and the input prompt looks fine. I tried multiple different sampler presets with no change. I'm using an AMD Vega 56 and I talked with another person on Discord with the same problem on an AMD card so it is probably an AMD-specific issue. I followed the oobabooga ROCm rentry guide in the README to set up textgen. One thing I noticed is the bitsandbytes package in the rocm fork is a little outdated, so could that be the issue?\n\n### Is there an existing issue for this?\n\n- [X] I have searched the existing issues\n\n### Reproduction\n\nSet up textgen webUI using instructions in the README and the ROCm installation guide in Arch Linux on a system with an AMD card. Clone exllama into the repositories folder and install the requirements. Test that exllama works with its on WebGUI. Load a model using exllama in textgen webUI, then generate output from any prompt. \n\n### Screenshot\n\n_No response_\n\n### Logs\n\n```shell\nNo log available\n```\n\n\n### System Info\n\n```shell\nArch Linux\r\nAMD 5600X\r\nAMD Vega 56\n```\n\n", "before_files": [{"content": "import os\nfrom pathlib import Path\nfrom typing import Any, Dict, Optional, Union\n\nimport torch\nfrom torch.nn import CrossEntropyLoss\nfrom transformers import GenerationConfig, PretrainedConfig, PreTrainedModel\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\nfrom modules import shared\nfrom modules.logging_colors import logger\n\ntry:\n    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig\nexcept:\n    logger.warning('Exllama module failed to load. Will attempt to load from repositories.')\n    try:\n        from modules.relative_imports import RelativeImport\n\n        with RelativeImport(\"repositories/exllama\"):\n            from model import ExLlama, ExLlamaCache, ExLlamaConfig\n    except:\n        logger.error(\"Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.\")\n        raise\n\n\nclass ExllamaHF(PreTrainedModel):\n    def __init__(self, config: ExLlamaConfig):\n        super().__init__(PretrainedConfig())\n        self.ex_config = config\n        self.ex_model = ExLlama(self.ex_config)\n        self.generation_config = GenerationConfig()\n        self.lora = None\n\n    def _validate_model_class(self):\n        pass\n\n    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):\n        pass\n\n    def prepare_inputs_for_generation(self, input_ids, **kwargs):\n        return {'input_ids': input_ids, **kwargs}\n\n    @property\n    def device(self) -> torch.device:\n        return torch.device(0)\n\n    def __call__(self, *args, **kwargs):\n        # TODO: Some decoding methods (such as Contrastive Search) may not work at this time\n        assert len(args) == 0, 'no *args should be passed to forward'\n        use_cache = kwargs.get('use_cache', True)\n        labels = kwargs.get('labels', None)\n        seq = kwargs['input_ids'][0].tolist()\n        cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None\n        if cache is None:\n            cache = ExLlamaCache(self.ex_model)\n            self.ex_model.forward(torch.tensor([seq[:-1]], dtype=torch.long), cache, preprocess_only=True, lora=self.lora)\n\n        logits = self.ex_model.forward(torch.tensor([seq[-1:]], dtype=torch.long), cache, lora=self.lora).to(kwargs['input_ids'].device)\n\n        loss = None\n        if labels is not None:\n            # Shift so that tokens < n predict n\n            shift_logits = logits[..., :-1, :].contiguous()\n            shift_labels = labels[..., 1:].contiguous()\n            # Flatten the tokens\n            loss_fct = CrossEntropyLoss()\n            shift_logits = shift_logits.view(-1, logits.shape[-1])\n            shift_labels = shift_labels.view(-1)\n            # Enable model parallelism\n            shift_labels = shift_labels.to(shift_logits.device)\n            loss = loss_fct(shift_logits, shift_labels)\n\n        return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None)\n\n    @classmethod\n    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):\n        assert len(model_args) == 0 and len(kwargs) == 0, \"extra args is currently not supported\"\n        if isinstance(pretrained_model_name_or_path, str):\n            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)\n\n        pretrained_model_name_or_path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)\n        config = ExLlamaConfig(pretrained_model_name_or_path / 'config.json')\n\n        # from 'oobabooga/text-generation-webui/modules/exllama.py'\n        weight_path = None\n        for ext in ['.safetensors', '.pt', '.bin']:\n            found = list(pretrained_model_name_or_path.glob(f\"*{ext}\"))\n            if len(found) > 0:\n                weight_path = found[-1]\n                break\n        assert weight_path is not None, f'could not find weight in \"{pretrained_model_name_or_path}\"'\n\n        config.model_path = str(weight_path)\n        config.max_seq_len = shared.args.max_seq_len\n        config.compress_pos_emb = shared.args.compress_pos_emb\n        if shared.args.gpu_split:\n            config.set_auto_map(shared.args.gpu_split)\n            config.gpu_peer_fix = True\n\n        # This slowes down a bit but align better with autogptq generation.\n        # TODO: Should give user choice to tune the exllama config\n        # config.fused_attn = False\n        # config.fused_mlp_thd = 0\n\n        return ExllamaHF(config)\n", "path": "modules/exllama_hf.py"}, {"content": "import sys\nfrom pathlib import Path\n\nfrom modules import shared\nfrom modules.logging_colors import logger\n\ntry:\n    from exllama.generator import ExLlamaGenerator\n    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig\n    from exllama.tokenizer import ExLlamaTokenizer\nexcept:\n    logger.warning('Exllama module failed to load. Will attempt to load from repositories.')\n    try:\n        from modules.relative_imports import RelativeImport\n\n        with RelativeImport(\"repositories/exllama\"):\n            from generator import ExLlamaGenerator\n            from model import ExLlama, ExLlamaCache, ExLlamaConfig\n            from tokenizer import ExLlamaTokenizer\n    except:\n        logger.error(\"Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.\")\n        raise\n\n\nclass ExllamaModel:\n    def __init__(self):\n        pass\n\n    @classmethod\n    def from_pretrained(self, path_to_model):\n\n        path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)\n        tokenizer_model_path = path_to_model / \"tokenizer.model\"\n        model_config_path = path_to_model / \"config.json\"\n\n        # Find the model checkpoint\n        model_path = None\n        for ext in ['.safetensors', '.pt', '.bin']:\n            found = list(path_to_model.glob(f\"*{ext}\"))\n            if len(found) > 0:\n                if len(found) > 1:\n                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')\n\n                model_path = found[-1]\n                break\n\n        config = ExLlamaConfig(str(model_config_path))\n        config.model_path = str(model_path)\n        config.max_seq_len = shared.args.max_seq_len\n        config.compress_pos_emb = shared.args.compress_pos_emb\n        if shared.args.gpu_split:\n            config.set_auto_map(shared.args.gpu_split)\n            config.gpu_peer_fix = True\n\n        model = ExLlama(config)\n        tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))\n        cache = ExLlamaCache(model)\n        generator = ExLlamaGenerator(model, tokenizer, cache)\n\n        result = self()\n        result.config = config\n        result.model = model\n        result.cache = cache\n        result.tokenizer = tokenizer\n        result.generator = generator\n        return result, result\n\n    def generate_with_streaming(self, prompt, state):\n        self.generator.settings.temperature = state['temperature']\n        self.generator.settings.top_p = state['top_p']\n        self.generator.settings.top_k = state['top_k']\n        self.generator.settings.typical = state['typical_p']\n        self.generator.settings.token_repetition_penalty_max = state['repetition_penalty']\n        self.generator.settings.token_repetition_penalty_sustain = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']\n        if state['ban_eos_token']:\n            self.generator.disallow_tokens([self.tokenizer.eos_token_id])\n        else:\n            self.generator.disallow_tokens(None)\n\n        self.generator.end_beam_search()\n        ids = self.generator.tokenizer.encode(prompt)\n        self.generator.gen_begin_reuse(ids)\n        initial_len = self.generator.sequence[0].shape[0]\n        has_leading_space = False\n        for i in range(state['max_new_tokens']):\n            token = self.generator.gen_single_token()\n            if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('\u2581'):\n                has_leading_space = True\n\n            decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])\n            if has_leading_space:\n                decoded_text = ' ' + decoded_text\n\n            yield decoded_text\n            if token.item() == self.generator.tokenizer.eos_token_id or shared.stop_everything:\n                break\n\n    def generate(self, prompt, state):\n        output = ''\n        for output in self.generate_with_streaming(prompt, state):\n            pass\n\n        return output\n\n    def encode(self, string, **kwargs):\n        return self.tokenizer.encode(string)\n", "path": "modules/exllama.py"}], "after_files": [{"content": "import os\nfrom pathlib import Path\nfrom typing import Any, Dict, Optional, Union\n\nimport torch\nfrom torch.nn import CrossEntropyLoss\nfrom transformers import GenerationConfig, PretrainedConfig, PreTrainedModel\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\nfrom modules import shared\nfrom modules.logging_colors import logger\n\ntry:\n    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig\nexcept:\n    logger.warning('Exllama module failed to load. Will attempt to load from repositories.')\n    try:\n        from modules.relative_imports import RelativeImport\n\n        with RelativeImport(\"repositories/exllama\"):\n            from model import ExLlama, ExLlamaCache, ExLlamaConfig\n    except:\n        logger.error(\"Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.\")\n        raise\n\n\nclass ExllamaHF(PreTrainedModel):\n    def __init__(self, config: ExLlamaConfig):\n        super().__init__(PretrainedConfig())\n        self.ex_config = config\n        self.ex_model = ExLlama(self.ex_config)\n        self.generation_config = GenerationConfig()\n        self.lora = None\n\n    def _validate_model_class(self):\n        pass\n\n    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):\n        pass\n\n    def prepare_inputs_for_generation(self, input_ids, **kwargs):\n        return {'input_ids': input_ids, **kwargs}\n\n    @property\n    def device(self) -> torch.device:\n        return torch.device(0)\n\n    def __call__(self, *args, **kwargs):\n        # TODO: Some decoding methods (such as Contrastive Search) may not work at this time\n        assert len(args) == 0, 'no *args should be passed to forward'\n        use_cache = kwargs.get('use_cache', True)\n        labels = kwargs.get('labels', None)\n        seq = kwargs['input_ids'][0].tolist()\n        cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None\n        if cache is None:\n            cache = ExLlamaCache(self.ex_model)\n            self.ex_model.forward(torch.tensor([seq[:-1]], dtype=torch.long), cache, preprocess_only=True, lora=self.lora)\n\n        logits = self.ex_model.forward(torch.tensor([seq[-1:]], dtype=torch.long), cache, lora=self.lora).to(kwargs['input_ids'].device)\n\n        loss = None\n        if labels is not None:\n            # Shift so that tokens < n predict n\n            shift_logits = logits[..., :-1, :].contiguous()\n            shift_labels = labels[..., 1:].contiguous()\n            # Flatten the tokens\n            loss_fct = CrossEntropyLoss()\n            shift_logits = shift_logits.view(-1, logits.shape[-1])\n            shift_labels = shift_labels.view(-1)\n            # Enable model parallelism\n            shift_labels = shift_labels.to(shift_logits.device)\n            loss = loss_fct(shift_logits, shift_labels)\n\n        return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None)\n\n    @classmethod\n    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):\n        assert len(model_args) == 0 and len(kwargs) == 0, \"extra args is currently not supported\"\n        if isinstance(pretrained_model_name_or_path, str):\n            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)\n\n        pretrained_model_name_or_path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)\n        config = ExLlamaConfig(pretrained_model_name_or_path / 'config.json')\n\n        # from 'oobabooga/text-generation-webui/modules/exllama.py'\n        weight_path = None\n        for ext in ['.safetensors', '.pt', '.bin']:\n            found = list(pretrained_model_name_or_path.glob(f\"*{ext}\"))\n            if len(found) > 0:\n                weight_path = found[-1]\n                break\n        assert weight_path is not None, f'could not find weight in \"{pretrained_model_name_or_path}\"'\n\n        config.model_path = str(weight_path)\n        config.max_seq_len = shared.args.max_seq_len\n        config.compress_pos_emb = shared.args.compress_pos_emb\n        if shared.args.gpu_split:\n            config.set_auto_map(shared.args.gpu_split)\n            config.gpu_peer_fix = True\n        if torch.version.hip:\n            config.rmsnorm_no_half2 = True\n            config.rope_no_half2 = True\n            config.matmul_no_half2 = True\n            config.silu_no_half2 = True\n\n        # This slowes down a bit but align better with autogptq generation.\n        # TODO: Should give user choice to tune the exllama config\n        # config.fused_attn = False\n        # config.fused_mlp_thd = 0\n\n        return ExllamaHF(config)\n", "path": "modules/exllama_hf.py"}, {"content": "import sys\nfrom pathlib import Path\n\nfrom torch import version as torch_version\n\nfrom modules import shared\nfrom modules.logging_colors import logger\n\ntry:\n    from exllama.generator import ExLlamaGenerator\n    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig\n    from exllama.tokenizer import ExLlamaTokenizer\nexcept:\n    logger.warning('Exllama module failed to load. Will attempt to load from repositories.')\n    try:\n        from modules.relative_imports import RelativeImport\n\n        with RelativeImport(\"repositories/exllama\"):\n            from generator import ExLlamaGenerator\n            from model import ExLlama, ExLlamaCache, ExLlamaConfig\n            from tokenizer import ExLlamaTokenizer\n    except:\n        logger.error(\"Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.\")\n        raise\n\n\nclass ExllamaModel:\n    def __init__(self):\n        pass\n\n    @classmethod\n    def from_pretrained(self, path_to_model):\n\n        path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)\n        tokenizer_model_path = path_to_model / \"tokenizer.model\"\n        model_config_path = path_to_model / \"config.json\"\n\n        # Find the model checkpoint\n        model_path = None\n        for ext in ['.safetensors', '.pt', '.bin']:\n            found = list(path_to_model.glob(f\"*{ext}\"))\n            if len(found) > 0:\n                if len(found) > 1:\n                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')\n\n                model_path = found[-1]\n                break\n\n        config = ExLlamaConfig(str(model_config_path))\n        config.model_path = str(model_path)\n        config.max_seq_len = shared.args.max_seq_len\n        config.compress_pos_emb = shared.args.compress_pos_emb\n        if shared.args.gpu_split:\n            config.set_auto_map(shared.args.gpu_split)\n            config.gpu_peer_fix = True\n        if torch_version.hip:\n            config.rmsnorm_no_half2 = True\n            config.rope_no_half2 = True\n            config.matmul_no_half2 = True\n            config.silu_no_half2 = True\n\n\n        model = ExLlama(config)\n        tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))\n        cache = ExLlamaCache(model)\n        generator = ExLlamaGenerator(model, tokenizer, cache)\n\n        result = self()\n        result.config = config\n        result.model = model\n        result.cache = cache\n        result.tokenizer = tokenizer\n        result.generator = generator\n        return result, result\n\n    def generate_with_streaming(self, prompt, state):\n        self.generator.settings.temperature = state['temperature']\n        self.generator.settings.top_p = state['top_p']\n        self.generator.settings.top_k = state['top_k']\n        self.generator.settings.typical = state['typical_p']\n        self.generator.settings.token_repetition_penalty_max = state['repetition_penalty']\n        if state['ban_eos_token']:\n            self.generator.disallow_tokens([self.tokenizer.eos_token_id])\n        else:\n            self.generator.disallow_tokens(None)\n\n        self.generator.end_beam_search()\n        ids = self.generator.tokenizer.encode(prompt)\n        self.generator.gen_begin_reuse(ids)\n        initial_len = self.generator.sequence[0].shape[0]\n        has_leading_space = False\n        for i in range(state['max_new_tokens']):\n            token = self.generator.gen_single_token()\n            if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('\u2581'):\n                has_leading_space = True\n\n            decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])\n            if has_leading_space:\n                decoded_text = ' ' + decoded_text\n\n            yield decoded_text\n            if token.item() == self.generator.tokenizer.eos_token_id or shared.stop_everything:\n                break\n\n    def generate(self, prompt, state):\n        output = ''\n        for output in self.generate_with_streaming(prompt, state):\n            pass\n\n        return output\n\n    def encode(self, string, **kwargs):\n        return self.tokenizer.encode(string)\n", "path": "modules/exllama.py"}]}