Support for DBRX models

[maziyarpanahi]

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

[OlivierDehaene]

On it. We will have a PR ready tomorrow.

[hanlint]

Let us know if we can help!

[OlivierDehaene]

#1685 adds dbrx support.
Usage:

curl 127.0.0.1:3000/generate_stream \
    -X POST \
    -d '{"inputs":"<|im_start|> user\n Who is the President of the USA? <|im_end|> \n <|im_start|> assistant\n","parameters":{"max_new_tokens":50}}' \
    -H 'Content-Type: application/json'

[OlivierDehaene]

For now, it only supports eetq, bitsandbytes, bitsandbytes-nf4 and bitsandbytes-fp4 quantization methods.

[L3o-pold]

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.

[RonanKMcGovern]

Using eetq gives an error that the datatype must be float16 or float32.

[maziyarpanahi]

I can confirm the 16bit works very well! Thanks.

[RonanKMcGovern]

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

[L3o-pold]

Are you running on Docker?

[wagstaffjfriel]

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,
)

[RonanKMcGovern]

Are you running on Docker?

Yes, this is using the latest TGI docker image.

[L3o-pold]

@RonanKMcGovern can you share command args you are using?