diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index ca2f9908966..b398de589f6 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -188,6 +188,7 @@ def run_simple(args): sampling_kwargs=sampling_kwargs, speculative_algorithm=args.speculative_algorithm, draft_model_dir=args.draft_model_dir, + draft_quantization=args.draft_quantization, speculative_num_steps=args.draft_length, speculative_num_draft_tokens=args.block_size, tensor_parallel_size=args.tp_size, @@ -305,7 +306,7 @@ def run_simple(args): type=str, required=False, default="EAGLE3", - choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "NONE"], + choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "DSPARK", "NONE"], help="Speculative algorithm to use", ) parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory") @@ -316,6 +317,14 @@ def run_simple(args): default=None, help="Path to the draft model directory", ) + parser.add_argument( + "--draft_quantization", + type=str, + required=False, + default=None, + help="Quantization method of the draft checkpoint (e.g. modelopt, modelopt_fp4). " + "Read from the draft's config.json when omitted; set it only to override.", + ) parser.add_argument( "--runtime_params", type=str, diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index 24062399cb8..0aa76de491b 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -14,6 +14,8 @@ # limitations under the License. import asyncio +import json +import os import time from .base import Model @@ -114,14 +116,48 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs "model": kwargs.get("draft_model_dir"), "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 8), } + elif kwargs.get("speculative_algorithm") == "DSPARK": + # Match draft sampling to the target's verify mode: a greedy target with a + # probabilistic draft (or the reverse) crushes acceptance at temp > 0. + temperature = sampling_kwargs.get("temperature", 1.0) + specdec = { + "method": "dspark", + "model": kwargs.get("draft_model_dir"), + "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 7), + "draft_sample_method": kwargs.get("dspark_draft_sample_method") + or ("greedy" if temperature == 0 else "probabilistic"), + } elif kwargs.get("speculative_algorithm") == "NONE": specdec = None + # vLLM copies the target's quantization onto the draft, so a quantized drafter under + # a bf16 target is built as bf16 and dies loading the packed weights. Read the format + # from the drafter's own config.json; --draft_quantization overrides. + if specdec is not None and specdec.get("model"): + draft_quantization = kwargs.get("draft_quantization") + if draft_quantization is None: + draft_config = os.path.join(specdec["model"], "config.json") + if os.path.isfile(draft_config): + with open(draft_config) as f: + quant_config = json.load(f).get("quantization_config") or {} + draft_quantization = quant_config.get("quant_method") + if draft_quantization: + specdec["quantization"] = draft_quantization + print(f"Draft model quantization: {draft_quantization}") + if specdec is None: num_speculative_tokens = 1 else: num_speculative_tokens = specdec.get("num_speculative_tokens", 3) + # DSpark's block-parallel draft can outgrow the pre-allocated workspace during + # CUDA-graph capture; acceptance length is unaffected by graph capture, so skip it. + # SPECDEC_ENFORCE_EAGER=1 forces the same for other algorithms. + enforce_eager = ( + kwargs.get("speculative_algorithm") == "DSPARK" + or os.environ.get("SPECDEC_ENFORCE_EAGER") == "1" + ) + engine_args = AsyncEngineArgs( model=model_dir, tokenizer=kwargs.get("tokenizer_path"), @@ -133,8 +169,25 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs max_num_seqs=max_concurrent_requests * num_speculative_tokens, skip_tokenizer_init=False, async_scheduling=kwargs.get("async_scheduling", True), - enforce_eager=False, + enforce_eager=enforce_eager, max_model_len=kwargs.get("max_model_len"), + # Engine knobs a model card may pin, passed through from + # --runtime_params engine_args.. Only keys the caller actually set are + # forwarded, so vLLM keeps its own defaults otherwise. Hybrid Mamba models + # need these: on Nemotron-3.5-Lightning the SSM-cache settings decide whether + # the first draft token is accepted, and leaving them at vLLM's defaults costs + # ~65% of acceptance length while every later position looks normal. + **{ + key: kwargs[key] + for key in ( + "mamba_backend", + "mamba_ssm_cache_dtype", + "mamba_cache_mode", + "mamba_cache_philox_rounds", + "enable_mamba_cache_stochastic_rounding", + ) + if kwargs.get(key) is not None + }, ) self.engine_args = engine_args self.model = AsyncLLM.from_engine_args(engine_args) diff --git a/examples/speculative_decoding/scripts/quantize_drafter.py b/examples/speculative_decoding/scripts/quantize_drafter.py new file mode 100644 index 00000000000..6bb41e7a5a1 --- /dev/null +++ b/examples/speculative_decoding/scripts/quantize_drafter.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Calibration-free PTQ for a speculative-decoding drafter. + +Every scale is derived from the weights, so this needs no dataset and no forward pass, and +never imports the drafter's modeling code: each 2-D weight is wrapped in a throwaway +``nn.Linear`` under its checkpoint name, and ModelOpt's usual ``quantizer_name`` patterns +select over those names. Works for any drafter layout (DSpark / DFlash / EAGLE3 / Medusa), +including exported ones that ship no importable model class. + +``fp8`` and ``nvfp4`` need a static activation amax, normally measured on calibration data; +a fixed ``input_scale`` of 1.0 is applied instead. Acceptance length is governed by +clipping rather than resolution, and AL sits on a flat plateau from input_scale ~0.3 to 4.0 +(Qwen3-8B + DSpark, MT-Bench: +0.1% for FP8, -3.9% for NVFP4), so the scale only has to be +big enough. AWQ is not offered: ``awq_lite`` silently degrades to RTN without a +``forward_loop``. + +Example: + python quantize_drafter.py \ + --drafter_path nvidia/MiniMax-M3-DSpark \ + --qformat fp8 \ + --export_path ./MiniMax-M3-DSpark-FP8 +""" + +import argparse +import copy +import json +import shutil +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors.torch import load_file, save_file + +import modelopt.torch.quantization as mtq +from modelopt.recipe.presets import QUANT_CFG_CHOICES +from modelopt.torch.export.quant_utils import ( + get_activation_scaling_factor, + get_quant_config, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + to_quantized_weight, +) +from modelopt.torch.quantization.config import need_calibration +from modelopt.torch.quantization.utils import is_quantized_linear + +# INT8/INT4 are absent on purpose: they quantize cleanly but vLLM's ModelOpt backend +# cannot serve them. +SUPPORTED_QFORMATS = [ + "w4a16_nvfp4", + "nvfp4", + "fp8", + "fp8_pc_pt", +] + +# All 2-D, so the flat view treats them as GEMMs, but none is one: markov_w1/embed_tokens +# are embeddings (ModelOpt's presets skip these via `parent_class`, which the flat view +# cannot see), and confidence_head has a single output whose per-channel scale is 0-dim. +# All tiny. `fc` is left to the presets; `lm_head` is excluded by the preset itself. +DEFAULT_EXCLUDE = ["*markov_head*", "*confidence_head*", "*embed_tokens*"] + +# Sidecars carried over to the export, and -- with the weights and config.json -- the only +# files fetched when --drafter_path is a repo id rather than a local directory. +SIDECAR_FILES = ("tokenizer.json", "tokenizer_config.json", "generation_config.json") + +# The amax that yields input_scale 1.0 for FP8. NVFP4 divides by 6*448, so the same amax +# records as 0.1667 there; both mean the same activation range. +FP8_E4M3_MAX = 448.0 +STATIC_ACT_AMAX = FP8_E4M3_MAX + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--drafter_path", required=True, help="HF repo id or local dir of the drafter checkpoint." + ) + parser.add_argument("--export_path", required=True, help="Output directory.") + parser.add_argument( + "--qformat", + default="w4a16_nvfp4", + choices=SUPPORTED_QFORMATS, + help="Quantization format. All are calibration-free; fp8 and nvfp4 additionally " + "quantize activations, using a fixed input_scale of 1.0.", + ) + parser.add_argument( + "--dtype", + default="bfloat16", + choices=["bfloat16", "float16", "float32"], + help="Compute dtype the weights are cast to before quantizing.", + ) + parser.add_argument( + "--exclude", + nargs="*", + default=[], + metavar="PATTERN", + help="Extra fnmatch patterns to leave unquantized, in `quantizer_name` form. " + f"Appended to the defaults ({' '.join(DEFAULT_EXCLUDE)}), which always apply.", + ) + parser.add_argument( + "--quantize_lm_head", + action="store_true", + help="Also quantize lm_head -- the largest drafter tensor, but it feeds the " + "acceptance test directly, so measure AL first.", + ) + return parser.parse_args() + + +def auto_map_modules(config: dict) -> set[str]: + """Module basenames referenced by a config's ``auto_map``, e.g. {"modeling_x"}. + + Values are either ``"modeling_x.XModel"`` or, for tokenizers, a list whose entries may + be null (``[null, "tokenization_x.XTokenizerFast"]``); a ``repo--`` prefix points at + another repository and is not a local file. + """ + modules = set() + for value in (config.get("auto_map") or {}).values(): + for ref in value if isinstance(value, list) else [value]: + if isinstance(ref, str) and "." in ref: + modules.add(ref.split("--")[-1].rsplit(".", 1)[0]) + return modules + + +def load_drafter(drafter_path: str) -> tuple[Path, dict[str, torch.Tensor]]: + """Resolve a local dir or HF repo id to (dir, state_dict).""" + local_dir = Path(drafter_path) + if not local_dir.is_dir(): + from huggingface_hub import snapshot_download + + # A drafter that ships custom modeling code references it from auto_map, and the + # export carries that config verbatim -- so those .py files have to be fetched too + # or the exported references dangle. + local_dir = Path( + snapshot_download( + drafter_path, + allow_patterns=["*.safetensors", "config.json", "*.py", *SIDECAR_FILES], + ) + ) + + shards = sorted(local_dir.glob("*.safetensors")) + assert shards, f"No .safetensors found under {local_dir}" + state_dict: dict[str, torch.Tensor] = {} + for shard in shards: + state_dict.update(load_file(shard)) + return local_dir, state_dict + + +def build_linear_view(state_dict: dict[str, torch.Tensor], dtype: torch.dtype) -> nn.Module: + """Expose every 2-D weight as an nn.Linear whose module name is its checkpoint key. + + Nested ModuleDicts so ``named_modules()`` reproduces the dotted checkpoint keys, which + is what ``quantizer_name`` patterns match against. + """ + root = nn.ModuleDict() + for key, weight in state_dict.items(): + if weight.dim() != 2 or not key.endswith(".weight"): + continue + *parents, leaf = key[: -len(".weight")].split(".") + node = root + for part in parents: + if part not in node: + node[part] = nn.ModuleDict() + node = node[part] + out_features, in_features = weight.shape + # On meta, so nn.Linear skips allocating and randomly initializing a weight that + # the next line replaces anyway. + with torch.device("meta"): + linear = nn.Linear(in_features, out_features, bias=False, dtype=dtype) + linear.weight = nn.Parameter(weight.to(dtype), requires_grad=False) + node[leaf] = linear + return root + + +def set_static_activation_amax(root: nn.Module, amax: float = STATIC_ACT_AMAX) -> int: + """Give every static ``input_quantizer`` the same fixed amax. Returns how many were set. + + Skips dynamic quantizers and any that already have an amax, so it composes as a + fallback rather than an overwrite. + """ + count = 0 + for _, module in root.named_modules(): + if not is_quantized_linear(module): + continue + input_quantizer = getattr(module, "input_quantizer", None) + if input_quantizer is None or not input_quantizer.is_enabled: + continue + if getattr(input_quantizer, "_dynamic", False): + continue + if getattr(input_quantizer, "amax", None) is not None: + continue + # Keep amax in fp32, as ModelOpt does everywhere else -- casting to the weight + # dtype would round a measured amax through bf16's 8-bit mantissa. + input_quantizer.amax = torch.tensor(amax, dtype=torch.float32) + count += 1 + return count + + +def resolve_activation_scales(root: nn.Module, quant_cfg: dict) -> None: + """Establish activation scales for a format that quantizes activations statically. + + The single place deciding where a static amax comes from: real calibration would call + ``mtq.calibrate`` here, ahead of the fixed fallback. + """ + if not need_calibration(quant_cfg): + return + n = set_static_activation_amax(root) + print(f"Set {n} static activation amax values (fixed, input_scale 1.0) -- not calibrated.") + + +def build_quant_cfg(qformat: str, exclude: list[str], quantize_lm_head: bool) -> dict: + """Take the shipped preset and layer the drafter-specific exclusions on top.""" + quant_cfg = copy.deepcopy(QUANT_CFG_CHOICES[qformat]) + if quantize_lm_head: + # The preset disables *all* of lm_head's quantizers. Re-enabling only the weight + # one would leave a W+A format exporting lm_head with no input_scale while the + # config still advertises it as fully quantized, which a runtime fails to load. + for quantizer in ("weight_quantizer", "input_quantizer"): + quant_cfg["quant_cfg"].append( + {"quantizer_name": f"*lm_head*{quantizer}", "enable": True} + ) + for pattern in DEFAULT_EXCLUDE + exclude: + quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) + return quant_cfg + + +def export_quantized_state_dict( + root: nn.Module, state_dict: dict[str, torch.Tensor], dtype: torch.dtype +) -> dict[str, torch.Tensor]: + """Pack each quantized weight and emit it alongside its scales. + + Unified-HF naming (``w.weight_scale`` etc). Untouched tensors carry through in ``dtype``. + """ + export_sd = {k: v.to(dtype) for k, v in state_dict.items()} + for name, module in root.named_modules(): + if not is_quantized_linear(module) or not module.weight_quantizer.is_enabled: + continue + quantization = get_quantization_format(module) + assert quantization is not None, f"{name}: enabled quantizer resolved to no format" + weight_scale = get_weight_scaling_factor(module) + weight_scale_2 = get_weight_scaling_factor_2(module) + # The packing helpers index the scale as ``scale[:, None]``, which a 0-dim scale + # cannot satisfy. One row of weights, so leave it in ``dtype``. + if weight_scale is not None and weight_scale.dim() == 0 and module.weight.shape[0] == 1: + print(f"Skipping {name}: single-output projection, per-channel scale is scalar") + continue + export_sd[f"{name}.weight"] = to_quantized_weight( + module.weight, + weight_scale, + quantization, + weight_scale_2, + get_weight_block_size(module), + ) + export_sd[f"{name}.weight_scale"] = weight_scale + if weight_scale_2 is not None: + export_sd[f"{name}.weight_scale_2"] = weight_scale_2 + # Without this the runtime has no activation scale and the format silently degrades. + activation_scale = get_activation_scaling_factor(module) + if activation_scale is not None: + export_sd[f"{name}.input_scale"] = activation_scale + return export_sd + + +def main(): + args = parse_args() + dtype = getattr(torch, args.dtype) + + source_dir, state_dict = load_drafter(args.drafter_path) + root = build_linear_view(state_dict, dtype) + + quant_cfg = build_quant_cfg(args.qformat, args.exclude, args.quantize_lm_head) + + mtq.quantize(root, quant_cfg) # no forward_loop: scales come from the weights + resolve_activation_scales(root, quant_cfg) + + mtq.print_quant_summary(root) + + export_sd = export_quantized_state_dict(root, state_dict, dtype) + + export_dir = Path(args.export_path) + export_dir.mkdir(parents=True, exist_ok=True) + save_file(export_sd, export_dir / "model.safetensors", metadata={"format": "pt"}) + + config = json.loads((source_dir / "config.json").read_text()) + hf_quant_config = get_quant_config(root) + # ``get_quant_config`` only knows the linear view, so tensors it never saw (norms, 1-D + # weights) are missing and a loader walking the checkpoint expects a scale for them. + quantized = { + name + for name, module in root.named_modules() + if is_quantized_linear(module) + and f"{name}.weight" in export_sd + and f"{name}.weight_scale" in export_sd + } + unquantized = sorted( + key[: -len(".weight")] + for key in state_dict + if key.endswith(".weight") and key[: -len(".weight")] not in quantized + ) + exclude_modules = hf_quant_config["quantization"].get("exclude_modules", []) + for name in unquantized: + if name not in exclude_modules: + exclude_modules.append(name) + # Runtimes match against their own module prefix, which is nested relative to the + # checkpoint key (vLLM builds the draft's ``fc`` at ``model.fc``). + wildcard = f"*{name}" + if wildcard not in exclude_modules: + exclude_modules.append(wildcard) + # Runtimes fuse sibling projections into one layer whose name is in no checkpoint key, + # so excluding only the parts would leave the fused layer quantized. + for fused, parts in ( + ("qkv_proj", ("q_proj", "k_proj", "v_proj")), + ("gate_up_proj", ("gate_proj", "up_proj")), + ): + if all(any(p in name for name in exclude_modules) for p in parts): + alias = f"*{fused}" + if alias not in exclude_modules: + exclude_modules.append(alias) + hf_quant_config["quantization"]["exclude_modules"] = exclude_modules + config["quantization_config"] = dict(hf_quant_config["quantization"]) + # ModelOpt names the format ``quant_algo``; vLLM reads ``quant_method`` and treats its + # absence as unquantized, splitting NVFP4 off into its own backend. Emit both. + quant_algo = str(hf_quant_config["quantization"].get("quant_algo") or "") + config["quantization_config"].setdefault( + "quant_method", "modelopt_fp4" if "NVFP4" in quant_algo.upper() else "modelopt" + ) + # Same list, second key: the flat ``quantization_config`` in config.json is read for + # ``ignore``, not ``exclude_modules``. + config["quantization_config"]["ignore"] = list(exclude_modules) + config["torch_dtype"] = args.dtype + (export_dir / "config.json").write_text(json.dumps(config, indent=2)) + (export_dir / "hf_quant_config.json").write_text(json.dumps(hf_quant_config, indent=2)) + + for extra in SIDECAR_FILES: + if (source_dir / extra).is_file(): + shutil.copy2(source_dir / extra, export_dir / extra) + + # A drafter that ships custom modeling code points at it from auto_map; the export + # carries that config verbatim, so the .py files have to come along or the reference + # dangles. (The DFlash/DSpark exports have no auto_map -- this is for the ones that do.) + for module in auto_map_modules(config): + source_py = source_dir / f"{module}.py" + if source_py.is_file(): + shutil.copy2(source_py, export_dir / source_py.name) + + before = sum(v.numel() * v.element_size() for v in state_dict.values()) + after = sum(v.numel() * v.element_size() for v in export_sd.values()) + print(f"\n{args.qformat}: {before / 2**30:.2f} GiB -> {after / 2**30:.2f} GiB") + print(f"Exported to {export_dir}") + + +if __name__ == "__main__": + main() diff --git a/tools/launcher/common/specdec/quantize_drafter.sh b/tools/launcher/common/specdec/quantize_drafter.sh new file mode 100644 index 00000000000..ddaac78ed00 --- /dev/null +++ b/tools/launcher/common/specdec/quantize_drafter.sh @@ -0,0 +1,47 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +# Calibration-free PTQ for an exported speculative-decoding drafter. +# +# Required env vars: +# DRAFTER_CKPT — exported drafter path, or a training output_dir to auto-detect under + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +source ${SCRIPT_DIR}/../service_utils.sh + +trap 'error_handler $0 $LINENO' ERR + +################################################################################################### + +DRAFTER="${DRAFTER_CKPT}" +# Training writes exported-checkpoint-/ under output_dir; take the newest. Only for a +# local directory -- anything else (an HF repo id) is passed through for the script to +# resolve. -V sorts numerically, so checkpoint-1000 beats checkpoint-900. +if [ -d "${DRAFTER}" ] && [ ! -f "${DRAFTER}/config.json" ]; then + latest=$(find "${DRAFTER}" -maxdepth 1 -mindepth 1 -type d \ + -name 'exported-checkpoint-*' -printf '%f\n' 2>/dev/null | sort -V | tail -1) + if [ -z "${latest}" ]; then + echo "ERROR: ${DRAFTER} is not a checkpoint and holds no exported-checkpoint-* directory." + exit 1 + fi + DRAFTER="${DRAFTER}/${latest}" + echo "Auto-detected drafter: ${DRAFTER}" +fi + +python modules/Model-Optimizer/examples/speculative_decoding/scripts/quantize_drafter.py \ + --drafter_path "${DRAFTER}" \ + "$@" diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml new file mode 100644 index 00000000000..0e87e97f3b7 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml @@ -0,0 +1,73 @@ +# Calibration-free NVFP4 PTQ of a DSpark drafter for Qwen3-8B. +# +# Takes an exported drafter (train one with hf_streaming_dspark.yaml, or point +# drafter at a published checkpoint) and quantizes it weight-only to NVFP4, then +# measures acceptance length so the cost is visible. No calibration data needed: +# every scale comes from the weights. +# +# 2-step pipeline: +# task_0: Quantize the drafter (CPU-only, ~1 min for an 8B-class draft) +# task_1: Benchmark acceptance length on MT-Bench via vLLM +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml --yes + +job_name: Qwen3-8B_DSpark_PTQ_nvfp4 +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + # Exported drafter to quantize. /scratchspace/export is where the streaming + # examples leave theirs; a plain checkpoint dir or HF repo id also works. + draft_model: /scratchspace/export + + # Step 1: Quantize. w4a16_nvfp4 keeps activations in bf16; use --qformat nvfp4 + # for weight+activation (fixed input_scale 1.0, ~3.9% AL on Qwen3-8B). + # + # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family + # drafters build their fused context-KV projection by reading qkv_proj.weight + # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. + # `fc` is optional -- quantizing it saves ~3% more size for ~0.7% AL on + # Qwen3-8B (3.0186 vs 3.0392); add '*fc*' to keep it in bf16. + task_0: + script: common/specdec/quantize_drafter.sh + args: + - --qformat w4a16_nvfp4 + - --export_path /scratchspace/export_quantized + - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' + environment: + - DRAFTER_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 2: Acceptance length on MT-Bench. Compare against the same run with + # --draft_model_dir <> to see what quantization cost. + task_1: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export_quantized + # DSPARK/DFLASH read --block_size, not --draft_length; must match the + # drafter's block_size (7 for dspark_qwen3_8b_block7). + - --block_size 7 + - --output_length 4096 + - --engine VLLM + - --tp_size 1 + - --ep_size 1 + - --speculative_algorithm DSPARK + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:nightly diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json new file mode 100644 index 00000000000..b4d22256dd7 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json @@ -0,0 +1,8 @@ +{ + "engine_args": { + "mamba_backend": "flashinfer", + "mamba_ssm_cache_dtype": "float16", + "enable_mamba_cache_stochastic_rounding": true, + "mamba_cache_philox_rounds": 5 + } +} diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml new file mode 100644 index 00000000000..76275297cd8 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml @@ -0,0 +1,88 @@ +# Calibration-free NVFP4 PTQ of the DSpark drafter for Nemotron-3.5-Lightning-30B-A3B. +# +# Same pipeline as examples/Qwen/Qwen3-8B/hf_dspark_ptq_nvfp4.yaml, against a hybrid +# Mamba-MoE target and the published DSpark drafter. Quantizes weight-only to NVFP4, +# then measures acceptance length so the cost is visible. No calibration data needed: +# every scale comes from the weights. +# +# 2-step pipeline: +# task_0: Quantize the drafter (CPU-only, well under a minute for this 0.97B draft) +# task_1: Benchmark acceptance length on MT-Bench via vLLM +# +# Usage: +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/hf_dspark_ptq_nvfp4.yaml --yes + +job_name: Nemotron-3.5-Lightning-30B-A3B_DSpark_PTQ_nvfp4 +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + draft_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark + + # Step 1: Quantize. w4a16_nvfp4 keeps activations in bf16; use --qformat nvfp4 + # for weight+activation (fixed input_scale 1.0). + # + # The q/k/v exclusions are mandatory, not a tuning choice: DFlash-family + # drafters build their fused context-KV projection by reading qkv_proj.weight + # raw, which cannot be a packed tensor. o_proj and the MLP still quantize. + # `fc` is optional -- quantizing it saves ~4% more size for ~1.3% AL on this + # model (4.2334 vs 4.2899); add '*fc*' to keep it in bf16. + # + # This drafter has has_lm_head=false (it shares the target's), so + # --quantize_lm_head does not apply. embed_tokens is 37% of the checkpoint but + # is excluded by default: it is an nn.Embedding the drafter inherits. + task_0: + script: common/specdec/quantize_drafter.sh + args: + - --qformat w4a16_nvfp4 + - --export_path /scratchspace/export_quantized + - --exclude '*q_proj*' '*k_proj*' '*v_proj*' '*qkv_proj*' + environment: + - DRAFTER_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 2: Acceptance length on MT-Bench. Compare against the same run with + # --draft_model_dir <> to see what quantization cost. + # + # This target is a hybrid Mamba-MoE model, so it differs from the Qwen3 example + # in three ways. It needs TP8. Startup is dominated by Mamba2 kernel warmup + # (~6 min before the engine is ready), not weight loading. And it needs the + # mamba engine settings from the model card (runtime_params below): without + # them the first draft token is rejected ~88% of the time and acceptance length + # collapses from ~4.3 to ~1.5, while every later position stays normal -- so it + # looks like a bad drafter rather than a serving misconfiguration. + task_1: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export_quantized + # DSPARK reads --block_size, not --draft_length; must match the drafter's + # block_size, which is 8 for this checkpoint. + - --block_size 8 + - --output_length 4096 + - --engine VLLM + - --tp_size 8 + - --ep_size 1 + - --speculative_algorithm DSPARK + - --trust_remote_code + - --temperature 0 + - --runtime_params modules/Model-Optimizer/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/engine_args.json + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 8 + environment: + - HF_MODEL_CKPT: <> + # The model card sets this on every Nemotron-3.5 serve command. + - VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: vllm/vllm-openai:nightly