diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a23d9eede88..658d1e5c851 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -53,11 +53,11 @@ Changelog **Backward Breaking Changes** +- Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``. - The Megatron-Core DeepSeek-V4 indexer (``CSAIndexer``) is now a quantization module and persists its quantizer state in the checkpoint as ``indexer._extra_state``. A DeepSeek-V4 model quantized with an earlier release resumes from its ``torch_dist`` checkpoint only with a non-strict load (``--dist-ckpt-strictness log_unexpected`` in Megatron-LM) until it is saved again. - ``examples/hf_ptq`` no longer detects MTP layers by name. Weights the loader could not place -- an MTP head, an auxiliary tower -- are identified from Transformers' own accounting: the model is loaded with ``from_pretrained(..., output_loading_info=True)`` and the reported ``unexpected_keys`` (present in the checkpoint, not in the model's architecture) are recorded on the model and carried into the export unchanged. Everything the loader *did* place goes through the normal export path. This removes ``load_mtp_weights``, ``mtp_layer_prefixes_from_checkpoint`` and their support matrix of MTP storage conventions, along with ``_add_mtp_exclusions`` and the pre-quantization ``enable: False`` entries ``hf_ptq`` appended to the recipe's ``quant_cfg``. Two consequences: MTP layers now follow the recipe like any other module instead of being force-excluded by the script -- matching ``examples/megatron_bridge``, which has no MTP-specific code at all -- and ``quantization_config.ignore`` can no longer claim a layer is unquantized that the export in fact quantized. Recipes importing ``configs/ptq/units/default_disabled_quantizers`` still disable ``mtp.*``, so their behaviour is unchanged; a recipe omitting that unit will now quantize an MTP the model actually built. - - ``examples/hf_ptq --vllm_fakequant_export`` now raises ``NotImplementedError`` when the checkpoint holds weights the model has no parameter for and a shard actually provides them (an MTP head, an auxiliary tower). The fake-quant exporter writes only model-backed state, so it would otherwise drop those weights silently -- and a fake-quant checkpoint is evaluated, where a missing head changes the score rather than failing loudly. Use the unified HF export, which carries them through. Buffers Transformers recomputes are not weights to lose: ``*.inv_freq`` is skipped even when a shard provides it, since older Llama/Mistral-lineage conversions do list it in the index and refusing an export over it would reject checkpoints that export correctly today. The check runs immediately after the model loads, not at export time, so an incompatible run fails before calibration rather than after it. - The ``modelopt.onnx.quantization.graph_utils`` module has been removed with no compatibility shim; update direct imports using this migration map: @@ -81,7 +81,6 @@ Changelog ``get_layer_precision_mapping``, ``get_resize_scales``, ``print_stat``, ``remove_partial_input_qdq``, ``should_quantize_to_8bit``, and ``validate_8bit_layers``. - - Layerwise calibration now uses prior-layer QDQ activations by default (``layerwise.get_qdq_activations_from_prev_layer=True``). Set it to ``False`` to preserve full-precision activations for subsequent layers (the default behavior for diff --git a/examples/alpamayo/qad.py b/examples/alpamayo/qad.py index fac41af6643..c6bd930969e 100644 --- a/examples/alpamayo/qad.py +++ b/examples/alpamayo/qad.py @@ -337,12 +337,7 @@ def main(): "per_device_eval_batch_size": 1, "gradient_accumulation_steps": args.grad_accum, "learning_rate": args.lr, - # 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, + "warmup_steps": args.warmup_ratio, # a warmup_steps below 1 is a ratio "max_steps": args.max_steps, "logging_steps": args.logging_steps, "save_steps": args.save_steps, diff --git a/examples/gpt-oss/configs/sft_full.yaml b/examples/gpt-oss/configs/sft_full.yaml index 34732956a0b..780b493e555 100644 --- a/examples/gpt-oss/configs/sft_full.yaml +++ b/examples/gpt-oss/configs/sft_full.yaml @@ -16,7 +16,7 @@ per_device_train_batch_size: 2 per_device_eval_batch_size: 2 gradient_accumulation_steps: 2 max_length: 4096 -warmup_steps: 0.03 # use warmup_ratio if using transformers<5.0 +warmup_steps: 0.03 lr_scheduler_type: cosine_with_min_lr lr_scheduler_kwargs: min_lr_rate: 0.1 diff --git a/examples/gpt-oss/configs/sft_lora.yaml b/examples/gpt-oss/configs/sft_lora.yaml index 9d298ba9a15..663c19743ef 100644 --- a/examples/gpt-oss/configs/sft_lora.yaml +++ b/examples/gpt-oss/configs/sft_lora.yaml @@ -21,7 +21,7 @@ lora_alpha: 16 lora_dropout: 0.0 lora_target_modules: all-linear max_length: 4096 -warmup_steps: 0.03 # use warmup_ratio if using transformers<5.0 +warmup_steps: 0.03 lr_scheduler_type: cosine_with_min_lr lr_scheduler_kwargs: min_lr_rate: 0.1 diff --git a/examples/gpt-oss/qat-finetune-transformers.ipynb b/examples/gpt-oss/qat-finetune-transformers.ipynb index 42226b29821..58dba84cb6a 100644 --- a/examples/gpt-oss/qat-finetune-transformers.ipynb +++ b/examples/gpt-oss/qat-finetune-transformers.ipynb @@ -207,7 +207,7 @@ " per_device_eval_batch_size=1,\n", " gradient_accumulation_steps=2,\n", " max_length=4096,\n", - " warmup_steps=0.03, # use warmup_ratio if using transformers<5.0\n", + " warmup_steps=0.03,\n", " eval_strategy=\"steps\",\n", " eval_on_start=True,\n", " logging_steps=10,\n", diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 59b31adbf1d..e10a99805e3 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -135,7 +135,7 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http > *6.Some models currently support export to HF format only.* \ > *7.[PTQ for DeepSeek](../deepseek/README.md)* \ > *8.GLM-4.7 has MTP (Multi-Token Prediction) layers that are automatically loaded and excluded from quantization.* \ -> *9.Running Whisper model with transformers>=5.0 requires [torchcodec](https://github.com/meta-pytorch/torchcodec?tab=readme-ov-file#installing-cuda-enabled-torchcodec) and other system packages (e.g. ffmpeg).* \ +> *9.Running Whisper model requires [torchcodec](https://github.com/meta-pytorch/torchcodec?tab=readme-ov-file#installing-cuda-enabled-torchcodec) and other system packages (e.g. ffmpeg).* \ > *10.GPT-OSS ships with native MXFP4 weights; NVFP4 export is produced via the closed-form `--cast_mxfp4_to_nvfp4` cast (see [MXFP4 → NVFP4 cast](#mxfp4--nvfp4-cast-for-gpt-oss)).* \ > *11.Vision-language model (VLM): by default, only the language model is quantized while the vision encoder is kept in high precision. Pass `--vlm` to the shell script (see [VLM quantization](#vlm-quantization)).* \ > *12.For VLMs, `int8_smoothquant` only supports TensorRT-LLM checkpoint export and is not compatible with the TensorRT-LLM torch backend.* \ diff --git a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb index f293eda7d43..6e65d2dcc9b 100644 --- a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb +++ b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb @@ -274,7 +274,7 @@ " per_device_eval_batch_size=1,\n", " gradient_accumulation_steps=2,\n", " max_length=4096,\n", - " warmup_steps=0.03, # use warmup_ratio if using transformers<5.0\n", + " warmup_steps=0.03,\n", " eval_strategy=\"steps\",\n", " eval_on_start=True,\n", " logging_steps=50,\n", diff --git a/examples/puzzletron/requirements.txt b/examples/puzzletron/requirements.txt index 6f4c94e08a8..9d6118eb938 100644 --- a/examples/puzzletron/requirements.txt +++ b/examples/puzzletron/requirements.txt @@ -1,4 +1,2 @@ math-verify ray -# Likely works for transformers v5 also, but we need to test it -transformers<5.0 diff --git a/modelopt/torch/__init__.py b/modelopt/torch/__init__.py index b540457724d..4898e20f8ea 100644 --- a/modelopt/torch/__init__.py +++ b/modelopt/torch/__init__.py @@ -16,7 +16,6 @@ """Model optimization and deployment subpackage for torch.""" import importlib -import importlib.util import warnings as _warnings from packaging.version import Version as _Version @@ -49,22 +48,13 @@ try: from transformers import __version__ as _transformers_version - if _Version(_transformers_version) < _Version("4.57") or _Version( + if _Version(_transformers_version) < _Version("5.5") or _Version( _transformers_version ) >= _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.", ) - - # CVE-2026-4372: transformers<5.3 + the optional `kernels` package allows remote code execution - # when loading untrusted models via the Hub kernel-download path. Fixed in 5.3.0. Only warn if - # `kernels` is actually installed, since ModelOpt's core install never pulls it in. - if _Version(_transformers_version) < _Version("5.3") and importlib.util.find_spec("kernels"): - _warnings.warn( - f"transformers {_transformers_version} with the `kernels` package is affected by" - ' CVE-2026-4372; consider `pip install -U "transformers>=5.3"` when loading untrusted models.', - ) except ImportError: pass diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index bb428d8e180..8bd9e6cc3f8 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -98,7 +98,7 @@ def _tied_quantized_modules(model: nn.Module) -> list[str]: Grouped by name, which survives offload: a ``data_ptr`` grouping sees nothing when the weights are on meta and would pass vacuously. Falls back to ``data_ptr`` when the model - publishes no map (transformers < 5). + publishes no map (e.g. not a ``PreTrainedModel``). """ tied_map = TiedWeightMap(model) groups: dict[str, list[str]] = {} diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index dc3acfbca89..d6363584cbc 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -245,11 +245,11 @@ class TiedWeightMap: """ def __init__(self, model: nn.Module) -> None: - """Source the tie map from HF's ``all_tied_weights_keys`` (transformers >=5.0). + """Source the tie map from HF's ``all_tied_weights_keys``. HF's ``{target: source}`` == our ``{alias: canonical}``, resolved at load, config-gated, - ``torch.equal``-pruned, and name-based so it survives FSDP shard / offload. Absent on - transformers <5.0 -> empty map (the ``data_ptr`` backstop in postprocess is the net). + ``torch.equal``-pruned, and name-based so it survives FSDP shard / offload. Absent on a + non-``PreTrainedModel`` -> empty map (the ``data_ptr`` backstop in postprocess is the net). """ all_tied = getattr(model, "all_tied_weights_keys", None) # Warn whenever a tie is declared (embedding tie or any ``_tied_weights_keys`` entry, e.g. @@ -259,12 +259,8 @@ def __init__(self, model: nn.Module) -> None: ) or bool(getattr(model, "_tied_weights_keys", None)) if all_tied is None and declares_tie: warnings.warn( - "This model may contain tied/shared weights, but deduplicating them on export " - "requires transformers>=5.0 (it uses model.all_tied_weights_keys, which is only " - "supported in newer versions). On older versions the exported checkpoint may keep " - "duplicate copies of the tied weights (larger files), and tied weights may not be " - "deduplicated correctly during export. Upgrade to transformers>=5.0 for correct " - "tied-weight export." + "This model declares tied weights but has no model.all_tied_weights_keys, so the " + "exported checkpoint may keep duplicate copies of the tied weights." ) # Drop any self-entry (alias == canonical): HF should not emit one, but a target==source # pair would schedule the kept canonical for deletion, so filter it out defensively. diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index ef2233116a1..7f692194fa6 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -185,17 +185,9 @@ def name_shards_and_write_index( def _parse_shard_size(size: int | str) -> int: """Convert a shard-size string (e.g. ``"10GB"``, ``"500MB"``) to bytes. - Mirrors transformers' ``convert_file_size_to_int``, which reads ``GB``/``MB``/``KB`` - as decimal and only ``GiB``/``MiB``/``KiB`` as binary. That helper was removed from - ``transformers.utils`` in transformers 5.x, so the fallback below is the live path - there, not a rarely-taken branch. + Mirrors transformers 4.x's ``convert_file_size_to_int`` (removed in 5.x), which reads + ``GB``/``MB``/``KB`` as decimal and only ``GiB``/``MiB``/``KiB`` as binary. """ - try: - from transformers.utils import convert_file_size_to_int - - return convert_file_size_to_int(size) - except ImportError: - pass if isinstance(size, int): return size s = size.strip().upper() @@ -257,9 +249,9 @@ def _undeclared_tied_aliases(model: nn.Module) -> set[str]: goes, so the shared address is gone by then. Take the same information up front, off the live model, and drop by name instead. - Needed on transformers <5.0, where ``all_tied_weights_keys`` does not exist and the tie map - is empty; without this the alias ships as a second full copy of the same weight. Walk order - matches ``state_dict()``, so the tensor kept here is the one the resident path keeps. + Covers ties missing from ``all_tied_weights_keys`` (or a model without it); without this the + alias ships as a second full copy of the same weight. Walk order matches ``state_dict()``, so + the tensor kept here is the one the resident path keeps. """ seen: dict[int, str] = {} aliases: set[str] = set() diff --git a/modelopt/torch/models/README.md b/modelopt/torch/models/README.md index f1ffcf6a54a..4411218e0a9 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -21,7 +21,7 @@ __all__: list[str] = [] register( ModelSpec( model_type="qwen3_moe", - min_transformers_version="4.57", + min_transformers_version="5.5", # modelopt policy, not a model fact: this model is validated for grouped export. export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( @@ -71,7 +71,7 @@ come from, and they must agree: | `min_transformers_version` | Earliest `transformers` release whose definitions match this spec; `None` for a `remote_code` model | Clamp `min_transformers_version` at the repo's minimum supported transformers (`tf_min` -in `noxfile.py`, currently `4.57`) — a model older than the floor records the floor, +in `noxfile.py`, currently `5.5`) — a model older than the floor records the floor, since nothing older is installed or tested. A model added later records its own release. The pair is what lets `test_specs_vs_transformers.py` assert rather than skip: it can tell a model that is legitimately absent on an older transformers from a spec that no diff --git a/modelopt/torch/models/deepseek_v3/specs.py b/modelopt/torch/models/deepseek_v3/specs.py index f4d50f9388a..3a774ecb931 100644 --- a/modelopt/torch/models/deepseek_v3/specs.py +++ b/modelopt/torch/models/deepseek_v3/specs.py @@ -28,15 +28,14 @@ # structural router+experts test in is_moe matches. block_names below is what puts it on # the MoE path at all. # -# The experts container changed upstream. transformers 4.57 builds an nn.ModuleList of -# DeepseekV3MLP, which is the naming recorded here; 5.x replaces it with a fused, 3-D -# DeepseekV3Experts holding gate_up_proj/down_proj. The fused layout resolves through -# the structural first-projection check in get_expert_linear_names and never consults -# this naming -- the same split Mixtral has across the same releases. +# The naming recorded here is for an nn.ModuleList of DeepseekV3MLP experts (remote code and +# transformers 4.x). transformers 5.x uses a fused, 3-D DeepseekV3Experts holding +# gate_up_proj/down_proj, which resolves through the structural first-projection check in +# get_expert_linear_names and never consults this naming -- the same split as Mixtral. register( ModelSpec( model_type="deepseek_v3", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( block_names=("DeepseekV3MoE",), diff --git a/modelopt/torch/models/gemma/specs.py b/modelopt/torch/models/gemma/specs.py index ec85399fe60..04a44e67779 100644 --- a/modelopt/torch/models/gemma/specs.py +++ b/modelopt/torch/models/gemma/specs.py @@ -24,7 +24,7 @@ register( ModelSpec( model_type="gemma", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(weight_plus_one_norm_names=("GemmaRMSNorm",)), ) ) diff --git a/modelopt/torch/models/gemma2/specs.py b/modelopt/torch/models/gemma2/specs.py index 7a8f857f010..633dd441541 100644 --- a/modelopt/torch/models/gemma2/specs.py +++ b/modelopt/torch/models/gemma2/specs.py @@ -24,7 +24,7 @@ register( ModelSpec( model_type="gemma2", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(weight_plus_one_norm_names=("Gemma2RMSNorm",)), ) ) diff --git a/modelopt/torch/models/gemma3/specs.py b/modelopt/torch/models/gemma3/specs.py index ecb840c443f..11eba7f19f9 100644 --- a/modelopt/torch/models/gemma3/specs.py +++ b/modelopt/torch/models/gemma3/specs.py @@ -24,7 +24,7 @@ register( ModelSpec( model_type="gemma3", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(weight_plus_one_norm_names=("Gemma3RMSNorm",)), ) ) diff --git a/modelopt/torch/models/gpt_oss/specs.py b/modelopt/torch/models/gpt_oss/specs.py index 34583b4a4df..be8419781c7 100644 --- a/modelopt/torch/models/gpt_oss/specs.py +++ b/modelopt/torch/models/gpt_oss/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="gpt_oss", - min_transformers_version="4.57", + min_transformers_version="5.5", moe_spec=MoESpec( # GPT-OSS fuses gate and up into a single gate_up_proj. # transformers names the block GptOssMLP; GptOssMoE is kept for the diff --git a/modelopt/torch/models/llama/specs.py b/modelopt/torch/models/llama/specs.py index 0d62d44feda..8b517ff1740 100644 --- a/modelopt/torch/models/llama/specs.py +++ b/modelopt/torch/models/llama/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="llama", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec( # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/torch/models/minimax/specs.py b/modelopt/torch/models/minimax/specs.py index 59c944f44e5..29ae0706653 100644 --- a/modelopt/torch/models/minimax/specs.py +++ b/modelopt/torch/models/minimax/specs.py @@ -30,7 +30,7 @@ register( ModelSpec( model_type="minimax", - min_transformers_version="4.57", + min_transformers_version="5.5", moe_spec=MoESpec( block_names=("MiniMaxSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), diff --git a/modelopt/torch/models/mixtral/specs.py b/modelopt/torch/models/mixtral/specs.py index f15909ba6e5..a0c3d11ee52 100644 --- a/modelopt/torch/models/mixtral/specs.py +++ b/modelopt/torch/models/mixtral/specs.py @@ -24,7 +24,7 @@ register( ModelSpec( model_type="mixtral", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( block_names=("MixtralSparseMoeBlock",), diff --git a/modelopt/torch/models/nemotron/specs.py b/modelopt/torch/models/nemotron/specs.py index 3b797601733..0bc32a0a210 100644 --- a/modelopt/torch/models/nemotron/specs.py +++ b/modelopt/torch/models/nemotron/specs.py @@ -24,7 +24,7 @@ register( ModelSpec( model_type="nemotron", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec( weight_plus_one_norm_names=("LayerNorm1P", "NemotronLayerNorm1P"), ), diff --git a/modelopt/torch/models/nemotron_h/specs.py b/modelopt/torch/models/nemotron_h/specs.py index fe651930153..7198c2b31b3 100644 --- a/modelopt/torch/models/nemotron_h/specs.py +++ b/modelopt/torch/models/nemotron_h/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="nemotron_h", - min_transformers_version="5.3", + min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). diff --git a/modelopt/torch/models/phimoe/specs.py b/modelopt/torch/models/phimoe/specs.py index a8937afcb6a..10ce40174d4 100644 --- a/modelopt/torch/models/phimoe/specs.py +++ b/modelopt/torch/models/phimoe/specs.py @@ -30,7 +30,7 @@ register( ModelSpec( model_type="phimoe", - min_transformers_version="4.57", + min_transformers_version="5.5", moe_spec=MoESpec( block_names=("PhimoeSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), diff --git a/modelopt/torch/models/qwen2_moe/specs.py b/modelopt/torch/models/qwen2_moe/specs.py index f491fb5e6dd..832c8ead8de 100644 --- a/modelopt/torch/models/qwen2_moe/specs.py +++ b/modelopt/torch/models/qwen2_moe/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="qwen2_moe", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( block_names=("Qwen2MoeSparseMoeBlock",), diff --git a/modelopt/torch/models/qwen3/specs.py b/modelopt/torch/models/qwen3/specs.py index f6bca9577e1..c803373a464 100644 --- a/modelopt/torch/models/qwen3/specs.py +++ b/modelopt/torch/models/qwen3/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="qwen3", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec( # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/torch/models/qwen3_5_moe/specs.py b/modelopt/torch/models/qwen3_5_moe/specs.py index 0982d43daa0..7146bcd91a8 100644 --- a/modelopt/torch/models/qwen3_5_moe/specs.py +++ b/modelopt/torch/models/qwen3_5_moe/specs.py @@ -28,7 +28,7 @@ register( ModelSpec( model_type="qwen3_5_moe", - min_transformers_version="5.2", + min_transformers_version="5.5", moe_spec=MoESpec( block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/torch/models/qwen3_moe/specs.py b/modelopt/torch/models/qwen3_moe/specs.py index 7bface2731a..c8b9bcd3e5b 100644 --- a/modelopt/torch/models/qwen3_moe/specs.py +++ b/modelopt/torch/models/qwen3_moe/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="qwen3_moe", - min_transformers_version="4.57", + min_transformers_version="5.5", moe_spec=MoESpec( block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/torch/models/qwen3_next/specs.py b/modelopt/torch/models/qwen3_next/specs.py index 222a1884bc8..31282a5d391 100644 --- a/modelopt/torch/models/qwen3_next/specs.py +++ b/modelopt/torch/models/qwen3_next/specs.py @@ -22,7 +22,7 @@ register( ModelSpec( model_type="qwen3_next", - min_transformers_version="4.57", + min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), moe_spec=MoESpec( block_names=("Qwen3NextSparseMoeBlock",), diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py index db077487c00..19b2e357631 100644 --- a/modelopt/torch/opt/plugins/huggingface.py +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -19,7 +19,7 @@ import os import threading import warnings -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from typing import Any import torch @@ -72,7 +72,7 @@ def _get_modelopt_state_path(model_name_or_path: str) -> str: @contextmanager -def _patch_model_init_for_modelopt(cls, model_path, extra_context=None): +def _patch_model_init_for_modelopt(cls, model_path): """Patch for `cls.init` method to restore ModelOpt state after `init`.""" # Note: Keeping original config in local as the package will be shared among threads added_original_init = False @@ -90,8 +90,7 @@ def new_init_fn(self, *args, **kwargs): modelopt_state_path = _get_modelopt_state_path(model_path) _original__init__(self, *args, **kwargs) if os.path.isfile(modelopt_state_path): - with extra_context() if extra_context else nullcontext(): - restore_from_modelopt_state(self, modelopt_state_path=modelopt_state_path) + restore_from_modelopt_state(self, modelopt_state_path=modelopt_state_path) print_rank_0(f"Restored ModelOpt state from {modelopt_state_path}") diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 250505dba91..4f54a59248c 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -30,7 +30,6 @@ import torch import torch.nn as nn import transformers -from packaging.version import Version from transformers import HfArgumentParser, PreTrainedModel, Trainer, TrainerCallback from transformers import modeling_utils as tf_modeling_utils @@ -55,10 +54,6 @@ "ModelOptTrainerArguments", ] -# transformers 5.0 changed `_tied_weights_keys` from a list to a {target: source} dict and added -# the `load_config` parameter to `_load_state_dict_into_zero3_model`. -_TRANSFORMERS_GE_5_0 = Version(transformers.__version__) >= Version("5.0") - def is_liger_available(): try: @@ -68,31 +63,6 @@ def is_liger_available(): return True -@contextmanager -def _undo_torch_init_override_by_transformers(): - if not hasattr(tf_modeling_utils, "TORCH_INIT_FUNCTIONS"): - yield - return - # transformers override weight initialization during model instantiation for faster loading; - # this leads to a secondary bug causing fx symbolic tracing to fail (torch does not allow - # overriding torch.nn.init functions - fx tracing asserts that this does not happen and fails) - # correct fx symbolic tracing is needed for NAS/Pruned model restoration - # lets restore the original init functions before modelopt restore so that tracing works during nas restore - # weight initialization is anyways done, so this wont affect performance - modelopt_reverted_torch_init_funcs = {} - for name, init_func in tf_modeling_utils.TORCH_INIT_FUNCTIONS.items(): - torch_init_func = getattr(torch.nn.init, name) - # Check if the init function has been overridden by transformers - if id(torch_init_func) != id(init_func): - modelopt_reverted_torch_init_funcs[name] = torch_init_func - setattr(torch.nn.init, name, init_func) - - yield - - for name, init_func in modelopt_reverted_torch_init_funcs.items(): - setattr(torch.nn.init, name, init_func) - - def _restore_qtensor_wrappers(model, model_path): """Re-wrap QTensorWrapper weights that were replaced during HF weight loading. @@ -146,9 +116,7 @@ def _restore_qtensor_wrappers(model, model_path): def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): """Patch for `cls.from_pretrained` method to restore ModelOpt state.""" - with _patch_model_init_for_modelopt( - cls, pretrained_model_name_or_path, extra_context=_undo_torch_init_override_by_transformers - ): + with _patch_model_init_for_modelopt(cls, pretrained_model_name_or_path): model = types.MethodType(cls._modelopt_cache["from_pretrained"].__func__, cls)( pretrained_model_name_or_path, *args, **kwargs ) @@ -160,9 +128,7 @@ def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs) def _new_from_config(cls, /, config, **kwargs): """Patch for `cls.from_config` method to restore ModelOpt state.""" - with _patch_model_init_for_modelopt( - cls, config._name_or_path, extra_context=_undo_torch_init_override_by_transformers - ): + with _patch_model_init_for_modelopt(cls, config._name_or_path): model = types.MethodType(cls._modelopt_cache["_from_config"].__func__, cls)( config, **kwargs ) @@ -184,10 +150,6 @@ def _legacy_tied_weights_keys_as_dict(model: nn.Module): the dedup patterns ``_get_tied_weight_keys`` is expected to return. The original attribute is restored on exit so the shim stays invisible to the rest of the model's lifetime. """ - if not _TRANSFORMERS_GE_5_0: - yield - return - patched = [] try: for module in model.modules(): @@ -222,9 +184,7 @@ def _load_params_and_buffers_into_zero3_model(model_to_load, state_dict, load_co buffer_state_dict = {k: v for k, v in state_dict.items() if k in buffer_names} model_to_load.load_state_dict(buffer_state_dict, strict=False) cached_fn = tf_modeling_utils._modelopt_cache["_load_state_dict_into_zero3_model"] - if _TRANSFORMERS_GE_5_0: - return cached_fn(model_to_load, state_dict, load_config) - return cached_fn(model_to_load, state_dict) + return cached_fn(model_to_load, state_dict, load_config) pretrained_model_patch_methods = [ @@ -312,20 +272,10 @@ 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") + """Map ``warmup_ratio``, removed in transformers 5.15, to ``warmup_steps`` (below 1 it is a ratio).""" + dests = {a.dest for a in actions} + if "warmup_ratio" in config and "warmup_ratio" not in dests and "warmup_steps" in dests: + config.setdefault("warmup_steps", config.pop("warmup_ratio")) class ModelOptArgParser(HfArgumentParser): diff --git a/modelopt/torch/prune/importance_hooks/expert_removal_hooks.py b/modelopt/torch/prune/importance_hooks/expert_removal_hooks.py index 2d0a9ad4c52..1b1dc989e16 100644 --- a/modelopt/torch/prune/importance_hooks/expert_removal_hooks.py +++ b/modelopt/torch/prune/importance_hooks/expert_removal_hooks.py @@ -20,8 +20,6 @@ from typing import TYPE_CHECKING import torch -import transformers -from packaging.version import Version from torch import nn from .base_hooks import ForwardHook @@ -363,10 +361,7 @@ def get_router_logits_and_routed_experts( Based on Qwen3VLMoeSparseMoe forward pass. """ orig_shape = hidden_states.shape - # Use hidden_states.shape[-1] instead of self.moe.hidden_size for transformers v5 compatibility - hidden_size = ( - self.moe.hidden_size if hasattr(self.moe, "hidden_size") else hidden_states.shape[-1] - ) + hidden_size = hidden_states.shape[-1] # Flatten to (num_tokens, hidden_size) for processing hidden_states_flat = hidden_states.reshape(-1, hidden_size) @@ -384,19 +379,8 @@ def get_router_logits_and_routed_experts( routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) routing_weights = routing_weights.to(hidden_states_flat.dtype) - if Version(transformers.__version__) >= Version("5.0"): - # transformers 5.x: grouped_mm_experts_forward expects - # (hidden_states_flat 2D, top_k_index, top_k_weights) - routed_out = self.moe.experts(hidden_states_flat, router_indices, routing_weights) - else: - # transformers 4.x: loop-based experts expects - # (hidden_states_3d 3D, routing_weights_full, router_indices) - batch_size = orig_shape[0] if hidden_states.ndim == 3 else 1 - hidden_states_3d = hidden_states_flat.reshape(batch_size, -1, hidden_size) - router_weights = torch.zeros( - router_logits.shape, dtype=routing_weights.dtype, device=router_logits.device - ).scatter_(1, router_indices, routing_weights) - routed_out = self.moe.experts(hidden_states_3d, router_weights, router_indices) + # Fused experts take (hidden_states_flat 2D, top_k_index, top_k_weights) + routed_out = self.moe.experts(hidden_states_flat, router_indices, routing_weights) # Return in same shape as input routed_out = routed_out.reshape(*orig_shape) diff --git a/modelopt/torch/puzzletron/anymodel/models/__init__.py b/modelopt/torch/puzzletron/anymodel/models/__init__.py index 850115720f6..af4966aa27a 100644 --- a/modelopt/torch/puzzletron/anymodel/models/__init__.py +++ b/modelopt/torch/puzzletron/anymodel/models/__init__.py @@ -14,7 +14,6 @@ # limitations under the License. from packaging.version import Version as _Version -from transformers import __version__ as _transformers_version # Import models to trigger factory registration from .gpt_oss import * diff --git a/modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py b/modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py index aa967d40a87..f2e5683a605 100644 --- a/modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py +++ b/modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py @@ -34,7 +34,6 @@ # Production models use MXFP4 quantized MoE with combined tensors # (gate_up_proj_blocks, down_proj_blocks), which is not yet supported. from ....pruning.pruning_mixin import PruningMixIn -from ....utils.dummy_modules import DummyBlock from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size @@ -51,14 +50,6 @@ class GptOssModelDescriptor(ModelDescriptor): _DECODER_LAYER_CLS: Type[nn.Module] = None - @classmethod - def create_dummy_block(cls, original_layer: GptOssDecoderLayer, block_index: int) -> nn.Module: - dummy_block = DummyBlock(block_index=block_index) - # Required by `GptOssModel.forward` in transformers<5.4 - if hasattr(original_layer, "attention_type"): - dummy_block.attention_type = original_layer.attention_type - return dummy_block - @staticmethod def decoder_layer_cls(): """Get the decoder layer class for GPT-OSS models. diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py b/modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py index 46588abd19e..77c62fe0b21 100644 --- a/modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py +++ b/modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py @@ -24,7 +24,6 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer, Qwen2RotaryEmbedding from ....block_config import BlockConfig -from ....utils.dummy_modules import DummyBlock from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size from ..llama.llama_model_descriptor import LlamaFFNIntermediateLayerDescriptor @@ -40,18 +39,6 @@ class Qwen2ModelDescriptor(ModelDescriptor): def decoder_layer_cls(): return Qwen2DecoderLayer - @classmethod - def create_dummy_block(cls, original_layer: nn.Module, block_index: int) -> nn.Module: - """Create a dummy block that preserves Qwen2-specific attributes like attention_type. - - Qwen2's forward pass accesses decoder_layer.attention_type for attention mask selection. - """ - dummy = DummyBlock(block_index=block_index) - # Copy attention_type from original layer (required by Qwen2's forward pass) - if hasattr(original_layer, "attention_type"): - dummy.attention_type = original_layer.attention_type - return dummy - @staticmethod def block_config_to_layer_overrides(block_config: BlockConfig): return { diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py b/modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py index 30ffd28d0c0..792e384edc6 100644 --- a/modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py @@ -19,7 +19,6 @@ from dataclasses import dataclass, field from typing import Dict, List -from torch import nn from transformers.models.qwen3.modeling_qwen3 import ( Qwen3DecoderLayer, Qwen3ForCausalLM, @@ -29,7 +28,6 @@ from ....block_config import BlockConfig from ....pruning.ffn_intermediate_pruning_mixin import FFNIntermediateLayerDescriptor from ....pruning.kv_heads_pruning_mixin import KVHeadsLayerDescriptor -from ....utils.dummy_modules import DummyBlock from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size @@ -46,18 +44,6 @@ class Qwen3ModelDescriptor(ModelDescriptor): def decoder_layer_cls(): return Qwen3DecoderLayer - @classmethod - def create_dummy_block(cls, original_layer: nn.Module, block_index: int) -> nn.Module: - """Create a dummy block that preserves Qwen3-specific attributes like attention_type. - - Qwen3's forward pass accesses decoder_layer.attention_type for attention mask selection. - """ - dummy = DummyBlock(block_index=block_index) - # Copy attention_type from original layer (required by Qwen3's forward pass) - if hasattr(original_layer, "attention_type"): - dummy.attention_type = original_layer.attention_type - return dummy - @staticmethod def block_config_to_layer_overrides(block_config: BlockConfig): return { diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 1524172cfec..7d21944fe12 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -26,9 +26,12 @@ import torch import torch.nn as nn import transformers -from packaging import version from torch import Tensor from torch.nn.functional import linear +from transformers.integrations.finegrained_fp8 import FP8Linear +from transformers.models.falcon.modeling_falcon import FalconLinear +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts +from transformers.models.llama4.modeling_llama4 import Llama4TextExperts from transformers.models.t5.modeling_t5 import T5Attention from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_FA_AVAILABLE @@ -74,11 +77,9 @@ __all__ = ["register_hf_attentions_on_the_fly"] -TRANSFORMERS_VERSION_GE_5_0 = version.parse(transformers.__version__) >= version.parse("5.0.0") - class _QuantAttention(QuantModule): - """Attention class for KV Cache quantization compatible with new_attention_interface in transformers >= 4.48.0.""" + """Attention class for KV Cache quantization of modules using transformers' attention interface.""" def _setup(self): self.q_bmm_quantizer = TensorQuantizer() @@ -297,12 +298,12 @@ def _quantized_attention( return attn_out.contiguous(), None def forward(self, *args, **kwargs): - """Forward method for KV cache quantization compatible with new_attention_interface in transformers >= 4.48.0. + """Forward method for KV cache quantization of modules using transformers' attention interface. The forward method is used to patch the attention interface with _quantized_attention. Once output tensors are generated, it restores the original attention interface. """ - # In transformers>=5.0 some attention classes (e.g. BertAttention) no longer store + # Some attention classes (e.g. BertAttention) do not store # `self.config` directly; fall back to searching child modules for a config attribute. _config = getattr(self, "config", None) if _config is None: @@ -352,9 +353,8 @@ def _is_eager_attention(): @staticmethod def is_compatible_attention(attn): - # The new_attention_interface is only available in transformers >= 4.48.0 - # In addition, the new attention interface is not available for some models such as T5 - # Hence lets do a crude check here to see if the attention module is using the new_attention_interface + # Some models (e.g. T5 before transformers 5.15) do not use the attention interface, + # hence lets do a crude check here to see if the attention module is using it # This is not foolproof but should work for most cases module = inspect.getmodule(attn) return getattr(module, "ALL_ATTENTION_FUNCTIONS", None) is not None @@ -469,7 +469,7 @@ def register_hf_attentions_on_the_fly(model): if registered_attn_module or not attention_cls: return - # this is the case for models that do not use the new_attention_interface or transformers version < 4.48.0 + # this is the case for models that do not use the attention interface # Register the attention class for KV Cache quantization success = any(register_attention_for_kv_quant(cls) for cls in attention_cls) if not success: @@ -486,16 +486,16 @@ class HFParallelLinear(torch.nn.Linear, DynamicModule): shard = None def _setup(self): - # transformers<5.0 and >=5.16 keep the sharded weight as a DTensor + # transformers>=5.16 keeps 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 - # transformers>=5.16: keep HF's TP wrapper on its local-tensor path (also in training), + # 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 + else: # transformers<5.16: 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) @@ -545,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, >=5.16 + if isinstance(self.weight, torch.distributed.tensor.DTensor): # transformers>=5.16 assert self.weight.placements == self.shard, ( f"Received unexpected shard {self.weight.placements} for {self}" ) @@ -556,7 +556,7 @@ def enable_weight_access_and_writeback(self): yield finally: self.weight = weight - else: # transformers 5.0-5.15: weights are already plain Parameters + else: # transformers<5.16: weights are already plain Parameters yield @@ -664,8 +664,8 @@ def iter_weights_for_calibration(self): class _QuantSparseSequentialMoe(QuantModule): """Quantization wrapper for HuggingFace sparse MoE blocks. - This base class is for Sequential MoEs (i.e each experts are implemented as standalone modules). - Transformers>=5.0 has batched experts, no per-expert quantizers. + This base class is for Sequential MoEs (i.e each experts are implemented as standalone modules), + e.g. remote-code models; native transformers MoEs use fused experts instead. Supports ``layer_sync_moe_local_experts_amax`` to sync input quantizer amax across experts. @@ -680,19 +680,18 @@ def _setup(self): self._moe_calib_experts_ratio = None self._token_counting_initialized = False + def _num_experts(self) -> int: + """Expert count from the gate, the block or the experts container (0 if unknown).""" + for obj in [getattr(self, "gate", None), self, getattr(self, "experts", None)]: + for attr in ("num_experts", "n_routed_experts"): + if getattr(obj, attr, None): + return getattr(obj, attr) + return len(self.experts) if hasattr(self.experts, "__len__") else 0 + def _init_token_counting(self): """Lazy-init token counting infra (buffer + gate hook). Called once from forward.""" self._token_counting_initialized = True - num_experts = 0 - for obj in [getattr(self, "gate", None), self, getattr(self, "experts", None)]: - if obj is not None: - for attr in ("num_experts", "n_routed_experts"): - if hasattr(obj, attr): - num_experts = getattr(obj, attr) - break - if num_experts: - break - + num_experts = self._num_experts() if num_experts == 0: warnings.warn( f"{self.__class__.__name__}: could not resolve num_experts; " @@ -714,10 +713,10 @@ def _gate_forward_hook(self, module, input, output): return with torch.no_grad(): if isinstance(output, tuple) and len(output) >= 3: - # v5.x TopKRouter: returns (logits, scores, indices) + # TopKRouter: returns (logits, scores, indices) indices = output[2] else: - # v4.x nn.Linear gate: returns logits tensor + # nn.Linear-style gate (e.g. remote code): returns logits tensor logits = output if not isinstance(output, tuple) else output[0] top_k = self.gate.top_k if hasattr(self.gate, "top_k") else self.top_k _, indices = torch.topk(logits.float(), top_k, dim=-1) @@ -733,42 +732,21 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # During calibration, forward all tokens to a larger fraction of experts to improve # calibration coverage, then re-run with the original top_k for actual outputs. if is_calib: - # Skip counting when all experts are calibrated (ratio == 1.0). - self._count_expert_tokens = self._moe_calib_experts_ratio < 1.0 - if self._count_expert_tokens and not self._token_counting_initialized: + # Skip counting when all experts are calibrated (ratio == 1.0). Init first: it resets + # the counting flag, which would otherwise drop the first batch's counts. + count_tokens = self._moe_calib_experts_ratio < 1.0 + if count_tokens and not self._token_counting_initialized: self._init_token_counting() - if TRANSFORMERS_VERSION_GE_5_0: - assert hasattr(self, "gate") and hasattr(self.gate, "top_k") - original_top_k = self.gate.top_k - self.gate.top_k = max( - original_top_k, round(self.gate.num_experts * self._moe_calib_experts_ratio) - ) - super().forward(hidden_states) - self.gate.top_k = original_top_k - else: - # Path for transformers<5.0 - if hasattr(self, "gate") and hasattr(self.gate, "top_k"): - top_k_owner = self.gate - else: - top_k_owner = self - original_top_k = top_k_owner.top_k - if hasattr(self, "num_experts"): - top_k_owner.top_k = max( - original_top_k, round(self.num_experts * self._moe_calib_experts_ratio) - ) - elif hasattr(self, "experts"): - num_experts = ( - self.experts.num_experts - if hasattr(self.experts, "num_experts") - else len(self.experts) - ) - top_k_owner.top_k = max( - original_top_k, - round(num_experts * self._moe_calib_experts_ratio), - ) - else: - raise ValueError(f"Could not find num_experts in module {self}") + self._count_expert_tokens = count_tokens + # top_k lives on the router gate, or on the block itself (e.g. remote-code MoEs) + top_k_owner = self.gate if hasattr(getattr(self, "gate", None), "top_k") else self + original_top_k = top_k_owner.top_k + top_k_owner.top_k = max( + original_top_k, round(self._num_experts() * self._moe_calib_experts_ratio) + ) + try: super().forward(hidden_states) + finally: top_k_owner.top_k = original_top_k self._count_expert_tokens = False @@ -813,94 +791,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return next_states -class _QuantQwen3VLMoeTextExperts(QuantModule): - """Quantized wrapper for the pre-transformers-5.12 ``Qwen3VLMoeTextExperts`` layout. - - That layout stores ``gate_up_proj`` as (num_experts, hidden_size, 2*expert_dim) and runs - the experts through ``torch.bmm``/``@``, so it is unrolled into ``nn.Linear`` modules here. - transformers>=5.12 moved this module to the standard fused layout handled by - :class:`_QuantFusedExperts`; see the registration site below. - """ - - def _setup(self): - """Modify the Qwen3VLMoeTextExperts by using nn.Linear layers.""" - from accelerate import init_empty_weights - - dtype, device = self.gate_up_proj.dtype, self.gate_up_proj.device - - def _copy_weight(module, weight): - module.to_empty(device=device) - with torch.no_grad(): - module.weight.data = weight.detach().data.to(dtype=dtype, device=device) - - # The attribute name was changed from `intermediate_size` to `intermediate_dim` in - # https://github.com/huggingface/transformers/commit/0642963ba13f2dae0596fe489415569e1d91fbda - if hasattr(self, "intermediate_size"): - expert_dim = self.intermediate_size - elif hasattr(self, "intermediate_dim"): - expert_dim = self.intermediate_dim - else: - raise AttributeError("Could not find intermediate dimension size in model") - - with init_empty_weights(): - gate_proj = nn.ModuleList( - [ - nn.Linear(self.hidden_size, expert_dim, bias=False) - for _ in range(self.num_experts) - ] - ) - up_proj = nn.ModuleList( - [ - nn.Linear(self.hidden_size, expert_dim, bias=False) - for _ in range(self.num_experts) - ] - ) - down_proj = nn.ModuleList( - [ - nn.Linear(expert_dim, self.hidden_size, bias=False) - for _ in range(self.num_experts) - ] - ) - - for idx in range(self.num_experts): - _copy_weight(gate_proj[idx], self.gate_up_proj[idx, :, :expert_dim].T) - _copy_weight(up_proj[idx], self.gate_up_proj[idx, :, expert_dim:].T) - _copy_weight(down_proj[idx], self.down_proj[idx, :].T) - - delattr(self, "gate_up_proj") - delattr(self, "down_proj") - self.gate_proj = gate_proj - self.up_proj = up_proj - self.down_proj = down_proj - - def forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - router_indices: torch.Tensor, - ) -> torch.Tensor: - batch_size = hidden_states.shape[0] - hidden_states = hidden_states.reshape(-1, self.hidden_size) - next_states = torch.zeros_like(hidden_states) - with torch.no_grad(): - expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=self.num_experts) - expert_mask = expert_mask.permute(2, 1, 0) - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - for expert_idx in expert_hit: - with torch.no_grad(): - _, token_idx = torch.where(expert_mask[expert_idx[0]]) - current_state = hidden_states[token_idx] - gate = self.gate_proj[expert_idx](current_state) - up = self.up_proj[expert_idx](current_state) - gated_output = up * self.act_fn(gate) - out = self.down_proj[expert_idx](gated_output) - weighted_output = out * routing_weights[token_idx, expert_idx, None] - next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype)) - next_states = next_states.view(batch_size, -1, self.hidden_size) - - return next_states - - def _get_fused_expert_intermediate_dim(module): """Resolve the intermediate (expert) dimension from a fused expert module. @@ -1325,23 +1215,13 @@ def unpack_weight(self): del self.weight_scale_inv -try: - from transformers.models.llama4.modeling_llama4 import Llama4TextExperts - - if Llama4TextExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})( - _QuantLlama4TextExperts - ) -except ImportError: - pass - -try: - from transformers.models.falcon.modeling_falcon import FalconLinear +if Llama4TextExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})( + _QuantLlama4TextExperts + ) - if FalconLinear not in QuantModuleRegistry: - QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear) -except ImportError: - pass +if FalconLinear not in QuantModuleRegistry: + QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear) try: from compressed_tensors.linear.compressed_linear import CompressedLinear @@ -1353,37 +1233,8 @@ def unpack_weight(self): except ImportError: pass -try: - from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts - - # transformers>=5.12 rewrote Qwen3VLMoeTextExperts onto the standard - # ``@use_experts_implementation`` fused layout: ``hidden_size``/``expert_dim`` became - # ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` was transposed to - # (num_experts, 2*intermediate_dim, hidden_dim), and the forward now calls ``F.linear`` - # twice per expert. ``_QuantQwen3VLMoeTextExperts`` only understands the older layout, - # so registering it against the new one crashes on ``self.hidden_size`` (nvbug 6518551). - # The decorator sets ``_apply_gate`` on the class; use it to detect the new layout and - # leave those modules to ``register_fused_experts_on_the_fly``, which claims them with - # the generic ``_QuantFusedExperts``. The old layout must stay explicitly registered: - # it is structurally indistinguishable from a generic fused-experts module, yet its - # forward uses ``torch.bmm``/``@`` rather than ``F.linear``, so the generic wrapper - # would silently quantize nothing. - if Qwen3VLMoeTextExperts not in QuantModuleRegistry and not hasattr( - Qwen3VLMoeTextExperts, "_apply_gate" - ): - QuantModuleRegistry.register({Qwen3VLMoeTextExperts: "hf.Qwen3VLMoeTextExperts"})( - _QuantQwen3VLMoeTextExperts - ) -except ImportError: - pass - -try: - from transformers.integrations.finegrained_fp8 import FP8Linear - - if FP8Linear not in QuantModuleRegistry: - QuantModuleRegistry.register({FP8Linear: "hf.FP8Linear"})(_QuantFP8Linear) -except ImportError: - pass +if FP8Linear not in QuantModuleRegistry: + QuantModuleRegistry.register({FP8Linear: "hf.FP8Linear"})(_QuantFP8Linear) class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): @@ -1472,13 +1323,8 @@ def forward( return super().forward(hidden_states, router_indices, routing_weights) -try: - from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts - - if GptOssExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) -except ImportError: - pass +if GptOssExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) def register_falcon_linears_on_the_fly(model): @@ -1514,7 +1360,7 @@ def _is_sparse_sequaential_moe_block(module): return False if not hasattr(module.experts, "__iter__"): - # transformers>=5.0 has batched experts, no per-expert quantizers + # Fused (batched) experts, no per-expert quantizers return False # Primary: gate sub-module has topk/top_k + num_experts (standard TopKRouter pattern) @@ -1523,7 +1369,7 @@ def _is_sparse_sequaential_moe_block(module): if hasattr(gate, "top_k") and _has_num_experts(gate): return True - # Fallback: top_k + num_experts on the block itself (older transformers, e.g. v4.x Qwen3Next) + # Fallback: top_k + num_experts on the block itself (e.g. remote-code MoE blocks) if hasattr(module, "top_k"): if not _has_num_experts(module) and hasattr(module.experts, "__len__"): module.num_experts = len(module.experts) diff --git a/modelopt/torch/speculative/eagle/default_config.py b/modelopt/torch/speculative/eagle/default_config.py index dfe280b855c..f139a608346 100644 --- a/modelopt/torch/speculative/eagle/default_config.py +++ b/modelopt/torch/speculative/eagle/default_config.py @@ -19,8 +19,8 @@ "hidden_act": "silu", "torch_dtype": "bfloat16", "position_embedding_type": "rope", - # rope_theta is set both inside rope_scaling and at the top level for cross-version - # compatibility: transformers 5.x reads it from rope_scaling, while 4.x reads it top-level. + # rope_theta is set both inside rope_scaling and at the top level: transformers reads it from + # rope_scaling, while the Megatron EAGLE plugin reads it top-level. "rope_scaling": {"rope_type": "default", "rope_theta": 10000}, "rope_theta": 10000, "num_hidden_layers": 1, @@ -79,8 +79,8 @@ "qk_nope_head_dim": 128, "qk_rope_head_dim": 64, "rms_norm_eps": 0.00001, - # rope_theta is set both inside rope_scaling and at the top level for cross-version - # compatibility: transformers 5.x reads it from rope_scaling, while 4.x reads it top-level. + # rope_theta is set both inside rope_scaling and at the top level: transformers reads it from + # rope_scaling, while the Megatron EAGLE plugin reads it top-level. "rope_scaling": { "beta_fast": 1.0, "beta_slow": 1.0, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 084040a5e06..fce0b496f26 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -73,11 +73,10 @@ import logging from pathlib import Path -from typing import Any, NamedTuple +from typing import NamedTuple import torch import torch.nn.functional as F -import transformers from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -103,7 +102,6 @@ __all__ = ["HFDFlashModel"] -_QWEN3_VL_MROPE_WORKAROUND_VERSION = "5.3.0" _MULTIMODAL_FORWARD_KWARGS = frozenset( { "pixel_values", @@ -127,30 +125,6 @@ def _multimodal_forward_kwargs(model_kwargs: dict) -> dict: } -def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor: - """Return the per-frame video grid representation used by Qwen3-VL RoPE. - - Qwen3-VL's video processor emits one ``[T, H, W]`` row per source video, but - its rendered prompt contains a separate visual-token group for every temporal - frame. Transformers 5.3's ``get_rope_index`` consumes one grid row per - rendered group, while the vision encoder still requires the original one-row- - per-video representation. This helper is therefore used *only* for mRoPE - position construction; callers must keep the original tensor for the model - forward. - """ - if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3: - raise ValueError( - "Qwen3-VL video_grid_thw must have shape [num_videos, 3], got " - f"{tuple(video_grid_thw.shape)}." - ) - if torch.any(video_grid_thw[:, 0] <= 0): - raise ValueError("Qwen3-VL video_grid_thw temporal lengths must be positive.") - - expanded_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - expanded_grid_thw[:, 0] = 1 - return expanded_grid_thw - - def _dpace_position_weights( confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None ) -> torch.Tensor: @@ -249,125 +223,6 @@ def _base_llm_config(self): or self.config ) - def _qwen3_vl_position_ids( - self, - input_ids, - attention_mask, - position_ids, - past_key_values, - inputs_embeds, - model_kwargs, - ): - """Precompute Qwen3-VL mRoPE positions for Transformers 5.3.0 batches. - - The video encoder consumes one grid row per source video, whereas mRoPE - consumes one row per rendered temporal-frame group. Calling the - top-level model with the original video grid makes the two contracts - conflict. Construct the mRoPE positions with a frame-expanded copy, - then pass the original grid to the vision encoder in ``forward``. - - Transformers 5.4.0 performs this frame expansion in ``get_rope_index`` - itself; only 5.3.0 needs the external workaround. See - https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py - - Prefer ``get_rope_index`` over ``compute_3d_position_ids``. The latter - writes ``rope_deltas`` into the base model even though DFlash training - never supplies a cache. Keeping this calculation side-effect free is - important when the frozen target is reused for consecutive training - batches or validation. - """ - model_type = str(getattr(self.config, "model_type", "")) - if ( - position_ids is not None - or not model_type.startswith("qwen3_vl") - # Cached decoding uses the base model's rope_deltas path. DFlash - # training has no cache and is the only path that needs the - # frame-expanded construction below. - or past_key_values is not None - ): - return position_ids - - image_grid_thw = model_kwargs.get("image_grid_thw") - video_grid_thw = model_kwargs.get("video_grid_thw") - if not isinstance(image_grid_thw, torch.Tensor) and not isinstance( - video_grid_thw, torch.Tensor - ): - return position_ids - - if transformers.__version__ != _QWEN3_VL_MROPE_WORKAROUND_VERSION: - if transformers.__version__.startswith("5.3."): - raise RuntimeError( - "Qwen3-VL DFlash mRoPE supports Transformers 5.3.0 or >=5.4.0; " - f"got {transformers.__version__}. A 5.3.x patch release may already " - "expand video_grid_thw internally." - ) - return position_ids - - mm_token_type_ids = model_kwargs.get("mm_token_type_ids") - backbone = getattr(self, "model", None) - # Probed dynamically: which one exists depends on the Transformers version. - get_rope_index: Any = getattr(backbone, "get_rope_index", None) - compute_position_ids: Any = getattr(backbone, "compute_3d_position_ids", None) - if ( - not isinstance(mm_token_type_ids, torch.Tensor) - or input_ids is None - or (not callable(get_rope_index) and not callable(compute_position_ids)) - ): - raise ValueError( - "Qwen3-VL DFlash training requires input_ids, mm_token_type_ids, and " - "a Qwen3-VL model with get_rope_index or compute_3d_position_ids. " - "Use the Qwen3-VL AutoProcessor without dropping mm_token_type_ids." - ) - - if mm_token_type_ids.shape != input_ids.shape: - raise ValueError( - "Qwen3-VL mm_token_type_ids must have the same shape as input_ids, got " - f"{tuple(mm_token_type_ids.shape)} and {tuple(input_ids.shape)}." - ) - - rope_video_grid_thw = video_grid_thw - if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0: - video_token_mask = mm_token_type_ids == 2 - if isinstance(attention_mask, torch.Tensor): - video_token_mask = video_token_mask & attention_mask.bool() - video_group_starts = video_token_mask.clone() - video_group_starts[:, 1:] &= ~video_token_mask[:, :-1] - expected_video_groups = int(video_grid_thw[:, 0].sum()) - actual_video_groups = int(video_group_starts.sum()) - if actual_video_groups != expected_video_groups: - raise ValueError( - "Qwen3-VL video frame groups do not match video_grid_thw: " - f"expected {expected_video_groups}, found {actual_video_groups}." - ) - rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw) - - rope_kwargs = { - "input_ids": input_ids, - "image_grid_thw": image_grid_thw, - "video_grid_thw": rope_video_grid_thw, - "attention_mask": attention_mask, - "mm_token_type_ids": mm_token_type_ids, - } - if callable(get_rope_index): - position_ids, _ = get_rope_index(**rope_kwargs) - else: - position_ids = compute_position_ids( - **rope_kwargs, - inputs_embeds=inputs_embeds, - past_key_values=past_key_values, - ) - - expected_shape = (3, *input_ids.shape) - valid_position_ids = isinstance(position_ids, torch.Tensor) and ( - tuple(position_ids.shape) == expected_shape - ) - if not valid_position_ids: - raise RuntimeError( - "Qwen3-VL produced invalid mRoPE position ids: expected shape " - f"{expected_shape}, got {getattr(position_ids, 'shape', None)}." - ) - return position_ids - def _find_base_model_parts(self): """Locate base model submodules (backbone, embeddings, lm_head) by probing known paths. @@ -982,16 +837,6 @@ def forward( - Label alignment: position k predicts token at anchor+k - Optional loss decay weighting """ - if self.training: - position_ids = self._qwen3_vl_position_ids( - input_ids, - attention_mask, - position_ids, - past_key_values, - inputs_embeds, - kwargs, - ) - if not self.training: if self.dflash_offline: raise RuntimeError( diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index 67adab385d2..94399912983 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -357,7 +357,7 @@ def modify( arch_config = config.eagle_architecture_config # Populate base-model-dependent fields before constructing PretrainedConfig, - # since transformers >=5.4 validates rope_scaling during __init__. + # since transformers validates rope_scaling during __init__. arch_config["hidden_size"] = self._base_llm_config.hidden_size arch_config["vocab_size"] = self._base_llm_config.vocab_size arch_config["max_position_embeddings"] = self._base_llm_config.max_position_embeddings @@ -366,7 +366,7 @@ def modify( rope_scaling["rope_theta"] = arch_config["rope_theta"] # Use the base model's config class so fields like max_position_embeddings are declared - # before transformers>=5.5 rope standardization runs in __post_init__. + # before transformers' rope standardization runs in __post_init__. base_config_cls = type(self._base_llm_config) self.eagle_config = base_config_cls.from_dict(arch_config) self.eagle_config.eagle_decoder_type = self.eagle_decoder_type diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index c419d24588f..4db1f4b0f76 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -775,7 +775,7 @@ def _warn_overrides_on_fake_base(): model_name_or_path, config=model_config if pass_config else None, trust_remote_code=trust_remote_code, - torch_dtype=dtype, + dtype=dtype, device_map=device_map, **extra, ) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1efe02e5748..cdd4e36c791 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -30,12 +30,8 @@ from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModelForCausalLM - -try: - from transformers.conversion_mapping import get_model_conversion_mapping - from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key -except ImportError: # transformers<5 has no weight-conversion engine - get_model_conversion_mapping = rename_source_key = WeightConverter = dot_natural_key = None +from transformers.conversion_mapping import get_model_conversion_mapping +from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key from modelopt.torch.utils.distributed import ( barrier, @@ -94,12 +90,12 @@ def _promote_non_dtensor_to_gpu(model: nn.Module, device: torch.device) -> None: def _conversion_plan(model: nn.Module) -> dict | None: """Transformers' own conversion mapping for ``model``, or ``None`` if nothing needs converting. - ``legacy_renames`` (``_checkpoint_conversion_mapping``) covers transformers<5; on 5+ the + ``legacy_renames`` covers a remote-code model's ``_checkpoint_conversion_mapping``; the ``renamings``/``converters`` from HF's engine drive renaming + MoE weight fusion directly. """ legacy_renames = dict(getattr(model, "_checkpoint_conversion_mapping", None) or {}) renamings, converters = [], [] - for entry in get_model_conversion_mapping(model) if get_model_conversion_mapping else []: + for entry in get_model_conversion_mapping(model): (converters if isinstance(entry, WeightConverter) else renamings).append(entry) if not (legacy_renames or renamings or converters): return None @@ -119,8 +115,6 @@ def _resolve_target(plan: dict, key: str) -> tuple[str, str | None]: """ for old, new in plan["legacy_renames"].items(): key = re.sub(old, new, key) - if rename_source_key is None: # transformers<5: legacy renames only, no converters - return key, None return rename_source_key( key, plan["renamings"], plan["converters"], plan["prefix"], plan["meta_state_dict"] ) @@ -128,8 +122,6 @@ def _resolve_target(plan: dict, key: str) -> tuple[str, str | None]: def _convert_keys(plan: dict, state: dict) -> dict: """Rename 1:1 keys and fuse per-expert keys by driving transformers' own conversion ops.""" - if rename_source_key is None: # transformers<5: legacy renames only, no fusion - return {_resolve_target(plan, k)[0]: v for k, v in state.items()} result: dict = {} collected: dict = {} # target -> (converter, {source_pattern: [(sort_key, tensor)]}) for key in sorted(state, key=dot_natural_key): @@ -172,12 +164,12 @@ def build_meta_causal_lm( elif attn_implementation is not None: # Honor the override even when the caller passed in a pre-fetched config. hf_config._attn_implementation = attn_implementation - dtype = getattr(hf_config, "torch_dtype", None) or torch.bfloat16 + dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 from accelerate import init_empty_weights # only real callers of this function need it with init_empty_weights(include_buffers=False): model = AutoModelForCausalLM.from_config( - hf_config, torch_dtype=dtype, trust_remote_code=trust_remote_code + hf_config, dtype=dtype, trust_remote_code=trust_remote_code ) model.eval() return model @@ -395,8 +387,8 @@ def parallel_load_and_prepare_fsdp2( module_to_name = {m: n for n, m in model.named_modules()} layer_prefixes = [module_to_name[layer] + "." for layer in decoder_layers] - # transformers>=5 fuses/renames checkpoint keys so they no longer match param names 1:1 - # (None => the pre-5.x identity path). + # transformers fuses/renames checkpoint keys so they no longer match param names 1:1 + # (None => nothing to convert). plan = _conversion_plan(model) # Valid targets; keys converting to anything else are aux weights (e.g. an MTP head) we skip. diff --git a/noxfile.py b/noxfile.py index 49ac9ed3b05..afa96010ac4 100644 --- a/noxfile.py +++ b/noxfile.py @@ -47,12 +47,7 @@ # ``.[all,dev-test]`` install to constrain that env). TRANSFORMERS_VERSIONS = { "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 - # models silently misroute to the LLM path on export. Likewise ``datasets>=5.1`` requires - # ``huggingface_hub>=1.31`` and fails to import against the downgraded hub. - "tf_min": ("transformers~=4.57.0", "diffusers<0.40", "datasets<5.1"), + "tf_min": ("transformers~=5.5.0",), } diff --git a/plugins/modelopt/skills/ptq/references/unsupported-models.md b/plugins/modelopt/skills/ptq/references/unsupported-models.md index 67563baa8c0..4ed7959c71b 100644 --- a/plugins/modelopt/skills/ptq/references/unsupported-models.md +++ b/plugins/modelopt/skills/ptq/references/unsupported-models.md @@ -135,8 +135,8 @@ class QuantCustomModule(OriginalModule): **Most MoE models are auto-detected** — ModelOpt handles two common patterns automatically: -- **transformers >= 5.0**: Unified fused experts (`gate_up_proj` + `down_proj` 3D tensors) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. -- **transformers < 5.0**: Sequential per-expert `nn.Linear` with `gate` + `experts` → auto-detected by `register_sparse_moe_on_the_fly`. +- **Fused experts** (`gate_up_proj` + `down_proj` 3D tensors, all native transformers MoEs) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. +- **Sequential per-expert `nn.Linear`** with `gate` + `experts` (e.g. remote-code models) → auto-detected by `register_sparse_moe_on_the_fly`. **Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in the plugin (`modelopt/torch/quantization/plugins/huggingface.py`): diff --git a/pyproject.toml b/pyproject.toml index 61ef7bebb44..cc95e6bca7c 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.19", # Should match modelopt/torch/__init__.py and noxfile.py + "transformers>=5.5,<5.19", # Should match modelopt/torch/__init__.py and noxfile.py "wonderwords", ] diff --git a/tests/_test_utils/examples/models.py b/tests/_test_utils/examples/models.py index dd1d8755f84..525ef3d94ad 100644 --- a/tests/_test_utils/examples/models.py +++ b/tests/_test_utils/examples/models.py @@ -38,11 +38,6 @@ def _select_path(remote_id: str, local_id: str) -> str: local_id="TinyLlama-1.1B-Chat-v1.0-x2-MoE", ) -WHISPER_PATH = _select_path( - remote_id="openai/whisper-tiny", - local_id="whisper-tiny", -) - TINY_LLAMA_PATH = _select_path( remote_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", local_id="TinyLlama-1.1B-Chat-v1.0", diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 2282e55d81a..7e3f275f924 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -13,16 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import pytest -import transformers from _test_utils.examples.hf_ptq_utils import PTQCommand -from _test_utils.examples.models import ( - BART_PATH, - MIXTRAL_PATH, - T5_PATH, - TINY_LLAMA_PATH, - WHISPER_PATH, -) -from packaging.version import Version +from _test_utils.examples.models import BART_PATH, MIXTRAL_PATH, T5_PATH, TINY_LLAMA_PATH @pytest.mark.parametrize( @@ -52,22 +44,6 @@ def test_ptq_mixtral(command): command.run(MIXTRAL_PATH) -@pytest.mark.skipif( - Version(transformers.__version__) >= Version("5.0"), - reason="Whisper requires torchcodec and other system packages for transformers>=5.0", -) -@pytest.mark.parametrize( - "command", - [ - # Auto-batch-size computation seems to take >10mins for Whisper hence using a fixed batch size - PTQCommand(quant="fp8", calib_batch_size=16, calib_dataset="peoples_speech", min_sm=89), - ], - ids=PTQCommand.param_str, -) -def test_ptq_whisper(command): - command.run(WHISPER_PATH) - - @pytest.mark.parametrize( "command", [ diff --git a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py index 1f4fd6c37bc..4c9638edd99 100644 --- a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py +++ b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py @@ -16,7 +16,6 @@ import pytest import torch -import transformers from _test_utils.torch.quantization.offload import make_cpu_offloaded_model, make_layerwise_cfg from _test_utils.torch.transformers_models import create_tiny_llama_dir, create_tiny_qwen3_moe_dir from transformers import AutoModelForCausalLM @@ -222,9 +221,9 @@ def test_hf_vllm_export_tiny_llama(tmp_path, quant_cfg): _test_hf_vllm_export(tmp_path, quant_cfg, tiny_model_dir) -@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.INT4_AWQ_CFG]) +@pytest.mark.parametrize( + "quant_cfg", [mtq.FP8_DEFAULT_CFG] +) # INT4_AWQ_CFG is unsupported on fused MoE def test_hf_vllm_export_tiny_qwen3_moe(tmp_path, quant_cfg): - if quant_cfg == mtq.INT4_AWQ_CFG and transformers.__version__.startswith("5."): - pytest.skip("INT4_AWQ_CFG is not supported for Qwen3 MoE in transformers > 5.x") tiny_model_dir = create_tiny_qwen3_moe_dir(tmp_path, num_hidden_layers=2) _test_hf_vllm_export(tmp_path, quant_cfg, tiny_model_dir) diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 58baa7defae..18846c9844d 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -102,14 +102,13 @@ def test_streaming_export_drops_tied_alias_and_writes_config(tmp_path): def test_streaming_export_drops_tied_alias_without_hf_tie_map(tmp_path, monkeypatch): """The alias must be dropped even when HF declares no tie map. - ``all_tied_weights_keys`` only exists on transformers>=5.0, so at our floor the map is - empty and nothing declares the tie by name. The resident path still catches it with an - address pass over the whole state dict; the streaming path has to take the same - information off the live model, before it copies each tensor to host. + Without ``all_tied_weights_keys`` (e.g. a tie HF never declared) nothing names the tie. The + resident path still catches it with an address pass over the whole state dict; the streaming + path has to take the same information off the live model, before it copies each tensor to host. """ model = _tiny_quantized_llama(tie=True) assert model.lm_head.weight is model.model.embed_tokens.weight, "fixture is not tied" - # Simulate the transformers floor: the attribute is simply not there. + # Hide HF's tie map: the attribute is simply not there. monkeypatch.setattr( type(model), "all_tied_weights_keys", property(lambda self: None), raising=False ) 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 03c4e48fbab..20f789dc078 100644 --- a/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py +++ b/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py @@ -42,9 +42,9 @@ class _NewWarmupArgs(ModelOptHFArguments): class _OldWarmupArgs(ModelOptHFArguments): - """transformers 4.x: an int ``warmup_steps`` plus ``warmup_ratio``.""" + """transformers < 5.15: a float ``warmup_steps`` plus ``warmup_ratio``.""" - warmup_steps: int = field(default=0) + warmup_steps: float = field(default=0) warmup_ratio: float = field(default=0.0) @@ -53,8 +53,7 @@ class _OldWarmupArgs(ModelOptHFArguments): [ (_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_steps: 0.05\n", {"warmup_steps": 0.05, "warmup_ratio": 0.0}), (_OldWarmupArgs, "warmup_ratio: 0.05\n", {"warmup_steps": 0, "warmup_ratio": 0.05}), ], ) diff --git a/tests/unit/torch/opt/plugins/test_transformers_save_load.py b/tests/unit/torch/opt/plugins/test_transformers_save_load.py index e8b3ed22151..1e526fc09c0 100644 --- a/tests/unit/torch/opt/plugins/test_transformers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_transformers_save_load.py @@ -27,10 +27,7 @@ from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForCausalLM, LlamaForCausalLM -from modelopt.torch.opt.plugins.transformers import ( - _TRANSFORMERS_GE_5_0, - _legacy_tied_weights_keys_as_dict, -) +from modelopt.torch.opt.plugins.transformers import _legacy_tied_weights_keys_as_dict @pytest.mark.parametrize("model_cls", [LlamaForCausalLM, AutoModelForCausalLM]) @@ -84,7 +81,6 @@ def test_save_pretrained_with_legacy_tied_weights_keys(tmp_path, tie_word_embedd tf_modelopt_state_and_output_tester(model, model_test) -@pytest.mark.skipif(not _TRANSFORMERS_GE_5_0, reason="list-style keys are native to transformers 4") def test_legacy_tied_weights_keys_as_dict_restores_class_attribute(): """The shim must not leave an instance attribute shadowing the class declaration.""" diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index e7b8fb69d61..22a86971a6d 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -85,10 +85,10 @@ def test_kv_quant_hf(model_getter, attn_cls): original_is_compatible_attention = None if attn_cls is not None: - # Test case for transformers < 4.48 + # Test case for attention modules without transformers' attention interface # This needs: # 1) replace the attention class with the test attention class - # 2) set _QuantAttention.is_compatible_attention output to False to fall back to the transformers < 4.48 support + # 2) set _QuantAttention.is_compatible_attention output to False to fall back to the AST-based support for name, module in model_test.named_modules(): if name.endswith(attention_module): if original_is_compatible_attention is None: diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 73db66a539a..7b029b393c5 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1478,8 +1478,8 @@ def _make_qwen3_vl_moe_experts(): num_key_value_heads=2, vocab_size=128, ) - # The fused forward of transformers>=5.12 dispatches on this; ``eager`` is the only - # backend that routes through ``F.linear`` and therefore through the quantizer hooks. + # The fused forward dispatches on this; ``eager`` is the only backend that routes + # through ``F.linear`` and therefore through the quantizer hooks. config._experts_implementation = "eager" experts = Qwen3VLMoeTextExperts(config) # Weights are created with ``torch.empty``; fill them so comparisons are meaningful. @@ -1490,36 +1490,18 @@ def _make_qwen3_vl_moe_experts(): return experts -def _qwen3_vl_moe_is_new_layout(): - """True when transformers>=5.12 moved the experts onto the generic fused layout.""" - from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts - - return hasattr(Qwen3VLMoeTextExperts, "_apply_gate") - - def _qwen3_vl_moe_forward_args(): - """Routing inputs for the installed layout, sized so every expert is hit. - - The two layouts disagree on both argument order and routing-weight shape, and the - pre-5.12 eval-mode forward computes a dense weighted sum over all experts. Zeroing the - weights outside the top-k keeps that dense path equal to the sparse per-expert loop. - """ + """Routing inputs sized so every expert is hit.""" seq_len = QWEN_NUM_EXPERTS // QWEN_TOP_K torch.manual_seed(0) hidden_states = torch.randn(seq_len, QWEN_HIDDEN_DIM) router_indices = torch.arange(QWEN_NUM_EXPERTS, dtype=torch.long).reshape(seq_len, QWEN_TOP_K) top_k_weights = torch.softmax(torch.randn(seq_len, QWEN_TOP_K), dim=-1) - - if _qwen3_vl_moe_is_new_layout(): - return hidden_states, router_indices, top_k_weights - - routing_weights = torch.zeros(seq_len, QWEN_NUM_EXPERTS) - routing_weights.scatter_(1, router_indices, top_k_weights) - return hidden_states.unsqueeze(0), routing_weights, router_indices + return hidden_states, router_indices, top_k_weights class TestQwen3VLMoeTextExperts: - """The registered wrapper must match the installed transformers layout (nvbug 6518551).""" + """Qwen3-VL MoE text experts must quantize through the generic fused-experts wrapper.""" @staticmethod def _experts_type(): @@ -1527,21 +1509,11 @@ def _experts_type(): return Qwen3VLMoeTextExperts - def test_registration_matches_installed_layout(self): - """transformers>=5.12 experts must be left to the generic fused-experts wrapper.""" - from modelopt.torch.quantization.plugins.huggingface import _QuantQwen3VLMoeTextExperts - - registered = QuantModuleRegistry.get(self._experts_type()) - if _qwen3_vl_moe_is_new_layout(): - # Statically registering the legacy wrapper here would shadow on-the-fly - # detection and crash on ``self.hidden_size`` during conversion. - assert registered is None - assert _fused_experts_wrapper_class(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts - else: - # The pre-5.12 forward uses ``torch.bmm``, which the generic wrapper cannot - # intercept, so the explicit registration must stay in place. - assert registered is not None - assert issubclass(registered, _QuantQwen3VLMoeTextExperts) + def test_left_to_generic_fused_experts_wrapper(self): + """The experts must be left to the generic fused-experts wrapper.""" + # A static registration would shadow on-the-fly detection. + assert QuantModuleRegistry.get(self._experts_type()) is None + assert _fused_experts_wrapper_class(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts def test_convert_and_forward_matches_reference(self): """Conversion must succeed and stay numerically transparent before calibration.""" diff --git a/tests/unit/torch/quantization/plugins/test_peft.py b/tests/unit/torch/quantization/plugins/test_peft.py index a3ef95cf6c1..ba69579e539 100644 --- a/tests/unit/torch/quantization/plugins/test_peft.py +++ b/tests/unit/torch/quantization/plugins/test_peft.py @@ -111,14 +111,16 @@ def forward_loop(model): outputs_peft = peft_model(input_ids).logits outputs_full = model_full(input_ids).logits - assert torch.allclose(outputs_peft, outputs_full, atol=1e-2, rtol=1e-2) + # The full model quantizes the merged weight, PEFT only the base one: allow for that INT8 gap. + tol = {"atol": 3e-2, "rtol": 1e-2} + assert torch.allclose(outputs_peft, outputs_full, **tol) peft_model.save_pretrained(tmp_path / "peft_model") peft_loaded = PeftModel.from_pretrained(model_original, tmp_path / "peft_model") outputs_peft_loaded = peft_loaded(input_ids).logits - assert torch.allclose(outputs_peft_loaded, outputs_full, atol=1e-2, rtol=1e-2) + assert torch.allclose(outputs_peft_loaded, outputs_full, **tol) model_after_peft_merge = peft_loaded.merge_and_unload() outputs_after_peft_merge = model_after_peft_merge(input_ids).logits - assert torch.allclose(outputs_after_peft_merge, outputs_full, atol=1e-2, rtol=1e-2) + assert torch.allclose(outputs_after_peft_merge, outputs_full, **tol) diff --git a/tests/unit/torch/quantization/plugins/test_sparse_sequential_moe.py b/tests/unit/torch/quantization/plugins/test_sparse_sequential_moe.py index 636a43bad55..e6d2740abbe 100644 --- a/tests/unit/torch/quantization/plugins/test_sparse_sequential_moe.py +++ b/tests/unit/torch/quantization/plugins/test_sparse_sequential_moe.py @@ -20,19 +20,12 @@ import pytest import torch import torch.nn as nn -from packaging.version import Version pytest.importorskip("transformers") -if Version(torch.__version__) < Version("2.9"): - pytest.skip("torch 2.8 grouped_mm is CUDA-only", allow_module_level=True) - -from _test_utils.torch.transformers_models import get_tiny_qwen3_moe - import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer from modelopt.torch.quantization.plugins.huggingface import ( - TRANSFORMERS_VERSION_GE_5_0, _is_sparse_sequaential_moe_block, register_sparse_moe_on_the_fly, ) @@ -42,7 +35,7 @@ # Helpers: lightweight mock modules for _is_sparse_sequaential_moe_block # --------------------------------------------------------------------------- class _FakeGateWithRouter(nn.Module): - """Mimics a v5.x TopKRouter gate with top_k and num_experts.""" + """Mimics a router gate with top_k and num_experts that returns logits.""" def __init__(self, top_k=2, num_experts=4): super().__init__() @@ -100,6 +93,37 @@ def forward(self, hidden_states): return out +class _RoutedExpertsGate(nn.Module): + """A router gate that names its expert count ``n_routed_experts`` and returns logits.""" + + def __init__(self, top_k=2, num_experts=4): + super().__init__() + self.top_k = top_k + self.n_routed_experts = num_experts + self.linear = nn.Linear(8, num_experts) + + def forward(self, x): + return self.linear(x) + + +class _MoEBlockRoutedExpertsGate(nn.Module): + """Matches the primary detection path through gate.top_k + gate.n_routed_experts.""" + + def __init__(self, num_experts=4, top_k=2): + super().__init__() + self.gate = _RoutedExpertsGate(top_k=top_k, num_experts=num_experts) + self.experts = nn.ModuleList([nn.Linear(8, 8) for _ in range(num_experts)]) + + def forward(self, hidden_states): + _, selected = torch.topk(self.gate(hidden_states), self.gate.top_k, dim=-1) + out = torch.zeros_like(hidden_states) + for i, expert in enumerate(self.experts): + mask = (selected == i).any(dim=-1) + if mask.any(): + out[mask] += expert(hidden_states[mask]) + return out + + # --------------------------------------------------------------------------- # Tests for _is_sparse_sequaential_moe_block # --------------------------------------------------------------------------- @@ -175,164 +199,120 @@ def test_n_routed_experts_accepted(self): # --------------------------------------------------------------------------- # Tests for _QuantSparseSequentialMoe # --------------------------------------------------------------------------- -@pytest.mark.skipif(TRANSFORMERS_VERSION_GE_5_0, reason="Transformers v5 has stacked MoE") class TestQuantSparseSequentialMoe: - """Tests for _QuantSparseSequentialMoe using a real tiny Qwen3Moe model.""" + """Tests for _QuantSparseSequentialMoe on a sequential (per-expert ``nn.Linear``) MoE block.""" @staticmethod - def _get_moe_block(model): - """Return the first MoE block from the model.""" - for module in model.modules(): - if _is_sparse_sequaential_moe_block(module): - return module - raise RuntimeError("No MoE block found in model") + def _convert(block): + if QuantModuleRegistry.get(type(block)) is None: + register_sparse_moe_on_the_fly(block) + return QuantModuleRegistry.convert(block) def test_register_sparse_moe_on_the_fly(self): - model = get_tiny_qwen3_moe() - moe_block = self._get_moe_block(model) - moe_type = type(moe_block) - - if QuantModuleRegistry.get(moe_type) is not None: - pytest.skip("MoE type already registered (upstream change)") - - register_sparse_moe_on_the_fly(model) - assert QuantModuleRegistry.get(moe_type) is not None + block = _MoEBlockWithGateRouter() + register_sparse_moe_on_the_fly(block) + assert QuantModuleRegistry.get(type(block)) is not None def test_setup_config_knobs_default(self): """_setup should only initialize config knobs, no buffer or hook.""" - model = get_tiny_qwen3_moe() - moe_block = self._get_moe_block(model) - if QuantModuleRegistry.get(type(moe_block)) is None: - register_sparse_moe_on_the_fly(model) - - converted = QuantModuleRegistry.convert(moe_block) + converted = self._convert(_MoEBlockWithGateRouter()) assert converted._moe_calib_experts_ratio is None assert not hasattr(converted, "expert_token_count") def test_forward_default_config_passthrough(self): """With default config (both features off), forward should be a direct pass-through.""" - model = get_tiny_qwen3_moe() - moe_block = self._get_moe_block(model) - if QuantModuleRegistry.get(type(moe_block)) is None: - register_sparse_moe_on_the_fly(model) + block = _MoEBlockWithGateRouter() + ref_block = copy.deepcopy(block) + converted = self._convert(block) - ref_block = self._get_moe_block(get_tiny_qwen3_moe()) - ref_block.load_state_dict(moe_block.state_dict()) - converted = QuantModuleRegistry.convert(moe_block) - - x = torch.randn(1, 4, 32, dtype=ref_block.gate.weight.dtype) + x = torch.randn(4, 8) with torch.no_grad(): - out_ref = ref_block(x) - out_test = converted(x) - - if isinstance(out_ref, tuple): - out_ref = out_ref[0] - if isinstance(out_test, tuple): - out_test = out_test[0] - assert torch.allclose(out_ref, out_test, atol=1e-5) + assert torch.allclose(ref_block(x), converted(x), atol=1e-5) assert not hasattr(converted, "expert_token_count") def test_forward_calib_restores_top_k(self): """After calibration forward with moe_calib_experts_ratio, top_k should be restored.""" - model = get_tiny_qwen3_moe() - moe_block = self._get_moe_block(model) - if QuantModuleRegistry.get(type(moe_block)) is None: - register_sparse_moe_on_the_fly(model) - - if TRANSFORMERS_VERSION_GE_5_0: - original_top_k = moe_block.gate.top_k - else: - original_top_k = moe_block.top_k - - converted = QuantModuleRegistry.convert(moe_block) + converted = self._convert(_MoEBlockWithGateRouter(top_k=2)) converted._moe_calib_experts_ratio = 1.0 + converted.experts[0]._if_calib = True # simulate calibration mode - # Simulate calibration mode - for m in converted.experts.modules(): - if hasattr(m, "_if_calib"): - m._if_calib = True - break - - x = torch.randn(1, 4, 32, dtype=converted.gate.weight.dtype) with torch.no_grad(): - converted(x) - - if TRANSFORMERS_VERSION_GE_5_0: - assert converted.gate.top_k == original_top_k - else: - assert converted.top_k == original_top_k + converted(torch.randn(4, 8)) + assert converted.gate.top_k == 2 def test_token_counting_lazy_init(self): """When moe_calib_experts_ratio > 0, token counting infra is lazy-inited.""" - model = get_tiny_qwen3_moe() - moe_block = self._get_moe_block(model) - if QuantModuleRegistry.get(type(moe_block)) is None: - register_sparse_moe_on_the_fly(model) - - converted = QuantModuleRegistry.convert(moe_block) + converted = self._convert(_MoEBlockWithGateRouter(num_experts=4, top_k=2)) converted._moe_calib_experts_ratio = 0.5 - assert not hasattr(converted, "expert_token_count") - # Simulate calibration mode so lazy-init triggers during forward - # Set _if_calib on an expert sub-module (not set by default since only the MoE - # block was converted, not the full model). - next(converted.experts.modules())._if_calib = True - - x = torch.randn(1, 4, 32, dtype=converted.gate.weight.dtype) + converted.experts[0]._if_calib = True # lazy-init triggers during a calibration forward with torch.no_grad(): - converted(x) + converted(torch.randn(4, 8)) + assert converted.expert_token_count.numel() == 4 - # Buffer and hook should now exist - assert hasattr(converted, "expert_token_count") - assert converted.expert_token_count.numel() > 0 - - # Manually enable counting and call gate to verify hook works + # Manually enable counting and call gate to verify the hook counts top_k per token converted._count_expert_tokens = True - if TRANSFORMERS_VERSION_GE_5_0: - hidden_size = converted.gate.weight.shape[1] - top_k = converted.gate.top_k - else: - hidden_size = converted.gate.in_features - top_k = converted.top_k if hasattr(converted, "top_k") else converted.gate.top_k - converted.expert_token_count.zero_() - tokens = torch.randn(8, hidden_size, dtype=converted.gate.weight.dtype) with torch.no_grad(): - converted.gate(tokens) - assert converted.expert_token_count.sum().item() == 8 * top_k + converted.gate(torch.randn(8, 8)) + assert converted.expert_token_count.sum().item() == 8 * 2 + + @pytest.mark.parametrize( + ("block_cls", "top_k_owner"), + [ + (_MoEBlockWithGateRouter, "gate"), + (_MoEBlockFallback, "block"), # top_k + num_experts on the block, nn.Linear gate + (_MoEBlockRoutedExpertsGate, "gate"), # gate.n_routed_experts instead of num_experts + ], + ) + def test_calib_forward_supports_detected_layouts(self, block_cls, top_k_owner): + """Every layout detection accepts must also calibrate with token forcing.""" + converted = self._convert(block_cls(num_experts=4, top_k=1)) + owner = converted.gate if top_k_owner == "gate" else converted + converted._moe_calib_experts_ratio = 0.75 + converted.experts[0]._if_calib = True + + with torch.no_grad(): + converted(torch.randn(16, 8)) + assert owner.top_k == 1 # restored after the forced pass + # The forced pass routed each token to round(4 * 0.75) = 3 experts. + assert converted.expert_token_count.sum().item() == 16 * 3 -@pytest.mark.skipif(TRANSFORMERS_VERSION_GE_5_0, reason="Transformers v5 has stacked MoE") -def test_qwen3_sequential_moe_quantize_with_token_forcing_and_counting(): - """End-to-end: mtq.quantize a Qwen3MoE with INT8 + moe_calib_experts_ratio + token counting.""" - model = get_tiny_qwen3_moe() +class _TinySequentialMoEModel(nn.Module): + def __init__(self, block): + super().__init__() + self.block = block + + def forward(self, x): + return self.block(x) - # Verify detection - moe_found = any(_is_sparse_sequaential_moe_block(m) for m in model.modules()) - assert moe_found, "Qwen3MoE should be detected as a sparse MoE block" +@pytest.mark.parametrize("block_cls", [_MoEBlockWithGateRouter, _MoEBlockFallback]) +def test_quantize_sequential_moe_with_token_forcing(block_cls): + """End-to-end: mtq.quantize calibrates every expert of a sequential MoE via token forcing.""" + torch.manual_seed(0) + model = _TinySequentialMoEModel(block_cls(num_experts=4, top_k=1)) + calib_data = [torch.randn(32, 8) for _ in range(2)] quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) - quant_cfg["algorithm"] = { - "method": "max", - "moe_calib_experts_ratio": 0.5, - } - - def calib_fn(model): - x = model.dummy_inputs["input_ids"] - for _ in range(2): - model(x) - - mtq.quantize(model, quant_cfg, calib_fn) - - # Verify token counting worked - for name, module in model.named_modules(): - if hasattr(module, "expert_token_count") and module.expert_token_count.numel() > 0: - assert (module.expert_token_count > 0).all(), ( - f"Not all experts received tokens in {name}: {module.expert_token_count}" - ) - - # Verify model still runs + quant_cfg["algorithm"] = {"method": "max", "moe_calib_experts_ratio": 0.75} + + def forward_loop(m): + for x in calib_data: + m(x) + + mtq.quantize(model, quant_cfg, forward_loop) + + block = model.block + assert block._moe_calib_experts_ratio == 0.75 + # Each calibration batch routed every token to round(4 * 0.75) = 3 experts. + assert block.expert_token_count.sum().item() == 2 * 32 * 3 + assert (block.expert_token_count > 0).all() + for name, module in block.experts.named_modules(): + if isinstance(module, TensorQuantizer) and module.is_enabled: + assert module.amax is not None, f"expert quantizer {name} was not calibrated" + owner = block.gate if hasattr(block.gate, "top_k") else block + assert owner.top_k == 1 # routing restored after calibration with torch.no_grad(): - out = model(model.dummy_inputs["input_ids"]) - assert out.logits is not None + assert torch.isfinite(model(torch.randn(4, 8))).all() diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index be959b22e63..99dace9ea9c 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -161,7 +161,7 @@ def from_pretrained(*args, **kwargs): assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None assert captured["args"] == ("qwen3-vl",) - assert captured["kwargs"]["torch_dtype"] == "auto" + assert captured["kwargs"]["dtype"] == "auto" class TestFakeBaseRopeTheta: diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 76050c0e568..8cf51d3daec 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -112,145 +112,6 @@ def test_convert_sets_mask_token_id(self): assert model.mask_token_id == 0 -def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch): - """Only mRoPE receives a per-frame video grid on Transformers 5.3.0.""" - original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]]) - expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long) - get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) - compute_position_ids = MagicMock() - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace( - get_rope_index=get_rope_index, - compute_3d_position_ids=compute_position_ids, - ), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") - - position_ids = HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 12, dtype=torch.long), - attention_mask=torch.ones(1, 12, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "video_grid_thw": original_grid, - "mm_token_type_ids": torch.tensor([[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]), - }, - ) - - assert position_ids is expected_position_ids - assert not compute_position_ids.called - assert torch.equal(original_grid, torch.tensor([[3, 4, 5], [2, 6, 7]])) - assert torch.equal( - get_rope_index.call_args.kwargs["video_grid_thw"], - torch.tensor([[1, 4, 5], [1, 4, 5], [1, 4, 5], [1, 6, 7], [1, 6, 7]]), - ) - - -def test_qwen3_vl_moe_transformers_530_position_ids_expand_video_grid(monkeypatch): - """Qwen3-VL family variants use the same 5.3.0 mRoPE workaround.""" - expected_position_ids = torch.ones(3, 1, 4, dtype=torch.long) - get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl_moe"), - model=SimpleNamespace(get_rope_index=get_rope_index), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") - - position_ids = HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 4, dtype=torch.long), - attention_mask=torch.ones(1, 4, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "video_grid_thw": torch.tensor([[2, 4, 4]]), - "mm_token_type_ids": torch.tensor([[2, 0, 2, 0]]), - }, - ) - - assert position_ids is expected_position_ids - assert torch.equal( - get_rope_index.call_args.kwargs["video_grid_thw"], - torch.tensor([[1, 4, 4], [1, 4, 4]]), - ) - - -def test_qwen3_vl_transformers_53_patch_release_raises(monkeypatch): - """Avoid double expansion when a 5.3 patch backports the upstream fix.""" - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace(get_rope_index=MagicMock()), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.1") - - with pytest.raises(RuntimeError, match=r"5\.3\.0 or >=5\.4\.0"): - HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 4, dtype=torch.long), - attention_mask=torch.ones(1, 4, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "video_grid_thw": torch.tensor([[1, 4, 4]]), - "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), - }, - ) - - -def test_qwen3_vl_transformers_54_uses_native_position_ids(monkeypatch): - """Transformers 5.4+ performs the grid expansion inside get_rope_index.""" - get_rope_index = MagicMock() - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace(get_rope_index=get_rope_index), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.4.0") - - position_ids = HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 4, dtype=torch.long), - attention_mask=torch.ones(1, 4, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "video_grid_thw": torch.tensor([[1, 4, 4]]), - "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), - }, - ) - - assert position_ids is None - assert not get_rope_index.called - - -def test_qwen3_vl_transformers_530_rejects_bad_video_frame_groups(monkeypatch): - """Fail before mRoPE construction when processor and video-grid contracts differ.""" - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace(get_rope_index=MagicMock()), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") - - with pytest.raises(ValueError, match="video frame groups"): - HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 4, dtype=torch.long), - attention_mask=torch.ones(1, 4, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "video_grid_thw": torch.tensor([[2, 4, 4]]), - "mm_token_type_ids": torch.tensor([[2, 2, 0, 0]]), - }, - ) - - def test_multimodal_forward_kwargs_exclude_non_model_inputs(): """Do not forward Trainer or collator-only fields to Hugging Face models.""" pixel_values = torch.ones(1) @@ -272,62 +133,6 @@ def test_multimodal_forward_kwargs_exclude_non_model_inputs(): assert forwarded["mm_token_type_ids"] is mm_token_type_ids -def test_eval_does_not_precompute_qwen3_vl_position_ids(monkeypatch): - """Evaluation delegates mRoPE construction to the base model and its cache.""" - model = get_tiny_llama(num_hidden_layers=4) - mtsp.convert(model, [("dflash", get_dflash_config())]) - precompute_position_ids = MagicMock() - monkeypatch.setattr(model, "_qwen3_vl_position_ids", precompute_position_ids) - - model.eval() - model(input_ids=torch.tensor([[1, 2, 3, 4]])) - - precompute_position_ids.assert_not_called() - - -def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch): - """Never silently fall back to one-dimensional positions for a visual batch.""" - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace(get_rope_index=MagicMock()), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") - - with pytest.raises(ValueError, match="mm_token_type_ids"): - HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 12, dtype=torch.long), - attention_mask=torch.ones(1, 12, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={"image_grid_thw": torch.tensor([[1, 4, 4]])}, - ) - - -def test_qwen3_vl_transformers_53_position_ids_reject_bad_mm_token_shape(monkeypatch): - """Keep processor-produced modality ids aligned with the padded text sequence.""" - fake_model = SimpleNamespace( - config=SimpleNamespace(model_type="qwen3_vl"), - model=SimpleNamespace(get_rope_index=MagicMock()), - ) - monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") - - with pytest.raises(ValueError, match="same shape as input_ids"): - HFDFlashModel._qwen3_vl_position_ids( - fake_model, - input_ids=torch.ones(1, 12, dtype=torch.long), - attention_mask=torch.ones(1, 12, dtype=torch.long), - position_ids=None, - past_key_values=None, - inputs_embeds=None, - model_kwargs={ - "image_grid_thw": torch.tensor([[1, 4, 4]]), - "mm_token_type_ids": torch.zeros(1, 11, dtype=torch.long), - }, - ) - - class TestDPaceWeights: """Test the D-PACE position-weighting objective (arXiv:2605.18810).""" diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 428caba0617..b4b22b74faf 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -17,7 +17,6 @@ import pytest import torch -from packaging.version import Version pytest.importorskip("accelerate") @@ -31,9 +30,7 @@ def _build_tiny_qwen3_moe(): """A tiny meta-init Qwen3-MoE (fused ``gate_up_proj`` + ``down_proj`` experts) for converter tests.""" - transformers = pytest.importorskip("transformers") - if Version(transformers.__version__) < Version("5.0"): - pytest.skip("multi-source fused-MoE conversion needs transformers>=5") + pytest.importorskip("transformers") from accelerate import init_empty_weights from transformers import AutoConfig, AutoModelForCausalLM