Exploration of Inference Speed of 🤗 Transformers Auto-Regressive Model

import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
import sagemaker
import numpy as np

# turn off transformers tqdm
from transformers.utils import logging
from transformers import BitsAndBytesConfig
import json
from sagemaker.huggingface import HuggingFaceModel

# Turn off logging and tqdm progress bars
logging.set_verbosity_error()
logging.disable_progress_bar()
os.environ["TOKENIZERS_PARALLELISM"] = "false"

model_id = "mistralai/Mistral-7B-v0.1"

# LM doesn't have a trained EOS token so it will go to max_length
generation_config = {
    "num_beams": 1,
    "max_new_tokens": 1000,
    "do_sample": False,
}
# This is the text that will start the LM generation
text = "The following is a extremely long 24 paragraph essay about the detailed plot of Romeo and Juliet:\n\n"
# Tokenizer for all models
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
inputs = tokenizer(
    text,
    return_tensors="pt",
    return_attention_mask=False,
).to("cuda")


# Help from https://deci.ai/blog/measure-inference-time-deep-neural-networks/
def run_model(model):
    model.eval()
    # INIT LOGGERS
    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
        enable_timing=True
    )
    repetitions = 50
    timings = np.zeros((repetitions, 1))
    # GPU-WARM-UP
    for _ in range(10):
        _ = model.generate(**inputs, max_length=35)

    with torch.inference_mode():
        for rep in range(repetitions):
            starter.record()
            outputs = model.generate(**inputs, **generation_config)
            ender.record()
            # WAIT FOR GPU SYNC
            torch.cuda.synchronize()
            curr_time = starter.elapsed_time(ender)
            assert len(outputs[0]) == generation_config["max_new_tokens"] + len(
                tokenizer(text)["input_ids"]
            )
            timings[rep] = curr_time

    mean_syn = np.sum(timings) / repetitions
    print(f"Average time: {mean_syn:0.2f} ms")
    tok_per_sec = generation_config["max_new_tokens"] / (mean_syn / 1000)
    print("Generate tok/s: ", round(tok_per_sec, 2))


# Worst case, GPU while tracking gradients
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="eager",
).to("cuda")
print("Running model while tracking gradients")

run_model(model, tracking_gradients=True)
model = None
torch.cuda.empty_cache()
# Worst case, GPU while tracking gradients
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="eager",
).to("cuda")
print("Running model without tracking gradients")
run_model(model)
model = None
torch.cuda.empty_cache()


# Compile the model with torch compile
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="eager",
).to("cuda")
model = torch.compile(model, mode="max-autotune", backend="inductor")
print("Running compiled model without tracking gradients")
run_model(model)
model = None
torch.cuda.empty_cache()

# Load model and use flash attention
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
).to("cuda")
print("Running model with flash attention")
run_model(model)
model = None
torch.cuda.empty_cache()

# Compile the flash attention model with torch compile
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
).to("cuda")
model = torch.compile(model, mode="max-autotune", backend="inductor")
print("Running model with compiled flash attention")
run_model(model)
model = None
torch.cuda.empty_cache()

# Load an 8bit quantized model
quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    bnb_8bit_compute_dtype=torch.float16,
)

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quant_config,
)
print("Running model with 8bit quantization")
run_model(model)
model = None
torch.cuda.empty_cache()

# load model in 4bit
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quant_config,
)
print("Running model with 4bit quantization")
run_model(model)
model = None
torch.cuda.empty_cache()

"""
Finished with local generation, now run generation on AWS Sagemaker with Text-Generation-Inference Server.
It's ok to use time.time() here because the server is running on a different machine and we don't have to worry about CPU/GPU sync
"""
# copied and edited from https://www.philschmid.de/sagemaker-deploy-mixtral
sess = sagemaker.Session()
sagemaker_session_bucket = sess.default_bucket()
role = sagemaker.get_execution_role()
sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

# Define Model and Endpoint configuration parameter
config = {
    "HF_MODEL_ID": model_id,  # model_id from hf.co/models
    "SM_NUM_GPUS": json.dumps(1),  # Number of GPU used per replica
    "MAX_INPUT_LENGTH": json.dumps(1000),  # Max length of input text
    "MAX_BATCH_PREFILL_TOKENS": json.dumps(
        2100
    ),  # Number of tokens for the prefill operation.
    "MAX_TOTAL_TOKENS": json.dumps(
        2100
    ),  # Max length of the generation (including input text)
    "MAX_BATCH_TOTAL_TOKENS": json.dumps(
        4000
    ),  # Limits the number of tokens that can be processed in parallel during the generation
}

# create HuggingFaceModel with the image uri
llm_model = HuggingFaceModel(hf_llm_image_ver="1.3.3", env=config)

# Deploy model to an endpoint
# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy
llm = llm_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    container_startup_health_check_timeout=600,
)

# Generation arguments
start = time.time()
chat = llm.predict({"inputs": text, "parameters": generation_config})
end = time.time()
tokens = tokenizer(chat[0]["generated_text"][len(text) :], return_tensors="pt")
tokens_per_sec_sagemaker_1 = (len(tokens[0]) - len(tokenizer(text))) / (end - start)
print("Generate tok/s of sagemaker model: ", round(tokens_per_sec_sagemaker_1, 2))
llm.delete_model()
llm.delete_endpoint()

config = {
    "HF_MODEL_ID": model_id,  # model_id from hf.co/models
    "SM_NUM_GPUS": json.dumps(4),  # Number of GPU used per replica
    "MAX_INPUT_LENGTH": json.dumps(1000),  # Max length of input text
    "MAX_BATCH_PREFILL_TOKENS": json.dumps(
        2100
    ),  # Number of tokens for the prefill operation.
    "MAX_TOTAL_TOKENS": json.dumps(
        2100
    ),  # Max length of the generation (including input text)
    "MAX_BATCH_TOTAL_TOKENS": json.dumps(
        4000
    ),  # Limits the number of tokens that can be processed in parallel during the generation
}
llm_model = HuggingFaceModel(hf_llm_image_ver="1.3.3", env=config)
llm = llm_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.12xlarge",
    container_startup_health_check_timeout=600,
)

# Generation arguments
start = time.time()
chat = llm.predict({"inputs": text, "parameters": generation_config})
end = time.time()
tokens = tokenizer(chat[0]["generated_text"][len(text) :], return_tensors="pt")
tokens_per_sec_sagemaker_2 = (len(tokens[0]) - len(tokenizer(text))) / (end - start)
print("Generate tok/s of sagemaker model: ", round(tokens_per_sec_sagemaker_2, 2))
llm.delete_model()
llm.delete_endpoint()