liger.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. # Copyright 2024 The HuggingFace Team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """
  15. Liger Kernel integration for applying optimized Triton kernels to transformer models.
  16. See https://github.com/linkedin/Liger-Kernel for details.
  17. """
  18. from ..modeling_utils import PreTrainedModel
  19. from ..trainer_utils import unwrap_peft_model
  20. from ..utils import is_liger_kernel_available, logging
  21. logger = logging.get_logger(__name__)
  22. def apply_liger_kernel(model, kernel_config):
  23. """
  24. Apply Liger Kernel optimizations to a model instance.
  25. Liger Kernel provides optimized Triton kernels for common transformer operations.
  26. This function patches the model in-place with those kernels.
  27. Args:
  28. model: The model to patch. Must be a `PreTrainedModel` or a PEFT wrapper around one.
  29. kernel_config: Kernel configuration.
  30. """
  31. if not is_liger_kernel_available():
  32. raise ImportError(
  33. "You have set `use_liger_kernel` to `True` but liger-kernel >= 0.3.0 is not available. "
  34. "Please install it with `pip install liger-kernel`"
  35. )
  36. from liger_kernel.transformers import _apply_liger_kernel_to_instance
  37. kernel_config = kernel_config or {}
  38. base_model = unwrap_peft_model(model)
  39. if isinstance(base_model, PreTrainedModel):
  40. _apply_liger_kernel_to_instance(model=base_model, **kernel_config)
  41. else:
  42. logger.warning("The model is not an instance of PreTrainedModel. No liger kernels will be applied.")