Chapter 7. Compressing language models with Red Hat AI Model Optimization Toolkit

Procedure

1. Open a shell prompt on the RHEL AI server.

2. Stop the Red Hat AI Inference Server service:

   [cloud-user@localhost ~]$ systemctl stop rhaiis

3. Create a working directory:

   [cloud-user@localhost ~]$ mkdir -p model-opt

4. Change permissions on the project folder and enter the folder:

   [cloud-user@localhost ~]$ chmod 775 model-opt && cd model-opt

5. Add the compression recipe Python script. For example, create the following example.py file that compresses the TinyLlama/TinyLlama-1.1B-Chat-v1.0 model in quantized FP8 format:

   from transformers import AutoModelForCausalLM, AutoTokenizer
   
   from llmcompressor import oneshot
   from llmcompressor.modifiers.quantization import QuantizationModifier
   from llmcompressor.utils import dispatch_for_generation
   
   import os
   
   MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
   
   model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
   tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
   
   # Configure the quantization algorithm and scheme
   recipe = QuantizationModifier(
       targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]
   )
   
   # Create log directory in a writable location
   LOG_DIR = "./sparse_logs"
   os.makedirs(LOG_DIR, exist_ok=True)
   
   # Apply quantization
   oneshot(model=model, recipe=recipe)
   
   # Confirm quantized model looks OK
   print("========== SAMPLE GENERATION ==============")
   dispatch_for_generation(model)
   input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to(
       model.device
   )
   output = model.generate(input_ids, max_new_tokens=20)
   print(tokenizer.decode(output[0]))
   print("==========================================")
   
   # Save to disk in compressed-tensors format
   SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-Dynamic"
   model.save_pretrained(SAVE_DIR)
   tokenizer.save_pretrained(SAVE_DIR)

6. Export your Hugging Face token.

   $ export HF=<YOUR_HUGGING_FACE_TOKEN>

7. Run the compression recipe using the Red Hat AI Model Optimization Toolkit container:

   [cloud-user@localhost ~]$ sudo podman run -it \
     -v ./model-opt:/opt/app-root/model-opt:z \
     -v /var/lib/rhaiis/models:/opt/app-root/models:z \
     --device nvidia.com/gpu=all \
     --workdir /opt/app-root/model-opt \
     -e HF_HOME=/opt/app-root/models \
     -e HF_TOKEN=$HF \
     --entrypoint python \
     registry.redhat.io/rhaiis/model-opt-cuda-rhel9:3.4.0-ea.1 \
     python example.py