Model-Optimizer/examples/speculative_decoding/example.ipynb

Synthesize data for speculative decoding training

The speculative decoding medule needs to learn to predict tokens from the base model. Therefore, we need to prepare the data generated from the base model. Note: if the target base model is a quantized version, the synthesized data should be generated using the quantized model.

First, quantize the base model (Llama-3.2-1B-Instruct) into FP8 and export to unified export format.

In [ ]:

!python llm_ptq/hf_ptq.py --pyt_ckpt_path meta-llama/Llama-3.2-1B-Instruct --qformat fp8 --batch_size 1 --export_path /tmp/llama3.2_1B_fp8 --export_fmt hf

Next, download the Daring-Anteater dataset.

In [ ]:

!git clone https://huggingface.co/datasets/nvidia/Daring-Anteater /tmp/Daring-Anteater

Then, lanuch an inference server that will run the quantized base model.

In [ ]:

!vllm serve /tmp/llama3.2_1B_fp8 --api-key token-abc123 --port 8000  --tensor-parallel-size 1 --quantization=modelopt

Open a new terminal and adapt the fine-tuning data by calling this server. Note: this may take a long time.

In [ ]:

!mkdir /tmp/finetune
!bash prepare_data.sh --data_path /tmp/Daring-Anteater/train.jsonl --output_path /tmp/finetune/data.jsonl --max_token 2048

Let's load the base model and convert it to EAGLE Model

In [ ]:

import transformers

import modelopt.torch.opt as mto
import modelopt.torch.speculative as mtsp

mto.enable_huggingface_checkpointing()

model = transformers.AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B-Instruct", torch_dtype="auto"
)
config = {
    "eagle_num_layers": 1,
    "use_input_layernorm_in_first_layer": True,
    "use_last_layernorm": False,
}
mtsp.convert(model, [("eagle", config)])

tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
tokenizer.pad_token_id = tokenizer.eos_token_id

Once synthesized data is ready, we can start training the eagle model.

In [ ]:

import json
from dataclasses import dataclass, field

from speculative_decoding.eagle_utils import DataCollatorWithPadding, LazySupervisedDataset
from transformers import Trainer

with open("/tmp/finetune/data.jsonl") as f:
    data_json = [json.loads(line) for line in f]
train_dataset = LazySupervisedDataset(data_json[: int(len(data_json) * 0.95)], tokenizer=tokenizer)
eval_dataset = LazySupervisedDataset(data_json[int(len(data_json) * 0.95) :], tokenizer=tokenizer)


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: str | None = field(default=None)
    model_max_length: int = field(
        default=4096,
        metadata={
            "help": (
                "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
            )
        },
    )
    dataloader_drop_last: bool = field(default=True)
    bf16: bool = field(default=True)


training_args = TrainingArguments(
    output_dir="/tmp/eagle_bf16",
    num_train_epochs=1.0,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
)
trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=DataCollatorWithPadding(),
)
trainer._move_model_to_device(model, trainer.args.device)

# Manually enable this to return loss in eval
trainer.can_return_loss = True
# Make sure label_smoother is None
assert trainer.label_smoother is None, "label_smoother is not supported in speculative decoding!"

trainer.train()
trainer.save_state()
trainer.save_model(training_args.output_dir)
tokenizer.save_pretrained(training_args.output_dir)

metrics = trainer.evaluate()
print(f"Evaluation results: \n{metrics}")

Now we have a EAGLE model in BF16 format. Next, we quantize this model into FP8 (PTQ).

In [ ]:

import modelopt.torch.quantization as mtq
import modelopt.torch.utils.dataset_utils as dataset_utils

mto.enable_huggingface_checkpointing()

model = transformers.AutoModelForCausalLM.from_pretrained("/tmp/eagle_bf16")
tokenizer = transformers.AutoTokenizer.from_pretrained("/tmp/eagle_bf16")

calib_dataloader = dataset_utils.get_dataset_dataloader(
    dataset_name="cnn_dailymail",
    tokenizer=tokenizer,
    batch_size=1,
    num_samples=512,
    device=model.device,
    include_labels=False,
)

quant_cfg = getattr(mtq, "FP8_DEFAULT_CFG")
quant_cfg["quant_cfg"]["*output_quantizer"] = {
    "num_bits": (4, 3),
    "axis": None,
    "enable": True,
}

calibrate_loop = dataset_utils.create_forward_loop(calib_dataloader, dataloader=calib_dataloader)
model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
mtq.print_quant_summary(model)

model.save_pretrained("/tmp/eagle_fp8_ptq")

To maintain the accuracy, we need to finetune the model (QAT).

In [ ]:

training_args.output_dir = "/tmp/eagle_fp8_qat"
trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=DataCollatorWithPadding(),
)
trainer._move_model_to_device(model, trainer.args.device)

# Manually enable this to return loss in eval
trainer.can_return_loss = True
# Make sure label_smoother is None
assert trainer.label_smoother is None, "label_smoother is not supported in speculative decoding!"

trainer.train()
trainer.save_state()
trainer.save_model(training_args.output_dir)
tokenizer.save_pretrained(training_args.output_dir)

metrics = trainer.evaluate()
print(f"Evaluation results: \n{metrics}")

To deploy this model, we need to first export it to a Unified checkpoint.

In [ ]:

from accelerate.hooks import remove_hook_from_module

from modelopt.torch.export import export_hf_checkpoint

# Move meta tensor back to device before exporting.
remove_hook_from_module(model, recurse=True)

export_hf_checkpoint(
    model,
    export_dir="/tmp/hf_ckpt",
)

Then convert the Unified ckeckpoint to TRTLLM checkpoint.

In [ ]:

!python TensorRT-LLM/examples/eagle/convert_checkpoint.py --model_dir /tmp/hf_ckpt --output_dir /tmp/trtllm_ckpt --num_eagle_layers 5 --max_non_leaves_per_layer 4 --max_draft_len 25 --dtype float16

Last, build a TensorRT-LLM engine.

In [ ]:

!trtllm-build --checkpoint_dir /tmp/trtllm_ckpt --output_dir /tmp/trtllm_engine --gemm_plugin float16 --use_paged_context_fmha enable --speculative_decoding_mode eagle  --max_batch_size 4

To run the EAGLE engine, please refer to TensorRT-LLM/examples/eagle:

In [ ]:

!python ../run.py --engine_dir /tmp/trtllm_engine \
                  --tokenizer_dir /tmp/eagle_fp8_qat \
                  --max_output_len=100 \
                  --eagle_choices="[[0],[1],[2],[3],[0,0],[0,1],[0,2],[1,0],[1,1],[2,0],[2,1],[3,0],[0,0,0],[0,0,1],[0,0,2],[0,1,0],[0,1,1],[0,2,0],[0,2,1],[1,0,0],[0,0,0,0],[0,0,0,1],[0,0,0,2],[0,0,0,0,0],[0,0,0,0,1]]" \
                  --temperature 1.0 \
                  --input_text "Once upon"