class_docstr: This class is a modification of the Transformer class implemented in
  the LLaMA repo. The class has been modified for training, since the
  original one just supports inference. The generate method was inspired by
  the generate function you can find in `llama.generation`.
code: def __init__(self, params: ModelArgs):
      super().__init__()
  
      self.params = params
      self.vocab_size = params.vocab_size
      self.n_layers = params.n_layers
      if params.use_fairscale:
          self.n_local_heads = (
              params.n_heads // fs_init.get_model_parallel_world_size()
          )
      else:
          self.n_local_heads = params.n_heads
  
      self.head_dim = params.dim // params.n_heads
      dim = (
          params.max_batch_size,
          params.max_seq_len,
          self.n_local_heads,
          self.head_dim,
      )
      self.cache_k = [torch.zeros(dim) for _ in range(self.n_layers)]
      self.cache_v = [torch.zeros(dim) for _ in range(self.n_layers)]
  
      if params.use_fairscale:
          self.tok_embeddings = ParallelEmbedding(
              params.vocab_size, params.dim, init_method=lambda x: x
          )
      else:
          self.tok_embeddings = nn.Embedding(
              params.vocab_size, params.dim
          )
  
      if params.froze_embeddings:
          for param in self.tok_embeddings.parameters():
              param.requires_grad = False
  
      self.layers = torch.nn.ModuleList()
      for layer_id in range(params.n_layers):
          self.layers.append(TransformerBlock(layer_id, params))
  
      self.norm = RMSNorm(params.dim, eps=params.norm_eps)
      if params.use_fairscale:
          self.output = ColumnParallelLinear(
              params.dim,
              params.vocab_size,
              bias=False,
              init_method=lambda x: x,
          )
      else:
          self.output = nn.Linear(
              params.dim, params.vocab_size, bias=False
          )
  
      # TODO: How too modify this for training?
      self.freqs_cis = precompute_freqs_cis(
          self.params.dim // self.params.n_heads,
          self.params.max_seq_len * 2,
      )
  
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 1 to create a new Transformer object from inherited base classes: nn with example: obj = Transformer(params)