diff --git a/docs/source/layers.md b/docs/source/layers.md index e1112041..c19c7899 100644 --- a/docs/source/layers.md +++ b/docs/source/layers.md @@ -27,6 +27,23 @@ The decorator does not change the behavior of the class -- it annotates the class with the given name (here `SiluAndMul`). The [`~kernels.kernelize`] function described below uses this name to look up kernels for the layer. +Sometimes you only want to kernelize a layer depending on some state inside +that layer. For instance, an MLP layer could support multiple activations, but +a kernel that you want to register only supports one particular activation. In +such cases, you can add a condition to a `use_kernel_forward_from_hub` decorator. +The layer will then only be kernelized when the condition holds. The condition +must be a callable that takes the instantiated layer and returns a `bool`. For +example: + +```python +@use_kernel_forward_from_hub( + "SwiGLUMLP", + condition=lambda module: module.config.hidden_act == "silu", +) +class MyMLP(nn.Module): + ... +``` + ### External layers An existing layer that does not (yet) have the [`~kernels.use_kernel_forward_from_hub`] diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py index 60d221d6..11188236 100644 --- a/kernels/src/kernels/layer/kernelize.py +++ b/kernels/src/kernels/layer/kernelize.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from copy import deepcopy from typing import TYPE_CHECKING @@ -13,6 +14,8 @@ import torch from torch import nn +logger = logging.getLogger(__name__) + def use_kernel_mapping( mapping: dict[ @@ -269,6 +272,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) if hasattr(module_class, "kernel_layer_name"): + cond = getattr(module_class, "kernel_condition", None) + if cond and not cond(module): + logger.info( + "Skipping kernelization for `%s` using `%s` due to kernel_condition.", + module_class.__name__, + module_class.kernel_layer_name, + ) + continue + kernelize_layer(module, mode=mode, device_type=device_type, use_fallback=use_fallback) return model diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 2e7aae20..74e205b6 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -272,10 +272,7 @@ def __str__(self) -> str: _CACHED_LAYER: dict[RepositoryProtocol, Type["nn.Module"]] = {} -def replace_kernel_forward_from_hub( - cls, - layer_name: str, -): +def replace_kernel_forward_from_hub(cls, layer_name: str, condition: Callable[["nn.Module"], bool] | None = None): """ Function that prepares a layer class to use kernels from the Hugging Face Hub. @@ -284,6 +281,15 @@ def replace_kernel_forward_from_hub( it is inherently fragile since the member variables and `forward` signature of such a layer can change. + Args: + layer_name (`str`): + The name of the layer to use for kernel lookup in registered mappings. + condition (`Callable[["nn.Module"], bool]`, *optional*): + Additional condition that is checked during kernelization. The callable + is passed the module instance and is evaluated for every instance of + the layer when [`~kernels.kernelize`] is called. If it returns `False`, + kernelization is skipped for that instance. + Example: ```python from kernels import replace_kernel_forward_from_hub @@ -293,9 +299,12 @@ def replace_kernel_forward_from_hub( ``` """ cls.kernel_layer_name = layer_name + # Wrap in `staticmethod` so that access through an instance does not bind + # it as a method (the condition takes the module as its only argument). + cls.kernel_condition = staticmethod(condition if condition is not None else lambda module: True) -def use_kernel_forward_from_hub(layer_name: str): +def use_kernel_forward_from_hub(layer_name: str, condition: Callable[["nn.Module"], bool] | None = None): """ Decorator factory that makes a layer extensible using the specified layer name. @@ -311,6 +320,11 @@ def use_kernel_forward_from_hub(layer_name: str): Args: layer_name (`str`): The name of the layer to use for kernel lookup in registered mappings. + condition (`Callable[["nn.Module"], bool]`, *optional*): + Additional condition that is checked during kernelization. The callable + is passed the module instance and is evaluated for every instance of + the layer when [`~kernels.kernelize`] is called. If it returns `False`, + kernelization is skipped for that instance. Returns: `Callable`: A decorator function that can be applied to layer classes. @@ -356,10 +370,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def decorator(ty): if inspect.isfunction(ty): Func = _create_func_module(ty) - replace_kernel_forward_from_hub(Func, layer_name) + replace_kernel_forward_from_hub(Func, layer_name, condition) return Func() elif inspect.isclass(ty): - replace_kernel_forward_from_hub(ty, layer_name) + replace_kernel_forward_from_hub(ty, layer_name, condition) return ty else: raise TypeError("@use_kernel_forward_from_hub can only be applied to classes or functions") diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index 699c4a11..4f28dc56 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -1,3 +1,4 @@ +import logging import sys from contextlib import nullcontext @@ -423,6 +424,84 @@ class SiluAndMulWithKernelFallback(SiluAndMul): kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE) +def test_kernel_condition_skips_kernelization(caplog): + @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: False) + class SiluAndMulConditionSkipped(SiluAndMul): + pass + + silu_and_mul = SiluAndMulConditionSkipped() + + with caplog.at_level(logging.INFO, logger="kernels.layer.kernelize"): + kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + + assert ( + "Skipping kernelization for `SiluAndMulConditionSkipped` using `SiluAndMulNonExisting` due to kernel_condition." + in caplog.text + ) + + # The forward was not replaced... + assert "forward" not in silu_and_mul.__dict__ + + # ... and the original implementation is still used. + X = torch.randn(32, 64) + silu_and_mul(X) + assert silu_and_mul.n_calls == 1 + + +def test_kernel_condition_holds(): + @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: True) + class SiluAndMulConditionHolds(SiluAndMul): + pass + + silu_and_mul = SiluAndMulConditionHolds() + + # The condition holds, so kernelization is attempted and fails because + # there is no kernel mapping for the layer. + with pytest.raises(ValueError, match="No layer mapping for `SiluAndMulNonExisting`"): + kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + + +def test_kernel_condition_receives_module_instance(): + received = [] + + @use_kernel_forward_from_hub( + "SiluAndMulNonExisting", + condition=lambda module: received.append(module) or False, + ) + class SiluAndMulConditionArg(SiluAndMul): + pass + + silu_and_mul = SiluAndMulConditionArg() + # No mapping, but uses fallback. + kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE) + + assert received == [silu_and_mul] + + +def test_kernel_condition_per_instance(): + @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: module.allow_kernel) + class SiluAndMulConditional(SiluAndMul): + def __init__(self, allow_kernel: bool): + super().__init__() + self.allow_kernel = allow_kernel + + skipped = SiluAndMulConditional(allow_kernel=False) + attempted = SiluAndMulConditional(allow_kernel=True) + + # Only the instance for which the condition holds is kernelized. + kernelize(skipped, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + with pytest.raises(ValueError, match="No layer mapping for `SiluAndMulNonExisting`"): + kernelize(attempted, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + + +def test_kernel_condition_defaults_to_true(): + # Without an explicit condition, kernelization is never skipped. Accessing + # the condition through an instance must not bind it as a method. + silu_and_mul = SiluAndMulWithKernel() + assert SiluAndMulWithKernel.kernel_condition(silu_and_mul) + assert silu_and_mul.kernel_condition(silu_and_mul) + + def test_local_layer_repo(device): # Fetch a kernel to the local cache. path = install_kernel("kernels-test/backward-marker-test", revision="main") @@ -601,7 +680,11 @@ def __init__(self, *args, **kwargs): self.foo = 42 def stub_repo(layer): - return LayerRepository(repo_id="kernels-test/nonexisting", layer_name=layer.__name__, revision="main") + return LayerRepository( + repo_id="kernels-test/nonexisting", + layer_name=layer.__name__, + revision="main", + ) with pytest.raises( TypeError,