mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-08 09:03:20 +08:00
h-guo18
a34d613d3c
## What does this PR do? **Type of change:** ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **Overview:** Main changes: - Refactored speculative decoding export logics into `class EagleExporter` to improve cohesion; - Separated speculative decoding export entrance with quantization export (`export_hf_checkpoint()`) due to their fundamental differences: - Quantization export base model's state_dict and config, while speculative decoding only export drafter's. - Most of the model-specific logics of quantization export (e.g. diffusers, vlms) are not needed for speculative decoding export. - Quantization export produce different format than speculative decoding checkpoint. (The former produce tokenizer config, generation config, e.t.c, while the later does not need. ) ## Usage <!-- You can potentially add a usage example below. --> To export an regular bf16 eagle checkpoint without quantization, the commands are the same: ```python python scripts/export_hf_checkpoint.py --model_path <x> --export_path <x> ``` To run PTQ on online-trained eagle checkpoint and export it: ```python python hf_ptq.py --pyt_ckpt_path <x> --qformat fp8 --export_path <x> ``` The above two commands will produce drafter ckpt for deployment, in the same foramt. ## Testing <!-- Mention how have you tested your change if applicable. --> Tested setting: - Base model: llama3.1-8b - Algorithms: eagle - Export path tested: - (Unquantized online ckpt) `python scripts/export_hf_checkpoint.py --model_path <x> --export_path <x>` - (PTQ) export `python hf_ptq.py --pyt_ckpt_path <x> --qformat fp8 --export_path <x>` - Tested deployment on vllm. Got normal AR. ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No <!--- If No, explain why. --> - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added export functionality for speculative decoding-optimized models * Support for multiple speculative decoding architectures with pre-configured deployment templates * Enhanced model export detection and automatic routing for optimized models * **Tests** * Updated export validation tests for speculative decoding models <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Export a HF checkpoint (with ModelOpt state) for deployment."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
import modelopt.torch.opt as mto
|
|
from modelopt.torch.export import export_speculative_decoding
|
|
from modelopt.torch.speculative.utils import load_vlm_or_llm_with_kwargs
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Export a HF checkpoint (with ModelOpt state) for deployment."
|
|
)
|
|
parser.add_argument("--model_path", type=str, default="Path of the trained checkpoint.")
|
|
parser.add_argument(
|
|
"--export_path", type=str, default="Destination directory for exported files."
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
mto.enable_huggingface_checkpointing()
|
|
|
|
args = parse_args()
|
|
_, model = load_vlm_or_llm_with_kwargs(args.model_path, torch_dtype="auto")
|
|
model.eval()
|
|
with torch.inference_mode():
|
|
export_speculative_decoding(
|
|
model,
|
|
export_dir=args.export_path,
|
|
)
|
|
print(f"Exported checkpoint to {args.export_path}")
|