code: def replace_llama_attn_with_flash_attn():
      transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (
          _prepare_decoder_attention_mask
      )
      transformers.models.llama.modeling_llama.LlamaAttention.forward = (
          forward
      )
  
file_path: ml-ferret/ferret/train/llama_flash_attn_monkey_patch.py
imports: from typing import List, Optional, Tuple
  import torch
  from torch import nn
  import transformers
  from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
  from einops import rearrange
  from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
  from flash_attn.bert_padding import unpad_input, pad_input
label_desc: Use image node_id 3 for calling a global function with example usage: replace_llama_attn_with_flash_attn() without return types