From 87b96f01fc4988840cbcd45160f54b3b0ea0ab7b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:58:52 +0000 Subject: [PATCH 1/4] Add require_transformers, the floor check for a feature transformers added Expert parallelism and the loaders that build on it rely on transformers APIs that only newer releases have, and an older release often runs them without an error while doing the wrong thing (transformers 5.5 shards no expert). moe.require_transformers(min_version, feature) raises naming the feature, the installed release and the one to upgrade to. transformers is imported where it is checked: this module imports without it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 19 +++++++++++++++++++ tests/unit/torch/utils/test_moe_a2a.py | 14 ++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 98c4bc2453c..f7452a49322 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -29,12 +29,31 @@ import torch import torch.distributed as dist +from packaging.version import parse from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single if TYPE_CHECKING: from torch import nn +def require_transformers(min_version: str, feature: str) -> None: + """Fail loudly if the installed transformers is older than ``min_version`` (the floor for ``feature``). + + Callers enforce the version where the needed APIs first exist, so a too-old transformers errors + clearly instead of silently mis-behaving (e.g. transformers 5.5 silently shards nothing for EP). + See each call site for the empirically-determined floor. + """ + # Imported here: this module imports without transformers. + import transformers + + have = transformers.__version__ + if parse(have) < parse(min_version): + raise RuntimeError( + f"{feature} requires transformers >= {min_version}, but {have} is installed. " + f"Upgrade with: pip install -U 'transformers>={min_version}'." + ) + + def _ep_config(hf_config): """Return the config that carries the EP plan / expert count. diff --git a/tests/unit/torch/utils/test_moe_a2a.py b/tests/unit/torch/utils/test_moe_a2a.py index 7a872183040..5d2498ae86d 100644 --- a/tests/unit/torch/utils/test_moe_a2a.py +++ b/tests/unit/torch/utils/test_moe_a2a.py @@ -30,6 +30,7 @@ import copy import warnings +from contextlib import nullcontext from datetime import timedelta from types import SimpleNamespace @@ -503,3 +504,16 @@ def test_routed_experts_are_the_plans_containers_under_any_prefix(plan, prefix): assert found == { name: modules[name] for name in [f"{prefix}layers.{i}.mlp.experts" for i in (0, 1)] } + + +@pytest.mark.parametrize(("installed", "refused"), [("5.6.2", True), ("5.7.0", False)]) +def test_a_transformers_below_a_features_floor_is_refused(monkeypatch, installed, refused): + transformers = pytest.importorskip("transformers") + monkeypatch.setattr(transformers, "__version__", installed) + + with ( + pytest.raises(RuntimeError, match=r"EP load requires transformers >= 5\.7, but 5\.6\.2 is") + if refused + else nullcontext() + ): + moe.require_transformers("5.7", "EP load") From 32d3b1d26ba3537ab143ea2966df48cc545f9610 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:59:07 +0000 Subject: [PATCH 2/4] Build a VLM without its LM head when the caller does not need it _auto_class takes require_lm_head, defaulting to True, as its second parameter: with it False, a model whose architecture reads as a VLM's, or an encoder-decoder model, builds through AutoModel, which leaves the head out. A causal LM keeps its head either way. _decide_distribution, its one caller, passes trust_remote_code by keyword. The class table gains the column and two headless rows. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 30 +++++++------ .../unit/torch/utils/test_model_load_utils.py | 44 ++++++++++++------- 2 files changed, 45 insertions(+), 29 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index babf7f11054..0591a92349f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -924,18 +924,20 @@ def parallel_load_and_prepare_fsdp2( return model -def _auto_class(hf_config, trust_remote_code: bool = False) -> Any: - """The auto class that builds ``hf_config``'s whole model, LM head included. - - A causal LM builds through ``AutoModelForCausalLM``. A VLM -- a nested text config, a vision - config, or a ``*ForConditionalGeneration`` architecture -- builds through the remote - ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` there, - its code is trusted and it registers no image-text class, else through - ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the - checkpoint registers one, else through ``AutoModelForCausalLM`` if that recognizes it (as it - does Phi-4-multimodal), else through ``AutoModel``, which leaves the head out. An - encoder-decoder model (``is_encoder_decoder``, whether or not it names its architecture, which - reads like a VLM's) builds through its seq2seq class. +def _auto_class(hf_config, require_lm_head: bool = True, trust_remote_code: bool = False) -> Any: + """The auto class that builds ``hf_config``'s whole model. + + A causal LM builds through ``AutoModelForCausalLM``, head included: ``require_lm_head`` only + applies to a model whose architecture reads as a VLM's -- a nested text config, a vision + config, or a ``*ForConditionalGeneration`` architecture -- or that is an encoder-decoder model + (``is_encoder_decoder``), whether or not it names its architecture. Such a model builds through + ``AutoModel``, which leaves its LM head out, unless ``require_lm_head``: then through the + remote ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` + there, its code is trusted and it registers no image-text class, else through its seq2seq + class if it is an encoder-decoder model, else through ``AutoModelForImageTextToText`` if that + recognizes the config or, its code trusted, the checkpoint registers one, else through + ``AutoModelForCausalLM`` if that recognizes it (as it does Phi-4-multimodal), else still + through ``AutoModel``. """ architectures = getattr(hf_config, "architectures", None) or [] is_vlm = ( @@ -945,6 +947,8 @@ def _auto_class(hf_config, trust_remote_code: bool = False) -> Any: ) if not is_vlm and not getattr(hf_config, "is_encoder_decoder", False): return AutoModelForCausalLM + if not require_lm_head: + return AutoModel auto_map = getattr(hf_config, "auto_map", None) or {} remote_with_head = ( trust_remote_code @@ -1027,7 +1031,7 @@ def _decide_distribution( trust_remote_code, None, hf_config, - auto_class=_auto_class(hf_config, trust_remote_code), + auto_class=_auto_class(hf_config, trust_remote_code=trust_remote_code), ) # init_empty_weights re-wraps every parameter it registers, which unties what from_config # tied: re-tie, so a tied embedding and LM head count once, as the loaded model holds them. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 2c225dafc14..b2aa0ab801a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -305,20 +305,28 @@ def _tiny_qwen2_5_vl(): @pytest.mark.parametrize( - ("config", "trust_remote_code", "expected"), + ("config", "require_lm_head", "trust_remote_code", "expected"), [ - (SimpleNamespace(architectures=["LlamaForCausalLM"]), False, AutoModelForCausalLM), - (_vlm_config(auto_map=_REMOTE_VLM), True, AutoModelForCausalLM), - (_vlm_config(auto_map=_REMOTE_VLM), False, AutoModel), - (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, AutoModelForImageTextToText), - (_tiny_qwen2_5_vl(), False, AutoModelForImageTextToText), - (Phi4MultimodalConfig(), False, AutoModelForCausalLM), - (_vlm_config(), False, AutoModel), - (T5Config(architectures=["T5ForConditionalGeneration"]), False, AutoModelForSeq2SeqLM), - (WhisperConfig(), False, AutoModelForSpeechSeq2Seq), + (SimpleNamespace(architectures=["LlamaForCausalLM"]), True, False, AutoModelForCausalLM), + (_tiny_qwen2_5_vl(), False, False, AutoModel), + (_vlm_config(auto_map=_REMOTE_VLM), True, True, AutoModelForCausalLM), + (_vlm_config(auto_map=_REMOTE_VLM), True, False, AutoModel), + (_vlm_config(auto_map=_REMOTE_IMAGE_TEXT_VLM), True, True, AutoModelForImageTextToText), + (_tiny_qwen2_5_vl(), True, False, AutoModelForImageTextToText), + (Phi4MultimodalConfig(), True, False, AutoModelForCausalLM), + (_vlm_config(), True, False, AutoModel), + ( + T5Config(architectures=["T5ForConditionalGeneration"]), + True, + False, + AutoModelForSeq2SeqLM, + ), + (WhisperConfig(), True, False, AutoModelForSpeechSeq2Seq), + (T5Config(architectures=["T5ForConditionalGeneration"]), False, False, AutoModel), ], ids=[ "causal-lm", + "vlm-backbone", "remote-vlm-with-head", "remote-vlm-untrusted", "remote-image-text-vlm", @@ -327,14 +335,18 @@ def _tiny_qwen2_5_vl(): "unknown-vlm", "encoder-decoder", "speech-encoder-decoder", + "encoder-decoder-backbone", ], ) -def test_the_class_that_builds_the_whole_model(config, trust_remote_code, expected): - """A VLM's AutoModel leaves the LM head out: it builds through the class that carries the head, - keeping the headless backbone only when no class recognizes the config. A remote class counts - only when the checkpoint's code is trusted. An encoder-decoder model's architecture reads like a - VLM's, but its seq2seq class carries its head, also for a config that names no architecture.""" - assert _auto_class(config, trust_remote_code) is expected +def test_the_class_that_builds_the_whole_model( + config, require_lm_head, trust_remote_code, expected +): + """A VLM's AutoModel leaves the LM head out: with the head required it builds through the class + that carries it, keeping the headless backbone only when no class recognizes the config. A remote + class counts only when the checkpoint's code is trusted. An encoder-decoder model's architecture + reads like a VLM's, but its seq2seq class carries its head, also for a config that names no + architecture.""" + assert _auto_class(config, require_lm_head, trust_remote_code=trust_remote_code) is expected # -------------------------------------------------------------------------------- the decision From e0b7b17ee1e0b2df2b5e5471afff91f2214279f4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:15:34 +0000 Subject: [PATCH 3/4] Build a headless model through AutoModel only when AutoModel can build it With require_lm_head=False, _auto_class returned AutoModel for every VLM-shaped config, before the remote-code checks. A trusted remote-code VLM whose auto_map registers only AutoModelForCausalLM has no AutoModel entry, and its config class is not in AutoModel's mapping, so building it raised "Unrecognized configuration class" although the same checkpoint builds with its head. AutoModel is now returned only for a native config or a trusted remote one that registers it; otherwise the model builds through the class that carries its head, as with require_lm_head=True. The docstring says the head is not required rather than absent, since some base models (Llama4's) keep it. Two table rows pin both sides: a trusted remote VLM without an AutoModel entry builds through AutoModelForCausalLM, and one with an entry through AutoModel. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 25 +++++++++++-------- .../unit/torch/utils/test_model_load_utils.py | 9 +++++++ 2 files changed, 24 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0591a92349f..b9cb961317b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -930,14 +930,16 @@ def _auto_class(hf_config, require_lm_head: bool = True, trust_remote_code: bool A causal LM builds through ``AutoModelForCausalLM``, head included: ``require_lm_head`` only applies to a model whose architecture reads as a VLM's -- a nested text config, a vision config, or a ``*ForConditionalGeneration`` architecture -- or that is an encoder-decoder model - (``is_encoder_decoder``), whether or not it names its architecture. Such a model builds through - ``AutoModel``, which leaves its LM head out, unless ``require_lm_head``: then through the - remote ``AutoModelForCausalLM`` when the checkpoint registers its ``*ForConditionalGeneration`` - there, its code is trusted and it registers no image-text class, else through its seq2seq - class if it is an encoder-decoder model, else through ``AutoModelForImageTextToText`` if that - recognizes the config or, its code trusted, the checkpoint registers one, else through - ``AutoModelForCausalLM`` if that recognizes it (as it does Phi-4-multimodal), else still - through ``AutoModel``. + (``is_encoder_decoder``), whether or not it names its architecture. Without ``require_lm_head`` + such a model builds through ``AutoModel``, its base model, when that can build it (the config + is a native one, or a trusted remote one registering ``AutoModel``): most base models leave the + head out, since it is not required, though some keep it (Llama4's). Otherwise it builds through + the remote ``AutoModelForCausalLM`` when the checkpoint registers its + ``*ForConditionalGeneration`` there, its code is trusted and it registers no image-text class, + else through its seq2seq class if it is an encoder-decoder model, else through + ``AutoModelForImageTextToText`` if that recognizes the config or, its code trusted, the + checkpoint registers one, else through ``AutoModelForCausalLM`` if that recognizes it (as it + does Phi-4-multimodal), else still through ``AutoModel``. """ architectures = getattr(hf_config, "architectures", None) or [] is_vlm = ( @@ -947,9 +949,12 @@ class if it is an encoder-decoder model, else through ``AutoModelForImageTextToT ) if not is_vlm and not getattr(hf_config, "is_encoder_decoder", False): return AutoModelForCausalLM - if not require_lm_head: - return AutoModel auto_map = getattr(hf_config, "auto_map", None) or {} + if not require_lm_head and ( + type(hf_config) in AutoModel._model_mapping + or (trust_remote_code and "AutoModel" in auto_map) + ): + return AutoModel remote_with_head = ( trust_remote_code and "conditionalgeneration" in str(auto_map.get("AutoModelForCausalLM", "")).lower() diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b2aa0ab801a..ca9e71b97f5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -323,6 +323,13 @@ def _tiny_qwen2_5_vl(): ), (WhisperConfig(), True, False, AutoModelForSpeechSeq2Seq), (T5Config(architectures=["T5ForConditionalGeneration"]), False, False, AutoModel), + (_vlm_config(auto_map=_REMOTE_VLM), False, True, AutoModelForCausalLM), + ( + _vlm_config(auto_map={**_REMOTE_VLM, "AutoModel": "m.KimiK25Model"}), + False, + True, + AutoModel, + ), ], ids=[ "causal-lm", @@ -336,6 +343,8 @@ def _tiny_qwen2_5_vl(): "encoder-decoder", "speech-encoder-decoder", "encoder-decoder-backbone", + "remote-vlm-without-a-base-class", + "remote-vlm-backbone", ], ) def test_the_class_that_builds_the_whole_model( From 158ec3425cd9f5035f37bc30deb8ffcc0524b3d2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:15:43 +0000 Subject: [PATCH 4/4] Say require_transformers refuses a pre-release of its floor on purpose A pre-release of the floor release (5.7.0.dev0 for 5.7) is refused, as PEP 440 orders it below the release; it can predate the feature the floor stands for. The docstring now says so, and drops a pointer to call sites this helper does not have yet. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index f7452a49322..563c28fea2d 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -41,7 +41,8 @@ def require_transformers(min_version: str, feature: str) -> None: Callers enforce the version where the needed APIs first exist, so a too-old transformers errors clearly instead of silently mis-behaving (e.g. transformers 5.5 silently shards nothing for EP). - See each call site for the empirically-determined floor. + A pre-release of the floor release (``5.7.0.dev0`` for ``5.7``) is refused too: it can predate + the feature the floor stands for. """ # Imported here: this module imports without transformers. import transformers