From e37d3b2457dbdcad77bc0ca3e75517cfc8f7fe70 Mon Sep 17 00:00:00 2001 From: Aman Karmani Date: Mon, 4 Sep 2023 04:04:34 +0000 Subject: [PATCH] use flash_attn xentropy when available --- src/axolotl/monkeypatch/llama_attn_hijack_flash.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py index cb0aa3fe6f..5da24373b0 100644 --- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py +++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py @@ -2,7 +2,9 @@ # copied from https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py +import logging import warnings +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -33,6 +35,9 @@ ) +LOG = logging.getLogger("axolotl") + + def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False): transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = ( # pylint: disable=protected-access _prepare_decoder_attention_mask @@ -43,6 +48,15 @@ def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False): transformers.models.llama.modeling_llama.LlamaModel.forward = ( llama_model_forward ) + try: + from flash_attn.losses.cross_entropy import CrossEntropyLoss + + LOG.info("patching with flash_attention.losses.cross_entropy") + transformers.models.llama.modeling_llama.CrossEntropyLoss = partial( + CrossEntropyLoss, inplace_backward=True + ) + except ImportError: + pass # Disable the transformation of the attention mask in LlamaModel as the flash attention