From ca88c66fd20961acf474b8fda14053b2156673a5 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:22:49 -0700 Subject: [PATCH 1/3] Drop version checks for long-unsupported dependency releases Remove branches for onnxruntime<1.21, Transformer Engine<2.0, megatron-core<=0.9, TensorRT-LLM<1.1.0rc2 and vLLM<=0.11. Every supported install or CI container is well past these versions. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/vllm_serve/vllm_serve_fakequant.py | 10 +--------- examples/vllm_serve/vllm_serve_sparse_attn.py | 8 +------- modelopt/deploy/llm/generate.py | 13 +------------ modelopt/onnx/quantization/ort_utils.py | 10 ---------- .../quantization/plugins/transformer_engine.py | 13 ++----------- .../torch/speculative/plugins/megatron_eagle.py | 13 ++----------- tests/unit/onnx/autocast/test_nodeclassifier.py | 14 -------------- 7 files changed, 7 insertions(+), 74 deletions(-) diff --git a/examples/vllm_serve/vllm_serve_fakequant.py b/examples/vllm_serve/vllm_serve_fakequant.py index 2892c850adf..b34715ee295 100644 --- a/examples/vllm_serve/vllm_serve_fakequant.py +++ b/examples/vllm_serve/vllm_serve_fakequant.py @@ -61,9 +61,8 @@ import sys from pathlib import Path -import vllm -from packaging import version from vllm.entrypoints.cli.main import main as vllm_main +from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_mlflow_utils import MLFLOW_ENV_VARS, add_mlflow_args, resolve_mlflow_args @@ -74,13 +73,6 @@ def _is_missing_entrypoint(error: ModuleNotFoundError, entrypoint: str) -> bool: ) -vllm_version = version.parse(vllm.__version__) -if vllm_version <= version.parse("0.11.0"): - from vllm.utils import FlexibleArgumentParser -else: - from vllm.utils.argparse_utils import FlexibleArgumentParser - - # Env vars to copy from the driver to Ray workers (must match fakequant_worker / vllm_ptq_utils). # The MLflow ones are settled by resolve_mlflow_args() below, after this list is published: # Ray reads the values when it creates the actors, so naming them here is enough. diff --git a/examples/vllm_serve/vllm_serve_sparse_attn.py b/examples/vllm_serve/vllm_serve_sparse_attn.py index c50c18c7146..6bc793057db 100644 --- a/examples/vllm_serve/vllm_serve_sparse_attn.py +++ b/examples/vllm_serve/vllm_serve_sparse_attn.py @@ -33,19 +33,13 @@ from pathlib import Path import uvloop -import vllm -from packaging import version from vllm.entrypoints.openai.api_server import run_server +from vllm.utils.argparse_utils import FlexibleArgumentParser -vllm_version = version.parse(vllm.__version__) try: # vLLM 0.29 moved the serve argument parser out of ``vllm.entrypoints.openai`` from vllm.entrypoints.launchers.cli_args import make_arg_parser except ImportError: from vllm.entrypoints.openai.cli_args import make_arg_parser -if vllm_version <= version.parse("0.11.0"): - from vllm.utils import FlexibleArgumentParser -else: - from vllm.utils.argparse_utils import FlexibleArgumentParser def main(): diff --git a/modelopt/deploy/llm/generate.py b/modelopt/deploy/llm/generate.py index 39306504137..ca6a310c0fa 100644 --- a/modelopt/deploy/llm/generate.py +++ b/modelopt/deploy/llm/generate.py @@ -21,9 +21,7 @@ from pathlib import Path from typing import Any -import tensorrt_llm import torch -from packaging.version import parse as parse_version from tensorrt_llm import SamplingParams try: @@ -141,10 +139,6 @@ def _find_max_position_embeddings(cfg: dict) -> int | None: enable_padding=True, ) - self._support_context_logits_and_stop_words = parse_version( - tensorrt_llm.__version__ - ) >= parse_version("1.1.0rc2") - super().__init__( backend="pytorch", model=checkpoint_dir, @@ -172,7 +166,7 @@ def max_beam_width(self): @property def gather_context_logits(self): """Returns whether the context_logits can be returned from the LLM instance.""" - return self._support_context_logits_and_stop_words + return True def _generate( self, @@ -184,8 +178,6 @@ def _generate( ): assert temperature >= 0.0, "Temperature must be greater than 0.0." - if not self._support_context_logits_and_stop_words: - stop_words = None beam_width = self.max_beam_width kwargs = _sanitize_temperature_and_top_p(temperature, top_p) sampling_config = SamplingParams( @@ -285,9 +277,6 @@ def generate_context_logits( Returns: a tensor list of the context_logits. """ - assert self._support_context_logits_and_stop_words, ( - "Context logits are not supported with the current tensorrt_llm version." - ) assert not self._enable_kv_cache_reuse, ( "Context logits require enable_kv_cache_reuse=False: with KV cache prefix reuse, " "shared-prefix requests only return logits for the recomputed suffix, producing " diff --git a/modelopt/onnx/quantization/ort_utils.py b/modelopt/onnx/quantization/ort_utils.py index 354a9e1566c..bdca64dc015 100755 --- a/modelopt/onnx/quantization/ort_utils.py +++ b/modelopt/onnx/quantization/ort_utils.py @@ -436,16 +436,6 @@ def _make_trt_ep_first_choice(calibration_eps, trt_plugins): "See https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#data-dependant-shape-dds-ops" ) calibration_eps.remove("trt") - try: - assert Version(ort.__version__) >= Version("1.21.0") - except AssertionError: - logger.warning( - "This model contains DDS ops, please upgrade your ORT version to 1.21.0+ for full support with " - "TRT EP backend. The reason for that is because DDS ops require TRT 10.6+, and earlier versions " - "of ORT were compiled with TRT 10.4 or lower. Disabling it for now. " - "See https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#data-dependant-shape-dds-ops" - ) - calibration_eps.remove("trt") if has_custom_op: if "trt" in calibration_eps: logger.warning( diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index c7517eb559c..c9fe5d4fe94 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -24,15 +24,12 @@ import transformer_engine.pytorch.module.grouped_linear as te_grouped_linear import transformer_engine.pytorch.module.layernorm_linear as te_layernorm_linear import transformer_engine.pytorch.module.linear as te_linear -from packaging.version import Version from modelopt.torch.quantization.utils import replace_function from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from .custom import _ParallelLinear -_TE_VERSION = Version(te.__version__) - _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV = "MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP" @@ -318,17 +315,11 @@ class _QuantLayerNormLinearFunc(torch.autograd.Function): @staticmethod def _get_original_gemm(): - if Version("2.0") <= _TE_VERSION: - return te_layernorm_linear.general_gemm - else: - return te_layernorm_linear.tex.gemm + return te_layernorm_linear.general_gemm @staticmethod def _gemm_replace_args(): - if Version("2.0") <= _TE_VERSION: - return (te_layernorm_linear, "general_gemm") - else: - return (te_layernorm_linear.tex, "gemm") + return (te_layernorm_linear, "general_gemm") @staticmethod def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs): diff --git a/modelopt/torch/speculative/plugins/megatron_eagle.py b/modelopt/torch/speculative/plugins/megatron_eagle.py index 65e68d8f280..e35f654f7c7 100644 --- a/modelopt/torch/speculative/plugins/megatron_eagle.py +++ b/modelopt/torch/speculative/plugins/megatron_eagle.py @@ -20,7 +20,6 @@ import warnings from contextlib import contextmanager -import megatron.core import torch import torch.nn.functional as F from megatron.core import InferenceParams, tensor_parallel @@ -54,7 +53,6 @@ from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.transformer.utils import sharded_state_dict_default from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint -from packaging.version import Version from ..eagle.conversion import EagleDMRegistry from ..eagle.eagle_model import EagleModel @@ -130,11 +128,6 @@ def dict_to_config( return config -def mcore_version_higher_than(target_version: str): - """Check if megatron-core is greater than this version.""" - return Version(megatron.core.__version__) > Version(target_version) - - def logits_kld_loss(logits, gt_logits, mapping=None): """KL Divergence loss using ground truth logits.""" gathered_logits = gather_from_tensor_model_parallel_region(logits) @@ -937,8 +930,6 @@ def _base_model_forward( # decoder will get hidden_states from decoder.input_tensor decoder_input = None - extra_kwargs = {"packed_seq_params": None} if mcore_version_higher_than("0.9.0") else {} - rotary_pos_emb = None yarn_mscale = 1.0 if self.config.multi_latent_attention: @@ -950,7 +941,7 @@ def _base_model_forward( self.decoder, decoder_input, self.config, - **extra_kwargs, + packed_seq_params=None, ) rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len) elif self.position_embedding_type == "yarn": @@ -959,7 +950,7 @@ def _base_model_forward( None, decoder_input, self.config, - **extra_kwargs, + packed_seq_params=None, ) rotary_pos_emb, yarn_mscale = self.rotary_pos_emb(rotary_seq_len) else: diff --git a/tests/unit/onnx/autocast/test_nodeclassifier.py b/tests/unit/onnx/autocast/test_nodeclassifier.py index 002189b6bd3..b5fdb46e9ea 100644 --- a/tests/unit/onnx/autocast/test_nodeclassifier.py +++ b/tests/unit/onnx/autocast/test_nodeclassifier.py @@ -20,8 +20,6 @@ import numpy as np import pytest from onnx import TensorProto, helper, numpy_helper -from onnxruntime import __version__ as ort_version -from packaging.version import Version from modelopt.onnx.autocast.logging_config import configure_logging from modelopt.onnx.autocast.nodeclassifier import ( @@ -233,9 +231,6 @@ def test_depth_of_reduction_rule(): assert rule.check(reduce_init_node) is True -@pytest.mark.skipif( - Version(ort_version) < Version("1.21.0"), reason="WAR: Requires onnxruntime>=1.21.0" -) def test_node_classifier(test_model): node_to_init_map = {key: [] for key in ["add_node", "mul_node"]} @@ -273,9 +268,6 @@ def test_node_classifier(test_model): assert not set(fp16_nodes).intersection(set(fp32_nodes)) -@pytest.mark.skipif( - Version(ort_version) < Version("1.21.0"), reason="WAR: Requires onnxruntime>=1.21.0" -) def test_node_classifier_custom_rule(test_model): node_to_init_map = { "add_node": [numpy_helper.from_array(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))], @@ -302,9 +294,6 @@ def _check_inner(self, node): assert not set(fp16_nodes).intersection(set(fp32_nodes)) -@pytest.mark.skipif( - Version(ort_version) < Version("1.21.0"), reason="WAR: Requires onnxruntime>=1.21.0" -) def test_node_classifier_op_types_to_exclude(test_model): node_to_init_map = { "add_node": [numpy_helper.from_array(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))], @@ -515,9 +504,6 @@ def test_depth_of_reduction_rule_with_tensor_stats(): assert rule.check(small_conv_node) is False -@pytest.mark.skipif( - Version(ort_version) < Version("1.21.0"), reason="WAR: Requires onnxruntime>=1.21.0" -) def test_node_classifier_with_multi_batch_calibration(test_model): """Test NodeClassifier with multi-batch calibration data.""" node_to_init_map = {key: [] for key in ["add_node", "mul_node"]} From 91d8ff13e22de1d72653732cd17662dbbd614744 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:03:28 -0700 Subject: [PATCH 2/3] Inline the TE GEMM lookup and keep the TensorRT-LLM 1.0 capability flag TensorRT-LLM 1.0 still imports this module (CudaGraphConfig exists there), so the <1.1.0rc2 context-logits/stop-words gate is not dead. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/deploy/llm/generate.py | 13 ++++++++++- .../plugins/transformer_engine.py | 22 +++++++------------ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/modelopt/deploy/llm/generate.py b/modelopt/deploy/llm/generate.py index ca6a310c0fa..39306504137 100644 --- a/modelopt/deploy/llm/generate.py +++ b/modelopt/deploy/llm/generate.py @@ -21,7 +21,9 @@ from pathlib import Path from typing import Any +import tensorrt_llm import torch +from packaging.version import parse as parse_version from tensorrt_llm import SamplingParams try: @@ -139,6 +141,10 @@ def _find_max_position_embeddings(cfg: dict) -> int | None: enable_padding=True, ) + self._support_context_logits_and_stop_words = parse_version( + tensorrt_llm.__version__ + ) >= parse_version("1.1.0rc2") + super().__init__( backend="pytorch", model=checkpoint_dir, @@ -166,7 +172,7 @@ def max_beam_width(self): @property def gather_context_logits(self): """Returns whether the context_logits can be returned from the LLM instance.""" - return True + return self._support_context_logits_and_stop_words def _generate( self, @@ -178,6 +184,8 @@ def _generate( ): assert temperature >= 0.0, "Temperature must be greater than 0.0." + if not self._support_context_logits_and_stop_words: + stop_words = None beam_width = self.max_beam_width kwargs = _sanitize_temperature_and_top_p(temperature, top_p) sampling_config = SamplingParams( @@ -277,6 +285,9 @@ def generate_context_logits( Returns: a tensor list of the context_logits. """ + assert self._support_context_logits_and_stop_words, ( + "Context logits are not supported with the current tensorrt_llm version." + ) assert not self._enable_kv_cache_reuse, ( "Context logits require enable_kv_cache_reuse=False: with KV cache prefix reuse, " "shared-prefix requests only return logits for the recomputed suffix, producing " diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index c9fe5d4fe94..32c9a3cffac 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -313,14 +313,6 @@ def te_grouped_quantized_linear_fn(package, func_name, self, *args): class _QuantLayerNormLinearFunc(torch.autograd.Function): """Patched version of _LayerNormLinear to quantize the input to the GEMM operation.""" - @staticmethod - def _get_original_gemm(): - return te_layernorm_linear.general_gemm - - @staticmethod - def _gemm_replace_args(): - return (te_layernorm_linear, "general_gemm") - @staticmethod def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs): input_quantizer, weight_quantizer = _QuantLayerNormLinearFunc.modelopt_quantizers @@ -331,15 +323,16 @@ def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs): # We need to recompute the quantized input for the backward pass, so we save the input_quantizer ctx.modelopt_input_quantizer = input_quantizer - original_gemm = _QuantLayerNormLinearFunc._get_original_gemm() + original_gemm = te_layernorm_linear.general_gemm def _patched_general_gemm(weight, input, *gemm_args, **gemm_kwargs): qinput = input_quantizer(input) return original_gemm(weight, qinput, *gemm_args, **gemm_kwargs) with replace_function( - *_QuantLayerNormLinearFunc._gemm_replace_args(), - _patched_general_gemm, # type: ignore[call-arg] + te_layernorm_linear, + "general_gemm", + _patched_general_gemm, ): outputs = te_layernorm_linear._og_LayerNormLinear.forward( ctx, inp, ln_weight, ln_bias, qweight, *args, **kwargs @@ -357,7 +350,7 @@ def backward(ctx, *grad_outputs): """ gemm_call_counter = {"count": 0} - original_gemm = _QuantLayerNormLinearFunc._get_original_gemm() + original_gemm = te_layernorm_linear.general_gemm def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs): # The first time, gemm is used for dgrad calculation @@ -376,8 +369,9 @@ def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs): return original_gemm(qinput, b, *gemm_args, **gemm_kwargs) with replace_function( - *_QuantLayerNormLinearFunc._gemm_replace_args(), - _patched_general_gemm, # type: ignore[call-arg] + te_layernorm_linear, + "general_gemm", + _patched_general_gemm, ): # During backward, the patch does not exist; autograd will automatically use # _QuantLayerNormLinearFunc.backward From c8867ebc50c15bd1e695bba0a87b5c3af6fe15a3 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:27:42 -0700 Subject: [PATCH 3/3] Document Transformer Engine 2.0 as the minimum for TE quantization Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + docs/source/getting_started/_installation_for_Linux.rst | 3 +++ 2 files changed, 4 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5b8b221cca7..721f8ebc81d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -48,6 +48,7 @@ Changelog - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. +- Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly. diff --git a/docs/source/getting_started/_installation_for_Linux.rst b/docs/source/getting_started/_installation_for_Linux.rst index 8f858f72f75..369a417d727 100644 --- a/docs/source/getting_started/_installation_for_Linux.rst +++ b/docs/source/getting_started/_installation_for_Linux.rst @@ -20,6 +20,9 @@ Latest Model Optimizer (``nvidia-modelopt``) currently has the following system +-------------------------+-----------------------------+ | TensorRT-LLM (Optional) | >=1.0 | +-------------------------+-----------------------------+ +| Transformer Engine | >=2.0 | +| (Optional) | | ++-------------------------+-----------------------------+ | ONNX Runtime (Optional) | See CUDA notes below | +-------------------------+-----------------------------+ | TensorRT (Optional) | >=10.0 |