code: 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
      )
      transformers.models.llama.modeling_llama.LlamaAttention.forward = (
          flashattn_forward
      )
      if packed:
          transformers.models.llama.modeling_llama.LlamaDecoderLayer = (
              LlamaDecoderLayer
          )
          transformers.models.llama.modeling_llama.LlamaModel.forward = (
              llama_model_forward
          )
  
      # skip only if explicitly disabled
      if cross_entropy:
          try:
              from flash_attn.losses.cross_entropy import (
                  CrossEntropyLoss,
              )
  
              LOG.info("patching with flash_attn.losses.cross_entropy")
              transformers.models.llama.modeling_llama.CrossEntropyLoss = partial(
                  CrossEntropyLoss, inplace_backward=True
              )
          except ImportError:
              LOG.info(
                  "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
  
              class LlamaRMSNorm(RMSNorm):
                  """Patched LLamaRMSNorm"""
  
                  def __init__(self, hidden_size, eps=1e-6):
                      super().__init__(hidden_size, eps=eps)
  
              LOG.info("patching with flash_attn.ops.rms_norm")
              transformers.models.llama.modeling_llama.LlamaRMSNorm = (
                  LlamaRMSNorm
              )
          except ImportError:
              LOG.info(
                  "optimized flash-attention RMSNorm not found (run `pip install 'git+https://github.com/Dao-AILab/flash-attention.git#egg=dropout_layer_norm&subdirectory=csrc/layer_norm'`)"
              )
  
imports: import logging
  import warnings
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import transformers
  from einops import rearrange
  from flash_attn.bert_padding import pad_input, unpad_input
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.modeling_llama import LlamaAttention
  from transformers.models.llama.modeling_llama import LlamaDecoderLayer
  from transformers.models.llama.modeling_llama import LlamaMLP, apply_rotary_pos_emb, repeat_kv
  from xformers.ops import SwiGLU
  from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids, set_module_name
label_desc: Use image node_id 3 for calling a global function with example usage: replace_llama_attn_with_flash_attn(packed, cross_entropy, rms_norm) without return types