skip some flash attn patches unless explicitly enabled (#643)

* skip some flash attn patches if explicitly disabled

* make the other patches optional

README.md

@@ -636,6 +636,8 @@ flash_optimum:
 xformers_attention:
 # whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
 flash_attention:
+flash_attn_cross_entropy:  # Whether to use flash-attention cross entropy implementation - advanced use only
+flash_attn_rms_norm:  # Whether to use flash-attention rms norm implementation - advanced use only
 # whether to use scaled-dot-product attention
 # https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 sdp_attention:

src/axolotl/monkeypatch/llama_attn_hijack_flash.py

@@ -38,7 +38,11 @@ except ImportError:
 LOG = logging.getLogger("axolotl")
 
 
-def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False):
+def replace_llama_attn_with_flash_attn(
+    packed: Optional[bool] = False,
+    cross_entropy: Optional[bool] = False,
+    rms_norm: Optional[bool] = False,
+):
     transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (  # pylint: disable=protected-access
         _prepare_decoder_attention_mask
     )
@@ -49,6 +53,8 @@ def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False):
             llama_model_forward
         )
 
+    # skip only if explicitly disabled
+    if cross_entropy:
         try:
             from flash_attn.losses.cross_entropy import CrossEntropyLoss
 
@@ -61,6 +67,8 @@ def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False):
                 "optimized flash-attention CrossEntropyLoss not found (run `pip install 'git+https://github.com/Dao-AILab/flash-attention.git#egg=xentropy_cuda_lib&subdirectory=csrc/xentropy'`)"
             )
 
+    # skip only if explicitly disabled
+    if rms_norm:
         try:
             from flash_attn.ops.rms_norm import RMSNorm
 

src/axolotl/utils/models.py

@@ -121,7 +121,11 @@ def load_model(
             )
 
             LOG.info("patching with flash attention for sample packing")
-            replace_llama_attn_with_flash_attn(packed=cfg.sample_packing)
+            replace_llama_attn_with_flash_attn(
+                packed=cfg.sample_packing,
+                cross_entropy=cfg.flash_attn_cross_entropy,
+                rms_norm=cfg.flash_attn_rms_norm,
+            )
     elif cfg.is_llama_derived_model and cfg.xformers_attention:
         from axolotl.monkeypatch.llama_attn_hijack_xformers import (
             hijack_llama_attention,