diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 76f31c28d90..a23d9eede88 100755 --- a/CHANGELOG.rst +++ b/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** diff --git a/docs/source/guides/_customized_model_quantization.rst b/docs/source/guides/_customized_model_quantization.rst index b2fa879ecd4..d89a94d7adb 100644 --- a/docs/source/guides/_customized_model_quantization.rst +++ b/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 `_ MoE. +Its `Llama4TextExperts `_ 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 `_ is an MoE model with customized MoE linear implementation. The MoE layer in DBRX is implemented as a `DbrxExperts `_ 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 `_. +#. Export the quantized model with :meth:`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 `_. -#. Export the quantized model with :meth:`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 + ) diff --git a/examples/alpamayo/qad.py b/examples/alpamayo/qad.py index 77c55cbd331..fac41af6643 100644 --- a/examples/alpamayo/qad.py +++ b/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, diff --git a/examples/alpamayo/quantize.py b/examples/alpamayo/quantize.py index 6a7b3d56f61..7d7624aac40 100644 --- a/examples/alpamayo/quantize.py +++ b/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()) diff --git a/examples/gpt-oss/requirements.txt b/examples/gpt-oss/requirements.txt index 03bc8cc59f9..200632c72f2 100644 --- a/examples/gpt-oss/requirements.txt +++ b/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 diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 4bdec9b8679..851ab22138a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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 diff --git a/examples/llm_qat/configs/train/finetune.yaml b/examples/llm_qat/configs/train/finetune.yaml index 5b3168a0f9b..86d283ce99d 100644 --- a/examples/llm_qat/configs/train/finetune.yaml +++ b/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 diff --git a/examples/llm_qat/configs/train/qad_nvfp4.yaml b/examples/llm_qat/configs/train/qad_nvfp4.yaml index cab093ae6dc..9453b616ee5 100644 --- a/examples/llm_qat/configs/train/qad_nvfp4.yaml +++ b/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 diff --git a/examples/llm_qat/configs/train/qad_scale_only.yaml b/examples/llm_qat/configs/train/qad_scale_only.yaml index 39c999d02e1..5b2adef98a9 100644 --- a/examples/llm_qat/configs/train/qad_scale_only.yaml +++ b/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 diff --git a/examples/llm_qat/configs/train/qad_with_learnt_amax.yaml b/examples/llm_qat/configs/train/qad_with_learnt_amax.yaml index 1923849f588..e284e67238d 100644 --- a/examples/llm_qat/configs/train/qad_with_learnt_amax.yaml +++ b/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 diff --git a/examples/llm_qat/configs/train/qat_nvfp4.yaml b/examples/llm_qat/configs/train/qat_nvfp4.yaml index 0afc7e3d1de..22ee839ec1b 100644 --- a/examples/llm_qat/configs/train/qat_nvfp4.yaml +++ b/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 diff --git a/examples/llm_qat/configs/train/qlora_nvfp4.yaml b/examples/llm_qat/configs/train/qlora_nvfp4.yaml index 346e846bb80..162cf0f8521 100644 --- a/examples/llm_qat/configs/train/qlora_nvfp4.yaml +++ b/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 diff --git a/modelopt/torch/__init__.py b/modelopt/torch/__init__.py index 5eaf1463584..b540457724d 100644 --- a/modelopt/torch/__init__.py +++ b/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.", diff --git a/modelopt/torch/distill/plugins/huggingface.py b/modelopt/torch/distill/plugins/huggingface.py index a6ecdd0d7c0..cb88c489e79 100644 --- a/modelopt/torch/distill/plugins/huggingface.py +++ b/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 diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 550fb0dbd6a..7c526d8255f 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/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.""" diff --git a/modelopt/torch/export/trtllm/decoder_type.py b/modelopt/torch/export/trtllm/decoder_type.py index ef117aa000e..c4d6ac1b4aa 100644 --- a/modelopt/torch/export/trtllm/decoder_type.py +++ b/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", diff --git a/modelopt/torch/export/trtllm/hf_config_map.py b/modelopt/torch/export/trtllm/hf_config_map.py index 377e1887b5d..31c1e3c6654 100644 --- a/modelopt/torch/export/trtllm/hf_config_map.py +++ b/modelopt/torch/export/trtllm/hf_config_map.py @@ -66,8 +66,8 @@ ( ["num_local_experts", "moe_num_experts"], "moe_num_experts", - ), # Mixture of Experts (Mixtral, DBRX) - (["num_experts_per_tok", "moe_top_k"], "moe_top_k"), # Mixture of Experts (Mixtral, DBRX) + ), # Mixture of Experts (Mixtral) + (["num_experts_per_tok", "moe_top_k"], "moe_top_k"), # Mixture of Experts (Mixtral) (["model_type"], "qwen_type"), # qwen (["lru_width"], "rnn_hidden_size"), # Recurrent Gemma (["embeddings_scale_by_sqrt_dim"], "emb_scale_by_sqrt_dim"), # Recurrent Gemma @@ -76,7 +76,7 @@ (["final_logit_softcapping"], "final_logit_softcapping"), # Gemma2 (["attn_logit_softcapping"], "attn_logit_softcapping"), # Gemma2 (["query_pre_attn_scalar"], "query_pre_attn_scalar"), # Gemma2 - (["clip_qkv"], "clip_qkv"), # DBRX + (["clip_qkv"], "clip_qkv"), # MPT (["rope_scaling"], "rope_scaling"), # Llama 3.1 (["cross_attention_layers"], "cross_attention_layers"), # Mllama (["vision_output_dim"], "vision_output_dim"), # Mllama diff --git a/modelopt/torch/export/trtllm/layer_utils.py b/modelopt/torch/export/trtllm/layer_utils.py index 342f1d99a43..8f02ad7298f 100755 --- a/modelopt/torch/export/trtllm/layer_utils.py +++ b/modelopt/torch/export/trtllm/layer_utils.py @@ -812,21 +812,6 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer): return config -def _get_dbrx_expert(experts: nn.Module, export_id: int, linear_name: str): - # DBRX experts layout is: - # experts: - # w1[0] - # w1[1] - # ... - # w2[0] - # w2[1] - # ... - # v1[0] - # v1[1] - # ... - return getattr(experts, linear_name)[export_id] - - def _build_stacked_linear(experts: nn.Module, module_name, linear_type, num_experts, expert_getter): config = LinearConfig(linear_type=linear_type) @@ -991,7 +976,7 @@ def build_stacked_experts( def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig: """Builds the MOE config for the module.""" assert is_moe(module) - assert decoder_type in ["llama", "dbrx", "phi3", "deepseek", "qwen"] + assert decoder_type in ["llama", "phi3", "deepseek", "qwen"] config = MOEConfig() @@ -1003,8 +988,6 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig: else: # For huggingface model config.router = build_linear_config(module.gate, LINEAR_ROW) - elif decoder_type == "dbrx": - config.router = build_linear_config(module.router.layer, LINEAR_ROW) elif decoder_type == "deepseek": # Build a linear on the fly. router_linear = nn.Linear(module.gate.gating_dim, module.gate.n_routed_experts, bias=None) @@ -1059,13 +1042,6 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig: len(module.experts), _get_expert_attr, ) - elif decoder_type == "dbrx": - experts.fc, experts.proj = build_stacked_experts( - module.experts.mlp, - ["w1_linear", "w2_linear", "v1_linear"], - len(module.experts.mlp.w1_linear), - _get_dbrx_expert, - ) elif decoder_type in ["deepseek", "qwen"]: experts.fc, experts.proj = build_stacked_experts( module.experts, @@ -1078,7 +1054,7 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig: config.experts = experts - # activation for mixtral and dbrx + # activation for mixtral config.hidden_act = "swiglu" return config @@ -1220,12 +1196,7 @@ def build_decoder_config( # Set all config fields in modelopt from HF config _set_layer_config_from_metaconfig(config, model_metadata_config) - if type(module).__name__ == "DbrxBlock": - # Flatten DBRX attention and ffn - module_layers = {} - module_layers.update(dict(getattr(module, "norm_attn_norm").named_children())) - module_layers.update({"ffn": module.ffn}) - elif decoder_type == "t5": + if decoder_type == "t5": # Combine two modules (T5LayerSelfAttention, T5LayerFF) / three modules # ((T5LayerSelfAttention, T5LayerCrossAttention, T5LayerFF)) of T5 model # (depending on whether it's encoder / decoder) into one decoder layer @@ -1536,10 +1507,6 @@ def update_experts_avg_prequant_scale(experts: nn.Module): get_func = _get_expert_attr num_experts = len(experts.experts) experts = experts.experts - elif "dbrx" in type(experts).__name__.lower(): - get_func = _get_dbrx_expert - num_experts = len(getattr(experts.experts.mlp, experts_linear_names[0])) - experts = experts.experts.mlp else: raise NotImplementedError("MoE model not supported") @@ -1568,8 +1535,6 @@ def get_experts_linear_names(model: torch.nn.Module): """Returns linear layer names based on decoder type for MoE models.""" if "mixtral" in type(model).__name__.lower(): return ["w1", "w2", "w3"] - elif "dbrx" in type(model).__name__.lower(): - return ["w1_linear", "w2_linear", "v1_linear"] else: raise NotImplementedError("MoE model not supported") diff --git a/modelopt/torch/export/trtllm/model_config.py b/modelopt/torch/export/trtllm/model_config.py index 0182ce01d72..c3d630f9fc9 100755 --- a/modelopt/torch/export/trtllm/model_config.py +++ b/modelopt/torch/export/trtllm/model_config.py @@ -489,7 +489,7 @@ class DecoderLayerConfig: attn_logit_softcapping: float = 0 query_pre_attn_scalar: float = 0 - # DBRX + # MPT clip_qkv: int = 0 # T5, Mllama diff --git a/modelopt/torch/export/trtllm/model_config_export.py b/modelopt/torch/export/trtllm/model_config_export.py index e79b6ee49d3..6965fd57c7b 100644 --- a/modelopt/torch/export/trtllm/model_config_export.py +++ b/modelopt/torch/export/trtllm/model_config_export.py @@ -216,7 +216,7 @@ def _torch_to_tensorrt_llm_checkpoint( if hasattr(model, "model") and hasattr(model.model, "alibi_mask"): model_metadata_config["alibi"] = True - # For MPT, DBRX + # For MPT for config_key in ["attn_config", "ffn_config"]: config_value = model_metadata_config.get(config_key, None) if config_value: diff --git a/modelopt/torch/export/trtllm/tensorrt_llm_utils.py b/modelopt/torch/export/trtllm/tensorrt_llm_utils.py index 04e0b7526d8..78943690d72 100755 --- a/modelopt/torch/export/trtllm/tensorrt_llm_utils.py +++ b/modelopt/torch/export/trtllm/tensorrt_llm_utils.py @@ -319,8 +319,6 @@ def convert_to_tensorrt_llm_config( config["eos_token_id"] = model_config.eos_token_id config["bos_token_id"] = model_config.bos_token_id config["pad_token_id"] = model_config.pad_token_id - elif decoder_type == "dbrx": - config["clip_qkv"] = first_attention_decoder_config.clip_qkv elif decoder_type == "mllama": num_layers = config["num_hidden_layers"] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 67e4ab83917..21f36e95444 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -319,7 +319,6 @@ def collect_shared_input_modules( def _input_hook(module, input, output): """Update dictionary with list of all modules that share the same input.""" if len(input) > 0 and isinstance(input[0], torch.Tensor): - # TODO: Handle DBRX MoE case input_to_linear[input[0]].append(module) def _output_hook(module, input, output): @@ -511,7 +510,6 @@ def _llm_dummy_forward(model: torch.nn.Module) -> None: def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): """Group modules that take the same input and register shared parameters in module.""" - # TODO: Handle DBRX MoE quantization_format = get_quantization_format(model) model_hf_type = hf_model_type(model) module_names = set() diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index c835e6b63e4..fef532fde08 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -67,7 +67,6 @@ # Importing the model packages registers every spec as a side effect. from . import ( # isort: skip arctic, - dbrx, deepseek, deepseek_v3, deepseek_v4, diff --git a/modelopt/torch/models/dbrx/__init__.py b/modelopt/torch/models/dbrx/__init__.py deleted file mode 100644 index 774a5a4e201..00000000000 --- a/modelopt/torch/models/dbrx/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - - -"""Per-model specs for HF model type ``dbrx``.""" - -from . import specs diff --git a/modelopt/torch/models/dbrx/specs.py b/modelopt/torch/models/dbrx/specs.py deleted file mode 100644 index b114abaeba3..00000000000 --- a/modelopt/torch/models/dbrx/specs.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""DBRX specs (HF model type ``dbrx``).""" - -from ..specs import ModelSpec, MoESpec, register - -__all__: list[str] = [] - -# Expert names refer to the quantized layout: _QuantDbrxExpertGLU rewrites the fused -# w1/v1/w2 parameters into per-expert w1_linear/v1_linear/w2_linear ModuleLists on -# experts.mlp (see modelopt/torch/quantization/plugins/huggingface.py). -register( - ModelSpec( - model_type="dbrx", - min_transformers_version="4.57", - moe_spec=MoESpec( - block_names=("DbrxFFN",), - expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), - ), - ) -) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 516d127c07a..aa9110223ee 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -110,7 +110,7 @@ class MoESpec(SpecSection): expert_linear_names: tuple[str, ...] | None = None """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``. - For layouts modelopt rewrites (e.g. quantized DBRX), these are the names on the + For layouts modelopt rewrites into per-expert modules, these are the names on the rewritten module.""" fused_expert_names: bool = False @@ -125,7 +125,7 @@ class MoESpec(SpecSection): """The (gate, up) pair among ``expert_linear_names`` that serving engines fuse into a single ``gate_up_proj``, e.g. ``("gate_proj", "up_proj")`` or ``("w1", "w3")``. ``None`` for non-gated experts (NemotronH) and already-fused - layouts (GptOss, DBRX).""" + layouts (GptOss).""" @property def gate_up_pairs(self) -> tuple[tuple[str, str], ...]: diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 6665fdcaedc..250505dba91 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -311,6 +311,23 @@ class ModelOptTrainerArguments(ModelOptHFArguments): ) +def _translate_warmup_keys(config: dict, actions) -> None: + """Accept either warmup spelling across transformers versions. + + transformers 5.15 removed ``warmup_ratio``; from 5.0, a ``warmup_steps`` below 1 is a ratio. + """ + types = {a.dest: a.type for a in actions} + if "warmup_steps" not in types: + return + if "warmup_ratio" in config and "warmup_ratio" not in types: + ratio = config.pop("warmup_ratio") + config.setdefault("warmup_steps", ratio) + elif "warmup_ratio" in types and types["warmup_steps"] is int: + steps = config.get("warmup_steps") + if isinstance(steps, float) and 0 < steps < 1: + config["warmup_ratio"] = config.pop("warmup_steps") + + class ModelOptArgParser(HfArgumentParser): """HfArgumentParser with ``--config`` YAML support and ``--generate_docs`` for ARGUMENTS.md.""" @@ -347,6 +364,7 @@ def parse_args_into_dataclasses(self, args=None, **kwargs): with open(config_path, encoding="utf-8") as f: config = yaml.safe_load(f) if config: + _translate_warmup_keys(config, self._actions) known_by_parser = {a.dest for a in self._actions} all_modelopt_fields = self._all_modelopt_fields() applicable = {} @@ -791,6 +809,9 @@ def _sharded_liger_compute(self, fn): def compute_loss(self, model, inputs, return_outputs=False, **kwargs): """Compute loss, patching lm_head to identity when using liger fused loss.""" if self.use_liger_kernel: + # The fused loss is computed here from the hidden states, so drop the ``skip_logits`` + # that transformers>=5.15 adds for eval: Liger would then need the popped labels. + inputs = {k: v for k, v in inputs.items() if k != "skip_logits"} with self._liger_identity_lm_head(): return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs) return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs) diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_model_descriptor.py b/modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_model_descriptor.py index a0f9c95c6ce..587c62e98b3 100644 --- a/modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_model_descriptor.py +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_model_descriptor.py @@ -15,6 +15,7 @@ # mypy: ignore-errors +import inspect import re from dataclasses import dataclass, field from typing import Dict, List @@ -105,8 +106,13 @@ def init_rotary_embedding(model, runtime): model.config.vision_config if hasattr(model.config, "vision_config") else None ) if vision_config is not None: - head_dim = vision_config.hidden_size // vision_config.num_heads - model.model.visual.rotary_pos_emb = Qwen3VLMoeVisionRotaryEmbedding(head_dim // 2).to( + # transformers 5.17 builds the vision RoPE from the vision config instead of its dim. + if "config" in inspect.signature(Qwen3VLMoeVisionRotaryEmbedding).parameters: + rotary_pos_emb = Qwen3VLMoeVisionRotaryEmbedding(vision_config) + else: + head_dim = vision_config.hidden_size // vision_config.num_heads + rotary_pos_emb = Qwen3VLMoeVisionRotaryEmbedding(head_dim // 2) + model.model.visual.rotary_pos_emb = rotary_pos_emb.to( device=runtime.device, dtype=runtime.dtype ) diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index caa6ef0eacb..32fec050a9f 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -688,7 +688,6 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): r"^(.*?\.mlp\.experts\.local_experts)\.\d+\.(linear_fc1|linear_fc2)$", r"^(.*?)\.(gate_proj|up_proj)$", # gate_proj, up_proj for llama like models r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts - r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts # Qwen3.5/3.6 hybrid linear_attn: vLLM fuses (in_proj_qkv, in_proj_z) # into ``in_proj_qkvz`` and (in_proj_a, in_proj_b) into ``in_proj_ba`` and # requires fused shards to share quant_algo. Two callables (not one @@ -1723,7 +1722,6 @@ class _AutoQuantizeBackwardScoringSearcher(_AutoQuantizeBaseSearcher): r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", r"^(.*?)\.(\d+\.(w1|w2|w3))$", - r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", ] _custom_support: list[tuple[Callable, Callable, Callable]] = [] diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index ab0d2975563..1524172cfec 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -486,26 +486,43 @@ class HFParallelLinear(torch.nn.Linear, DynamicModule): shard = None def _setup(self): - if isinstance(self.weight, torch.distributed.tensor.DTensor): # transformers<5.0 + # transformers<5.0 and >=5.16 keep the sharded weight as a DTensor + if isinstance(self.weight, torch.distributed.tensor.DTensor): assert self.weight.placements == self.shard, ( f"Received unexpected shard {self.weight.placements} for {self}" ) device_mesh = self.weight.device_mesh - else: # transformers>=5.0: weights are plain Parameters, mesh is on the module + # transformers>=5.16: keep HF's TP wrapper on its local-tensor path (also in training), + # so the quantized linear sees this rank's plain weight shard. + self._hf_quantized_needs_local_tp = True + else: # transformers 5.0-5.15: weights are plain Parameters, mesh is on the module device_mesh = self._hf_device_mesh tp_group = device_mesh.get_group() self._parallel_state = ParallelState(data_parallel_group=-1, tensor_parallel_group=tp_group) @classmethod - def is_compatible(cls, linear) -> bool: + def is_compatible(cls, linear, tp_mesh=None) -> bool: if not isinstance(linear, torch.nn.Linear): return False - if not hasattr(linear, "_hf_tp_plan"): - return False - return linear._hf_tp_plan in cls.supported_hf_tp_plans + if hasattr(linear, "_hf_tp_plan"): # transformers<5.16 + return linear._hf_tp_plan in cls.supported_hf_tp_plans + # transformers>=5.16 marks nothing on the module: the weight is a DTensor on the model's TP + # mesh, sharded along the TP style's dim. Matching the mesh keeps FSDP2 shards out. + weight = linear.weight + return ( + tp_mesh is not None + and isinstance(weight, torch.distributed.tensor.DTensor) + and weight.device_mesh == tp_mesh + and tuple(weight.placements) == cls.shard + ) - # This is hack for now, otherwise DMRegistry treats this class same as nn.Linear + # Defined so DMRegistry does not treat this class the same as nn.Linear. def forward(self, x): + # transformers>=5.16 installs its TP input/output transforms as an instance forward, which + # DynamicModule.convert keeps as _forward_pre_dm; it wraps the unconverted nn.Linear.forward. + pre_fwd = getattr(self, "_forward_pre_dm", None) + if pre_fwd is not None: + return pre_fwd(x) return super().forward(x) @@ -528,7 +545,7 @@ def fold_weight(self, keep_attrs: bool = False): @contextmanager def enable_weight_access_and_writeback(self): - if isinstance(self.weight, torch.distributed.tensor.DTensor): # transformers<5.0 + if isinstance(self.weight, torch.distributed.tensor.DTensor): # transformers<5.0, >=5.16 assert self.weight.placements == self.shard, ( f"Received unexpected shard {self.weight.placements} for {self}" ) @@ -539,7 +556,7 @@ def enable_weight_access_and_writeback(self): yield finally: self.weight = weight - else: # transformers>=5.0: weights are already plain Parameters + else: # transformers 5.0-5.15: weights are already plain Parameters yield @@ -560,11 +577,24 @@ def convert_hf_parallel_linears_on_the_fly(model): This method converts them to `HFColumnParallelLinear` and `HFRowParallelLinear` so that they can be treated as TP sharded layers and not like regular nn.Linear layers. """ - for name, module in model.named_modules(): - if HFColumnParallelLinear.is_compatible(module): + # transformers>=5.16 keeps the TP mesh on the PreTrainedModel, which may sit inside a wrapper. + tp_mesh = next( + (m._device_mesh for m in model.modules() if getattr(m, "_device_mesh", None) is not None), + None, + ) + converted = False + for module in model.modules(): + if HFColumnParallelLinear.is_compatible(module, tp_mesh): HFColumnParallelLinear.convert(module) - elif HFRowParallelLinear.is_compatible(module): + converted = True + elif HFRowParallelLinear.is_compatible(module, tp_mesh): HFRowParallelLinear.convert(module) + converted = True + if tp_mesh is not None and not converted: + warnings.warn( + "Found a transformers tensor-parallel mesh but no TP-sharded linear layer; they will be" + " quantized as unsharded layers." + ) if transformers.pytorch_utils.Conv1D not in QuantModuleRegistry: @@ -783,108 +813,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return next_states -# For more information on DbrxExpert, see https://github.com/huggingface/transformers/blob/dcdda532/src/transformers/models/dbrx/modeling_dbrx.py#L756 -class _QuantDbrxExperts(QuantModule): - def _setup(self): - """Modify the DbrxExpert.""" - # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU - - def forward( - self, - x: torch.Tensor, - top_experts: torch.LongTensor, - top_weights: torch.Tensor, - ) -> 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.num_experts).permute( - 2, 1, 0 - ) - for expert_idx in range(self.num_experts): - topk_idx, token_idx = torch.where(expert_mask[expert_idx]) - if token_idx.shape[0] == 0: - continue - - 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] - ) - - out.index_add_(0, token_idx, expert_out) - - out = out.reshape(bsz, q_len, hidden_size) - return out - - -class _QuantDbrxExpertGLU(QuantModule): - 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()) - - # In transformers 5.0, DbrxExpertGLU.forward uses raw matmul: x @ w1[i] where - # w1[i] has shape (ffn_hidden_size, hidden_size). To match via F.linear (which - # computes x @ W.T), we store weights transposed: W = w1[i].T. - self.w1_linear = nn.ModuleList( - [ - nn.Linear(self.ffn_hidden_size, self.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).transpose( - 1, 2 - ), - ) - delattr(self, "w1") - - self.v1_linear = nn.ModuleList( - [ - nn.Linear(self.ffn_hidden_size, self.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).transpose( - 1, 2 - ), - ) - delattr(self, "v1") - - # w2: down_proj uses intermediate.matmul(w2[i].t()) = F.linear(intermediate, w2[i]) - # so W = w2[i] directly (no extra transpose needed). - self.w2_linear = nn.ModuleList( - [ - nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False) - for _ in range(self.moe_num_experts) - ] - ) - _copy_weights( - self.w2_linear, - self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), - ) - 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) - - class _QuantQwen3VLMoeTextExperts(QuantModule): """Quantized wrapper for the pre-transformers-5.12 ``Qwen3VLMoeTextExperts`` layout. @@ -1171,27 +1099,6 @@ def _is_quant_fused_experts_module(module): return isinstance(module, _QuantFusedExperts) -class _QuantDbrxFFN(_QuantSparseSequentialMoe): - @property - def num_experts(self): - return self.router.moe_num_experts - - @property - def top_k(self): - # In older transformers, top_k was stored on DbrxRouter as moe_top_k. - # In transformers 5.0, DbrxFFN stores it as a plain attribute (top_k). - if hasattr(self.router, "moe_top_k"): - return self.router.moe_top_k - return self.__dict__.get("top_k", 1) - - @top_k.setter - def top_k(self, value): - if hasattr(self.router, "moe_top_k"): - self.router.moe_top_k = value - else: - self.__dict__["top_k"] = value - - @contextmanager def patch_compressed_linear_loading(): """Context manager that patches CompressedLinear to survive custom ``_init_weights`` calls. @@ -1428,20 +1335,6 @@ def unpack_weight(self): except ImportError: pass -try: - from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN - - if DbrxExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxExperts: "hf.DbrxExperts"})(_QuantDbrxExperts) - - if DbrxExpertGLU not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxExpertGLU: "hf.DbrxExpertGLU"})(_QuantDbrxExpertGLU) - - if DbrxFFN not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxFFN: "hf.DbrxFFN"})(_QuantDbrxFFN) -except ImportError: - pass - try: from transformers.models.falcon.modeling_falcon import FalconLinear @@ -1588,18 +1481,6 @@ def forward( pass -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__ == "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) - - def register_falcon_linears_on_the_fly(model): """Register Falcon linear modules as a QUANT_MODULE. @@ -1691,7 +1572,7 @@ def _fused_experts_wrapper_class(module): * non-gated (``_QuantNonGatedFusedExperts``): a 3-D ``up_proj`` with no ``gate_proj`` and no ``gate_up_proj``. Matches NemotronH ``NemotronHExperts``. - Returns ``None`` for non-standard layouts (DBRX, GptOss, GraniteMoE, + Returns ``None`` for non-standard layouts (GptOss, GraniteMoE, Llama4TextExperts) which have their own explicit registrations. ``act_fn`` is not required: these wrappers only intercept the two ``F.linear`` @@ -2134,7 +2015,6 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: CUSTOM_MODEL_PLUGINS.update( [ register_falcon_linears_on_the_fly, - register_dbrx_moe_on_the_fly, register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 885d6f9546a..63f349dd10a 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -635,7 +635,9 @@ def enable_weight_access_and_writeback( """ if _get_enclosing_fsdp_module(module, root_model, names) is not None: context = fsdp2_weight_access_and_writeback_context(module, root_model, writeback, names) - elif is_quantized_parallel_linear(module) and hasattr(module, "_hf_tp_plan"): + elif is_quantized_parallel_linear(module) and hasattr( + module, "enable_weight_access_and_writeback" + ): # HF transformers TP sharded linear layer context = module.enable_weight_access_and_writeback() elif hasattr(module, "_hf_hook"): @@ -666,7 +668,9 @@ def requires_weight_materialization(module, root_model, names: "ModuleNames | No return False if _get_enclosing_fsdp_module(module, root_model, names) is not None: return True - if is_quantized_parallel_linear(module) and hasattr(module, "_hf_tp_plan"): + if is_quantized_parallel_linear(module) and hasattr( + module, "enable_weight_access_and_writeback" + ): return True hook = getattr(module, "_hf_hook", None) if hook is None: diff --git a/noxfile.py b/noxfile.py index 19ca05fe1d8..49ac9ed3b05 100644 --- a/noxfile.py +++ b/noxfile.py @@ -46,7 +46,7 @@ # Extra install pins applied per transformers matrix entry (installed after the base # ``.[all,dev-test]`` install to constrain that env). TRANSFORMERS_VERSIONS = { - "tf_latest": ("transformers~=5.14.0",), + "tf_latest": ("transformers~=5.18.0",), # transformers 4.57 caps ``huggingface_hub<1.0``, but ``diffusers>=0.40`` requires # ``huggingface_hub>=1.23``. Bound diffusers to a hub<1.0-compatible release so this env # stays internally consistent; otherwise diffusers' pipeline import fails and diffusers diff --git a/plugins/modelopt/skills/ptq/references/unsupported-models.md b/plugins/modelopt/skills/ptq/references/unsupported-models.md index 2d536af6aa3..67563baa8c0 100644 --- a/plugins/modelopt/skills/ptq/references/unsupported-models.md +++ b/plugins/modelopt/skills/ptq/references/unsupported-models.md @@ -142,9 +142,8 @@ class QuantCustomModule(OriginalModule): | MoE design | Strategy | Plugin example | | --- | --- | --- | -| Fused weights + `torch.bmm` | Add `TensorQuantizer` around bmm | `_QuantLlama4TextExperts` | +| Fused weights + `torch.bmm` | Add `TensorQuantizer` around bmm | `_QuantLlama4TextExperts` (walkthrough in `docs/source/guides/_customized_model_quantization.rst`) | | Fused weights + functional interception | Intercept matmul ops | `_QuantGptOssExperts` | -| Fused 2D weights (experts stacked in rows) | Two-level expansion | `_QuantDbrxExpertGLU` | | Fused weights + `forward(x, expert_id)` | Expand + reconstruct on export | `_QuantMoELinear` (Step3.5) | For the full guide, see `examples/hf_ptq/README.md`. diff --git a/pyproject.toml b/pyproject.toml index 1404766235d..61ef7bebb44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,7 +88,7 @@ hf = [ "peft>=0.17.0", "sentencepiece>=0.2.1", # Also implicitly used in test_unified_export_megatron, test_vllm_fakequant_megatron_export "tiktoken", - "transformers>=4.57,<5.15", # Should match modelopt/torch/__init__.py and noxfile.py + "transformers>=4.57,<5.19", # Should match modelopt/torch/__init__.py and noxfile.py "wonderwords", ] diff --git a/tests/gpu/torch/quantization/plugins/test_transformers_tp.py b/tests/gpu/torch/quantization/plugins/test_transformers_tp.py index 1a8a821b972..514647d8d01 100644 --- a/tests/gpu/torch/quantization/plugins/test_transformers_tp.py +++ b/tests/gpu/torch/quantization/plugins/test_transformers_tp.py @@ -15,6 +15,7 @@ from functools import partial +import pytest import torch from _test_utils.torch.transformers_models import create_tiny_llama_dir from transformers import AutoModelForCausalLM @@ -22,10 +23,30 @@ import modelopt.torch.quantization as mtq -def _test_transformers_tp(model_path, rank, size): +class _Wrapper(torch.nn.Module): + """A module holding the HF model, as trainers and PEFT do.""" + + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, input_ids): + return self.model(input_ids) + + +def _test_transformers_tp(model_path, wrapped, rank, size): model_tp = AutoModelForCausalLM.from_pretrained(model_path, tp_plan="auto") input_ids = torch.randint(0, model_tp.config.vocab_size, (10, 512), device=f"cuda:{rank}") - mtq.quantize(model_tp, mtq.NVFP4_AWQ_LITE_CFG, lambda model: model(input_ids)) + to_quantize = _Wrapper(model_tp) if wrapped else model_tp + mtq.quantize(to_quantize, mtq.NVFP4_AWQ_LITE_CFG, lambda model: model(input_ids)) + # Every TP-sharded decoder linear must be quantized as a TP-aware layer, on every transformers + # TP layout (module-level plan attribute before 5.16, DTensor weights after). + decoder_linears = [m for m in model_tp.model.layers.modules() if isinstance(m, torch.nn.Linear)] + assert decoder_linears + assert all( + getattr(m, "_is_column_parallel", False) or getattr(m, "_is_row_parallel", False) + for m in decoder_linears + ) outputs_ref = model_tp(input_ids) # Test that the model forward pass works mtq.fold_weight(model_tp) @@ -33,6 +54,7 @@ def _test_transformers_tp(model_path, rank, size): assert torch.allclose(outputs_ref.logits, outputs_test.logits, atol=1e-4) -def test_transformers_tp(need_2_gpus, dist_workers, tmp_path): +@pytest.mark.parametrize("wrapped", [False, True]) +def test_transformers_tp(need_2_gpus, dist_workers, tmp_path, wrapped): model_path = create_tiny_llama_dir(tmp_path) - dist_workers.run(partial(_test_transformers_tp, model_path)) + dist_workers.run(partial(_test_transformers_tp, model_path, wrapped)) diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py index 4c93388ca08..1ca438bf33c 100644 --- a/tests/unit/torch/export/test_export_registry.py +++ b/tests/unit/torch/export/test_export_registry.py @@ -29,7 +29,6 @@ _export_quant_embedding, _export_quant_linear, _prepare_bmm_experts, - _prepare_dbrx_experts, _prepare_fused_experts, _prepare_iterable_experts, ) @@ -212,27 +211,6 @@ def test_builtin_dispatch_covers_all_handler_shapes(): assert PrepareMoEInputsRegistry.match(moe_linear) is None assert ExportModuleRegistry.match(_named_module("QuantMoELinear")) is None - # DBRX experts container: the usual generated class name... - dbrx = _named_module("QuantDbrxExperts", base_name="DbrxExperts") - assert PrepareMoEInputsRegistry.match(dbrx) is _prepare_dbrx_experts - assert ExportModuleRegistry.match(dbrx) is None - # ...and the _DMRegistryCls collision-fallback name, where only the mixin - # class name ("_QuantDbrxExperts") remains recognizable in the MRO. - fallback = _named_module( - "transformers_modules_modeling_dbrx_QuantDbrxExperts", base_name="_QuantDbrxExperts" - ) - assert PrepareMoEInputsRegistry.match(fallback) is _prepare_dbrx_experts - - # DBRX preparation remains type-specific even if a future DBRX variant gains - # plural quantizers; the independent export registry takes the fused path. - fused_dbrx = _named_module( - "QuantDbrxExperts", - base_name="DbrxExperts", - gate_up_proj_weight_quantizers=nn.ModuleList(), - ) - assert PrepareMoEInputsRegistry.match(fused_dbrx) is _prepare_dbrx_experts - assert ExportModuleRegistry.match(fused_dbrx) is _export_fused_experts_module - # Structural fused-experts matching wins over BMM name matching independently # in both registries. fused = _named_module( diff --git a/tests/unit/torch/models/test_model_specs.py b/tests/unit/torch/models/test_model_specs.py index 74bc6a09102..3e77844f91a 100644 --- a/tests/unit/torch/models/test_model_specs.py +++ b/tests/unit/torch/models/test_model_specs.py @@ -89,13 +89,6 @@ def test_get_expert_linear_names_raises_when_unmatched(): def test_get_expert_linear_names_from_specs(): # Arctic keeps the w1/w2/w3 naming it previously got from the engine default. assert get_expert_linear_names(_UnknownMoeBlock(), "arctic") == ["w1", "w2", "w3"] - # DBRX resolves the quantized per-expert ModuleList names (previously it fell - # through to the w1/w2/w3 default, which never existed on the quantized module). - assert get_expert_linear_names(_UnknownMoeBlock(), "dbrx") == [ - "w1_linear", - "w2_linear", - "v1_linear", - ] class NemotronHMOE(nn.Module): @@ -198,41 +191,21 @@ def __init__(self): assert get_experts_list(MixtralSparseMoeBlockFused(), "mixtral") == [] -def test_get_experts_list_still_rejects_unsupported_non_iterable_layouts(): - """Non-iterable is not by itself a reason to skip grouping. +def test_get_experts_list_rejects_specs_without_grouped_export(): + """A matching spec must still allow grouped export, whatever the experts' layout. - DBRX's experts container is not iterable either, but its per-expert linears exist - under ``experts.mlp`` -- grouping is possible, just not by this function, and its - spec says so. Skipping it would silently drop AWQ/SVDQuant resmoothing instead of - reporting an unsupported layout, so the fused shortcut is scoped to specs that - claim iterable experts. + Non-iterable experts are not by themselves a reason to skip grouping: that would + silently drop AWQ/SVDQuant resmoothing instead of reporting an unsupported layout. """ - class DbrxExperts(nn.Module): - def __init__(self): + class Qwen3_5MoeSparseMoeBlock(nn.Module): # noqa: N801 + def __init__(self, experts): super().__init__() - self.mlp = nn.Module() # per-expert linears hang off here, not off self + self.experts = experts - class DbrxFFN(nn.Module): - def __init__(self): - super().__init__() - self.experts = DbrxExperts() - - assert not hasattr(DbrxFFN().experts, "__iter__") - with pytest.raises(NotImplementedError): - get_experts_list(DbrxFFN(), "dbrx") - - -def test_get_experts_list_rejects_non_iterable_layouts(): - # DBRX matches a spec but is not an iterable-experts layout; grouped export - # must keep rejecting it (legacy behavior). - class DbrxFFN(nn.Module): - def __init__(self): - super().__init__() - self.experts = nn.ModuleList() - - with pytest.raises(NotImplementedError): - get_experts_list(DbrxFFN(), "dbrx") + for experts in (nn.Module(), nn.ModuleList()): + with pytest.raises(NotImplementedError): + get_experts_list(Qwen3_5MoeSparseMoeBlock(experts), "qwen3_5_moe") with pytest.raises(NotImplementedError): get_experts_list(_UnknownMoeBlock(), "some_unknown_model") @@ -242,15 +215,10 @@ class ArcticMoE(nn.Module): pass -class DbrxFFN(nn.Module): - pass - - def test_is_moe_matches_registered_non_standard_names(): # Non-standard MoE block names (no *SparseMoeBlock suffix, no router/experts # attributes) resolve through the model spec registry. assert is_moe(ArcticMoE()) - assert is_moe(DbrxFFN()) assert not is_moe(_UnknownMoeBlock()) @@ -554,7 +522,6 @@ def test_fused_experts_shortcut_ignores_unrelated_attributes(): # (model_type, block_names, expert_linear_names, fused_expert_names, gate_up_pair) EXPECTED_MOE_LAYOUTS = [ ("arctic", ("ArcticMoE",), ("w1", "w2", "w3"), False, ("w1", "w3")), - ("dbrx", ("DbrxFFN",), ("w1_linear", "w2_linear", "v1_linear"), False, None), ( "deepseek", ("DeepseekMoE",), @@ -744,7 +711,6 @@ def test_iterable_experts_matches_pre_refactor_support(): # The real HF root class name for each registered MoE model type. root_class_names = { "arctic": "ArcticForCausalLM", - "dbrx": "DbrxForCausalLM", "deepseek": "DeepseekForCausalLM", "deepseek_v3": "DeepseekV3ForCausalLM", "deepseek_v4": "DeepseekV4ForCausalLM", diff --git a/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py b/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py index 78a89a27ff9..03c4e48fbab 100644 --- a/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py +++ b/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py @@ -35,6 +35,39 @@ class _TrainArgs(ModelOptHFArguments): epochs: int = field(default=3, metadata={"help": "Number of epochs."}) +class _NewWarmupArgs(ModelOptHFArguments): + """transformers >= 5.15: only ``warmup_steps`` (a float; below 1 it is a ratio).""" + + warmup_steps: float = field(default=0) + + +class _OldWarmupArgs(ModelOptHFArguments): + """transformers 4.x: an int ``warmup_steps`` plus ``warmup_ratio``.""" + + warmup_steps: int = field(default=0) + warmup_ratio: float = field(default=0.0) + + +@pytest.mark.parametrize( + ("arg_type", "yaml", "expected"), + [ + (_NewWarmupArgs, "warmup_ratio: 0.05\n", {"warmup_steps": 0.05}), + (_NewWarmupArgs, "warmup_steps: 0.05\n", {"warmup_steps": 0.05}), + (_OldWarmupArgs, "warmup_steps: 0.05\n", {"warmup_steps": 0, "warmup_ratio": 0.05}), + (_OldWarmupArgs, "warmup_steps: 10\n", {"warmup_steps": 10, "warmup_ratio": 0.0}), + (_OldWarmupArgs, "warmup_ratio: 0.05\n", {"warmup_steps": 0, "warmup_ratio": 0.05}), + ], +) +def test_yaml_warmup_spelling_follows_transformers(tmp_path, arg_type, yaml, expected): + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml) + + (args,) = ModelOptArgParser((arg_type,)).parse_args_into_dataclasses( + args=["--config", str(config_file)] + ) + assert {k: getattr(args, k) for k in expected} == expected + + class TestModelOptArgParser: """Tests for ModelOptArgParser --config and --generate_docs features.""" diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 0da6cfbb1fc..189d1808fbe 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -47,8 +47,6 @@ import transformers from transformers import AutoModelForCausalLM, LlamaForCausalLM from transformers.integrations.finegrained_fp8 import FP8Linear -from transformers.models.dbrx.configuration_dbrx import DbrxConfig, DbrxFFNConfig -from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN class HFModel(nn.Module): @@ -152,53 +150,6 @@ def test_fp8_linear_per_tensor_dequant(monkeypatch): ) -@pytest.mark.skipif( - Version(transformers.__version__) < Version("5.0"), - reason="test_dbrx is not supported for transformers<5.0", -) -def test_dbrx(): - assert DbrxExperts in QuantModuleRegistry - assert DbrxExpertGLU in QuantModuleRegistry - - config = DbrxConfig( - ffn_config=DbrxFFNConfig(ffn_hidden_size=8, moe_num_experts=2, hidden_size=32), - hidden_size=32, - ) - - model_ref = DbrxFFN(config) - model_test = DbrxFFN(config) - with torch.no_grad(): - model_ref.experts.mlp.w1.copy_(torch.randn(16, 32)) - model_ref.experts.mlp.v1.copy_(torch.randn(16, 32)) - model_ref.experts.mlp.w2.copy_(torch.randn(16, 32)) - - model_test.load_state_dict(model_ref.state_dict()) - - mtq.replace_quant_module(model_test) - - expertglu_ref = model_ref.experts.mlp - expertglu_test = model_test.experts.mlp - - assert hasattr(expertglu_test, "w1_linear") and not hasattr(expertglu_test, "w1") - assert hasattr(expertglu_test, "v1_linear") and not hasattr(expertglu_test, "v1") - assert hasattr(expertglu_test, "w2_linear") and not hasattr(expertglu_test, "w2") - - # Weights are stored transposed (W = w1[i].T) to match F.linear semantics with - # transformers 5.0's raw matmul: x @ w1[i] = F.linear(x, w1[i].T) - assert torch.allclose( - torch.concat([m.weight.T for m in expertglu_test.w1_linear], dim=0), - expertglu_ref.w1, - ) - - mtq.set_quantizer_attributes_partial(model_test, "*", {"enable": False}) - - # In transformers 5.0, the FFN input dimension is ffn_hidden_size (not hidden_size) - x = torch.randn(1, 4, 8) - out_1 = model_ref(x) - out_2 = model_test(x) - assert torch.allclose(out_1[0], out_2[0]) - - @pytest.mark.skipif( not hasattr(transformers, "NemotronHConfig"), reason="NemotronH is not supported by this Transformers version",