class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
code: def __init__(
      self,
      config,
      init_method,
      output_layer_init_method,
      num_layers,
      hidden_size,
      ffn_hidden_size,
      num_attention_heads,
      apply_query_key_layer_scaling=True,
      kv_channels=None,
      layer_type=LayerType.encoder,  # it can be a list of types or single type
      self_attn_mask_type=AttnMaskType.padding,
      pre_process=True,
      post_process=True,
      precision=16,
      fp32_residual_connection=False,
      activations_checkpoint_method=None,
      activations_checkpoint_num_layers=None,
      layernorm_epsilon=1e-5,
      hidden_dropout=0.1,
      attention_dropout=0.1,
      ffn_dropout=0.0,
      drop_path_rate=0.0,
      bias_activation_fusion=True,
      bias_dropout_add_fusion=True,
      masked_softmax_fusion=True,
      persist_layer_norm=False,
      openai_gelu=False,
      onnx_safe=False,
      activation="gelu",
      model_type=ModelType.encoder_or_decoder,
      megatron_legacy=False,
      bias=True,
      chunk_size=64,
      normalization="layernorm",
      transformer_block_type="pre_ln",
      headscale=False,
      layer_number_offset=0,  # this is use only for attention norm_factor scaling
      activations_checkpoint_granularity=None,
      normalize_attention_scores=True,
      ub_tp_comm_overlap=False,
      use_flash_attention=False,
  ):
      kwargs = locals()
      for key in ["self", "__class__"]:
          kwargs.pop(key)
      self.drop_path_rate = kwargs.pop("drop_path_rate")
      super(ParallelVisionTransformer, self).__init__(**kwargs)
  
      self.num_layers = self.get_num_layers(num_layers)
  
      self.drop_path_rates = [
          rate.item()
          for rate in torch.linspace(
              0,
              self.drop_path_rate,
              self.num_layers
              * parallel_state.get_pipeline_model_parallel_world_size(),
          )
      ]
  
      # Rebuild with vision transformer layers.
      def build_layer(layer_number):
          if isinstance(layer_type, list):
              lt = layer_type[layer_number - 1]
          else:
              lt = layer_type
          return ParallelVisionTransformerLayer(
              config=config,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              layer_number=layer_number + layer_number_offset,
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              num_attention_heads=num_attention_heads,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              layer_type=lt,
              self_attn_mask_type=self_attn_mask_type,
              precision=precision,
              fp32_residual_connection=fp32_residual_connection,
              layernorm_epsilon=layernorm_epsilon,
              hidden_dropout=hidden_dropout,
              attention_dropout=attention_dropout,
              ffn_dropout=ffn_dropout,
              drop_path_rate=self.drop_path_rates[layer_number - 1],
              bias_activation_fusion=bias_activation_fusion,
              bias_dropout_add_fusion=bias_dropout_add_fusion,
              masked_softmax_fusion=masked_softmax_fusion,
              persist_layer_norm=persist_layer_norm,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              activation=activation,
              megatron_legacy=megatron_legacy,
              bias=bias,
              chunk_size=chunk_size,
              normalization=normalization,
              transformer_block_type=transformer_block_type,
              headscale=headscale,
              activations_checkpoint_granularity=activations_checkpoint_granularity,
              normalize_attention_scores=normalize_attention_scores,
              use_flash_attention=use_flash_attention,
          )
  
      if (
          parallel_state.get_virtual_pipeline_model_parallel_world_size()
          is not None
      ):
          assert (
              num_layers
              % parallel_state.get_virtual_pipeline_model_parallel_world_size()
              == 0
          ), (
              "num_layers_per_stage must be divisible by "
              "virtual_pipeline_model_parallel_size"
          )
  
          # self.model_type != ModelType.encoder_and_decoder
          assert (
              self.model_type.value != 2
          ), f"virtual pipeline parallel currently only supported for GPT"
  
          # Number of layers in each model chunk is the number of layers in the stage,
          # divided by the number of model chunks in a stage.
          self.num_layers = (
              self.num_layers
              // parallel_state.get_virtual_pipeline_model_parallel_world_size()
          )
          # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
          # layers to stages like (each list is a model chunk):
          # Stage 0: [0]  [2]  [4]  [6]
          # Stage 1: [1]  [3]  [5]  [7]
          # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
          # layers to stages like (each list is a model chunk):
          # Stage 0: [0, 1]  [4, 5]
          # Stage 1: [2, 3]  [6, 7]
          offset = parallel_state.get_virtual_pipeline_model_parallel_rank() * (
              num_layers
              // parallel_state.get_virtual_pipeline_model_parallel_world_size()
          ) + (
              parallel_state.get_pipeline_model_parallel_rank()
              * self.num_layers
          )
      else:
          # Each stage gets a contiguous set of layers.
          if (
              self.model_type == ModelType.encoder_and_decoder
              and parallel_state.get_pipeline_model_parallel_world_size()
              > 1
          ):
              pipeline_rank = (
                  parallel_state.get_pipeline_model_parallel_rank()
              )
              if layer_type == LayerType.encoder:
                  offset = pipeline_rank * self.num_layers
              else:
                  num_ranks_in_enc = (
                      parallel_state.get_pipeline_model_parallel_split_rank()
                  )
                  offset = (
                      pipeline_rank - num_ranks_in_enc
                  ) * self.num_layers
          else:
              offset = (
                  parallel_state.get_pipeline_model_parallel_rank()
                  * self.num_layers
              )
  
      self.layers = torch.nn.ModuleList(
          [build_layer(i + 1 + offset) for i in range(self.num_layers)]
      )
  
file_path: NeMo/nemo/collections/vision/modules/common/megatron/vision_transformer.py
imports: import torch
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.transformer import ParallelTransformer, ParallelTransformerLayer_
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
label_desc: Use image node_id 1 to create a new ParallelVisionTransformer object from inherited base classes: ParallelTransformer with example: obj = ParallelVisionTransformer(config, init_method, output_layer_init_method, num_layers, hidden_size, ffn_hidden_size, num_attention_heads, apply_query_key_layer_scaling, kv_channels, layer_type, _attn_mask_type, pre_process, post_process, precision, fp32_residual_connection, activations_checkpoint_method, activations_checkpoint_num_layers, layernorm_epsilon, hidden_dropout, attention_dropout, ffn_dropout, drop_path_rate, bias_activation_fusion, bias_dropout_add_fusion, masked_softmax_fusion, persist_layer_norm, openai_gelu, onnx_safe, activation, model_type, megatron_legacy, bias, chunk_size, normalization, transformer_block_type, headscale, layer_number_offset, activations_checkpoint_granularity, normalize_attention_scores, ub_tp_comm_overlap, use_flash_attention)