mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
Keval Morabia
53a2ddebab
- [x] Product Rename: TensorRT Model Optimizer to Model Optimizer (OMNIML-3033) - [x] Mention in Latest News section with date on the date of merging this PR (12/08) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
515 lines
16 KiB
Plaintext
515 lines
16 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Prepare Data\n",
|
|
"In this example, we use the Daring-Anteater dataset. For improved accuracy, please refer to the [Data Synthesis Section](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/speculative_decoding#optional-data-synthesis) in the README."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!git clone https://huggingface.co/datasets/nvidia/Daring-Anteater /tmp/Daring-Anteater"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Convert Model for Speculative Decoding\n",
|
|
"Here, we'll adapt our base model for speculative decoding by attaching a smaller EAGLE module. The upcoming code first loads meta-llama/Llama-3.2-1B as our base model and then configures the new draft module. To ensure compatibility, the draft module's dimensions must match the target model. Finally, the modelopt toolkit attaches this new, untrained module, leaving us with a combined model that is ready for the training phase later."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import transformers\n",
|
|
"\n",
|
|
"import modelopt.torch.opt as mto\n",
|
|
"import modelopt.torch.speculative as mtsp\n",
|
|
"from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG\n",
|
|
"\n",
|
|
"mto.enable_huggingface_checkpointing()\n",
|
|
"\n",
|
|
"# Load original HF model\n",
|
|
"base_model = \"meta-llama/Llama-3.2-1B\"\n",
|
|
"model = transformers.AutoModelForCausalLM.from_pretrained(\n",
|
|
" base_model, torch_dtype=\"auto\", device_map=\"cuda\"\n",
|
|
")\n",
|
|
"\n",
|
|
"# Read Default Config for EAGLE3\n",
|
|
"config = EAGLE3_DEFAULT_CFG[\"config\"]\n",
|
|
"\n",
|
|
"# Hidden size and vocab size must match base model\n",
|
|
"config[\"eagle_architecture_config\"].update(\n",
|
|
" {\n",
|
|
" \"hidden_size\": model.config.hidden_size,\n",
|
|
" \"vocab_size\": model.config.vocab_size,\n",
|
|
" \"draft_vocab_size\": model.config.vocab_size,\n",
|
|
" \"max_position_embeddings\": model.config.max_position_embeddings,\n",
|
|
" }\n",
|
|
")\n",
|
|
"\n",
|
|
"# Convert Model for eagle speculative decoding\n",
|
|
"mtsp.convert(model, [(\"eagle\", config)])\n",
|
|
"\n",
|
|
"# Prepare Tokenizer\n",
|
|
"tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, model_max_length=1024)\n",
|
|
"tokenizer.pad_token_id = tokenizer.eos_token_id\n",
|
|
"if tokenizer.chat_template is None:\n",
|
|
" tokenizer.chat_template = (\n",
|
|
" \"{%- for message in messages %}\"\n",
|
|
" \"{{- '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}\"\n",
|
|
" \"{%- endfor %}\"\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Train Draft module On Daring-Anteater\n",
|
|
"We will fine-tune the draft module on the Daring-Anteater dataset using the standard Hugging Face Trainer. Note that only the draft module's weights are updated during this process; the original target model remains frozen. After training, our speculative decoding model will be ready for export and deployment. Note that the time to train will be significantly dependent on the epochs (default=4) and the hardware being used."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"from dataclasses import dataclass, field\n",
|
|
"\n",
|
|
"from eagle_utils import DataCollatorWithPadding, LazySupervisedDataset\n",
|
|
"from transformers import Trainer\n",
|
|
"\n",
|
|
"with open(\"/tmp/Daring-Anteater/train.jsonl\") as f:\n",
|
|
" data_json = [json.loads(line) for line in f]\n",
|
|
"train_dataset = LazySupervisedDataset(data_json[: int(len(data_json) * 0.95)], tokenizer=tokenizer)\n",
|
|
"eval_dataset = LazySupervisedDataset(data_json[int(len(data_json) * 0.95) :], tokenizer=tokenizer)\n",
|
|
"\n",
|
|
"\n",
|
|
"@dataclass\n",
|
|
"class TrainingArguments(transformers.TrainingArguments):\n",
|
|
" dataloader_drop_last: bool = field(default=True)\n",
|
|
" bf16: bool = field(default=True)\n",
|
|
"\n",
|
|
"\n",
|
|
"training_args = TrainingArguments(\n",
|
|
" output_dir=\"/tmp/eagle_bf16\",\n",
|
|
" num_train_epochs=4,\n",
|
|
" per_device_train_batch_size=1,\n",
|
|
" per_device_eval_batch_size=1,\n",
|
|
")\n",
|
|
"trainer = Trainer(\n",
|
|
" model=model,\n",
|
|
" tokenizer=tokenizer,\n",
|
|
" args=training_args,\n",
|
|
" train_dataset=train_dataset,\n",
|
|
" eval_dataset=eval_dataset,\n",
|
|
" data_collator=DataCollatorWithPadding(),\n",
|
|
")\n",
|
|
"trainer._move_model_to_device(model, trainer.args.device)\n",
|
|
"\n",
|
|
"# Make sure label_smoother is None\n",
|
|
"assert trainer.label_smoother is None, \"label_smoother is not supported in speculative decoding!\"\n",
|
|
"\n",
|
|
"trainer.train()\n",
|
|
"trainer.save_state()\n",
|
|
"trainer.save_model(training_args.output_dir)\n",
|
|
"tokenizer.save_pretrained(training_args.output_dir)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Export Model Checkpoint\n",
|
|
"To deploy this model, we need to first export it to a Unified checkpoint."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from modelopt.torch.export import export_hf_checkpoint\n",
|
|
"\n",
|
|
"model.eval()\n",
|
|
"export_hf_checkpoint(\n",
|
|
" model,\n",
|
|
" export_dir=\"/tmp/hf_ckpt\",\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Deploying on TensorRT-LLM\n",
|
|
"\n",
|
|
"Here we show an example to deploy on TRT-LLM with `trtllm-serve` and [TRT-LLM container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release). See [Deployment](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/speculative_decoding#deployment) section for more info. \n",
|
|
"\n",
|
|
"First, we dump the `trtllm-serve` command and config file we need to `/tmp` folder."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"trtllm_serve_script = f\"\"\"trtllm-serve {base_model} \\\\\n",
|
|
" --host 0.0.0.0 \\\\\n",
|
|
" --port 8000 \\\\\n",
|
|
" --backend pytorch \\\\\n",
|
|
" --max_batch_size 32 \\\\\n",
|
|
" --max_num_tokens 8192 \\\\\n",
|
|
" --max_seq_len 8192 \\\\\n",
|
|
" --extra_llm_api_options /tmp/extra-llm-api-config.yml\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"extra_llm_api_config = \"\"\"enable_attention_dp: false\n",
|
|
"disable_overlap_scheduler: true\n",
|
|
"enable_autotuner: false\n",
|
|
"\n",
|
|
"cuda_graph_config:\n",
|
|
" max_batch_size: 1\n",
|
|
"\n",
|
|
"speculative_config:\n",
|
|
" decoding_type: Eagle\n",
|
|
" max_draft_len: 3\n",
|
|
" speculative_model_dir: /tmp/hf_ckpt\n",
|
|
"\n",
|
|
"kv_cache_config:\n",
|
|
" enable_block_reuse: false\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"# Dump the two scripts into /tmp\n",
|
|
"with open(\"/tmp/trtllm_serve.sh\", \"w\") as f:\n",
|
|
" f.write(trtllm_serve_script)\n",
|
|
"\n",
|
|
"with open(\"/tmp/extra-llm-api-config.yml\", \"w\") as f:\n",
|
|
" f.write(extra_llm_api_config)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Next, we start a TRT-LLM container in the background and run `trtllm-serve` inside it, using our exported checkpoint and the configuration scripts we just created:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import subprocess\n",
|
|
"import threading\n",
|
|
"\n",
|
|
"# Generate a unique container name so we can stop/remove it later\n",
|
|
"container_name = \"trtllm_serve_spec\"\n",
|
|
"\n",
|
|
"docker_cmd = [\n",
|
|
" \"docker\",\n",
|
|
" \"run\",\n",
|
|
" \"--rm\",\n",
|
|
" \"--net\",\n",
|
|
" \"host\",\n",
|
|
" \"--shm-size=2g\",\n",
|
|
" \"--ulimit\",\n",
|
|
" \"memlock=-1\",\n",
|
|
" \"--ulimit\",\n",
|
|
" \"stack=67108864\",\n",
|
|
" \"--gpus\",\n",
|
|
" \"all\",\n",
|
|
" \"-v\",\n",
|
|
" \"/tmp:/tmp\",\n",
|
|
" \"--name\",\n",
|
|
" container_name,\n",
|
|
" \"nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc2\",\n",
|
|
" \"bash\",\n",
|
|
" \"-c\",\n",
|
|
" \"bash /tmp/trtllm_serve.sh\",\n",
|
|
"]\n",
|
|
"\n",
|
|
"# print docker outputs\n",
|
|
"proc = subprocess.Popen(\n",
|
|
" docker_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"def stream_output(pipe):\n",
|
|
" for line in iter(pipe.readline, \"\"):\n",
|
|
" print(line, end=\"\")\n",
|
|
"\n",
|
|
"\n",
|
|
"# Use thread to print outputs\n",
|
|
"thread = threading.Thread(target=stream_output, args=(proc.stdout,))\n",
|
|
"thread.daemon = True\n",
|
|
"thread.start()\n",
|
|
"\n",
|
|
"print(\n",
|
|
" f\"Starting trtllm-serve in Docker (PID: {proc.pid}, container name: {container_name}) in the background:\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Please wait for the service to fully start inside the container. \n",
|
|
"Once you see the message `INFO: Application startup complete.`, you can proceed to send requests to the service:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"\n",
|
|
"import requests\n",
|
|
"\n",
|
|
"payload = {\n",
|
|
" \"model\": base_model,\n",
|
|
" \"messages\": [\n",
|
|
" {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
|
|
" {\"role\": \"user\", \"content\": \"Tell me about speculative decoding.\"},\n",
|
|
" ],\n",
|
|
" \"max_tokens\": 512,\n",
|
|
" \"temperature\": 0,\n",
|
|
" \"chat_template\": tokenizer.chat_template,\n",
|
|
"}\n",
|
|
"headers = {\"Content-Type\": \"application/json\", \"Accept\": \"application/json\"}\n",
|
|
"\n",
|
|
"response = requests.post(\n",
|
|
" \"http://localhost:8000/v1/chat/completions\", headers=headers, data=json.dumps(payload)\n",
|
|
")\n",
|
|
"output = response.json()\n",
|
|
"\n",
|
|
"print(output)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Finally, we clean up the container we created."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!docker rm -f trtllm_serve_spec"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Deploying on SGLang\n",
|
|
"Here, we deploy our trained model using SGLang. The following code defines the command needed to run the SGLang server with our specific configuration for speculative decoding."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# SGLang server launch command shell script\n",
|
|
"sglang_serve_script = f\"\"\"python3 -m sglang.launch_server \\\\\n",
|
|
" --model {base_model} \\\\\n",
|
|
" --host 0.0.0.0 \\\\\n",
|
|
" --port 30000 \\\\\n",
|
|
" --speculative-algorithm EAGLE3 \\\\\n",
|
|
" --speculative-eagle-topk 8 \\\\\n",
|
|
" --speculative-draft-model-path /tmp/hf_ckpt \\\\\n",
|
|
" --speculative-num-draft-tokens 3 \\\\\n",
|
|
" --speculative-num-steps 3 \\\\\n",
|
|
" --mem-fraction 0.6 \\\\\n",
|
|
" --cuda-graph-max-bs 2 \\\\\n",
|
|
" --dtype float16\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"with open(\"/tmp/sglang_serve.sh\", \"w\") as f:\n",
|
|
" f.write(sglang_serve_script)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Launch the SGLang server inside a Docker container as a background process."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"import subprocess\n",
|
|
"import threading\n",
|
|
"\n",
|
|
"container_name = \"sglang_serve_spec\"\n",
|
|
"home_dir = os.path.expanduser(\"~\")\n",
|
|
"hf_cache_dir = os.path.join(home_dir, \".cache\", \"huggingface\")\n",
|
|
"\n",
|
|
"# Ensure the Hugging Face cache directory exists. This directory should exist as ~/.cache/huggingface, when the model files for meta-llama/Llama-3.2-1B were downloaded earlier.\n",
|
|
"os.makedirs(hf_cache_dir, exist_ok=True)\n",
|
|
"\n",
|
|
"docker_cmd = [\n",
|
|
" \"docker\",\n",
|
|
" \"run\",\n",
|
|
" \"--rm\",\n",
|
|
" \"--net\",\n",
|
|
" \"host\",\n",
|
|
" \"--shm-size=32g\",\n",
|
|
" \"--gpus\",\n",
|
|
" \"all\",\n",
|
|
" \"-v\",\n",
|
|
" f\"{hf_cache_dir}:/root/.cache/huggingface\",\n",
|
|
" \"-v\",\n",
|
|
" \"/tmp:/tmp\",\n",
|
|
" \"--ipc=host\",\n",
|
|
" \"--name\",\n",
|
|
" container_name,\n",
|
|
" \"lmsysorg/sglang:latest\",\n",
|
|
" \"bash\",\n",
|
|
" \"-c\",\n",
|
|
" \"bash /tmp/sglang_serve.sh\",\n",
|
|
"]\n",
|
|
"\n",
|
|
"# Launch the Docker container\n",
|
|
"proc = subprocess.Popen(\n",
|
|
" docker_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"# Stream the process output\n",
|
|
"def stream_output(pipe):\n",
|
|
" for line in iter(pipe.readline, \"\"):\n",
|
|
" print(line, end=\"\")\n",
|
|
"\n",
|
|
"\n",
|
|
"# Use a thread to stream the output in without blocking the notebook\n",
|
|
"thread = threading.Thread(target=stream_output, args=(proc.stdout,))\n",
|
|
"thread.daemon = True\n",
|
|
"thread.start()\n",
|
|
"\n",
|
|
"print(\n",
|
|
" f\"Starting SGLang server in Docker (PID: {proc.pid}, container name: {container_name}) in the background:\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"As with TRT-LLM, please wait for the service to fully start inside the container. \n",
|
|
"Once you see the message `INFO: Application startup complete.`, you can proceed to send requests to the service:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"\n",
|
|
"import requests\n",
|
|
"\n",
|
|
"payload = {\n",
|
|
" \"model\": base_model,\n",
|
|
" \"messages\": [\n",
|
|
" {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
|
|
" {\"role\": \"user\", \"content\": \"Tell me about speculative decoding.\"},\n",
|
|
" ],\n",
|
|
" \"max_tokens\": 512,\n",
|
|
" \"temperature\": 0,\n",
|
|
"}\n",
|
|
"headers = {\"Content-Type\": \"application/json\", \"Accept\": \"application/json\"}\n",
|
|
"\n",
|
|
"# Send request to the SGLang server\n",
|
|
"response = requests.post(\n",
|
|
" \"http://localhost:30000/v1/chat/completions\", headers=headers, data=json.dumps(payload)\n",
|
|
")\n",
|
|
"output = response.json()\n",
|
|
"\n",
|
|
"print(output)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Clean up the container"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!docker rm -f sglang_serve_spec"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Deploying on vLLM (Coming Soon)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"While vLLM is another extremely popular, high-performance inference server, direct support for speculative decoding with this demo notebook is still under active development. This notebook will be updated once deployment is possible."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "modelopt+vllm",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.0"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|