docstr: ### Generator to create layers
  
  The layers are generated in the same order as checkpoints.
  
  It gives `None` when a layer is not available; we use the layer indices as NeoX and there are two
  transformation layers we don't need in our implementation.
  
  :param n_vocab: is the number of tokens in the vocabulary
  :param n_hidden: is the number of features in the embeddings
  :param n_layers: is the number of transformer layers
  :param n_heads: is the number of attention heads
  :param filter_layers: are the set of layers to be used. All layers will be used if None.
      This is used to test smaller versions of the model with fewer layers
  :param is_clone_layers: specifies whether to clone the transformer layers (a bit faster)
  :param dtype: is the data type of the model
  :param device: is the device of the model
  :param is_llm_int8: specifies whether to use int8 quantization
  :param llm_int8_threshold: is the threshold $lpha$ used to separate outlier features
  :param is_flash_attention: specifies whether to use
      [FlashAttention](https://github.com/HazyResearch/flash-attention)
code: def __init__(
      self,
      *,
      n_vocab: int = 50_432,
      n_hidden: int = 6_144,
      n_layers: int = 44,
      n_heads: int = 64,
      filter_layers: Optional[Set] = None,
      is_clone_layers: bool = True,
      dtype: torch.dtype = torch.float,
      device: torch.device = torch.device("cpu"),
      is_llm_int8: bool = False,
      llm_int8_threshold: float = 6.0,
      is_flash_attention: bool = False,
  ):
      """
      ### Generator to create layers
  
      The layers are generated in the same order as checkpoints.
  
      It gives `None` when a layer is not available; we use the layer indices as NeoX and there are two
      transformation layers we don't need in our implementation.
  
      :param n_vocab: is the number of tokens in the vocabulary
      :param n_hidden: is the number of features in the embeddings
      :param n_layers: is the number of transformer layers
      :param n_heads: is the number of attention heads
      :param filter_layers: are the set of layers to be used. All layers will be used if None.
          This is used to test smaller versions of the model with fewer layers
      :param is_clone_layers: specifies whether to clone the transformer layers (a bit faster)
      :param dtype: is the data type of the model
      :param device: is the device of the model
      :param is_llm_int8: specifies whether to use int8 quantization
      :param llm_int8_threshold: is the threshold $\alpha$ used to separate outlier features
      :param is_flash_attention: specifies whether to use
          [FlashAttention](https://github.com/HazyResearch/flash-attention)
      """
      if filter_layers is None:
          filter_layers = set(range(n_layers + 3))
  
      self.n_vocab = n_vocab
      self.n_hidden = n_hidden
      self.n_layers = n_layers
      self.n_heads = n_heads
      self.filter_layers = filter_layers
      self.is_clone_layers = is_clone_layers
      self.dtype = dtype
      self.device = device
      self.is_llm_int8 = is_llm_int8
      self.llm_int8_threshold = llm_int8_threshold
      self.is_flash_attention = is_flash_attention
  
      self.pre_created_layers = dict(
          transformer_layer=None,
      )
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 1 to create a new LayerGenerator object  with example: obj = LayerGenerator()