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 sharded_state_dict(self, prefix=""):
      offset = self._get_layer_offset()
      num_layers = self.config.num_layers
  
      global_layer_offset = (
          self.layer_number - 1
      )  # self.layer_number starts at 1
      state_dict_prefix = f"{prefix}{global_layer_offset - offset}."  # module list index in TransformerBlock
      sharded_pp_offset = [
          (0, global_layer_offset, num_layers)
      ]  # PP sharding offset for ShardedTensors
  
      attn_state_dict = self.self_attention.sharded_state_dict(
          prefix=f"{state_dict_prefix}self_attention.",
          sharded_key_prefix=f"{prefix}self_attention.",
          sharded_offsets=sharded_pp_offset,
      )
  
      mlp_state_dict = self.mlp.sharded_state_dict(
          prefix=f"{state_dict_prefix}mlp.",
          sharded_key_prefix=f"{prefix}mlp.",
          sharded_offsets=sharded_pp_offset,
      )
  
      sharded_state_dict = {**mlp_state_dict, **attn_state_dict}
  
      return sharded_state_dict
  
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 4 for calling the TransformerLayer obj's underlying member method code with example usage: obj.sharded_state_dict(prefix) and returns: sharded_state_dict