LLMの​推論を​スピードアップ!TensorRT Model Optimizerで​投機的デコードを​試してみよう

from copy import deepcopy
from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.speculative as mtsp

# ベースモデルを読み込み
model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.bfloat16, device_map="cuda")

# EAGLEモジュールの設定をモデルに合わせて更新
eagle_config = deepcopy(mtsp.EAGLE1_DEFAULT_CFG)
eagle_config["config"]["eagle_architecture_config"].update({
    "hidden_size": model.config.hidden_size,
    "vocab_size": model.config.vocab_size,
    # ... その他モデルに合わせた設定
})
model = mtsp.convert(model, [("eagle", eagle_config["config"])])

import modelopt.torch.opt as mto
from transformers import Trainer, TrainingArguments

# EAGLEモジュールのみを学習対象に
for name, param in model.named_parameters():
    if "eagle" not in name:
        param.requires_grad = False

mto.enable_huggingface_checkpointing()

# 学習の実行
training_args = TrainingArguments(
    output_dir="eagle_tinyllama_v2",
    num_train_epochs=3,
    # ... その他学習パラメータ
)
trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("eagle_tinyllama_v2")

# チェックポイントへの変換
python3 -m tensorrt_llm.commands.convert_checkpoint \
    --model_dir eagle_tinyllama_v2_hf \
    --output_dir trtllm_eagle_ckpt \
    --dtype float16 \
    --speculative_decoding_mode eagle

# エンジンのビルド
trtllm-build \
    --checkpoint_dir trtllm_eagle_ckpt \
    --output_dir trtllm_eagle_engine \
    --gemm_plugin float16 \
    --max_batch_size 1 \
    --speculative_decoding_mode eagle \
    --max_draft_len 63

import tensorrt_llm.bindings.executor as tllm

# EAGLE用の設定を追加
eagle_config = tllm.EagleConfig(use_dynamic_tree=True, dynamic_tree_max_topK=10)
decoding_config = tllm.DecodingConfig(
    decoding_mode=tllm.DecodingMode.Eagle(),
    eagle_config=eagle_config,
)
executor_config = tllm.ExecutorConfig(max_beam_width=1, decoding_config=decoding_config)

# エンジンを読み込んで推論を実行
executor = tllm.Executor("trtllm_eagle_engine", tllm.ModelType.DECODER_ONLY, executor_config)
request = tllm.Request(input_ids, max_tokens=100, end_id=eos_id, pad_id=eos_id)
request_id = executor.enqueue_request(request)
responses = executor.await_responses(request_id)

Sea otters is a fascinating and unique species that has adapted to 
life in the Pacific Northwest. They are known for their unique fur, 
which is a combination of brown and white, and their ability to live 
in the water for extended periods.

ラッコは、太平洋岸北西部の環境に適応した、魅力的でユニークな種です。
茶色と白が混ざり合った独特の毛皮と、長時間水中で生活できる能力で知られています。