class_docstr: Fused QKV Attention layer for incrementally improved training efficiency
class_docstr_tok: ["Fused","QKV","Attention","layer","for","incrementally","improved","training","efficiency"]
code: def _post_training(self, model, name):
      q_proj, k_proj, v_proj = torch.split(
          self.qkv_proj.weight.data, self.out_features, dim=0
      )
  
      new_attn = LlamaAttention(self.config)
      new_attn.q_proj.weight.data = q_proj
      new_attn.k_proj.weight.data = k_proj
      new_attn.v_proj.weight.data = v_proj
      new_attn.o_proj.weight.data = self.o_proj.weight.data
  
      set_module_name(model, name, new_attn)
  
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 2 for calling the FusedAttention obj's underlying member method code with example usage: obj._post_training(model, name) without return types