Skip to content

GitLab

Unverified Commit c5de7cd8 authored by Nicolas Patry's avatar Nicolas Patry Committed by GitHub
Browse files

Add AWQ quantization inference support (#1019) (#1054) 

# Add AWQ quantization inference support

Fixes
https://github.com/huggingface/text-generation-inference/issues/781

This PR (partially) adds support for AWQ quantization for inference.
More information on AWQ [here](https://arxiv.org/abs/2306.00978). In
general, AWQ is faster and more accurate than GPTQ, which is currently
supported by TGI.

This PR installs 4-bit GEMM custom CUDA kernels released by AWQ authors
(in `requirements.txt`, just one line change).

Quick way to test this PR would be bring up TGI as follows:

```
text-generation-server download-weights abhinavkulkarni/codellama-CodeLlama-7b-Python-hf-w4-g128-awq

text-generation-launcher \
--huggingface-hub-cache ~/.cache/huggingface/hub/ \
--model-id abhinavkulkarni/codellama-CodeLlama-7b-Python-hf-w4-g128-awq \
--trust-remote-code --port 8080 \
--max-input-length 2048 --max-total-tokens 4096 --max-batch-prefill-tokens 4096 \
--quantize awq
```

Please note:
* This PR was tested with FlashAttention v2 and vLLM.
* This PR adds support for AWQ inference, not quantizing the models.
That needs to be done outside of TGI, instructions

[here](https://github.com/mit-han-lab/llm-awq/tree/f084f40bd996f3cf3a0633c1ad7d9d476c318aaa).
* This PR only adds support for `FlashLlama` models for now.
* Multi-GPU setup has not been tested. 
* No integration tests have been added so far, will add later if
maintainers are interested in this change.
* This PR can be tested on any of the models released

[here](https://huggingface.co/abhinavkulkarni?sort_models=downloads#models).

Please refer to the linked issue for benchmarks for

[abhinavkulkarni/meta-llama-Llama-2-7b-chat-hf-w4-g128-awq](https://huggingface.co/abhinavkulkarni/meta-llama-Llama-2-7b-chat-hf-w4-g128-awq)
vs

[TheBloke/Llama-2-7b-Chat-GPTQ](https://huggingface.co/TheBloke/Llama-2-7b-Chat-GPTQ).

Please note, AWQ has released faster (and in case of Llama, fused)
kernels for 4-bit GEMM, currently at the top of the `main` branch at
https://github.com/mit-han-lab/llm-awq, but this PR uses an older commit
that has been tested to work. We can switch to latest commit later on.

## Who can review?

@OlivierDehaene OR @Narsil

---------



# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet
though.

Once merged, your PR is going to appear in the release notes with the
title you set, so make sure it's a great title that fully reflects the
extent of your awesome contribution.

Then, please replace this with a description of the change and which
issue is fixed (if applicable). Please also include relevant motivation
and context. List any dependencies (if any) that are required for this
change.

Once you're done, someone will review your PR shortly (see the section
"Who can review?" below to tag some potential reviewers). They may
suggest changes to make the code even better. If no one reviewed your PR
after a week has passed, don't hesitate to post a new comment
@-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the
other checks if that's the case).
- [ ] Did you read the [contributor
guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the
[forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes?
Here are the
[documentation
guidelines](https://github.com/huggingface/transformers/tree/main/docs),
and
[here are tips on formatting
docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation

).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have
passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the
right person to tag with @


@OlivierDehaene OR @Narsil

 -->

---------
Co-authored-by: default avatarAbhinav M Kulkarni <abhinavkulkarni@gmail.com>
Co-authored-by: default avatarAbhinav Kulkarni <abhinav@concentric.ai>
parent fef36cea
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1650 additions and 302 deletions
Dockerfile
View file @ c5de7cd8
...@@ -111,22 +111,22 @@ RUN make build-flash-attention-v2 ...@@ -111,22 +111,22 @@ RUN make build-flash-attention-v2
# Build Transformers exllama kernels # Build Transformers exllama kernels
FROM kernel-builder as exllama-kernels-builder FROM kernel-builder as exllama-kernels-builder
WORKDIR /usr/src WORKDIR /usr/src
COPY server/exllama_kernels/ . COPY server/exllama_kernels/ .
# Build specific version of transformers # Build specific version of transformers
RUN TORCH_CUDA_ARCH_LIST="8.0;8.6+PTX" python setup.py build RUN TORCH_CUDA_ARCH_LIST="8.0;8.6+PTX" python setup.py build
# Build Transformers awq kernels
FROM kernel-builder as awq-kernels-builder
WORKDIR /usr/src
COPY server/Makefile-awq Makefile
# Build specific version of transformers
RUN TORCH_CUDA_ARCH_LIST="8.0;8.6+PTX" make build-awq
# Build Transformers CUDA kernels # Build Transformers CUDA kernels
FROM kernel-builder as custom-kernels-builder FROM kernel-builder as custom-kernels-builder
WORKDIR /usr/src WORKDIR /usr/src
COPY server/custom_kernels/ . COPY server/custom_kernels/ .
# Build specific version of transformers # Build specific version of transformers
RUN python setup.py build RUN python setup.py build
...@@ -175,6 +175,8 @@ COPY --from=flash-att-v2-builder /usr/src/flash-attention-v2/build/lib.linux-x86 ...@@ -175,6 +175,8 @@ COPY --from=flash-att-v2-builder /usr/src/flash-attention-v2/build/lib.linux-x86
COPY --from=custom-kernels-builder /usr/src/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages COPY --from=custom-kernels-builder /usr/src/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
# Copy build artifacts from exllama kernels builder # Copy build artifacts from exllama kernels builder
COPY --from=exllama-kernels-builder /usr/src/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages COPY --from=exllama-kernels-builder /usr/src/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
# Copy build artifacts from awq kernels builder
COPY --from=awq-kernels-builder /usr/src/llm-awq/awq/kernels/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
# Copy builds artifacts from vllm builder # Copy builds artifacts from vllm builder
COPY --from=vllm-builder /usr/src/vllm/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages COPY --from=vllm-builder /usr/src/vllm/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages
......
docs/source/basic_tutorials/preparing_model.md
View file @ c5de7cd8
...@@ -4,7 +4,7 @@ Text Generation Inference improves the model in several aspects. ...@@ -4,7 +4,7 @@ Text Generation Inference improves the model in several aspects.
## Quantization ## Quantization
TGI supports [bits-and-bytes](https://github.com/TimDettmers/bitsandbytes#bitsandbytes) and [GPT-Q](https://arxiv.org/abs/2210.17323) quantization. To speed up inference with quantization, simply set `quantize` flag to `bitsandbytes` or `gptq` depending on the quantization technique you wish to use. When using GPT-Q quantization, you need to point to one of the models [here](https://huggingface.co/models?search=gptq). To get more information about quantization, please refer to (./conceptual/quantization.md) TGI supports [bits-and-bytes](https://github.com/TimDettmers/bitsandbytes#bitsandbytes), [GPT-Q](https://arxiv.org/abs/2210.17323) and [AWQ](https://arxiv.org/abs/2306.00978) quantization. To speed up inference with quantization, simply set `quantize` flag to `bitsandbytes`, `gptq` or `awq` depending on the quantization technique you wish to use. When using GPT-Q quantization, you need to point to one of the models [here](https://huggingface.co/models?search=gptq) when using AWQ quantization, you need to point to one of the models [here](https://huggingface.co/models?search=awq). To get more information about quantization, please refer to (./conceptual/quantization.md)
## RoPE Scaling ## RoPE Scaling
......
integration-tests/models/__snapshots__/test_flash_awq/test_flash_llama_awq.json 0 → 100644
View file @ c5de7cd8
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.703125,
"text": "What"
},
{
"id": 338,
"logprob": -1.4765625,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8583984,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7548828,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9306641,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4550781,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.5732422,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5761719,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5888672,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026504517,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4287109,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15856934,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62646484,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
}
integration-tests/models/__snapshots__/test_flash_awq/test_flash_llama_awq_all_params.json 0 → 100644
View file @ c5de7cd8
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 338,
"logprob": -9.0859375,
"text": "is"
},
{
"id": 21784,
"logprob": -10.90625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -2.65625,
"text": "Learning"
},
{
"id": 29973,
"logprob": -4.8085938,
"text": "?"
}
],
"seed": 0,
"tokens": [
{
"id": 13,
"logprob": -0.19958496,
"special": false,
"text": "\n"
},
{
"id": 4013,
"logprob": -2.203125,
"special": false,
"text": "This"
},
{
"id": 1139,
"logprob": -0.23693848,
"special": false,
"text": " question"
},
{
"id": 756,
"logprob": 0.0,
"special": false,
"text": " has"
},
{
"id": 1063,
"logprob": -0.076538086,
"special": false,
"text": " been"
},
{
"id": 4433,
"logprob": 0.0,
"special": false,
"text": " asked"
},
{
"id": 1784,
"logprob": -1.1367188,
"special": false,
"text": " many"
},
{
"id": 3064,
"logprob": 0.0,
"special": false,
"text": " times"
},
{
"id": 322,
"logprob": -1.7460938,
"special": false,
"text": " and"
},
{
"id": 306,
"logprob": 0.0,
"special": false,
"text": " I"
}
],
"top_tokens": null
},
"generated_text": "What is Deep Learning?\nThis question has been asked many times and I"
}
integration-tests/models/__snapshots__/test_flash_awq/test_flash_llama_awq_load.json 0 → 100644
View file @ c5de7cd8
[
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.703125,
"text": "What"
},
{
"id": 338,
"logprob": -1.4765625,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8652344,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7548828,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9306641,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4550781,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.5732422,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5761719,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5888672,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026504517,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4287109,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15856934,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62646484,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.703125,
"text": "What"
},
{
"id": 338,
"logprob": -1.4765625,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8583984,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7548828,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9306641,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4550781,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.5732422,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5761719,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5888672,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026504517,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4287109,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15856934,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62646484,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.703125,
"text": "What"
},
{
"id": 338,
"logprob": -1.4765625,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8652344,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7548828,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9306641,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4550781,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.5732422,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5761719,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5888672,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026504517,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4287109,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15856934,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62646484,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.703125,
"text": "What"
},
{
"id": 338,
"logprob": -1.4765625,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8652344,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7548828,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9306641,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4550781,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.5732422,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5761719,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5888672,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026504517,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4287109,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15856934,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62646484,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
}
]
integration-tests/models/__snapshots__/test_flash_awq_sharded/test_flash_llama_awq_load_sharded.json 0 → 100644
View file @ c5de7cd8
[
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.6914062,
"text": "What"
},
{
"id": 338,
"logprob": -1.4746094,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8623047,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7558594,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9228516,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4609375,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.57177734,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5722656,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5859375,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.02633667,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4335938,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15991211,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62060547,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.6914062,
"text": "What"
},
{
"id": 338,
"logprob": -1.4746094,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8623047,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7558594,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9228516,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4609375,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.57177734,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5722656,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5859375,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.02633667,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4335938,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15991211,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62060547,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.6914062,
"text": "What"
},
{
"id": 338,
"logprob": -1.4746094,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8623047,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7558594,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9228516,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4609375,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.57177734,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5722656,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5859375,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.02633667,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4335938,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15991211,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62060547,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
},
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.6914062,
"text": "What"
},
{
"id": 338,
"logprob": -1.4746094,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8623047,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7558594,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9228516,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4609375,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.57177734,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5722656,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5859375,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.02633667,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4335938,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.15991211,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17456055,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62060547,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
}
]
integration-tests/models/__snapshots__/test_flash_awq_sharded/test_flash_llama_awq_sharded.json 0 → 100644
View file @ c5de7cd8
{
"details": {
"best_of_sequences": null,
"finish_reason": "length",
"generated_tokens": 10,
"prefill": [
{
"id": 1,
"logprob": null,
"text": "<s>"
},
{
"id": 1724,
"logprob": -7.6914062,
"text": "What"
},
{
"id": 338,
"logprob": -1.4746094,
"text": "is"
},
{
"id": 21784,
"logprob": -9.390625,
"text": "Deep"
},
{
"id": 29257,
"logprob": -1.8623047,
"text": "Learning"
},
{
"id": 29973,
"logprob": -0.7558594,
"text": "?"
}
],
"seed": null,
"tokens": [
{
"id": 13,
"logprob": -1.9228516,
"special": false,
"text": "\n"
},
{
"id": 5618,
"logprob": -2.4609375,
"special": false,
"text": "What"
},
{
"id": 338,
"logprob": -0.57177734,
"special": false,
"text": " is"
},
{
"id": 278,
"logprob": -1.5722656,
"special": false,
"text": " the"
},
{
"id": 4328,
"logprob": -1.5927734,
"special": false,
"text": " difference"
},
{
"id": 1546,
"logprob": -0.026428223,
"special": false,
"text": " between"
},
{
"id": 21784,
"logprob": -1.4267578,
"special": false,
"text": " Deep"
},
{
"id": 29257,
"logprob": -0.16015625,
"special": false,
"text": " Learning"
},
{
"id": 322,
"logprob": -0.17382812,
"special": false,
"text": " and"
},
{
"id": 6189,
"logprob": -0.62060547,
"special": false,
"text": " Machine"
}
],
"top_tokens": null
},
"generated_text": "\nWhat is the difference between Deep Learning and Machine"
}
integration-tests/models/test_flash_awq.py 0 → 100644
View file @ c5de7cd8
import pytest
@pytest.fixture(scope="module")
def flash_llama_awq_handle(launcher):
with launcher("abhinavkulkarni/codellama-CodeLlama-7b-Python-hf-w4-g128-awq", num_shard=1, quantize="awq") as handle:
yield handle
@pytest.fixture(scope="module")
async def flash_llama_awq(flash_llama_awq_handle):
await flash_llama_awq_handle.health(300)
return flash_llama_awq_handle.client
@pytest.mark.asyncio
@pytest.mark.private
async def test_flash_llama_awq(flash_llama_awq, response_snapshot):
response = await flash_llama_awq.generate(
"What is Deep Learning?", max_new_tokens=10, decoder_input_details=True
)
assert response.details.generated_tokens == 10
assert response.generated_text == "\nWhat is the difference between Deep Learning and Machine"
assert response == response_snapshot
@pytest.mark.asyncio
@pytest.mark.private
async def test_flash_llama_awq_all_params(flash_llama_awq, response_snapshot):
response = await flash_llama_awq.generate(
"What is Deep Learning?",
max_new_tokens=10,
repetition_penalty=1.2,
return_full_text=True,
temperature=0.5,
top_p=0.9,
top_k=10,
truncate=5,
typical_p=0.9,
watermark=True,
decoder_input_details=True,
seed=0,
)
assert response.details.generated_tokens == 10
assert response == response_snapshot
@pytest.mark.asyncio
@pytest.mark.private
async def test_flash_llama_awq_load(
flash_llama_awq, generate_load, response_snapshot
):
responses = await generate_load(
flash_llama_awq, "What is Deep Learning?", max_new_tokens=10, n=4
)
assert len(responses) == 4
assert all([r.generated_text == "\nWhat is the difference between Deep Learning and Machine" for r in responses])
assert responses == response_snapshot
integration-tests/models/test_flash_awq_sharded.py 0 → 100644
View file @ c5de7cd8
import pytest
@pytest.fixture(scope="module")
def flash_llama_awq_handle_sharded(launcher):
with launcher("abhinavkulkarni/codellama-CodeLlama-7b-Python-hf-w4-g128-awq", num_shard=2, quantize="awq") as handle:
yield handle
@pytest.fixture(scope="module")
async def flash_llama_awq_sharded(flash_llama_awq_handle_sharded):
await flash_llama_awq_handle_sharded.health(300)
return flash_llama_awq_handle_sharded.client
@pytest.mark.asyncio
@pytest.mark.private
async def test_flash_llama_awq_sharded(flash_llama_awq_sharded, response_snapshot):
response = await flash_llama_awq_sharded.generate(
"What is Deep Learning?", max_new_tokens=10, decoder_input_details=True
)
assert response.details.generated_tokens == 10
assert response.generated_text == "\nWhat is the difference between Deep Learning and Machine"
assert response == response_snapshot
@pytest.mark.asyncio
@pytest.mark.private
async def test_flash_llama_awq_load_sharded(
flash_llama_awq_sharded, generate_load, response_snapshot
):
responses = await generate_load(
flash_llama_awq_sharded, "What is Deep Learning?", max_new_tokens=10, n=4
)
assert len(responses) == 4
assert all([r.generated_text == "\nWhat is the difference between Deep Learning and Machine" for r in responses])
assert responses == response_snapshot
launcher/src/main.rs
View file @ c5de7cd8
...@@ -25,6 +25,7 @@ enum Quantization { ...@@ -25,6 +25,7 @@ enum Quantization {
BitsandbytesNF4, BitsandbytesNF4,
BitsandbytesFP4, BitsandbytesFP4,
Gptq, Gptq,
Awq,
} }
impl std::fmt::Display for Quantization { impl std::fmt::Display for Quantization {
...@@ -43,6 +44,9 @@ impl std::fmt::Display for Quantization { ...@@ -43,6 +44,9 @@ impl std::fmt::Display for Quantization {
Quantization::Gptq => { Quantization::Gptq => {
write!(f, "gptq") write!(f, "gptq")
} }
Quantization::Awq => {
write!(f, "awq")
}
} }
} }
} }
......
server/.gitignore
View file @ c5de7cd8
...@@ -159,3 +159,4 @@ safetensors ...@@ -159,3 +159,4 @@ safetensors
flash-attention/ flash-attention/
flash-attention-v2/ flash-attention-v2/
vllm/ vllm/
llm-awq/
server/Makefile
View file @ c5de7cd8
include Makefile-flash-att include Makefile-flash-att
include Makefile-flash-att-v2 include Makefile-flash-att-v2
include Makefile-vllm include Makefile-vllm
include Makefile-awq
unit-tests: unit-tests:
pytest -s -vv -m "not private" tests pytest -s -vv -m "not private" tests
......
server/Makefile-awq 0 → 100644
View file @ c5de7cd8
awq_commit := f084f40bd996f3cf3a0633c1ad7d9d476c318aaa
awq:
rm -rf llm-awq
git clone https://github.com/mit-han-lab/llm-awq
build-awq: awq
cd llm-awq/ && git fetch && git checkout $(awq_commit)
cd llm-awq/awq/kernels && python setup.py build
install-awq: build-awq
pip uninstall awq_inference_engine -y || true
cd llm-awq/awq/kernels && python setup.py install
server/poetry.lock
View file @ c5de7cd8
This diff is collapsed.
server/pyproject.toml
View file @ c5de7cd8
...@@ -54,5 +54,7 @@ priority = "explicit" ...@@ -54,5 +54,7 @@ priority = "explicit"
markers = ["private: marks tests as requiring an admin hf token (deselect with '-m \"not private\"')"] markers = ["private: marks tests as requiring an admin hf token (deselect with '-m \"not private\"')"]
[build-system] [build-system]
requires = ["poetry-core>=1.0.0"] requires = [
"poetry-core>=1.0.0",
]
build-backend = "poetry.core.masonry.api" build-backend = "poetry.core.masonry.api"
server/text_generation_server/cli.py
View file @ c5de7cd8
...@@ -17,6 +17,7 @@ class Quantization(str, Enum): ...@@ -17,6 +17,7 @@ class Quantization(str, Enum):
bitsandbytes_nf4 = "bitsandbytes-nf4" bitsandbytes_nf4 = "bitsandbytes-nf4"
bitsandbytes_fp4 = "bitsandbytes-fp4" bitsandbytes_fp4 = "bitsandbytes-fp4"
gptq = "gptq" gptq = "gptq"
awq = "awq"
class Dtype(str, Enum): class Dtype(str, Enum):
......
server/text_generation_server/models/__init__.py
View file @ c5de7cd8
...@@ -268,6 +268,10 @@ def get_model( ...@@ -268,6 +268,10 @@ def get_model(
raise ValueError( raise ValueError(
"gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`" "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
) )
if quantize == "awq":
raise ValueError(
"awq quantization is not supported for AutoModel"
)
elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"): elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
raise ValueError( raise ValueError(
"4bit quantization is not supported for AutoModel" "4bit quantization is not supported for AutoModel"
......
server/text_generation_server/models/flash_llama.py
View file @ c5de7cd8
...@@ -62,7 +62,7 @@ class FlashLlama(FlashCausalLM): ...@@ -62,7 +62,7 @@ class FlashLlama(FlashCausalLM):
filenames = weight_files(model_id, revision=revision, extension=".safetensors") filenames = weight_files(model_id, revision=revision, extension=".safetensors")
weights = Weights(filenames, device, dtype, process_group=self.process_group) weights = Weights(filenames, device, dtype, process_group=self.process_group)
if config.quantize == "gptq": if config.quantize in ["gptq", "awq"]:
weights._set_gptq_params(model_id) weights._set_gptq_params(model_id)
model = FlashLlamaForCausalLM(config, weights) model = FlashLlamaForCausalLM(config, weights)
......
server/text_generation_server/utils/awq/quantize/qmodule.py 0 → 100644
View file @ c5de7cd8
# Copied logic from https://github.com/mit-han-lab/llm-awq/blob/f084f40bd996f3cf3a0633c1ad7d9d476c318aaa/awq/quantize/qmodule.py
import math
import torch
import torch.nn as nn
import awq_inference_engine # with CUDA kernels
# class ScaledActivation(nn.Module):
# def __init__(self, module, scales):
# super().__init__()
# self.act = module
# self.scales = nn.Parameter(scales.data)
#
# def forward(self, x):
# return self.act(x) / self.scales.view(1, 1, -1).to(x.device)
class WQLinear(nn.Module):
def __init__(self, w_bit, group_size, qweight, qzeros, scales, bias):
super().__init__()
if w_bit not in [4]:
raise NotImplementedError("Only 4-bit are supported for now.")
self.in_features = qweight.shape[0]
self.out_features = qweight.shape[1] * 32 // w_bit
self.w_bit = w_bit
self.group_size = group_size if group_size != -1 else self.in_features
# quick sanity check (make sure aligment)
assert self.in_features % self.group_size == 0
assert self.out_features % (32 // self.w_bit) == 0
self.qweight = qweight
self.qzeros = qzeros
self.scales = scales
if bias:
self.bias = bias
else:
self.bias = None
@torch.no_grad()
def forward(self, x):
out_shape = x.shape[:-1] + (self.out_features, )
out = awq_inference_engine.gemm_forward_cuda(x.reshape(-1, x.shape[-1]), self.qweight, self.scales, self.qzeros, 8)
out = out + self.bias if self.bias is not None else out
return out.reshape(out_shape)
server/text_generation_server/utils/layers.py
View file @ c5de7cd8
...@@ -18,6 +18,13 @@ from accelerate import init_empty_weights ...@@ -18,6 +18,13 @@ from accelerate import init_empty_weights
from text_generation_server.utils.gptq.quant_linear import QuantLinear from text_generation_server.utils.gptq.quant_linear import QuantLinear
HAS_AWQ = True
try:
from text_generation_server.utils.awq.quantize.qmodule import WQLinear
except ImportError:
HAS_AWQ = False
try: try:
major, _minor = torch.cuda.get_device_capability() major, _minor = torch.cuda.get_device_capability()
except Exception: except Exception:
...@@ -248,6 +255,14 @@ def get_linear(weight, bias, quantize): ...@@ -248,6 +255,14 @@ def get_linear(weight, bias, quantize):
bits, bits,
groupsize, groupsize,
) )
elif quantize == "awq":
try:
qweight, qzeros, scales, _, bits, groupsize, _ = weight
except Exception:
raise NotImplementedError(
f"The passed weight is not `awq` compatible, loader needs to be updated."
)
linear = WQLinear(w_bit=bits, group_size=groupsize, qweight=qweight, qzeros=qzeros, scales=scales, bias=bias is not None)
else: else:
raise NotImplementedError(f"Quantization `{quantize}` is not implemented yet.") raise NotImplementedError(f"Quantization `{quantize}` is not implemented yet.")
return linear return linear
...@@ -283,8 +298,8 @@ class TensorParallelHead(SuperLayer): ...@@ -283,8 +298,8 @@ class TensorParallelHead(SuperLayer):
weight = weights.get_tensor(f"{prefix}.weight") weight = weights.get_tensor(f"{prefix}.weight")
should_gather = False should_gather = False
# GPTQ doesn't quantize heads (nor embeddings) # GPTQ and AWQ don't quantize heads (nor embeddings)
if config.quantize == "gptq": if config.quantize in ["gptq", "awq"]:
quantize = None quantize = None
else: else:
quantize = config.quantize quantize = config.quantize
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment