class_docstr: This class is a modification of the Attention class implemented in
  the LLaMA repo. The class has been modified for training, since the
  original one just supports inference.
code: def forward(
      self,
      x: torch.Tensor,
      kv_mask: torch.Tensor,
      freqs_cis: torch.Tensor,
      cache_k: Optional[torch.Tensor] = None,
      cache_v: Optional[torch.Tensor] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
      start_pos = 0  # Temporary
  
      bsz, seqlen, _ = x.shape
      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
  
      xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
      xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
      xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
  
      xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
  
      # Modified code to allow training, caching is not good for training
      if (cache_k is None and cache_v is not None) or (
          cache_k is not None and cache_v is None
      ):
          raise ValueError("cache_k is None while cache_v is not None")
      if cache_k is None:
          keys = xk
          values = xv
      else:
          cache_k.to(xk.device)
          cache_v.to(xv.device)
          cache_k[
              :bsz, start_pos : start_pos + seqlen
          ] = xk  # noqa E203
          cache_v[
              :bsz, start_pos : start_pos + seqlen
          ] = xv  # noqa E203
          keys = self.cache_k[:bsz, : start_pos + seqlen]  # noqa E203
          values = self.cache_v[:bsz, : start_pos + seqlen]  # noqa E203
  
      xq = xq.transpose(1, 2)
      keys = keys.transpose(1, 2)
      values = values.transpose(1, 2)
      scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(
          self.head_dim
      )
      if kv_mask is not None:
          scores = scores + kv_mask
      scores = F.softmax(scores.float(), dim=-1).type_as(xq)
      output = torch.matmul(scores, values)
      output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
      if cache_k is None:
          return self.wo(output), None, None
      else:
          return self.wo(output), self.cache_k, self.cache_v
  
imports: import json
  import math
  import os
  from dataclasses import dataclass
  from pathlib import Path
  from typing import Tuple, List, Union, Optional
  import deepspeed
  import torch
  import torch.distributed
  import torch.nn.functional
  import fairscale.nn.model_parallel.initialize
  from fairscale.nn.model_parallel.initialize import initialize_model_parallel
  from fairscale.nn.model_parallel.layers import ParallelEmbedding, RowParallelLinear, ColumnParallelLinear
  from torch import nn
  from transformers import AutoTokenizer
  from llama import Tokenizer
  from llama.generation import sample_top_p
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, kv_mask, freqs_cis, cache_k, cache_v) and returns: self,  None,  None, self,  self,  self