Repository: huggingface/text-generation-inference

Issue:
Support for DBRX models
### Feature request

support new DBRX model(s): https://huggingface.co/databricks/dbrx-instruct

### Motivation

It is the new state-of-the-art Open LLM

### Your contribution

I have the resources to test any PR with the 132B model to confirm


Discussion:
On it. We will have a PR ready tomorrow. 
Let us know if we can help!
#1685 adds dbrx support.
Usage: 

```shell
curl 127.0.0.1:3000/generate_stream \
    -X POST \
    -d '{"inputs":"
 user\n Who is the President of the USA? 
 \n 
 assistant\n","parameters":{"max_new_tokens":50}}' \
    -H 'Content-Type: application/json'
```
For now, it only supports eetq, bitsandbytes, bitsandbytes-nf4 and bitsandbytes-fp4 quantization methods.
Fails on boot on docker version 1.4.5 with:

```
--model-id=databricks/dbrx-instruct
--revision=3b5d968eab47b0cb5b075fd984612b63f92841c2
--trust-remote-code
```

```
2024-03-29T23:39:50.846327Z ERROR text_generation_launcher: Error when initializing model
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 41, in __init__
    tokenizer = GPT2TokenizerFast.from_pretrained(
  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2086, in from_pretrained
    return cls._from_pretrained(
  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2120, in _from_pretrained
    slow_tokenizer = (cls.slow_tokenizer_class)._from_pretrained(
  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2325, in _from_pretrained
    tokenizer = cls(*init_inputs, **init_kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/gpt2/tokenization_gpt2.py", line 181, in __init__
    with open(vocab_file, encoding="utf-8") as vocab_handle:
TypeError: expected str, bytes or os.PathLike object, not NoneType

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 52, in __init__
    tokenizer = AutoTokenizer.from_pretrained(
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py", line 818, in from_pretrained
    tokenizer_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 489, in get_class_from_dynamic_module
    final_module = get_cached_module_file(
  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 315, in get_cached_module_file
    modules_needed = check_imports(resolved_module_file)
  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 180, in check_imports
    raise ImportError(
ImportError: This modeling file requires the following packages that were not found in your environment: tiktoken. Run `pip install tiktoken`

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())
  File "/opt/conda/lib/python3.10/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.10/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/cli.py", line 89, in serve
    server.serve(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 235, in serve
    asyncio.run(
  File "/opt/conda/lib/python3.10/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 636, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
> File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 196, in serve_inner
    model = get_model(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/__init__.py", line 388, in get_model
    return FlashDbrx(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 63, in __init__
    tokenizer = GPT2TokenizerFast.from_pretrained(
  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2070, in from_pretrained
    raise EnvironmentError(
OSError: Can't load tokenizer for 'Xenova/dbrx-instruct-tokenizer'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure 'Xenova/dbrx-instruct-tokenizer' is the correct path to a directory containing all relevant files for a GPT2TokenizerFast tokenizer.

2024-03-29T23:39:51.479717Z ERROR shard-manager: text_generation_launcher: Shard complete standard error output:

The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. 
The tokenizer class you load from this checkpoint is 'TiktokenTokenizerWrapper'. 
The class this function is called from is 'GPT2Tokenizer'.
Traceback (most recent call last):

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 41, in __init__
    tokenizer = GPT2TokenizerFast.from_pretrained(

  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2086, in from_pretrained
    return cls._from_pretrained(

  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2120, in _from_pretrained
    slow_tokenizer = (cls.slow_tokenizer_class)._from_pretrained(

  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2325, in _from_pretrained
    tokenizer = cls(*init_inputs, **init_kwargs)

  File "/opt/conda/lib/python3.10/site-packages/transformers/models/gpt2/tokenization_gpt2.py", line 181, in __init__
    with open(vocab_file, encoding="utf-8") as vocab_handle:

TypeError: expected str, bytes or os.PathLike object, not NoneType


During handling of the above exception, another exception occurred:


Traceback (most recent call last):

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 52, in __init__
    tokenizer = AutoTokenizer.from_pretrained(

  File "/opt/conda/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py", line 818, in from_pretrained
    tokenizer_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)

  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 489, in get_class_from_dynamic_module
    final_module = get_cached_module_file(

  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 315, in get_cached_module_file
    modules_needed = check_imports(resolved_module_file)

  File "/opt/conda/lib/python3.10/site-packages/transformers/dynamic_module_utils.py", line 180, in check_imports
    raise ImportError(

ImportError: This modeling file requires the following packages that were not found in your environment: tiktoken. Run `pip install tiktoken`


During handling of the above exception, another exception occurred:


Traceback (most recent call last):

  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/cli.py", line 89, in serve
    server.serve(

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 235, in serve
    asyncio.run(

  File "/opt/conda/lib/python3.10/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)

  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 649, in run_until_complete
    return future.result()

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 196, in serve_inner
    model = get_model(

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/__init__.py", line 388, in get_model
    return FlashDbrx(

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_dbrx.py", line 63, in __init__
    tokenizer = GPT2TokenizerFast.from_pretrained(

  File "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2070, in from_pretrained
    raise EnvironmentError(

OSError: Can't load tokenizer for 'Xenova/dbrx-instruct-tokenizer'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure 'Xenova/dbrx-instruct-tokenizer' is the correct path to a directory containing all relevant files for a GPT2TokenizerFast tokenizer.
```
Using eetq gives an error that the datatype must be float16 or float32.
I can confirm the 16bit works very well! Thanks.
Yes, can confirm 16 bit runs fine on 4x A100 80 GB SXM.

- eetq gives an error that the datatype must be float16 or float32.
- 16 bit on 8x A6000 gives ""text_generation_launcher: Method Warmup encountered an error." when the final stage ." See [this issue](https://github.com/TrelisResearch/one-click-llms/issues/5)
Are you running on Docker?
Getting:
`ValueError("sharded is not supported for AutoModel")`

When deploying via SageMaker:
```
hub = {
  'HF_MODEL_ID': "databricks/dbrx-instruct", # model_id from hf.co/models
  'SM_NUM_GPUS': json.dumps(8), # Number of GPU used per replica
  'TOKEN':  'true',
  'TRUST_REMOTE_CODE':  'true',
    # 'REVISION': "e0bbb53cee412aba95f3b3fa4fc0265b1a0788b2",
    'HUGGING_FACE_HUB_TOKEN': "<token>",
  'MAX_INPUT_LENGTH': json.dumps(24000),  # Max length of input text
  'MAX_BATCH_PREFILL_TOKENS': json.dumps(32000),  # Number of tokens for the prefill operation.
  'MAX_TOTAL_TOKENS': json.dumps(32000),  # Max length of the generation (including input text)
  'MAX_BATCH_TOTAL_TOKENS': json.dumps(512000),  # Limits the number of tokens that can be processed in parallel during the generation
  # ,'HF_MODEL_QUANTIZE': "awq", # comment in to quantize not supported yet
}
role="arn:aws:iam::153172432509:role/Llama2_13B_Role"  # generic name but this is correct role 
llm_image ="763104351884.dkr.ecr.us-east-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.3.1-gpu-py310-cu121-ubuntu20.04-v1.0"

huggingface_model = HuggingFaceModel(
  image_uri =llm_image,
  env=hub,
  role=role, 
)
predictor = huggingface_model.deploy(
  initial_instance_count=1,
  instance_type="ml.g5.48xlarge", 
    endpoint_name=endpoint_name_var,
  container_startup_health_check_timeout=300,
)
```
> Are you running on Docker?

Yes, this is using the latest TGI docker image.
@RonanKMcGovern can you share command args you are using?
Sure

> @RonanKMcGovern can you share command args you are using?

Sure, it's this runpod template here: https://runpod.io/console/gpu-cloud?template=tlt1i1welu&ref=jmfkcdio

And the arguments for eetq would be:
```
--model-id databricks/dbrx-instruct --trust-remote-code --port 8080 --max-input-length 3000 --max-total-tokens 4096 --max-batch-prefill-tokens 4096 --quantize eetq
```
@RonanKMcGovern with same args I have an error message: `RuntimeError: Invalid datatype. Weight must be FP16 or FP32`
> @RonanKMcGovern with same args I have an error message: `RuntimeError: Invalid datatype. Weight must be FP16 or FP32`

Exactly, same as my note above on eetq
@RonanKMcGovern can you try:

```shell
--model-id databricks/dbrx-instruct --trust-remote-code --port 8080 --max-input-length 3000 --max-total-tokens 4096 --max-batch-prefill-tokens 4096 --quantize eetq --dtype float16
```

in the meantime? 
Note the added `--dtype float16` at the end.

I will patch the issue later today.
@OlivierDehaene `RuntimeError: Only 1 can be set between `dtype` and `quantize`, as they both decide how goes the final model.`

Resolving patch:
diff --git a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
--- a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
@@ -18,9 +18,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, Optional, Tuple
+
 import torch
 import torch.distributed
-
 from torch import nn
 from transformers.activations import ACT2FN
 from typing import Optional, List, Tuple
diff --git a/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py b/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
--- a/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
@@ -1,26 +1,21 @@
+from typing import List, Optional, Tuple
+
 import torch
 import torch.distributed
-
 from torch import nn
-from transformers.modeling_utils import PreTrainedModel
 from transformers.configuration_utils import PretrainedConfig
-from typing import Optional, List, Tuple
+from transformers.modeling_utils import PreTrainedModel
 
-from text_generation_server.utils import paged_attention, flash_attn
-from text_generation_server.utils.flash_attn import attention
 from text_generation_server.layers import (
-    TensorParallelRowLinear,
+    SpeculativeHead,
     TensorParallelColumnLinear,
     TensorParallelEmbedding,
-    SpeculativeHead,
+    TensorParallelRowLinear,
     get_linear,
 )
-from text_generation_server.layers.layernorm import (
-    FastLayerNorm,
-)
-from text_generation_server.layers.rotary import (
-    PositionRotaryEmbedding,
-)
+from text_generation_server.layers.layernorm import FastLayerNorm
+from text_generation_server.layers.rotary import PositionRotaryEmbedding
+from text_generation_server.utils import flash_attn, paged_attention
 
 
 def load_row(config, prefix: str, weights, bias: bool):
@@ -52,6 +47,7 @@ def __init__(
         hidden_size=64,
         num_hidden_layers=None,
         num_attention_heads=None,
+        num_ln_in_prallel_attention=None,
         layer_norm_epsilon=1e-5,
         initializer_range=0.02,
         use_cache=True,
@@ -65,6 +61,7 @@ def __init__(
         new_decoder_architecture=None,
         bias=False,
         parallel_attn=False,
+        rope_theta=10_000.0,
         **kwargs,
     ):
         if alibi:
@@ -75,6 +72,7 @@ def __init__(
         self.model_type = model_type
         self.alibi = False
         self.rotary = True
+        self.rope_theta = rope_theta
 
         self.vocab_size = vocab_size
         # Backward compatibility with n_embed kwarg
@@ -91,6 +89,7 @@ def __init__(
             else kwargs.pop("n_head", 8)
         )
         self.layer_norm_epsilon = layer_norm_epsilon
+        self.num_ln_in_parallel_attention = num_ln_in_prallel_attention
         self.initializer_range = initializer_range
         self.use_cache = use_cache
         self.hidden_dropout = hidden_dropout
@@ -132,9 +131,13 @@ def __init__(
         self.num_heads_kv = config.n_head_kv
         self.hidden_size = config.hidden_size
         self.head_size = self.hidden_size // self.num_heads
+        self.rope_theta = config.rope_theta
 
         self.rotary_emb = PositionRotaryEmbedding.static(
-            config=config, dim=self.head_size, base=10000.0, device=weights.device
+            config=config,
+            dim=self.head_size,
+            base=self.rope_theta,
+            device=weights.device,
         )
         self.softmax_scale = self.head_size ** (-0.5)
 
@@ -244,9 +247,13 @@ def __init__(
         self.hidden_size = hidden_size
         self.head_size = hidden_size // num_heads
         self.num_groups = num_groups
+        self.rope_theta = config.rope_theta
 
         self.rotary_emb = PositionRotaryEmbedding.static(
-            config=config, dim=self.head_size, base=10000.0, device=weights.device
+            config=config,
+            dim=self.head_size,
+            base=self.rope_theta,
+            device=weights.device,
         )
         self.softmax_scale = self.head_size ** (-0.5)
 
@@ -257,7 +264,7 @@ def __init__(
 
         if process_group.size() > self.num_groups:
             raise NotImplementedError(
-                f"Tensor Parallelism is not implemented for world_size > n groups"
+                "Tensor Parallelism is not implemented for world_size > n groups"
             )
         if self.num_groups % process_group.size() != 0:
             raise NotImplementedError(
@@ -459,29 +466,61 @@ def forward(
                 max_s,
             )
 
-            hidden_states, residual = self.post_attention_layernorm(
-                hidden_states, residual
-            )
+            if self.post_attention_layernorm is not None:
+                hidden_states, residual = self.post_attention_layernorm(
+                    hidden_states, residual
+                )
 
             mlp_output = self.mlp(hidden_states)
 
             return mlp_output, residual
 
 
+class FlashRWLayerNorm(nn.Module):
+    def __init__(self, config, prefix, weights):
+        super().__init__()
+        self.num_ln = config.num_ln_in_parallel_attn
+
+        if self.num_ln == 1:
+            self.input_ln = FastLayerNorm.load(
+                prefix=f"{prefix}.input_layernorm",
+                weights=weights,
+                eps=config.layer_norm_epsilon,
+            )
+        elif self.num_ln == 2:
+            self.ln_attn = FastLayerNorm.load(
+                prefix=f"{prefix}.ln_attn",
+                weights=weights,
+                eps=config.layer_norm_epsilon,
+            )
+            self.ln_mlp = FastLayerNorm.load(
+                prefix=f"{prefix}.ln_mlp",
+                weights=weights,
+                eps=config.layer_norm_epsilon,
+            )
+        else:
+            raise ValueError("Number of layer norms can either be 1 or 2.")
+
+    def forward(
+        self,
+        hidden_states,
+        residual,
+    ):
+        if self.num_ln == 1:
+            ln_hidden_states, residual = self.input_ln(hidden_states, residual)
+            return ln_hidden_states, ln_hidden_states, residual
+        elif self.num_ln == 2:
+            ln_attn, residual = self.ln_attn(hidden_states, residual)
+            ln_mlp, _ = self.ln_mlp(residual)
+            return ln_attn, ln_mlp, residual
+
+
 class FlashRWLargeLayer(nn.Module):
     def __init__(self, layer_id, config, weights):
         super().__init__()
         prefix = f"transformer.h.{layer_id}"
-        self.ln_attn = FastLayerNorm.load(
-            prefix=f"{prefix}.ln_attn",
-            weights=weights,
-            eps=config.layer_norm_epsilon,
-        )
-        self.ln_mlp = FastLayerNorm.load(
-            prefix=f"{prefix}.ln_mlp",
-            weights=weights,
-            eps=config.layer_norm_epsilon,
-        )
+
+        self.ln_layer = FlashRWLayerNorm(config, prefix, weights)
 
         self.self_attention = FlashRWLargeAttention(
             config,
@@ -507,8 +546,8 @@ def forward(
         input_lengths,
         max_s,
     ):
-        ln_attn, residual = self.ln_attn(hidden_states, residual)
-        ln_mlp, _ = self.ln_mlp(residual)
+        # Layer norm.
+        ln_attn, ln_mlp, residual = self.ln_layer(hidden_states, residual)
 
         # Self attention.
         attn_output = self.self_attention(


Test patch: