Support transformers 5.15-5.18 (<5.19) and drop DBRX

CHANGELOG.rst

@@ -100,6 +100,7 @@ Changelog
 - The TensorRT-LLM checkpoint export format is deprecated and will be removed in 0.49.0: ``export_tensorrt_llm_checkpoint`` and ``torch_to_tensorrt_llm_checkpoint`` now emit a ``DeprecationWarning`` on use. Use ``export_hf_checkpoint``, which exports a unified Hugging Face checkpoint deployable on TensorRT-LLM, vLLM and SGLang. Its implementation moved to ``modelopt.torch.export.trtllm``, so import those two functions from there and the ``ModelConfig`` dataclasses from ``modelopt.torch.export.trtllm.model_config``; both functions remain importable from ``modelopt.torch.export`` for this release only.
 - ``modelopt.torch.export.get_model_type`` and ``modelopt.torch.export.model_utils.MODEL_NAME_TO_TYPE`` return TensorRT-LLM model names and are deprecated for removal in 0.49.0 (they now emit a ``DeprecationWarning``); use the Hugging Face model type, ``model.config.model_type``, instead. ``export_tensorrt_llm_checkpoint`` now detects ``decoder_type`` when it is omitted, and ``examples/hf_ptq`` uses Hugging Face model types throughout.
 - The ``evaluation`` agent skill no longer supports GDPVal: its task recipe, example config and Apptainer SIF helper are removed, and the shared NeMo Gym machinery it carried now lives at ``references/gym.md`` with the launcher wrapper renamed ``scripts/nel-gym.sh``. GDPVal is an AA-suite member, so an "AA" request now generates the ``aa/`` tasks only -- report per-task scores rather than an aggregate compared against a published AA Index.
+- DBRX model support is removed: ModelOpt no longer quantizes or exports DBRX (``DbrxForCausalLM``) checkpoints. Use ModelOpt 0.47 to quantize DBRX, or write your own plugin as described in the customized-model quantization guide.
 - Deprecate ``metadata.recipe_type`` in recipe YAML, for both single-file recipes and a directory recipe's ``metadata.yml``. A recipe now says what kind it is with a ``# modelopt-schema:`` comment naming its schema class, or by delegating to a recipe that does (single-file recipes only -- a directory recipe has no body to delegate through); ``recipe_type`` is still read and still honoured, so a recipe outside this repo keeps working unchanged, but new recipes should leave it out and every recipe shipped here has been converted. Where both are present they must agree, and so must a recipe and the recipe it delegates to -- a disagreement is an error rather than a silent preference.
 
 **Bug Fixes**

docs/source/guides/_customized_model_quantization.rst

@@ -3,144 +3,72 @@ Guides to quantize a customized model from Hugging Face for TensorRT-LLM deploym
 ===================================================================================
 
 ModelOpt can usually quantize PyTorch models from the Hugging Face directly. By default, ModelOpt searches the PyTorch model and replaces the torch ``nn.Linear`` module with a quantized linear module.
+Fused MoE experts that store their weights as 3-D ``gate_up_proj`` / ``down_proj`` parameters and call ``F.linear`` (e.g. Mixtral, Qwen3-MoE, DeepSeek-V3 in transformers 5.x) are also detected and quantized automatically.
 
-If the model happens not using the ``nn.Linear`` for the linear layers, a customized Hugging Face plugin needs to be implemented to convert the model to use ``nn.Linear`` instead.
+If a model computes its linear ops some other way, a customized Hugging Face plugin is needed to insert the quantizers.
 
-The following is an example about how a customized Hugging Face model can be supported using modelopt:
+The following example shows how ModelOpt supports the `Llama 4 <https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct>`_ MoE.
+Its `Llama4TextExperts <https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama4/modeling_llama4.py>`_ module stores all experts in two ``nn.Parameter`` tensors, ``gate_up_proj`` of shape ``(num_experts, hidden_size, 2 * expert_dim)`` and ``down_proj`` of shape ``(num_experts, expert_dim, hidden_size)``, and applies them with ``torch.bmm``, so there is no ``nn.Linear`` for ModelOpt to replace.
 
-The `DBRX model <https://huggingface.co/databricks/dbrx-instruct>`_ is an MoE model with customized MoE linear implementation. The MoE layer in DBRX is implemented as a `DbrxExperts <https://github.com/databricks/dbrx/blob/main/model/modeling_dbrx.py>`_ module, where the three linear ops (w1, v1 and v2) are represented as ``nn.Parameter``. The linear op is forwarded as a pure ``matmul`` op.
+The plugin works as follows:
 
-As ModelOpt cannot detect these linear ops out-of-the-box, a HugggingFace plugin is implemented as the following:
+#. Define a ``_QuantLlama4TextExperts`` subclass of ``QuantModule`` and create the input and weight ``TensorQuantizer`` modules in ``_setup``.
+#. Re-implement ``forward`` with the same signature, quantizing the inputs and the weights of both ``bmm`` calls.
+   ModelOpt's per-channel and per-block quantization expect the input dimension last, but these weights are ``(num_experts, in_dim, out_dim)``, so they are quantized transposed.
+#. Register ``_QuantLlama4TextExperts`` to replace ``Llama4TextExperts`` from the ``transformers`` library.
+#. Quantize the model after the plugin is registered, for example with the `hf_ptq example <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/hf_ptq>`_.
+#. Export the quantized model with :meth:`export_hf_checkpoint <modelopt.torch.export.unified_export_hf.export_hf_checkpoint>`. A new module type may also need an export handler; see ``modelopt/torch/export/hf_export_handlers.py``.
+   If the customized model is not supported by TensorRT-LLM, add support in its PyTorch backend. See the :doc:`unified HF export guide <../deployment/3_unified_hf>` or :doc:`contact us <../support/1_contact>` for help.
 
-#. Define a customized ``_QuantDbrxExpertGLU`` as a ``DynamicModule`` with the same ``forward`` signature.
-#. Rewrite the linear ops (w1, v1 and v2) as a standard ``nn.Linear`` op, and re-implement the ``forward`` method.
-#. Register the new dynamic ``_QuantDbrxExperts`` to replace the ``DbrxExperts`` from the modeling_dbrx.py in the ``transformers`` library
-#. Try quantize the DBRX model after the plugin is implemented, feel free to follow the `hf_ptq example <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/hf_ptq>`_.
-#. Export the quantized model with :meth:`export_hf_checkpoint <modelopt.torch.export.unified_export_hf.export_hf_checkpoint>`. If the customized model is not supported by TensorRT-LLM, add support in its PyTorch backend and adapt the HF exporter if needed. See the :doc:`unified HF export guide <../deployment/3_unified_hf>` or :doc:`contact us <../support/1_contact>` for help.
-
-The following code snippet is excerpted from ``modelopt/torch/quantization/plugins/huggingface.py``
+The following code snippet is simplified from the plugin in ``modelopt/torch/quantization/plugins/huggingface.py``, which also handles weight-only calibration.
+For your own model, write a plugin like it in your own code and run it before quantizing.
 
 .. code-block:: python
 
-    from modelopt.torch.opt.dynamic import DynamicModule
-    from modelopt.torch.quantization.nn import QuantModuleRegistry
+    import torch
+    from transformers.models.llama4.modeling_llama4 import Llama4TextExperts
 
-    # For more information on DbrxExpert, see https://github.com/huggingface/transformers/blob/dcdda532/src/transformers/models/dbrx/modeling_dbrx.py#L756
-    class _QuantDbrxExperts(DynamicModule):
-        def _setup(self):
-            """Modify the DbrxExpert."""
-            # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU
-            pass
-
-        # forward method copied from the original dbrx repo - https://github.com/databricks/dbrx/blob/a3200393/model/modeling_dbrx.py#L795
-        def forward(
-            self,
-            x: torch.Tensor,
-            weights: torch.Tensor,
-            top_weights: torch.Tensor,
-            top_experts: torch.LongTensor,
-        ) -> torch.Tensor:
-            bsz, q_len, hidden_size = x.shape
-            x = x.view(-1, hidden_size)
-            out = torch.zeros_like(x)
-
-            expert_mask = nn.functional.one_hot(top_experts, num_classes=self.moe_num_experts).permute(
-                2, 1, 0
-            )
-            for expert_idx in range(0, self.moe_num_experts):
-                topk_idx, token_idx = torch.where(expert_mask[expert_idx])
-                if token_idx.shape[0] == 0:
-                    continue
+    from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer
 
-                token_list = token_idx.tolist()
-                topk_list = topk_idx.tolist()
 
-                expert_tokens = x[None, token_list].reshape(-1, hidden_size)
-                expert_out = (
-                    self.mlp(expert_tokens, expert_idx) * top_weights[token_list, topk_list, None]
-                )
+    class _TransposedQuantization(torch.autograd.Function):
+        """Quantize a (num_experts, in_dim, out_dim) weight with in_dim last, using a straight-through gradient."""
 
-                out.index_add_(0, token_idx, expert_out)
+        @staticmethod
+        def forward(ctx, inputs, quantizer):
+            return quantizer(inputs.transpose(-1, -2).contiguous()).transpose(-1, -2)
 
-            out = out.reshape(bsz, q_len, hidden_size)
-            return out
+        @staticmethod
+        def backward(ctx, grad_output):
+            return grad_output, None
 
 
-    class _QuantDbrxExpertGLU(DynamicModule):
-        def _setup(self):
-            """Modify the DbrxExpertGLU by using nn.Linear layers."""
-            dtype, device = self.w1.dtype, self.w1.device
-
-            def _copy_weights(modules, weights):
-                modules.to(dtype=dtype, device=device)
-                for expert_idx, module in enumerate(modules):
-                    with torch.no_grad():
-                        module.weight.copy_(weights[expert_idx].detach())
-
-            self.w1_linear = nn.ModuleList(
-                [
-                    nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False)
-                    for _ in range(self.moe_num_experts)
-                ]
-            )
-            _copy_weights(
-                self.w1_linear,
-                self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size),
-            )
-            delattr(self, "w1")
+    _transposed_quantize = _TransposedQuantization.apply
 
-            self.v1_linear = nn.ModuleList(
-                [
-                    nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False)
-                    for _ in range(self.moe_num_experts)
-                ]
-            )
-            _copy_weights(
-                self.v1_linear,
-                self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size),
-            )
-            delattr(self, "v1")
 
-            self.w2_linear = nn.ModuleList(
-                [
-                    nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False)
-                    for _ in range(self.moe_num_experts)
-                ]
+    class _QuantLlama4TextExperts(QuantModule):
+        def _setup(self):
+            self.gate_up_proj_input_quantizer = TensorQuantizer()
+            self.gate_up_proj_weight_quantizer = TensorQuantizer()
+            self.down_proj_input_quantizer = TensorQuantizer()
+            self.down_proj_weight_quantizer = TensorQuantizer()
+
+        # Same as Llama4TextExperts.forward, with quantizers on both bmm inputs and weights
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
+            gate_up = torch.bmm(
+                self.gate_up_proj_input_quantizer(hidden_states),
+                _transposed_quantize(self.gate_up_proj, self.gate_up_proj_weight_quantizer),
             )
-            _copy_weights(
-                self.w2_linear,
-                self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose(
-                    1, 2
-                ),
+            gate, up = gate_up.chunk(2, dim=-1)
+            next_states = torch.bmm(
+                self.down_proj_input_quantizer(up * self.act_fn(gate)),
+                _transposed_quantize(self.down_proj, self.down_proj_weight_quantizer),
             )
-            delattr(self, "w2")
-
-        def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:
-            x1 = self.w1_linear[expert_idx](x)
-            x2 = self.v1_linear[expert_idx](x)
-            x1 = self.activation_fn(x1)
-            x1 = x1 * x2
-            return self.w2_linear[expert_idx](x1)
-
-
-    if transformers.models.dbrx.modeling_dbrx.DbrxExperts not in QuantModuleRegistry:
-        QuantModuleRegistry.register(
-            {transformers.models.dbrx.modeling_dbrx.DbrxExperts: "hf.DbrxExperts"}
-        )(_QuantDbrxExperts)
-
-    if transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU not in QuantModuleRegistry:
-        QuantModuleRegistry.register(
-            {transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU: "hf.DbrxExpertGLU"}
-        )(_QuantDbrxExpertGLU)
-
+            return next_states.view(-1, self.hidden_size)
 
-    def register_dbrx_moe_on_the_fly(model):
-        """Register DBRX MoE modules as QUANT_MODULE.
 
-        The MoE class in DBRX is `transformers_modules.modeling_dbrx.DbrxExpertGLU`, which loads dynamically.
-        """
-        if type(model).__name__ in ["DbrxForCausalLM"]:
-            moe_type = type(model.transformer.blocks[0].ffn.experts.mlp)
-            # Create a QuantDbrxExpertGLU class on the fly
-            if QuantModuleRegistry.get(moe_type) is None:
-                QuantModuleRegistry.register({moe_type: moe_type.__name__})(_QuantDbrxExpertGLU)
+    if Llama4TextExperts not in QuantModuleRegistry:
+        QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})(
+            _QuantLlama4TextExperts
+        )

examples/alpamayo/qad.py

@@ -337,7 +337,12 @@ def main():
         "per_device_eval_batch_size": 1,
         "gradient_accumulation_steps": args.grad_accum,
         "learning_rate": args.lr,
-        "warmup_ratio": args.warmup_ratio,
+        # transformers 5.15 removed warmup_ratio; from 5.0, a warmup_steps below 1 is a ratio.
+        (
+            "warmup_ratio"
+            if "warmup_ratio" in transformers.TrainingArguments.__dataclass_fields__
+            else "warmup_steps"
+        ): args.warmup_ratio,
         "max_steps": args.max_steps,
         "logging_steps": args.logging_steps,
         "save_steps": args.save_steps,

examples/alpamayo/quantize.py

@@ -236,7 +236,11 @@ def _teacher_forced_flow_loss_forward(
         use_cache=True,
         **forward_kwargs,
     )
-    prompt_cache.crop(prefill_seq_len)
+    # Drop the diffusion tokens appended to the prompt cache. A negative count works on every
+    # transformers version; positive (absolute-size) crops are deprecated since 5.15.
+    appended = prompt_cache.get_seq_length() - prefill_seq_len
+    if appended > 0:
+        prompt_cache.crop(-appended)
     last_hidden = expert_out.last_hidden_state[:, -n_diffusion_tokens:]
     v_pred = self.action_out_proj(last_hidden).view(b, *self.action_space.get_action_space_dims())
 

examples/gpt-oss/requirements.txt

@@ -1,5 +1,5 @@
 kernels>=0.9.0,<0.13
-trackio<0.21
+trackio
 # transformers>=5.3 avoids CVE-2026-4372 (RCE via the `kernels` Hub download path, which this example installs)
 transformers>=5.3
 trl>=1.0

examples/hf_ptq/hf_ptq.py

@@ -1141,7 +1141,6 @@ def quantize_main(
                 is_nemotron_vl_model,
             )
         else:
-            assert model_type != "dbrx", f"Does not support export {model_type} without quantizaton"
             print(f"qformat: {args.qformat}. No quantization applied, export {device} model")
 
     # If asked, run the closed-form MXFP4 -> NVFP4 cast: read the source MXFP4

examples/llm_qat/configs/train/finetune.yaml

@@ -16,7 +16,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 use_liger_kernel: true
 seed: 42

examples/llm_qat/configs/train/qad_nvfp4.yaml

@@ -21,7 +21,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 use_liger_kernel: true
 manual_gc: true

examples/llm_qat/configs/train/qad_scale_only.yaml

@@ -28,7 +28,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 use_liger_kernel: true
 manual_gc: true

examples/llm_qat/configs/train/qad_with_learnt_amax.yaml

@@ -23,7 +23,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 use_liger_kernel: true
 manual_gc: true

examples/llm_qat/configs/train/qat_nvfp4.yaml

@@ -17,7 +17,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 use_liger_kernel: true
 manual_gc: true

examples/llm_qat/configs/train/qlora_nvfp4.yaml

@@ -19,7 +19,7 @@ per_device_train_batch_size: 2
 per_device_eval_batch_size: 2
 gradient_accumulation_steps: 2
 model_max_length: 8192
-warmup_ratio: 0.05
+warmup_steps: 0.05
 lr_scheduler_type: cosine
 gradient_checkpointing: true
 use_liger_kernel: true

modelopt/torch/__init__.py

@@ -51,7 +51,7 @@
 
     if _Version(_transformers_version) < _Version("4.57") or _Version(
         _transformers_version
-    ) >= _Version("5.15"):
+    ) >= _Version("5.19"):
         _warnings.warn(
             f"transformers {_transformers_version} is not tested with current version of modelopt and may cause issues."
             " Please install recommended version with `pip install -U nvidia-modelopt[hf]` if working with HF models.",

modelopt/torch/distill/plugins/huggingface.py

@@ -190,7 +190,9 @@ def _ds_gather(self, params):
     def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         """Train and evaluate on KD loss, with eval CE tracked as a metric."""
         self._ensure_teacher_prepared()
-        kd_inputs = {k: v for k, v in inputs.items() if k != "labels"}
+        # ``skip_logits`` (added for Liger eval by transformers>=5.15) would make the teacher's
+        # Liger forward demand labels; the fused KD loss needs its hidden states instead.
+        kd_inputs = {k: v for k, v in inputs.items() if k not in ("labels", "skip_logits")}
         labels = inputs.get("labels")
         is_training = model.training
 

modelopt/torch/export/hf_export_handlers.py

@@ -49,23 +49,6 @@ def _export_weight(
 # Preparation handlers are registered in the same precedence as the legacy MoE prepass.
 
 
-# Keyed on the mixin class name too: the generated class is normally named
-# "QuantDbrxExperts", but _DMRegistryCls falls back to a module-prefixed name on
-# collision, while "_QuantDbrxExperts" remains in the generated class's MRO.
-@PrepareMoEInputsRegistry.register("QuantDbrxExperts", "_QuantDbrxExperts")
-def _prepare_dbrx_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
-    """Fill missing input amax values for DBRX per-expert ModuleLists."""
-    experts_mlp = moe_module.experts.mlp
-    for linear_name in get_expert_linear_names(moe_module, ctx.model_type):
-        if hasattr(experts_mlp, linear_name):
-            linear_modulelist = getattr(experts_mlp, linear_name)
-            if hasattr(linear_modulelist, "__iter__"):
-                set_expert_quantizer_amax(
-                    modules=list(linear_modulelist),
-                    quantizer_attrs=["input_quantizer"],
-                )
-
-
 @PrepareMoEInputsRegistry.register(predicate=_has_fused_experts_quantizers)
 def _prepare_fused_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
     """Mark fused experts handled; their missing amax fallback occurs during export."""

modelopt/torch/export/trtllm/decoder_type.py

@@ -53,7 +53,6 @@
     "MixtralForCausalLM": "llama",
     "ArcticForCausalLM": "llama",
     "StarCoder": "gpt",
-    "Dbrx": "dbrx",
     "T5": "t5",
     "Bart": "bart",
     "GLM": "glm",