neuralmagic/TinyLlama-1.1B-Chat-v0.4-pruned50-quant-ds

TinyLlama 1.1B Chat 0.4 - DeepSparse

This repo contains model files for TinyLlama 1.1B Chat optimized for DeepSparse, a CPU inference runtime for sparse models.

This model was quantized and pruned with SparseGPT, using SparseML.

Inference

Install DeepSparse LLM for fast inference on CPUs:

pip install deepsparse-nightly[llm]

Run in a Python pipeline:

from deepsparse import TextGeneration

prompt = "How to make banana bread?"
formatted_prompt =  f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"

model = TextGeneration(model="hf:neuralmagic/TinyLlama-1.1B-Chat-v0.4-pruned50-quant-ds")
print(model(formatted_prompt, max_new_tokens=500).generations[0].text)

"""
Banana bread is a delicious and easy-to-make recipe that is sure to please. Here is a recipe for making banana bread:

Ingredients:

    For the Banana Bread:

    - 1 cup of sugar
    - 1 cup of flour
    - 1/2 cup of mashed bananas
    - 1/4 cup of milk
    - 1/2 cup of melted butter
    - 1/4 cup of baking powder
    - 1/4 cup of baking soda
    - 1/4 cup of eggs
    - 1/4 cup of milk
    - 1/4 cup of sugar


Instructions:

    1. Preheat the oven to 325°F (160°C).
    2. In a large bowl, combine the sugar and flour.
    3. In a separate bow, combine the mashed bananas, milk, butter, baking powder, baking soda, milk, sugar.
    4. Add the bananas and milk into the flour-sugar mixture.
    5. Pour the milk into the bowl of the flour-sugar mixture.
    6. Pour the baking powder into the bowl of the flour-sugar mixture.
    7. Pour the mashed bananas into the bowl of the flour-sugar mixture.
    8. Add the eggs into the bowl of the flour-sugar mixture.
    9. Stir the mixture until it becomes a dough.
    10. Grease a 9-inch (23 cm) square pan.
    11. Pour the mixture into the pan.
    12. Bake the banana bread in the oven for 40 minutes.
    13. Remove the banana bread from the oven and cool it.
    14. Cut the bread into 16 pieces.
    15. Make the glaze:
    16. Sprinkle the sugar over the bread.
    17. Bake the bread in the oven for 30 minutes.
"""

Prompt template

<|im_start|>user\n
{prompt}<|im_end|>\n
<|im_start|>assistant\n

Sparsification

For details on how this model was sparsified, see the recipe.yaml in this repo and follow the instructions below.

git clone https://github.com/neuralmagic/sparseml
pip install -e "sparseml[transformers]"
wget https://huggingface.co/neuralmagic/TinyLlama-1.1B-Chat-v0.4-pruned50-quant/raw/main/recipe.yaml # download recipe
python sparseml/src/sparseml/transformers/sparsification/obcq/obcq.py TinyLlama/TinyLlama-1.1B-Chat-v0.4 open_platypus --recipe recipe.yaml --save True
python sparseml/src/sparseml/transformers/sparsification/obcq/export.py --task text-generation --model_path obcq_deployment 
cp deployment/model.onnx deployment/model-orig.onnx

Run this kv-cache injection to speed up the model at inference by caching the Key and Value states:

import os
import onnx
from sparseml.exporters.kv_cache_injector import KeyValueCacheInjector
input_file = "deployment/model-orig.onnx"
output_file = "deployment/model.onnx"
model = onnx.load(input_file, load_external_data=False)
model = KeyValueCacheInjector(model_path=os.path.dirname(input_file)).apply(model)
onnx.save(model, output_file)
print(f"Modified model saved to: {output_file}")

Follow the instructions on our One Shot With SparseML page for a step-by-step guide for performing one-shot quantization of large language models.

Slack

For further support, and discussions on these models and AI in general, join Neural Magic's Slack Community