code: def test_sharded_state_dict(self, tp_pp):
      Utils.destroy_model_parallel()
      Utils.initialize_model_parallel(*tp_pp)
  
      model_parallel_cuda_manual_seed(123)
      transformer_config = TransformerConfig(
          num_layers=2,
          hidden_size=128,
          num_attention_heads=8,
          use_cpu_initialization=True,
      )
      parallel_transformer_layer = TransformerLayer(
          transformer_config,
          get_gpt_layer_with_transformer_engine_spec().submodules,
      )
  
      sharded_state_dict = (
          parallel_transformer_layer.sharded_state_dict()
      )
  
      extra_states = {
          k: v
          for k, v in sharded_state_dict.items()
          if k.endswith("extra_state")
      }
      sharded_tensors = {
          k: v
          for k, v in sharded_state_dict.items()
          if not k.endswith("extra_state")
      }
      assert all(
          isinstance(t, ShardedObject) for t in extra_states.values()
      )
      assert all(
          isinstance(t, ShardedTensor) for t in sharded_tensors.values()
      )
  
      # Test all local shapes
      tensor_local_shapes = {
          k: v.local_shape for k, v in sharded_tensors.items()
      }
      tp_size = parallel_state.get_tensor_model_parallel_world_size()
      assert tensor_local_shapes == get_tensor_shapes_for_tp(
          transformer_config, tp_size
      )
  
      # Test all global shapes. Prepend num layers in front of expected shapes
      tensor_global_shapes = {
          k: v.global_shape for k, v in sharded_tensors.items()
      }
      expected_global_shapes = {
          k: (transformer_config.num_layers, *v)
          for k, v in get_tensor_shapes_for_tp(
              transformer_config, 1
          ).items()
      }
      assert tensor_global_shapes == expected_global_shapes
  
      # Test ShardedTensor keys
      for state_dict_key, sh_ten in sharded_tensors.items():
          assert state_dict_key == f"0.{sh_ten.key}"
  
      Utils.destroy_model_parallel()
      Utils.initialize_model_parallel(1, 1)
  
file_path: megatron-lm/tests/unit_tests/transformer/test_transformer_layer.py
imports: import pytest
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 5 for calling the TestParallelTransformerLayer obj's underlying member method code with example usage: obj.test_sharded_state_dict(tp_pp) without return types