Drop version checks for long-unsupported dependency releases

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.

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                     |

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.

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():

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(

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"
 
 
@@ -316,20 +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():
-        if Version("2.0") <= _TE_VERSION:
-            return te_layernorm_linear.general_gemm
-        else:
-            return te_layernorm_linear.tex.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")
-
     @staticmethod
     def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs):
         input_quantizer, weight_quantizer = _QuantLayerNormLinearFunc.modelopt_quantizers
@@ -340,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
@@ -366,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
@@ -385,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

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:

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"]}