diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 98c4bc2453c..563c28fea2d 100644 --- a/modelopt/torch/utils/moe.py +++ b/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. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index babf7f11054..b9cb961317b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/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. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 2c225dafc14..ca9e71b97f5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/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 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")