diff --git a/conftest.py b/conftest.py index 41a90c2623cf..91a48a7a8e24 100644 --- a/conftest.py +++ b/conftest.py @@ -30,6 +30,7 @@ import _pytest import pytest +from transformers import PreTrainedModel from transformers.testing_utils import ( HfDoctestModule, HfDocTestParser, @@ -39,7 +40,7 @@ patch_testing_methods_to_collect_info, patch_torch_compile_force_graph, ) -from transformers.utils import enable_tf32 +from transformers.utils import ENV_VARS_TRUE_VALUES, enable_tf32 from transformers.utils.network_logging import register_network_debug_plugin @@ -213,6 +214,10 @@ def wrapper(*args, **kwargs): # they become normal warnings - i.e. the tests still need to test the current functionality warnings.simplefilter(action="ignore", category=FutureWarning) +# we patch `from_pretrained` temporarily to be controlled via env variables +_from_pretrained_kernels_patch = None +_original_from_pretrained = PreTrainedModel.from_pretrained.__func__ + def pytest_configure(config): # Shared directory for the read-only cache fallback events. The controller creates it and @@ -298,6 +303,32 @@ def pytest_configure(config): os.environ["DISABLE_SAFETENSORS_CONVERSION"] = "true" register_network_debug_plugin(config) + # Kernels related patching + global _from_pretrained_kernels_patch, _original_from_pretrained + os.environ["USE_HUB_KERNELS"] = "0" # Disable kernels for tests by default + + @classmethod + def from_pretrained(cls, *args, **kwargs): + # Manipulate default as per env variables + kwargs.setdefault("use_kernels", os.environ.get("USE_HUB_KERNELS", "NO").upper() in ENV_VARS_TRUE_VALUES) + return _original_from_pretrained(cls, *args, **kwargs) + + _from_pretrained_kernels_patch = mock.patch.object( + PreTrainedModel, + "from_pretrained", + from_pretrained, + ) + _from_pretrained_kernels_patch.start() + + +def pytest_unconfigure(config): + # Disable default kwargs per env for kernels activation + global _from_pretrained_kernels_patch + + if _from_pretrained_kernels_patch is not None: + _from_pretrained_kernels_patch.stop() + _from_pretrained_kernels_patch = None + def pytest_collection_modifyitems(items): for item in items: diff --git a/docker/transformers-all-latest-gpu/Dockerfile b/docker/transformers-all-latest-gpu/Dockerfile index 6e0909857677..d0ffa0f9ffde 100644 --- a/docker/transformers-all-latest-gpu/Dockerfile +++ b/docker/transformers-all-latest-gpu/Dockerfile @@ -143,9 +143,6 @@ RUN python3 -m pip install --no-cache-dir einops # For `VibeVoice` (added in PR #40546) RUN python3 -m pip install --no-cache-dir diffusers -# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs -RUN python3 -m pip uninstall -y kernels - # When installing in editable mode, `transformers` is not recognized as a package. # this line must be added in order for python to be aware of transformers. RUN cd transformers && python3 setup.py develop diff --git a/docker/transformers-pytorch-amd-gpu/Dockerfile b/docker/transformers-pytorch-amd-gpu/Dockerfile index 4fdfde3dd415..786fd98a5e38 100644 --- a/docker/transformers-pytorch-amd-gpu/Dockerfile +++ b/docker/transformers-pytorch-amd-gpu/Dockerfile @@ -42,9 +42,6 @@ RUN cd transformers && python3 setup.py develop # Remove nvml and nvidia-ml-py as it is not compatible with ROCm. apex is not tested on NVIDIA either. RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y -# `kernels` may causes many failing tests -RUN python3 -m pip uninstall -y kernels - # On ROCm, torchcodec is required to decode audio files. 0.10 matches the # torch 2.10 C++ ABI in the base image above; RUN python3 -m pip install --no-cache-dir "torchcodec==0.10" diff --git a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile index 2c087c47ea5a..267125104b5c 100644 --- a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile +++ b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile @@ -48,6 +48,3 @@ RUN python3 -c "from deepspeed.launcher.runner import main" # Remove nvml as it is not compatible with ROCm RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y - -# `kernels` may causes many failing tests -RUN python3 -m pip uninstall -y kernels diff --git a/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile index e78d7faf1db1..aef934e6c2d1 100644 --- a/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile +++ b/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile @@ -50,9 +50,6 @@ RUN python3 -m pip uninstall -y deepspeed # TODO: Find out why test fail. RUN DS_BUILD_CPU_ADAM=1 DS_BUILD_FUSED_ADAM=1 python3 -m pip install deepspeed --no-build-isolation --config-settings="--build-option=build_ext" --config-settings="--build-option=-j8" --no-cache -v --disable-pip-version-check 2>&1 -# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs -RUN python3 -m pip uninstall -y kernels - # When installing in editable mode, `transformers` is not recognized as a package. # this line must be added in order for python to be aware of transformers. RUN cd transformers && python3 setup.py develop diff --git a/docker/transformers-pytorch-deepspeed-nightly-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-nightly-gpu/Dockerfile index 5e35d639ab49..b4cc0d21a911 100644 --- a/docker/transformers-pytorch-deepspeed-nightly-gpu/Dockerfile +++ b/docker/transformers-pytorch-deepspeed-nightly-gpu/Dockerfile @@ -62,9 +62,6 @@ RUN python3 -m pip uninstall -y deepspeed #RUN git clone https://github.com/pytorch/TensorRT.git #RUN cd TensorRT/py && python3 setup.py install --fx-only -# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs -RUN python3 -m pip uninstall -y kernels - # When installing in editable mode, `transformers` is not recognized as a package. # this line must be added in order for python to be aware of transformers. RUN cd transformers && python3 setup.py develop diff --git a/docker/transformers-pytorch-gpu/Dockerfile b/docker/transformers-pytorch-gpu/Dockerfile index 61f65acb928b..05d479277b58 100644 --- a/docker/transformers-pytorch-gpu/Dockerfile +++ b/docker/transformers-pytorch-gpu/Dockerfile @@ -28,9 +28,6 @@ RUN [ ${#TORCH_AUDIO} -gt 0 ] && VERSION='torchaudio=='TORCH_AUDIO'.*' || VERSI RUN python3 -m pip install --no-cache-dir git+https://github.com/facebookresearch/detectron2.git pytesseract RUN python3 -m pip install -U "itsdangerous<2.1.0" -# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs -RUN python3 -m pip uninstall -y kernels - # When installing in editable mode, `transformers` is not recognized as a package. # this line must be added in order for python to be aware of transformers. RUN cd transformers && python3 setup.py develop diff --git a/docker/transformers-quantization-latest-gpu/Dockerfile b/docker/transformers-quantization-latest-gpu/Dockerfile index 5554a48e45d7..a43b8494ecdf 100755 --- a/docker/transformers-quantization-latest-gpu/Dockerfile +++ b/docker/transformers-quantization-latest-gpu/Dockerfile @@ -85,9 +85,6 @@ RUN python3 -m pip install --no-cache-dir torchao # Add transformers in editable mode RUN python3 -m pip install --no-cache-dir -e ./transformers[dev] -# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs -RUN python3 -m pip uninstall -y kernels - # Uninstall flash-attn installed by autoawq, it causes issues here : https://github.com/huggingface/transformers/actions/runs/15915442841/job/44892146131 RUN python3 -m pip uninstall -y flash-attn diff --git a/setup.py b/setup.py index fc24319b2b1e..ed888313a65f 100644 --- a/setup.py +++ b/setup.py @@ -175,7 +175,7 @@ def deps_list(*pkgs): extras = {} -extras["torch"] = deps_list("torch", "accelerate") +extras["torch"] = deps_list("torch", "accelerate", "kernels") extras["vision"] = deps_list("torchvision", "Pillow") extras["audio"] = deps_list("torchaudio", "librosa", "pyctcdecode", "phonemizer") if PYTHON_MINOR_VERSION < 13: @@ -186,7 +186,6 @@ def deps_list(*pkgs): "datasets", "ruff", "GitPython", "urllib3", "libcst", "rich", "ty", "tomli", "transformers-mlinter" ) extras["docs"] = deps_list("hf-doc-builder") -extras["kernels"] = deps_list("kernels") extras["sentencepiece"] = deps_list("sentencepiece", "protobuf") extras["tiktoken"] = deps_list("tiktoken", "blobfile") extras["mistral-common"] = deps_list("mistral-common[image]") @@ -249,7 +248,6 @@ def deps_list(*pkgs): + extras["vision"] + extras["audio"] + extras["video"] - + extras["kernels"] + extras["timm"] + extras["sentencepiece"] + extras["tiktoken"] diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index c1b0af1ce698..a59021deedec 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -1021,3 +1021,27 @@ def wrapped(self, *args, **kwargs): return wrapped return decorator + + +def device_map_uses_accelerator(device_map: dict | str | int | None) -> bool: + """Detect whether the device map contains any non-cpu device""" + if device_map is None: + return False + + # Single device + if not isinstance(device_map, dict): + if isinstance(device_map, int): + return True + device = torch.device(device_map) + return device.type not in {"cpu", "meta"} + + # Multi device + for target in device_map.values(): + if target == "disk": + continue + if isinstance(target, int): + return True + if torch.device(target).type not in {"cpu", "meta"}: + return True + + return False diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index b58adc35c063..847b178ee00a 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -21,12 +21,12 @@ from contextlib import contextmanager from pathlib import Path from types import ModuleType -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from ..conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping +from ..modeling_flash_attention_utils import FLASH_ATTN_KERNEL_VERSIONS from ..monkey_patching import register_patch_mapping from ..utils import ENV_VARS_TRUE_VALUES, logging -from ..utils.generic import is_flash_attention_requested from ..utils.import_utils import ( KERNELS_MAX_VERSION, KERNELS_MIN_VERSION, @@ -58,8 +58,17 @@ ) -_TRANSFORMERS_USE_HUB_KERNELS = os.environ.get("USE_HUB_KERNELS", "YES").upper() -_kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES +def _kernels_enabled() -> bool: + return os.environ.get("USE_HUB_KERNELS", "YES").upper() in ENV_VARS_TRUE_VALUES + + +def _ensure_kernels_enabled() -> None: + if not _kernels_enabled(): + value = os.environ.get("USE_HUB_KERNELS", "YES") + raise ValueError( + "Hub kernels are disabled through the environment variable " + f"`USE_HUB_KERNELS={value}`. Set `USE_HUB_KERNELS=YES` to enable them." + ) # Maps from func name to the internal module path @@ -98,37 +107,21 @@ ) from kernels import use_kernelized_func as _kernels_use_kernelized_func + # We change the strategy to prepare the metadata like kernels but the exchange + # at runtime can be determined by the env variable `USE_HUB_KERNELS` def use_kernel_forward_from_hub(layer_name: str): - if _kernels_enabled: - return _kernels_use_kernel_forward_from_hub(layer_name) - else: - logger.warning_once( - f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}" - ) - return lambda cls: cls + return _kernels_use_kernel_forward_from_hub(layer_name) def use_kernelized_func(module_names: list[Callable] | Callable): - if _kernels_enabled: - if isinstance(module_names, Callable): - module_names = [module_names] - return _kernels_use_kernelized_func(*module_names) - else: - logger.warning_once( - f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}" - ) - return lambda cls: cls + if isinstance(module_names, Callable): + module_names = [module_names] + return _kernels_use_kernelized_func(*module_names) def use_kernel_func_from_hub(layer_name: str): - if _kernels_enabled: - logger.warning_once( - "`use_kernel_func_from_hub` is deprecated in transformers v5.16 and will be removed in the future. Please use `use_kernel_forward_from_hub` instead." - ) - return _kernels_use_kernel_forward_from_hub(layer_name) - else: - logger.warning_once( - f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}" - ) - return lambda cls: cls + logger.warning_once( + "`use_kernel_func_from_hub` is deprecated in transformers v5.16 and will be removed in the future. Please use `use_kernel_forward_from_hub` instead." + ) + return _kernels_use_kernel_forward_from_hub(layer_name) # The default kernel mapping is built lazily (see `get_kernel_mapping_transformers`) so that simply # importing transformers (or `transformers.pipeline`) does not instantiate any `LayerRepository` / @@ -556,14 +549,13 @@ def get_kernel_mapping_transformers() -> dict: return _KERNEL_MAPPING_CACHE def register_kernel_mapping_transformers(mapping=None): + _ensure_kernels_enabled() if mapping is None: mapping = get_kernel_mapping_transformers() register_kernel_mapping(mapping) else: - _kernels_enabled = False - - # Stub to make decorators in transformers work when `kernels` + # Stub to make decorators int transformers work when `kernels` # is not installed. def use_kernel_forward_from_hub(*args, **kwargs): def decorator(cls): @@ -624,16 +616,6 @@ def register_kernel_mapping_transformers(*args, **kwargs): "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, } -# Flash attention version -> major version of its hub kernel repo. Flash attention flavors that are not -# listed here, and all other attention kernels, use `_DEFAULT_ATTN_KERNEL_VERSION`. -_FLASH_ATTN_KERNEL_VERSION_MAPPING: dict[int, int] = { - # v3 is the first version shipping the Torch stable ABI (CUDA/ROCm) and Torch 2.13 builds (incl. XPU) - 2: 3, - # FA4 is still in beta -> only v0 has been released - 4: 0, -} -_DEFAULT_ATTN_KERNEL_VERSION = 1 - _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {} @@ -647,10 +629,7 @@ def is_kernel(attn_implementation: str | None) -> bool: def get_attn_kernel_version(repo_id: str) -> int: """Return the major version of the hub kernel repo `repo_id` to load, e.g. `3` for `kernels-community/flash-attn2`.""" - for flash_attn_version, kernel_version in _FLASH_ATTN_KERNEL_VERSION_MAPPING.items(): - if is_flash_attention_requested(requested_attention_implementation=repo_id, version=flash_attn_version): - return kernel_version - return _DEFAULT_ATTN_KERNEL_VERSION + return FLASH_ATTN_KERNEL_VERSIONS.get(repo_id, 1) def load_and_register_attn_kernel( @@ -692,9 +671,7 @@ def load_and_register_attn_kernel( # create revision xor version rev = rev.strip() if rev else None - version = None - if rev is None: - version = get_attn_kernel_version(repo_id) + version = get_attn_kernel_version(repo_id) if rev is None else None # Load the kernel from hub try: @@ -736,7 +713,7 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ logger.warning_once(f"Kernel {kernel_name} not found in _HUB_KERNEL_MAPPING") mapping[kernel_name] = None return None - if is_kernels_available() and _kernels_enabled: + if is_kernels_available() and _kernels_enabled(): try: repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"] revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None) @@ -779,10 +756,11 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ return mapping[kernel_name] -def kernelize(model: "PreTrainedModel", mode: "Mode | None" = None): +def kernelize(model: "PreTrainedModel", mode: "Mode | None" = None, kernel_config: Optional["KernelConfig"] = None): """Temporarily register hidden kernel wrappers so `kernelize` can discover and replace them.""" if not is_kernels_available(): raise ImportError(_MISSING_KERNELS_MESSAGE) + _ensure_kernels_enabled() def get_device(device_type): if device_type == "cuda" and is_rocm_platform(): @@ -792,9 +770,10 @@ def get_device(device_type): mode = Mode.INFERENCE if not model.training else Mode.TRAINING if mode is None else mode device = get_device(model.device.type) - if model.kernel_config is not None: - inherit_mapping = not model.kernel_config.use_local_kernel - with use_kernel_mapping(model.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping): + used_kernel_config = kernel_config or model.kernel_config + if used_kernel_config is not None: + inherit_mapping = not used_kernel_config.use_local_kernel and used_kernel_config.inherit_mapping + with use_kernel_mapping(used_kernel_config.kernel_mapping, inherit_mapping=inherit_mapping): _kernels_kernelize(model, device=device, mode=mode) else: _kernels_kernelize(model, device=device, mode=mode) diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index f6b261c37665..de9c5b6bc14f 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -70,6 +70,14 @@ def is_flash_attn_available(): "flash_attention_4": "kernels-community/flash-attn4", } +FLASH_ATTN_KERNEL_VERSIONS = { + "kernels-community/flash-attn2": 3, + "kernels-community/flash-attn3": 1, + "kernels-community/vllm-flash-attn3": 1, + "kernels-community/aiter-flash-attn": 1, + "kernels-community/flash-attn4": 0, + "kernels-community/metal-flash-sdpa": 1, +} # Meta information on each mainline FA compatibility: # 1. The import structure and availability diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index ded5825dddc1..ca32925b26da 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -71,6 +71,7 @@ accelerate_disk_offload, accelerate_dispatch, check_and_set_device_map, + device_map_uses_accelerator, expand_device_map, get_device, load_offloaded_parameter, @@ -81,7 +82,7 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward -from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel, kernelize +from .integrations.hub_kernels import _kernels_enabled, allow_all_hub_kernels, is_kernel, kernelize from .integrations.moe import ALL_EXPERTS_FUNCTIONS from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward @@ -3894,6 +3895,30 @@ def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None else: self._use_kernels = False + def reset_kernels(self): + """ + Reset any kernelization applied on the model, i.e. use the torch implementation for any function or module + attached with kernels specific decorators (`use_kernel_forward_from_hub`, `use_kernelized_func`). + """ + if not is_kernels_available(): + return + + from kernels import Mode + + with warnings.catch_warnings(): + # Temporarily ignore user warnings as this is intentional from our side + warnings.filterwarnings( + "ignore", + message=r"\s*No kernel mapping found for layer .*", + category=UserWarning, + module=r"kernels\.layer\.layer", + ) + + # Force kernelization with an empty mapping to force the torch fallbacks in all cases + kernelize(self, mode=Mode.FALLBACK, kernel_config=KernelConfig(kernel_mapping={}, inherit_mapping=False)) + + self._use_kernels = False + @classmethod def from_pretrained( cls: type[SpecificPreTrainedModelType], @@ -4164,7 +4189,7 @@ def from_pretrained( tp_size = kwargs.pop("tp_size", None) trust_remote_code = kwargs.pop("trust_remote_code", None) allow_all_kernels = kwargs.pop("allow_all_kernels", False) - use_kernels = kwargs.pop("use_kernels", False) + use_kernels = kwargs.pop("use_kernels", True) kernel_config = kwargs.pop("kernel_config", None) key_mapping = kwargs.pop("key_mapping", None) @@ -4301,12 +4326,6 @@ def from_pretrained( "loaded from GGUF files." ) - if kernel_config is not None and not use_kernels: - logger.warning_once( - "A kernel_config was provided but use_kernels is False; setting use_kernels=True automatically. To suppress this warning, explicitly set use_kernels to True." - ) - use_kernels = True - checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files( pretrained_model_name_or_path=pretrained_model_name_or_path, variant=variant, @@ -4351,8 +4370,15 @@ def from_pretrained( register_fusion_patches(cls, config, fusion_config) + if use_kernels and not _kernels_enabled(): + logger.warning_once( + "Detected the usage of `use_kernels=True` but the env variable `USE_HUB_KERNELS` is set to turn kernels off. " + "We set `use_kernels=False` to follow the priority of the env varibale. Please consider passing `use_kernels=False` as well." + ) + use_kernels = False + # Kernel patches: single-layer replacement (stateful __init__) then fusions. - if kernel_config is not None and use_kernels: + if use_kernels and kernel_config is not None: from .integrations.hub_kernels import register_kernel_replacements_and_fusions # For remote kernels, we need to apply the context manager @@ -4389,6 +4415,12 @@ def from_pretrained( if device_map is not None: device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) + if device_map_uses_accelerator(device_map) and not use_kernels: + logger.warning_once( + "We detected the usage of an accelerator (e.g. cuda) but kernels is turned off (`use_kernels=False`). " + "Consider turning on kernels to gain maximum speedups!" + ) + # Finalize model weight initialization load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 3ee8ab29e667..d4655fce47bc 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -4822,6 +4822,72 @@ def new_save(self, save_directory, *args, **kwargs): PreTrainedModel.save_pretrained = original_save +def scoped_kernels(test): + """ + Decorator that treats a kernels test as isolated instance in which we remove the kernelization + side effects from the models (as they are attached on a class level). + + We guarantee this in any case even if the test fails. + """ + _MISSING = object() + + @require_kernels + @functools.wraps(test) + def wrapper(*args, **kwargs): + from kernels import use_kernel_mapping + from kernels.layer import layer as kernel_layer + + changed = {} + original_replace = kernel_layer._replace_forward + + def tracked_replace(module, layer): + changed.setdefault( + id(module), + (module, module.__dict__.get("forward", _MISSING)), + ) + original_replace(module, layer) + + with ( + use_kernel_mapping({}, inherit_mapping=True), + patch.object(kernel_layer, "_replace_forward", tracked_replace), + patch.dict("os.environ", {"USE_HUB_KERNELS": "YES"}), + ): + try: + return test(*args, **kwargs) + finally: + for module, forward in reversed(list(changed.values())): + if forward is _MISSING: + module.__dict__.pop("forward", None) + else: + module.__dict__["forward"] = forward + + # Mark to avoid class level usage duplication + wrapper._is_scoped_kernels = True + + return wrapper + + +def scoped_kernels_class(test_class): + """ + Applies `scoped_kernels` on each test function individually, i.e. each kernelize gets a fresh state. + """ + for name in dir(test_class): + if not name.startswith("test"): + continue + + test = getattr(test_class, name) + + if not callable(test): + continue + + if getattr(test, "_is_scoped_kernels", False): + continue + + setattr(test_class, name, scoped_kernels(test)) + + return test_class + + @contextmanager def preserve_module_forwards(model: "PreTrainedModel"): """ diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index 0b469286f669..58e807f511bd 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -101,10 +101,11 @@ class KernelConfig(PushToHubMixin): Kernel configuration class. This class is used to configure the kernel mapping for a model. """ - def __init__(self, kernel_mapping=None, use_local_kernel=False): + def __init__(self, kernel_mapping=None, use_local_kernel=False, inherit_mapping=True): self.kernel_mapping = kernel_mapping if kernel_mapping is not None else {} self.registered_layer_names = {} self.use_local_kernel = use_local_kernel + self.inherit_mapping = inherit_mapping def update_kernel( self, repo_id, registered_name, layer_name, device, mode, revision=None, version=1, trust_remote_code=False diff --git a/tests/alm_tester.py b/tests/alm_tester.py index 4c05751f564b..ddc5190bdba3 100644 --- a/tests/alm_tester.py +++ b/tests/alm_tester.py @@ -17,6 +17,8 @@ from inspect import signature from unittest.mock import patch +from transformers.testing_utils import scoped_kernels_class + from .multimodal_tester import MultiModalModelTest, MultiModalModelTester from .test_modeling_common import ( floats_tensor, @@ -148,6 +150,7 @@ def get_audio_config(self): return self.audio_config_class(**kwargs) +@scoped_kernels_class class ALMModelTest(MultiModalModelTest): """ Base test class for Audio-Language Models. diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py index 1ba5c3bf9a0d..1b8d23194ee4 100644 --- a/tests/causal_lm_tester.py +++ b/tests/causal_lm_tester.py @@ -25,6 +25,7 @@ is_flaky, require_flash_attn, require_torch_accelerator, + scoped_kernels_class, slow, ) @@ -270,6 +271,7 @@ def prepare_config_and_inputs_for_common(self): @require_torch +@scoped_kernels_class class CausalLMModelTest( ModelTesterMixin, GenerationTesterMixin, diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 507caf04f65e..076911de9190 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -43,6 +43,7 @@ require_kernels, require_rocm, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -56,8 +57,9 @@ import transformers.integrations.hub_kernels as hub_kernels_pkg -@require_kernels @slow +@require_torch_accelerator +@scoped_kernels class TestHubKernels(TestCasePlus): @classmethod def setUpClass(cls): @@ -398,7 +400,7 @@ def test_disable_hub_kernels(self): try: with patch.dict(os.environ, {"USE_HUB_KERNELS": "OFF"}): importlib.reload(hub_kernels_pkg) - self.assertFalse(hub_kernels_pkg._kernels_enabled) + self.assertFalse(hub_kernels_pkg._kernels_enabled()) finally: hub_kernels_pkg.__dict__.clear() hub_kernels_pkg.__dict__.update(original_state) @@ -411,7 +413,7 @@ def test_enable_hub_kernels(self): try: with patch.dict(os.environ, {"USE_HUB_KERNELS": "ON"}): importlib.reload(hub_kernels_pkg) - self.assertTrue(hub_kernels_pkg._kernels_enabled) + self.assertTrue(hub_kernels_pkg._kernels_enabled()) finally: hub_kernels_pkg.__dict__.clear() hub_kernels_pkg.__dict__.update(original_state) @@ -439,6 +441,8 @@ def test_is_kernel_regex(self): for s in invalid: self.assertFalse(is_kernel(s)) + @require_torch_accelerator + @scoped_kernels def test_lazy_load_kernel_success_and_cache(self): sentinel = types.ModuleType("sentinel_kernel_module") @@ -447,12 +451,19 @@ def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=Fals self.assertFalse(allow_all_kernels) return sentinel + patched_hub_mapping = copy.deepcopy(_HUB_KERNEL_MAPPING) + patched_hub_mapping["causal-conv1d"] = { + "repo_id": "kernels-community/causal-conv1d", + "version": 1, + } + patched_module_mapping = copy.copy(_KERNEL_MODULE_MAPPING) patched_module_mapping.pop("causal-conv1d", None) with patch.dict( lazy_load_kernel.__globals__, { + "_HUB_KERNEL_MAPPING": patched_hub_mapping, "_KERNEL_MODULE_MAPPING": patched_module_mapping, "get_kernel": fake_get_kernel, "ALLOW_ALL_KERNELS": False, @@ -473,6 +484,8 @@ def test_lazy_load_kernel_unknown(self): # Cleanup cache entry to avoid growth across tests _KERNEL_MODULE_MAPPING.pop(name, None) + @require_torch_accelerator + @scoped_kernels def test_lazy_load_kernel_version(self): name = "causal-conv1d" version_spec = ">=0.0.4,<0.1.0" @@ -652,7 +665,8 @@ def test_add_to_mapping_local_then_load(self): self.assertTrue(issubclass(layer_cls, torch.nn.Module)) -@require_kernels +@require_torch_accelerator +@scoped_kernels class TestUseKernelsLifecycle(TestCasePlus): @classmethod def setUpClass(cls): diff --git a/tests/models/bamba/test_modeling_bamba.py b/tests/models/bamba/test_modeling_bamba.py index 9c7c0e1640b4..3300293ec272 100644 --- a/tests/models/bamba/test_modeling_bamba.py +++ b/tests/models/bamba/test_modeling_bamba.py @@ -33,9 +33,9 @@ get_device_properties, require_deterministic_for_xpu, require_flash_attn, - require_kernels, require_torch, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -274,6 +274,10 @@ def create_and_check_mamba_chunked_prefill(self, config, input_ids, *args, devic model.to(device) model.eval() + # Enable kernels path + if device != "cpu": + model.use_kernels = True + input_ids = input_ids[:1].to(device) prefill_len = input_ids.shape[1] // 2 + 1 prompt = input_ids[:, :prefill_len] @@ -351,7 +355,7 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device=torch_device) diff --git a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py index 1836894231c9..5ef37c531703 100644 --- a/tests/models/deepseek_v3/test_modeling_deepseek_v3.py +++ b/tests/models/deepseek_v3/test_modeling_deepseek_v3.py @@ -24,6 +24,7 @@ require_torch, require_torch_accelerator, require_torch_large_accelerator, + scoped_kernels_class, slow, torch_device, ) @@ -213,6 +214,7 @@ def prepare_config_and_inputs_for_common(self): @require_torch +@scoped_kernels_class class DeepseekV3ModelTest( ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase, TensorParallelTesterMixin ): diff --git a/tests/models/falcon_h1/test_modeling_falcon_h1.py b/tests/models/falcon_h1/test_modeling_falcon_h1.py index 0e0670f8c2ad..e502b76f7ded 100644 --- a/tests/models/falcon_h1/test_modeling_falcon_h1.py +++ b/tests/models/falcon_h1/test_modeling_falcon_h1.py @@ -21,9 +21,9 @@ from transformers import DynamicCache, FalconH1Config, is_torch_available from transformers.testing_utils import ( Expectations, - require_kernels, require_torch, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -252,6 +252,10 @@ def create_and_check_mamba_chunked_prefill(self, config, input_ids, *args, devic model.to(device) model.eval() + # Enable kernels path + if device != "cpu": + model.use_kernels = True + input_ids = input_ids[:1].to(device) prefill_len = input_ids.shape[1] // 2 + 1 prompt = input_ids[:, :prefill_len] @@ -327,7 +331,7 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device=torch_device) diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 8e6605485196..b7d305b8ba20 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -36,11 +36,10 @@ get_cpu_ram_total_gib, is_kernels_available, require_deterministic_for_xpu, - require_kernels, require_torch, require_torch_accelerator, - require_torch_gpu, require_torch_multi_accelerator, + scoped_kernels, slow, torch_device, ) @@ -90,19 +89,9 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase): model_split_percents = [0.5, 0.6] model_tester_class = GptOssModelTester - @require_kernels - @require_torch_accelerator - def test_kernelize_does_not_crash(self): - """Regression test #45799 and #46619: `kernelize` should not crash with `use_kernelized_func` + `use_kernel_forward_from_hub`.""" - config, _ = self.model_tester.prepare_config_and_inputs_for_common() - model = GptOssModel(config).to(device=torch_device) - # This used to raise TypeError because apply_rotary_pos_emb was not wrapped as nn.Module - # and also because a stale `position_ids` arg made its signature mismatch the hub rotary kernel - model.set_use_kernels(True) - - @require_kernels @pytest.mark.flash_attn_test - @require_torch_gpu + @require_torch_accelerator + @scoped_kernels def test_default_flash_implementation_auto_correction(self): """ Tests that setting attn_implementation="flash_attention_2" during model initialization @@ -426,8 +415,8 @@ def run_distributed_test(quantized, model, kernels, attn_impl, mode): # Non-distributed test # ------------------------ @parameterized.expand(PARAMETERS) - @require_kernels @require_deterministic_for_xpu + @scoped_kernels def test_model_outputs(self, quantized, model, kernels, attn_impl, mode): if torch_device == "cpu": if attn_impl == "kernels-community/vllm-flash-attn3": diff --git a/tests/models/granitemoehybrid/test_modeling_granitemoehybrid.py b/tests/models/granitemoehybrid/test_modeling_granitemoehybrid.py index 9e7c2843bf46..391a8f936fc7 100644 --- a/tests/models/granitemoehybrid/test_modeling_granitemoehybrid.py +++ b/tests/models/granitemoehybrid/test_modeling_granitemoehybrid.py @@ -30,9 +30,9 @@ ) from transformers.testing_utils import ( require_flash_attn, - require_kernels, require_torch, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -130,7 +130,7 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_mamba_chunked_prefill(*config_and_inputs, device=torch_device) diff --git a/tests/models/jamba/test_modeling_jamba.py b/tests/models/jamba/test_modeling_jamba.py index 04c6ac8ec49f..ef9196d9f428 100644 --- a/tests/models/jamba/test_modeling_jamba.py +++ b/tests/models/jamba/test_modeling_jamba.py @@ -105,7 +105,7 @@ def __init__( num_attention_heads=2, num_key_value_heads=2, intermediate_size=40, - hidden_act="gelu", + hidden_act="silu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, diff --git a/tests/models/mamba2/test_modeling_mamba2.py b/tests/models/mamba2/test_modeling_mamba2.py index 329c7754c430..3eed04a17517 100644 --- a/tests/models/mamba2/test_modeling_mamba2.py +++ b/tests/models/mamba2/test_modeling_mamba2.py @@ -18,9 +18,9 @@ from transformers import AutoTokenizer, Mamba2Config, is_torch_available from transformers.testing_utils import ( Expectations, - require_kernels, require_torch, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -217,6 +217,10 @@ def create_and_check_mamba2_chunked_prefill(self, config, input_ids, *args, devi model.to(device) model.eval() + # Enable kernels path + if device != "cpu": + model.use_kernels = True + input_ids = input_ids[:1].to(device) prefill_len = input_ids.shape[1] // 2 + 1 prompt = input_ids[:, :prefill_len] @@ -250,8 +254,9 @@ def create_and_check_mamba2_slow_vs_fast_forward(self, config, input_ids, *args, model.gradient_checkpointing_enable() token_emb = model.embeddings(input_ids) - outputs_fast = model.layers[0].mixer.cuda_kernels_forward(token_emb) - outputs_slow = model.layers[0].mixer.torch_forward(token_emb) + outputs_slow = model.layers[0].mixer(token_emb) + model.use_kernels = True + outputs_fast = model.layers[0].mixer(token_emb) self.parent.assertTrue(torch.allclose(outputs_fast, outputs_slow, atol=1e-3, rtol=1e-3)) @@ -293,13 +298,13 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_mamba2_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_mamba2_chunked_prefill(*config_and_inputs, device=torch_device) @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_slow_vs_fast_forward(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_mamba2_slow_vs_fast_forward(*config_and_inputs) @@ -308,7 +313,7 @@ def test_mamba2_slow_vs_fast_forward(self): # creates a grouped SSD configuration in the mamba2 layers # See https://github.com/huggingface/transformers/pull/37533/ @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_slow_vs_fast_forward_grouped(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() config_and_inputs[0].n_groups //= 2 diff --git a/tests/models/nemotron_h/test_modeling_nemotron_h.py b/tests/models/nemotron_h/test_modeling_nemotron_h.py index 50874d5fd55a..931b09f20a4b 100644 --- a/tests/models/nemotron_h/test_modeling_nemotron_h.py +++ b/tests/models/nemotron_h/test_modeling_nemotron_h.py @@ -23,10 +23,10 @@ from transformers.testing_utils import ( require_bitsandbytes, require_flash_attn, - require_kernels, require_torch, require_torch_accelerator, require_torch_greater_or_equal, + scoped_kernels, slow, torch_device, ) @@ -332,6 +332,10 @@ def create_and_check_nemotron_h_chunked_prefill(self, config, input_ids, *args, model.to(device) model.eval() + # Enable kernels path + if device != "cpu": + model.use_kernels = True + input_ids = input_ids[:1].to(device) prefill_len = input_ids.shape[1] // 2 + 1 prompt = input_ids[:, :prefill_len] @@ -465,7 +469,7 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_nemotron_h_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_nemotron_h_chunked_prefill(*config_and_inputs, device=torch_device) diff --git a/tests/models/olmoe/test_modeling_olmoe.py b/tests/models/olmoe/test_modeling_olmoe.py index 65452ed7346c..777536bd89ec 100644 --- a/tests/models/olmoe/test_modeling_olmoe.py +++ b/tests/models/olmoe/test_modeling_olmoe.py @@ -21,6 +21,7 @@ from transformers.testing_utils import ( require_tokenizers, require_torch, + scoped_kernels_class, slow, torch_device, ) @@ -176,6 +177,7 @@ def prepare_config_and_inputs_for_common(self): @require_torch +@scoped_kernels_class class OlmoeModelTest( ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase ): diff --git a/tests/models/perceiver/test_modeling_perceiver.py b/tests/models/perceiver/test_modeling_perceiver.py index 7efcbe262c1c..497f60b0e5d1 100644 --- a/tests/models/perceiver/test_modeling_perceiver.py +++ b/tests/models/perceiver/test_modeling_perceiver.py @@ -27,8 +27,10 @@ from transformers.testing_utils import ( IS_ROCM_SYSTEM, require_torch, + require_torch_accelerator, require_torch_multi_gpu, require_vision, + scoped_kernels, slow, torch_device, ) @@ -816,6 +818,28 @@ def test_problem_types(self): loss.backward() + @require_torch_accelerator + @scoped_kernels + def test_kernels_can_run_without_crashing(self): + """Overriden to accomodate unique input preparation""" + + for model_class in self.all_model_classes: + with self.subTest(model_class=model_class.__name__): + config, inputs = self.model_tester.prepare_config_and_inputs_for_model_class(model_class) + + model = model_class(config).to(torch_device) + model.eval() + model.use_kernels = True + + prepared_inputs = self._prepare_for_class(inputs, model_class) + prepared_inputs = { + key: value.to(torch_device) if isinstance(value, torch.Tensor) else value + for key, value in prepared_inputs.items() + } + + with torch.no_grad(): + model(**prepared_inputs) + @require_torch_multi_gpu @unittest.skip( reason=( diff --git a/tests/models/wav2vec2/test_modeling_wav2vec2.py b/tests/models/wav2vec2/test_modeling_wav2vec2.py index 1d1fddad6051..76b46d259157 100644 --- a/tests/models/wav2vec2/test_modeling_wav2vec2.py +++ b/tests/models/wav2vec2/test_modeling_wav2vec2.py @@ -98,9 +98,9 @@ def _test_wav2vec2_with_lm_invalid_pool(in_queue, out_queue, timeout): torch.tensor(sample["audio"]["array"]), 48_000, 16_000 ).numpy() - model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm").to( - torch_device - ) + model = Wav2Vec2ForCTC.from_pretrained( + "patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm", use_kernels=False + ).to(torch_device) processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm") input_values = processor(resampled_audio, return_tensors="pt").input_values diff --git a/tests/models/zamba2/test_modeling_zamba2.py b/tests/models/zamba2/test_modeling_zamba2.py index 4700b9aa2b9f..abd8f3308789 100644 --- a/tests/models/zamba2/test_modeling_zamba2.py +++ b/tests/models/zamba2/test_modeling_zamba2.py @@ -24,9 +24,9 @@ Expectations, require_bitsandbytes, require_flash_attn, - require_kernels, require_torch, require_torch_accelerator, + scoped_kernels, slow, torch_device, ) @@ -274,6 +274,10 @@ def create_and_check_zamba2_chunked_prefill(self, config, input_ids, *args, devi model.to(device) model.eval() + # Enable kernels path + if device != "cpu": + model.use_kernels = True + input_ids = input_ids[:1].to(device) prefill_len = input_ids.shape[1] // 2 + 1 prompt = input_ids[:, :prefill_len] @@ -405,7 +409,7 @@ def test_mamba2_chunked_prefill_cpu(self): self.model_tester.create_and_check_zamba2_chunked_prefill(*config_and_inputs, device="cpu") @require_torch_accelerator - @require_kernels + @scoped_kernels def test_mamba2_chunked_prefill_torch_device(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_zamba2_chunked_prefill(*config_and_inputs, device=torch_device) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index f305b8d14967..eba74d5fc122 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -154,13 +154,18 @@ def _deterministic_init_model_dir(rank, config, dtype): yield model_dir -def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file): +def _fsdp_global_wrapper( + rank, test_name, func, func_args, func_kwargs, world_size, port, results_file, use_hub_kernels +): os.environ["WORLD_SIZE"] = str(world_size) os.environ["RANK"] = str(rank) os.environ["LOCAL_RANK"] = str(rank) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(port) + # Kernelize is propogated via envs so we pass it to the child as well + os.environ["USE_HUB_KERNELS"] = use_hub_kernels + _set_determinism(SEED) dist.init_process_group(backend=_get_distributed_backend(), rank=rank, world_size=world_size) _set_rank_device(rank) @@ -567,6 +572,7 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k config_class, config_dict = self._get_tiny_config() func_args = (config_class, config_dict, *test_args) + use_hub_kernels = os.environ.get("USE_HUB_KERNELS", "NO") results_file = tempfile.mktemp(suffix=".json") # port binding with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -576,7 +582,16 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k try: mp.spawn( _fsdp_global_wrapper, - args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file), + args=( + test_name, + test_impl, + func_args, + test_kwargs, + self.fsdp_nproc_per_node, + port, + results_file, + use_hub_kernels, + ), nprocs=self.fsdp_nproc_per_node, ) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 7b0e294fda2d..ca5bdff92e6e 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -46,7 +46,7 @@ from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import PrefixChange, WeightRenaming, process_target_pattern from transformers.integrations import HfDeepSpeedConfig -from transformers.integrations.deepgemm import _get_nvcc_version +from transformers.integrations.deepgemm import is_deepgemm_loadable from transformers.integrations.deepspeed import ( is_deepspeed_available, is_deepspeed_zero3_enabled, @@ -58,6 +58,7 @@ grouped_mm_experts_forward, sonicmoe_experts_forward, ) +from transformers.integrations.sonicmoe import is_sonicmoe_loadable from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, _get_tied_weight_keys from transformers.models.auto import get_values @@ -91,7 +92,6 @@ get_device_properties, hub_retry, is_flaky, - preserve_module_forwards, require_accelerate, require_bitsandbytes, require_deepspeed, @@ -109,6 +109,7 @@ rocm_has_sdpa_flash_backend, run_first, run_test_using_subprocess, + scoped_kernels, set_config_for_less_flaky_test, set_model_for_less_flaky_test, slow, @@ -119,7 +120,6 @@ GENERATION_CONFIG_NAME, SAFE_WEIGHTS_NAME, ModelOutput, - is_kernels_available, is_torch_bf16_available_on_device, is_torch_fp16_available_on_device, ) @@ -603,25 +603,14 @@ def _test_eager_matches_batched_and_grouped_inference(self, name, dtype): "grouped_mm": Mock(wraps=grouped_mm_experts_forward), } - if ( - dtype != torch.float32 - and is_kernels_available() - and torch.cuda.is_available() - and torch.cuda.get_device_capability() >= (9, 0) - ): - # we also need nvidia-cutlass-dsl and apache-tvm-ffi + # Kernels based implementations that need specific requirements, please see `is_xxx_loadable` + # for more information, e.g. which hardware, nvcc, etc. + if dtype != torch.float32 and is_sonicmoe_loadable(): mocks["sonicmoe"] = Mock(wraps=sonicmoe_experts_forward) implementations.append("sonicmoe") - nvcc_version = _get_nvcc_version() or (0, 0) - device_major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0 - # DeepGEMM ships kernels only for Hopper (SM90, needs nvcc 12.3+) and Blackwell (SM100, needs 12.9+). - if ( - dtype == torch.bfloat16 - and is_kernels_available() - and ((device_major == 9 and nvcc_version >= (12, 3)) or (device_major == 10 and nvcc_version >= (12, 9))) - ): - # DeepGEMM BF16 grouped forward requires Hopper+, a new-enough nvcc toolkit, and bf16 hidden states + # TODO(@Ilyas) deepgemm needs % 64 weights otherwise it is rejected at runtime + if dtype == torch.bfloat16 and is_deepgemm_loadable(): mocks["deepgemm"] = Mock(wraps=deepgemm_bf16_experts_forward) implementations.append("deepgemm") @@ -5814,19 +5803,38 @@ def test_format_of_can_record_outputs(self): is_valid_recorder = isinstance(recorder, (str, type, OutputRecorder)) self.assertTrue(is_valid_recorder, f"Invalid recorder: {recorder}") - @require_kernels @require_torch_accelerator + @scoped_kernels def test_kernels_can_load_without_crashing(self): """Check whether activating kernels leads to an (value) error""" config, _ = self.model_tester.prepare_config_and_inputs_for_common() for model_class in self.all_model_classes: model = model_class(config).to(torch_device) + # Most kernels are inference mode so ensure to enable as much as we can + model.eval() + + # Using kernels should not raise a `ValueError` + model.use_kernels = True + + @require_torch_accelerator + @scoped_kernels + def test_kernels_can_run_without_crashing(self): + """Check whether activating kernels and then running through some input leads to an (value) error""" + config, inputs = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config).to(torch_device) + # Most kernels are inference mode so ensure to enable as much as we can + model.eval() + model.use_kernels = True + + prepared_inputs = self._prepare_for_class(inputs, model_class) + prepared_inputs = { + k: v.to(torch_device) if isinstance(v, torch.Tensor) else v for k, v in prepared_inputs.items() + } - # `kernelize` mutates module-level singletons, so restore them to keep later tests kernel-free - with preserve_module_forwards(model): - # Using kernels should not raise a `ValueError` - model.use_kernels = True + model(**prepared_inputs) @parameterized.expand([("linear",), ("dynamic",), ("yarn",)]) def test_model_rope_scaling_from_config(self, scaling_type): diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 9b7f2e53387e..728a87c34fa2 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -101,8 +101,8 @@ def get_packed_grad_shard(grad, world_size, rank, dim): return grad.index_select(dim, torch.tensor(indices, device=grad.device)) -def _global_wrapper(rank, func, tp, port, backend, func_args, func_kwargs): - """Wrapper to set up distributed environment and run the test function.""" +def _global_wrapper(rank, func, tp, port, backend, func_args, func_kwargs, use_hub_kernels): + """Wrapper to set up (distributed) environment and run the test function.""" def setup_dist_env(rank, world_size, port): os.environ["WORLD_SIZE"] = str(world_size) @@ -111,6 +111,9 @@ def setup_dist_env(rank, world_size, port): os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(port) + # Kernelize is propogated via envs so we pass it to the child as well + os.environ["USE_HUB_KERNELS"] = use_hub_kernels + world_size = tp setup_dist_env(rank, world_size, port) @@ -123,14 +126,15 @@ def setup_dist_env(rank, world_size, port): def _init_distributed(tp: int, max_retries: int = 5, backend: str = "gloo"): - """Decorator to initialize distributed environment and spawn processes.""" + """Decorator to initialize (distributed) environment and spawn processes.""" def _init_distributed_inner(func): def wrapper(*args, **kwargs): world_size = tp for attempt in range(max_retries): port = _find_free_port() - spawn_args = (func, tp, port, backend, args, kwargs) + use_hub_kernels = os.environ.get("USE_HUB_KERNELS", "NO") + spawn_args = (func, tp, port, backend, args, kwargs, use_hub_kernels) try: mp.spawn(_global_wrapper, args=spawn_args, nprocs=world_size) return diff --git a/tests/utils/test_offline.py b/tests/utils/test_offline.py index 20f1690bb719..c8b6593de219 100644 --- a/tests/utils/test_offline.py +++ b/tests/utils/test_offline.py @@ -38,9 +38,9 @@ def test_offline_mode(self): run = """ mname = "hf-internal-testing/tiny-random-bert" BertConfig.from_pretrained(mname) -BertModel.from_pretrained(mname) +BertModel.from_pretrained(mname, use_kernels=False) BertTokenizer.from_pretrained(mname) -pipe = pipeline(task="fill-mask", model=mname) +pipe = pipeline(task="fill-mask", model=mname, model_kwargs={"use_kernels": False}) print("success") """ @@ -53,7 +53,7 @@ def offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled # Force fetching the files so that we can use the cache mname = "hf-internal-testing/tiny-random-bert" BertConfig.from_pretrained(mname) - BertModel.from_pretrained(mname) + BertModel.from_pretrained(mname, use_kernels=False) BertTokenizer.from_pretrained(mname) pipeline(task="fill-mask", model=mname) @@ -73,9 +73,9 @@ def test_offline_mode_no_internet(self): run = """ mname = "hf-internal-testing/tiny-random-bert" BertConfig.from_pretrained(mname) -BertModel.from_pretrained(mname) +BertModel.from_pretrained(mname, use_kernels=False) BertTokenizer.from_pretrained(mname) -pipe = pipeline(task="fill-mask", model=mname) +pipe = pipeline(task="fill-mask", model=mname, model_kwargs={"use_kernels": False}) print("success") """ @@ -88,7 +88,7 @@ def offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet") # Force fetching the files so that we can use the cache mname = "hf-internal-testing/tiny-random-bert" BertConfig.from_pretrained(mname) - BertModel.from_pretrained(mname) + BertModel.from_pretrained(mname, use_kernels=False) BertTokenizer.from_pretrained(mname) pipeline(task="fill-mask", model=mname) @@ -113,7 +113,7 @@ def test_offline_mode_sharded_checkpoint(self): run = """ mname = "hf-internal-testing/tiny-random-bert-sharded" BertConfig.from_pretrained(mname) -BertModel.from_pretrained(mname) +BertModel.from_pretrained(mname, use_kernels=False) print("success") """ @@ -143,7 +143,7 @@ def test_offline_mode_pipeline_exception(self): """ run = """ mname = "hf-internal-testing/tiny-random-bert" -pipe = pipeline(model=mname) +pipe = pipeline(model=mname, model_kwargs={"use_kernels": False}) """ mock = """ @@ -165,7 +165,7 @@ def test_offline_model_dynamic_model(self): """ run = """ mname = "hf-internal-testing/test_dynamic_model" -AutoModel.from_pretrained(mname, trust_remote_code=True) +AutoModel.from_pretrained(mname, trust_remote_code=True, use_kernels=False) print("success") """ @@ -206,6 +206,8 @@ def _execute_with_env(self, *commands: tuple[str, ...], should_fail: bool = Fals # Configure env new_env = self.get_env() new_env.update(env) + # Kernels always requires special treatment so we ignore it in general for each test here + new_env["USE_HUB_KERNELS"] = "0" # Run command result = subprocess.run(cmd, env=new_env, check=False, capture_output=True) diff --git a/tests/vlm_tester.py b/tests/vlm_tester.py index 05be8bdfa8f1..f84d2a1f4386 100644 --- a/tests/vlm_tester.py +++ b/tests/vlm_tester.py @@ -16,6 +16,8 @@ import unittest from inspect import signature +from transformers.testing_utils import scoped_kernels_class + from .multimodal_tester import MultiModalModelTest, MultiModalModelTester from .test_modeling_common import ( floats_tensor, @@ -119,6 +121,7 @@ def get_vision_config(self): return self.vision_config_class(**kwargs) +@scoped_kernels_class class VLMModelTest(MultiModalModelTest): """ Base test class for Vision-Language Models.