code: def __init__(self, args: ModelArgs):
      super().__init__()
      self.n_kv_heads = (
          args.n_heads if args.n_kv_heads is None else args.n_kv_heads
      )
      model_parallel_size = fs_init.get_model_parallel_world_size()
      self.n_local_heads = args.n_heads // model_parallel_size
      self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
      self.n_rep = self.n_local_heads // self.n_local_kv_heads
      self.head_dim = args.dim // args.n_heads
  
      self.wq = ColumnParallelLinear(
          args.dim,
          args.n_heads * self.head_dim,
          bias=False,
          gather_output=False,
          init_method=lambda x: x,
      )
      self.wk = ColumnParallelLinear(
          args.dim,
          self.n_kv_heads * self.head_dim,
          bias=False,
          gather_output=False,
          init_method=lambda x: x,
      )
      self.wv = ColumnParallelLinear(
          args.dim,
          self.n_kv_heads * self.head_dim,
          bias=False,
          gather_output=False,
          init_method=lambda x: x,
      )
      self.wo = RowParallelLinear(
          args.n_heads * self.head_dim,
          args.dim,
          bias=False,
          input_is_parallel=True,
          init_method=lambda x: x,
      )
  
      self.cache_k = torch.zeros(
          (
              args.max_batch_size,
              args.max_seq_len,
              self.n_local_kv_heads,
              self.head_dim,
          )
      ).to(device)
      self.cache_v = torch.zeros(
          (
              args.max_batch_size,
              args.max_seq_len,
              self.n_local_kv_heads,
              self.head_dim,
          )
      ).to(device)
  
imports: import math
  from dataclasses import dataclass
  from typing import Any, Optional, Tuple
  import fairscale.nn.model_parallel.initialize
  import torch
  import torch.nn.functional
  from fairscale.nn.model_parallel.layers import ColumnParallelLinear, ParallelEmbedding, RowParallelLinear
  from torch import nn
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(args)