[OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head

modelopt/torch/utils/moe.py

@@ -29,12 +29,32 @@
 
 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).
+    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
+
+    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.
 

modelopt/torch/utils/plugins/model_load_utils.py

@@ -924,18 +924,22 @@ 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
+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. 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 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.
+    does Phi-4-multimodal), else still through ``AutoModel``.
     """
     architectures = getattr(hf_config, "architectures", None) or []
     is_vlm = (
@@ -946,6 +950,11 @@ 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
     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()
@@ -1027,7 +1036,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.

tests/unit/torch/utils/test_model_load_utils.py

@@ -305,20 +305,35 @@ 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),
+        (_vlm_config(auto_map=_REMOTE_VLM), False, True, AutoModelForCausalLM),
+        (
+            _vlm_config(auto_map={**_REMOTE_VLM, "AutoModel": "m.KimiK25Model"}),
+            False,
+            True,
+            AutoModel,
+        ),
     ],
     ids=[
         "causal-lm",
+        "vlm-backbone",
         "remote-vlm-with-head",
         "remote-vlm-untrusted",
         "remote-image-text-vlm",
@@ -327,14 +342,20 @@ def _tiny_qwen2_5_vl():
         "unknown-vlm",
         "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(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

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")