Obtain ~2x speed-up in LLM inference with Medusa-1 on Amazon SageMaker AI

This weblog submit is co-written with Moran Beladev, Manos Stergiadis, and Ilya Gusev from Reserving.com.

Giant language fashions (LLMs) have revolutionized the sector of pure language processing with their capacity to grasp and generate humanlike textual content. Skilled on broad, generic datasets spanning a variety of matters and domains, LLMs use their parametric data to carry out more and more advanced and versatile duties throughout a number of enterprise use instances. Moreover, corporations are more and more investing sources in customizing LLMs by way of few-shot studying and fine-tuning to optimize their efficiency for specialised functions.

Nevertheless, the spectacular efficiency of LLMs comes at the price of important computational necessities, pushed by their massive variety of parameters and autoregressive decoding course of which is sequential in nature. This mix makes reaching low latency a problem to be used instances comparable to real-time textual content completion, simultaneous translation, or conversational voice assistants, the place subsecond response instances are important.

Researchers developed Medusa, a framework to hurry up LLM inference by including additional heads to foretell a number of tokens concurrently. This submit demonstrates how you can use Medusa-1, the primary model of the framework, to hurry up an LLM by fine-tuning it on Amazon SageMaker AI and confirms the pace up with deployment and a easy load check. Medusa-1 achieves an inference speedup of round two instances with out sacrificing mannequin high quality, with the precise enchancment various primarily based on mannequin dimension and information used. On this submit, we display its effectiveness with a 1.8 instances speedup noticed on a pattern dataset.

Introduction to Medusa and its advantages for LLM inference pace

LLMs generate textual content in a sequential method, which entails autoregressive sampling, with every new token conditional on the earlier ones. Producing Okay tokens necessitates Okay sequential executions of the mannequin. This token-by-token processing introduces an inherent latency and computational overhead as a result of the mannequin must carry out a separate ahead go for every new token within the output sequence. The next diagram from Function-Play with Giant Language Fashions illustrates this circulation.

Speculative decoding tackles this problem by utilizing a smaller, quicker draft mannequin to generate a number of potential token continuations in parallel, that are then verified by a bigger, extra correct goal mannequin. This parallelization accelerates textual content era whereas sustaining the standard of the goal mannequin as a result of the verification process is quicker than autoregressive token era. For an in depth clarification of the idea, discuss with the paper Accelerating Giant Language Mannequin Decoding with Speculative Sampling. The speculative decoding method will be applied utilizing the inference optimization toolkit on Amazon SageMaker Jumpstart.

The paper Medusa: Easy LLM Inference Acceleration Framework with A number of Decoding Heads launched Medusa as a substitute for speculative decoding. As a substitute of including a separate draft mannequin, it provides additional decoding heads to the LLM that generate candidate continuations concurrently. These candidates are then evaluated in parallel utilizing a tree-based consideration mechanism. This parallel processing reduces the variety of sequential steps wanted, resulting in quicker inference instances. The primary benefit of Medusa over speculative decoding is that it eliminates the necessity to purchase and preserve a separate draft mannequin whereas reaching greater speedups. For instance, when examined on the MT-Bench dataset, the paper experiences that Medusa-2 (the second model of Medusa) accelerates inference time by 2.8 instances. This outperforms speculative decoding, which solely manages to hurry up inference time by 1.5 instances on the identical dataset.

The Medusa framework at the moment helps Llama and Mistral fashions. Though it presents important pace enhancements, it does include a reminiscence trade-off (much like speculative decoding). As an example, including 5 Medusa heads to the 7-billion-parameter Mistral mannequin will increase the full parameter rely by 750 million (150 million per head), which suggests these extra parameters should be saved in GPU reminiscence, resulting in a better reminiscence requirement. Nevertheless, generally, this improve doesn’t necessitate switching to a better GPU reminiscence occasion. For instance, you’ll be able to nonetheless use an ml.g5.4xlarge occasion with 24 GB of GPU reminiscence to host your 7-billion-parameter Llama or Mistral mannequin with additional Medusa heads.

Coaching Medusa heads requires extra growth time and computational sources, which needs to be factored into undertaking planning and useful resource allocation. One other essential limitation to say is that the present framework, when deployed on an Amazon SageMaker AI endpoint, solely helps a batch dimension of 1—a configuration sometimes used for low-latency functions.

The next diagram from the unique Medusa paper authors’ FasterDecoding repository offers a visible Medusa framework overview.

There are two principal variants of Medusa:

Medusa-1 – Requires a two-stage method the place you first fine-tune your LLM after which add Medusa heads and prepare them on high of your frozen fine-tuned LLM
Medusa-2 – Launched later as an enchancment, fine-tunes each the extra heads and the spine LLM parameters collectively, enabling doubtlessly even additional latency speedups

The Medusa paper experiences that throughout fashions of various sizes, you’ll be able to obtain inference speedups of round two instances for Medusa-1 and round thrice for Medusa-2. With Medusa-1, the predictions are equivalent to these of the initially fine-tuned LLM. In distinction, with Medusa-2, we’d observe barely completely different outcomes in comparison with easy fine-tuning of the LLM as a result of each the heads and the spine LLM parameters are up to date collectively. On this submit, we give attention to Medusa-1.

Resolution overview

We cowl the next steps in our answer:

Stipulations
Load and put together the dataset
Advantageous-tune an LLM utilizing a SageMaker AI coaching job
Prepare Medusa heads on high of a frozen fine-tuned LLM utilizing a SageMaker AI coaching job
Deploy the fine-tuned LLM with Medusa heads on a SageMaker AI endpoint
Reveal LLM inference speedup

By following this answer, you’ll be able to speed up LLM inference in your functions, resulting in quicker response instances and improved consumer expertise.

Stipulations

To construct the answer your self, there are the next stipulations:

Load and put together the dataset

Now that you’ve cloned the GitHub repository and opened the medusa_1_train.ipynb pocket book, you’ll load and put together the dataset within the pocket book. We encourage you to learn this submit whereas operating the code within the pocket book. For this submit, we use a dataset known as sql-create-context, which incorporates samples of pure language directions, schema definitions and the corresponding SQL question. It incorporates 78,577 examples of pure language queries, SQL CREATE TABLE statements, and SQL queries answering the query utilizing the CREATE assertion as context. For demonstration functions, we choose 3,000 samples and cut up them into prepare, validation, and check units.

It’s essential to run the “Load and put together the dataset” part of the medusa_1_train.ipynb to organize the dataset for fine-tuning. We additionally included a knowledge exploration script to research the size of enter and output tokens. After information exploration, we put together the prepare, validation, and check units and add them to Amazon Easy Storage Service (Amazon S3).

Advantageous-tune an LLM utilizing SageMaker AI coaching job

We use the Zephyr 7B β mannequin as our spine LLM. Zephyr is a sequence of language fashions educated to behave as useful assistants, and Zephyr 7B β is a fine-tuned model of Mistral-7B-v0.1, educated on a mixture of publicly accessible and artificial datasets utilizing Direct Desire Optimization.

To launch a SageMaker AI coaching job, we have to use the PyTorch or Hugging Face estimator. SageMaker AI begins and manages all the required Amazon Elastic Compute Cloud (Amazon EC2) situations for us, provides the suitable containers, downloads information from our S3 bucket to the container and uploads and runs the required coaching script, in our case fine_tune_llm.py. We choose the hyperparameters primarily based on the QLoRA paper, however we encourage you to experiment with your personal combos. To expedite the execution of this code, we set the variety of epochs to 1. Nevertheless, for higher outcomes, it’s typically beneficial to set the variety of epochs to at the least 2 or 3.

from sagemaker.pytorch.estimator import PyTorch
from sagemaker.debugger import TensorBoardOutputConfig
import time
import os

def get_current_time():
    return time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())

def create_estimator(hyperparameters_dict, job_name, function, sess, train_scipt_path):
    metric=[
        {"Name": "loss", "Regex": r"'loss':s*([0-9.]+)"},
        {"Title": "epoch", "Regex": r"'epoch':s*([0-9.]+)"},
    ]

    tensorboard_s3_output_path = os.path.be part of(
       "s3://", sess.default_bucket(), job_name, 'tensorboard'
    )
    print("Tensorboard output path:", tensorboard_s3_output_path)

    tensorboard_output_config = TensorBoardOutputConfig(
        s3_output_path=tensorboard_s3_output_path,
        container_local_output_path=hyperparameters_dict['logging_dir']
    )
    estimator = PyTorch(
        sagemaker_session    = sess,
        entry_point          = train_scipt_path,    # prepare script
        source_dir="prepare",      # listing which incorporates all of the information wanted for coaching
        instance_type="ml.g5.4xlarge",   # situations sort used for the coaching job, "local_gpu" for native mode
        metric_definitions   = metric,
        instance_count       = 1,                 # the variety of situations used for coaching
        function                 = function,              # Iam function utilized in coaching job to entry AWS ressources, e.g. S3
        volume_size          = 300,               # the scale of the EBS quantity in GB
        framework_version      = '2.1.0',             # the pytorch_version model used within the coaching job
        py_version           = 'py310',           # the python model used within the coaching job
        hyperparameters      =  hyperparameters_dict,  # the hyperparameters handed to the coaching job
        disable_output_compression = True,        # not compress output to save lots of coaching time and value
        tensorboard_output_config = tensorboard_output_config
    )
    return estimator
    
# hyperparameters, that are handed into the coaching job
sft_hyperparameters = {
  ### SCRIPT PARAMETERS ###
  'train_dataset_path': '/choose/ml/enter/information/prepare/train_dataset.json', # path the place sagemaker will save coaching dataset
  'eval_dataset_path': '/choose/ml/enter/information/eval/eval_dataset.json', # path the place sagemaker will save analysis dataset
  'model_id': model_id,
  'max_seq_len': 256,                               # max sequence size for mannequin and packing of the dataset
  'use_qlora': True,                                 # use QLoRA mannequin
  ### TRAINING PARAMETERS ###
  'num_train_epochs': 1,                             # variety of coaching epochs
  'per_device_train_batch_size': 1,                  # batch dimension per gadget throughout coaching
  'gradient_accumulation_steps': 16,                  # variety of steps earlier than performing a backward/replace go
  'gradient_checkpointing': True,                    # use gradient checkpointing to save lots of reminiscence
  'optim': "adamw_8bit",                             # use fused adamw 8bit optimizer
  'logging_steps': 15,                               # log each 10 steps
  'save_strategy': "steps",                          # save checkpoint each epoch
  'save_steps': 15,
  'save_total_limit': 2,
  'eval_strategy': "steps",
  'eval_steps': 15,
  'learning_rate': 1e-4,                             # studying price, primarily based on QLoRA paper
  'bf16': True,                                      # use bfloat16 precision
  'max_grad_norm': 10,                              # max gradient norm primarily based on QLoRA paper
  'warmup_ratio': 0.03,                              # warmup ratio primarily based on QLoRA paper
  'lr_scheduler_type': "fixed",                   # use fixed studying price scheduler
  'output_dir': '/choose/ml/checkpoints/',              # Momentary output listing for mannequin checkpoints
  'merge_adapters': True,                            # merge LoRA adapters into mannequin for simpler deployment
  'report_to': "tensorboard",                        # report metrics to tensorboard
  'logging_dir': "/choose/ml/output/tensorboard"        # tensorboard logging listing
}
 
sft_job_name = f"sft-qlora-text-to-sql-{get_current_time()}"
information = {
    'prepare': train_dataset_path,
    'eval': eval_dataset_path
}

sft_estimator = create_estimator(sft_hyperparameters, sft_job_name, function, sess, "fine_tune_llm.py")

sft_estimator.match(job_name=sft_job_name, inputs=information, wait=False)

When our coaching job has accomplished efficiently after roughly 1 hour, we are able to use the fine-tuned mannequin artifact for the following step, coaching the Medusa heads on high of it. To visualise the coaching metrics in Tensorboard, you’ll be able to observe the steering on this documentation: Load and visualize output tensors utilizing the TensorBoard software

Prepare Medusa heads on high of frozen fine-tuned LLM utilizing a SageMaker AI coaching job

For coaching Medusa heads, we are able to reuse the capabilities beforehand talked about to launch the coaching job. We chosen hyperparameters primarily based on a mix of what the Medusa paper reported and what we discovered to be finest performing after a number of experiments. We set the variety of Medusa heads to five and used the 8-bit AdamW optimizer, as beneficial by the paper. For simplicity, we maintained a continuing studying price of 1e-4 with a continuing scheduler, much like the earlier fine-tuning step. Though the paper recommends an elevated studying price and a cosine scheduler, we discovered that our chosen mixture of hyperparameters carried out properly on this dataset. Nevertheless, we encourage you to experiment with your personal hyperparameter settings to doubtlessly obtain even higher outcomes.

# hyperparameters, that are handed into the coaching job
medusa_hyperparameters = {
  ### SCRIPT PARAMETERS ###
  'train_dataset_path': '/choose/ml/enter/information/prepare/train_dataset.json', # path the place sagemaker will save coaching dataset
  'eval_dataset_path': '/choose/ml/enter/information/eval/eval_dataset.json', # path the place sagemaker will save analysis dataset
  'model_path': '/choose/ml/enter/information/fine-tuned-model/',
  'max_seq_len': 256,                               # max sequence size for mannequin and packing of the dataset
  'medusa_num_heads': 5,
  ### TRAINING PARAMETERS ###
  'num_train_epochs': 3,                             # variety of coaching epochs
  'per_device_train_batch_size': 1,                  # batch dimension per gadget throughout coaching
  'gradient_accumulation_steps': 16,                  # variety of steps earlier than performing a backward/replace go
  'gradient_checkpointing': True,                    # use gradient checkpointing to save lots of reminiscence
  'optim': "adamw_8bit",                             # use fused adamw 8bit optimizer
  'logging_steps': 15,                               # log each 10 steps
  'save_strategy': "steps",                          # save checkpoint each epoch
  'save_steps': 15,
  'save_total_limit':2,
  'eval_strategy': "steps",
  'eval_steps': 15,
  'learning_rate': 1e-4,                             # studying price
  'bf16': True,                                      # use bfloat16 precision
  'max_grad_norm': 10,                              # max gradient norm primarily based on QLoRA paper
  'warmup_ratio': 0.03,                              # warmup ratio primarily based on QLoRA paper
  'lr_scheduler_type': "fixed",                   # use fixed studying price scheduler
  'output_dir': '/choose/ml/checkpoints/',              # Momentary output listing for mannequin checkpoints
  'report_to': "tensorboard",                        # report metrics to tensorboard
  'logging_dir': "/choose/ml/output/tensorboard"        # tensorboard logging listing
}

medusa_train_job_name = f"medusa-text-to-sql-{get_current_time()}"
information = {
    'prepare': train_dataset_path,
    'eval': eval_dataset_path,
    'fine-tuned-model': fine_tuned_model_path
}

medusa_estimator = create_estimator(medusa_hyperparameters, medusa_train_job_name, function, sess, "train_medusa_heads.py")

medusa_estimator.match(job_name=medusa_train_job_name, inputs=information, wait=False)

We discovered that after 3 epochs, the analysis lack of Medusa heads was converging, which will be noticed within the TensorBoard graph within the following picture.

Moreover the hyperparameters, the principle distinction is that we go train_medusa_heads.py because the coaching entrypoint, the place we first add Medusa heads, then freeze the fine-tuned LLM, and we create customized MedusaSFTTrainer class, which is a subclass of the transformers SFTTrainer.

# Add medusa heads and freeze base mannequin
add_medusa_heads(
    mannequin,
    medusa_num_heads=script_args.medusa_num_heads,
)
freeze_layers(mannequin)
mannequin.config.torch_dtype = torch_dtype
mannequin.config.use_cache = False

logger.information("Completed loading mannequin and medusa heads")

tokenizer = AutoTokenizer.from_pretrained(script_args.model_path, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token

################
# Coaching
################
coach = MedusaSFTTrainer(
    mannequin=mannequin,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    max_seq_length=script_args.max_seq_length,
    tokenizer=tokenizer,
    dataset_kwargs={
        "add_special_tokens": False,  # We template with particular tokens
        "append_concat_token": False,  # No want so as to add extra separator token
    },
    medusa_num_heads=script_args.medusa_num_heads,
    medusa_heads_coefficient=script_args.medusa_heads_coefficient,
    medusa_decay_coefficient=script_args.medusa_decay_coefficient,
    medusa_scheduler=script_args.medusa_scheduler,
    train_only_medusa_heads=script_args.train_only_medusa_heads,
    medusa_lr_multiplier=script_args.medusa_lr_multiplier
)
coach.prepare()

Within the add_medusa_heads() operate, we add the residual blocks of the Medusa heads, and in addition override the ahead go for our mannequin to verify to not prepare the frozen spine LLM:

def add_medusa_heads(
    mannequin,
    medusa_num_heads,
):
    """
    Args:
        mannequin (nn.Module): The bottom language mannequin for use.
        medusa_num_heads (int, elective): Variety of extra tokens to foretell
    """
    hidden_size = mannequin.lm_head.weight.form[-1]
    vocab_size = mannequin.lm_head.weight.form[0]
    mannequin.config.medusa_num_layers = 1
    mannequin.config.medusa_num_heads = medusa_num_heads
    mannequin.medusa_num_heads = medusa_num_heads
    # Create a listing of Medusa heads
    mannequin.medusa_heads = nn.ModuleList(
        [
            nn.Sequential(
                ResBlock(hidden_size),
                nn.Linear(hidden_size, vocab_size, bias=False),
            )
            for _ in range(medusa_num_heads)
        ]
    )

    # Guarantee medusa_head's dtype and gadget align with the base_model
    mannequin.medusa_heads.to(mannequin.dtype).to(mannequin.gadget)
    logger.information(f"Loading medusa heads in {str(mannequin.dtype)} to gadget {mannequin.gadget}")

    for i in vary(medusa_num_heads):
        # Initialize the weights of every medusa_head utilizing the bottom mannequin's weights
        mannequin.medusa_heads[i][-1].weight.information[:] = mannequin.lm_head.weight.information[:]

    def ahead(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Non-compulsory[torch.Tensor] = None,
        position_ids: Non-compulsory[torch.LongTensor] = None,
        past_key_values: Non-compulsory[List[torch.FloatTensor]] = None,
        inputs_embeds: Non-compulsory[torch.FloatTensor] = None,
        labels: Non-compulsory[torch.LongTensor] = None,
        use_cache: Non-compulsory[bool] = None,
        output_attentions: Non-compulsory[bool] = None,
        output_hidden_states: Non-compulsory[bool] = None,
        return_dict: Non-compulsory[bool] = None,
        train_only_medusa_heads: bool = False,
    ):
        """Ahead go of the MedusaModel.
        Returns:
            torch.Tensor: A tensor containing predictions from all Medusa heads.
            (Non-compulsory) Authentic predictions from the bottom mannequin's LM head.
        """
        maybe_grad = torch.no_grad() if train_only_medusa_heads else nullcontext()
        with maybe_grad:
            outputs = self.mannequin(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
            hidden_states = outputs[0]
            medusa_logits = [self.lm_head(hidden_states)]
        for i in vary(self.medusa_num_heads):
            medusa_logits.append(self.medusa_heads[i](hidden_states))
        return torch.stack(medusa_logits, dim=0)

    mannequin.ahead = varieties.MethodType(ahead, mannequin)

After the mannequin coaching is completed (which takes 1 hour), we put together the mannequin artefacts for deployment and add it to Amazon S3. Your ultimate mannequin artifact incorporates each the unique fine-tuned mannequin from the earlier step underneath the base-model prefix and the educated Medusa heads in a file named medusa_heads.safetensors.

Deploy the fine-tuned LLM with Medusa heads on a SageMaker AI endpoint

The Medusa framework is supported by the Textual content Era Inference (TGI) server. After coaching the LLM with Medusa heads, we deploy it to a SageMaker AI real-time endpoint utilizing the Hugging Face Inference Container arrange with TGI.

First, we create a SageMaker AI HuggingFaceModel object after which deploy the mannequin to an endpoint with the next operate:

import json
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri


def deploy_model(endpoint_name, instance_type, model_s3_path=None, hf_model_id=None):
    llm_image = get_huggingface_llm_image_uri(
      "huggingface",
      model="2.2.0",
      session=sess,
    )

    print(f"llm picture uri: {llm_image}")

    model_data = None
    if model_s3_path:
        model_data = {'S3DataSource': {'S3Uri': model_s3_path, 'S3DataType': 'S3Prefix', 'CompressionType': 'None'}}
        hf_model_id = "/choose/ml/mannequin"
    else:
        assert hf_model_id, "It's essential to present both pretrained HF mannequin id, or S3 mannequin information to deploy"
    config = {
      'HF_MODEL_ID': hf_model_id,  # path to the place sagemaker shops the mannequin
      'SM_NUM_GPUS': json.dumps(1),  # Variety of GPU used per reproduction
      'MAX_INPUT_LENGTH': json.dumps(1024),  # Max size of enter textual content
      'MAX_TOTAL_TOKENS': json.dumps(2048),  # Max size of the era (together with enter textual content)
    }

    llm_model = HuggingFaceModel(
      identify=endpoint_name,
      function=function,
      image_uri=llm_image,
      model_data=model_data,
      env=config
    )

    deployed_llm = llm_model.deploy(
      endpoint_name=endpoint_name,
      initial_instance_count=1,
      instance_type=instance_type,
      container_startup_health_check_timeout=300,
    )
    return deployed_llm

We deploy three LLMs on three SageMaker AI endpoints:

Base LLM which isn’t fine-tuned
The LLM that we fine-tuned
The fine-tuned LLM that additionally has educated Medusa heads

You possibly can deploy the three fashions in parallel by utilizing a operate that we included within the pocket book, or you’ll be able to deploy the fashions one after the other by operating the code under:

base_deployed_llm = deploy_model( f"base-{get_current_time()}", instance_type="ml.g5.4xlarge", model_s3_path=None, hf_model_id=model_id )
sft_deployed_llm = deploy_model( f"sft-{get_current_time()}", instance_type="ml.g5.4xlarge", model_s3_path=fine_tuned_model_path )
medusa_deployed_llm = deploy_model( f"medusa-{get_current_time()}", instance_type="ml.g5.4xlarge", model_s3_path=medusa_trained_model_path )

After the standing for every endpoint turns into InService, which ought to take round quarter-hour, we are able to invoke them for inference. We ship the next enter:

“You're a textual content to SQL question translator. Customers will ask you questions in English and you'll generate a SQL question primarily based on the offered SCHEMA. SCHEMA: CREATE TABLE table_name_32 (time VARCHAR, warmth VARCHAR, identify VARCHAR)
What was Liu Limin's time in warmth 4?“

We are able to observe the next responses:

The bottom LLM response incorporates additional phrases that aren’t wanted:

“To retrieve the time of Liu Limin in warmth 4 primarily based on the offered SCHEMA, we have to write a SQL question. For the reason that desk identify shouldn't be offered, let 's assume it is "my_table".  Assuming the desk has a main key (which isn't talked about within the SCHEMA), we are able to be part of the desk with itself to check the warmth and identify of Liu Limin in warmth 4 with all of the information within the desk.  This is the SQL question:  ``sql SELECT t1.time FROM my_table t1 JOIN my_table t2 ON t1.identify = t2.identify AND t2.warmth = 4 WHERE t1.identify="Liu Limin" AND t1.warmth <> 4; `  Clarification:  1. We're choosing the time` column from the primary desk (t1) within the FROM clause. 2. We're becoming a member of the desk with itself (my_table t1 JOIN my_table t2) to check the identify and warmth of Liu Limin in warmth 4 with all of the information within the desk. 3. We're filtering the outcomes utilizing the WHERE clause. We're choosing solely the information the place the identify is 'Liu Limin' and the warmth shouldn't be equal to 4 (i.e., not warmth 4). That is to make sure that we're choosing the time of Liu Limin in warmth 3.  Be aware: This question assumes that the desk has a novel main key. If the desk does not have a main key, it's possible you'll want so as to add extra situations to the JOIN and WHERE clauses to make sure that we're choosing the proper information.“

The fine-tuned LLM response is improved considerably, and incorporates solely the required output:

'SELECT time FROM table_name_32 WHERE warmth = 4 AND identify = "liu limin"'

The fine-tuned LLM with educated Medusa heads gives the very same response because the fine-tuned mannequin, demonstrating that Medusa-1, by design, maintains the output (high quality) of the unique mannequin:

'SELECT time FROM table_name_32 WHERE warmth = 4 AND identify = "liu limin"'

Reveal LLM inference speedup

To measure the inference pace enhancements, we examine the response instances of the deployed fine-tuned LLM and the fine-tuned LLM with Medusa heads on 450 check observations with the next code:

import time
import numpy as np
from tqdm import tqdm

def request(pattern, deployed_llm):
    immediate = tokenizer.apply_chat_template(pattern, tokenize=False, add_generation_prompt=True)
    outputs = deployed_llm.predict({
      "inputs": immediate,
      "parameters": {
        "max_new_tokens": 512,
        "do_sample": False,
        "return_full_text": False,
      }
    })
    return {"function": "assistant", "content material": outputs[0]["generated_text"].strip()}

def predict(deployed_llm, test_dataset):
    predicted_answers = []
    latencies = []

    for pattern in tqdm(test_dataset):
        start_time = time.time()
        predicted_answer = request(pattern["messages"][:2], deployed_llm)
        end_time = time.time()

        latency = end_time - start_time
        latencies.append(latency)
        predicted_answers.append(predicted_answer)

    # Calculate p90 and common latencies
    p90_latency = np.percentile(latencies, 90)
    avg_latency = np.imply(latencies)

    print(f"P90 Latency: {p90_latency:.2f} seconds")
    print(f"Common Latency: {avg_latency:.2f} seconds")

    return predicted_answers

First, we run predictions utilizing the fine-tuned LLM:

sft_predictions = predict(sft_deployed_llm, test_dataset)
P90 Latency: 1.28 seconds
Common Latency: 0.95 seconds

Then, we run predictions utilizing the fine-tuned LLM with Medusa heads:

medusa_predictions = predict(medusa_deployed_llm, test_dataset)
P90 Latency: 0.80 seconds
Common Latency: 0.53 seconds

The prediction runs ought to take round 8 and 4 minutes respectively. We are able to observe that the typical latency decreased from 950 to 530 milliseconds, which is an enchancment of 1.8 instances. You possibly can obtain even greater enhancements in case your dataset incorporates longer inputs and outputs. In our dataset, we solely had a median of 18 enter tokens and 30 output tokens.

We need to as soon as once more spotlight that, with this system, the output high quality is absolutely maintained, and all of the prediction outputs are the identical. The mannequin responses for the check set of 450 observations are the identical for each with Medusa heads and with out Medusa heads:

match_percentage = sum(a["content"] == b["content"] for a, b in zip(sft_predictions, medusa_predictions)) / len(sft_predictions) * 100
print(f"Predictions with the fine-tuned mannequin with medusa heads are the identical as with out medusa heads: {match_percentage:.2f}% of check set ")

Predictions with fine-tuned mannequin with medusa heads are the identical as with out medusa heads: 100.00% of check set

You may discover in your run that a number of observations aren’t precisely matching, and also you may get a 99% match attributable to small errors in floating level operations brought on by optimizations on GPUs.

Cleanup

On the finish of this experiment, don’t neglect to delete the SageMaker AI endpoints you created:

base_deployed_llm.delete_model()
base_deployed_llm.delete_endpoint()
sft_deployed_llm.delete_model()
sft_deployed_llm.delete_endpoint()
medusa_deployed_llm.delete_model()
medusa_deployed_llm.delete_endpoint()

Conclusion

On this submit, we demonstrated how you can fine-tune and deploy an LLM with Medusa heads utilizing the Medusa-1 method on Amazon SageMaker AI to speed up LLM inference. Through the use of this framework and SageMaker AI scalable infrastructure, we confirmed how you can obtain as much as twofold speedups in LLM inference whereas sustaining mannequin high quality. This answer is especially useful for functions requiring low-latency textual content era, comparable to customer support chat assistants, content material creation, and advice programs.

As a subsequent step, you’ll be able to discover fine-tuning your personal LLM with Medusa heads by yourself dataset and benchmark the outcomes in your particular use case, utilizing the offered GitHub repository.

In regards to the authors

Daniel Zagyva is a Senior ML Engineer at AWS Skilled Companies. He makes a speciality of growing scalable, production-grade machine studying options for AWS clients. His expertise extends throughout completely different areas, together with pure language processing, generative AI and machine studying operations.

Aleksandra Dokic is a Senior Information Scientist at AWS Skilled Companies. She enjoys supporting clients to construct progressive AI/ML options on AWS and he or she is happy about enterprise transformations by way of the ability of information.

Moran Beladev is a Senior ML Supervisor at Reserving.com. She is main the content material intelligence observe which is targeted on constructing, coaching and deploying content material fashions (pc imaginative and prescient, NLP and generative AI) utilizing essentially the most superior applied sciences and fashions. Moran can also be a PhD candidate, researching making use of NLP fashions on social graphs.

Manos Stergiadis is a Senior ML Scientist at Reserving.com. He makes a speciality of generative NLP and has expertise researching, implementing and deploying massive deep studying fashions at scale.

Ilya Gusev is a Senior Machine Studying Engineer at Reserving.com. He leads the event of the a number of LLM programs inside Reserving.com. His work focuses on constructing manufacturing ML programs that assist thousands and thousands of vacationers plan their journeys successfully.

Laurens van der Maas is a Machine Studying Engineer at AWS Skilled Companies. He works intently with clients constructing their machine studying options on AWS, makes a speciality of pure language processing, experimentation and accountable AI, and is keen about utilizing machine studying to drive significant change on the earth.