code: def replace_with_flash_attention(model: LlamaForCausalLM) -> None:
      for name, module in model.named_modules():
          if isinstance(module, LlamaAttention):
              module.forward = MethodType(attention_forward, module)
          if isinstance(module, LlamaModel):
              module._prepare_decoder_attention_mask = MethodType(
                  _prepare_decoder_attention_mask, module
              )
          if isinstance(module, LlamaRMSNorm):
              module.forward = MethodType(rms_norm_forward, module)
  
file_path: colossalai/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
imports: from types import MethodType
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  from einops import rearrange
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_kvpacked_func
  from flash_attn.ops.rms_norm import rms_norm
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaForCausalLM, LlamaModel, LlamaRMSNorm, apply_rotary_pos_emb, repeat_kv
  from colossalai.logging import get_dist_logger
label_desc: Use image node_id 4 for calling a global function with example usage: replace_with_flash_attention(model) without return types