class_docstr: A single transformer layer.
  
  Transformer layer takes input with size [s, b, h] and returns an
  output of the same size.
class_docstr_tok: ["A","single","transformer","layer",".","Transformer","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","an","output","of","the","same","size","."]
code: def _get_layer_offset(self):
      pipeline_rank = parallel_state.get_pipeline_model_parallel_rank()
  
      num_layers_per_pipeline_rank = (
          self.config.num_layers
          // parallel_state.get_pipeline_model_parallel_world_size()
      )
  
      if (
          parallel_state.get_virtual_pipeline_model_parallel_world_size()
          is not None
      ):
          vp_rank = (
              parallel_state.get_virtual_pipeline_model_parallel_rank()
          )
          vp_size = (
              parallel_state.get_virtual_pipeline_model_parallel_world_size()
          )
  
          total_num_layers = self.config.num_layers
          num_layers_per_virtual_rank = (
              num_layers_per_pipeline_rank // vp_size
          )
          total_virtual_chunks = total_num_layers // vp_size
          offset = vp_rank * total_virtual_chunks + (
              pipeline_rank * num_layers_per_virtual_rank
          )
  
      else:
          # Each stage gets a contiguous set of layers.
          if (
              parallel_state.get_pipeline_model_parallel_world_size()
              > 1
          ):
              offset = pipeline_rank * num_layers_per_pipeline_rank
          else:
              offset = 0
  
      return offset
  
file_path: megatron-lm/megatron/core/transformer/transformer_layer.py
imports: from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_viewless_tensor
label_desc: Use image node_id 2 for calling the TransformerLayer obj's underlying member method code with example usage: obj._get_layer_offset() and returns: offset