code: def test_single_process_save_load(self, tmp_path_dist_ckpt):
      Utils.initialize_model_parallel(1, 1)
  
      sharded_state_dict = {
          "sd_keyA": ShardedTensor.from_rank_offsets(
              "keyA", torch.ones(2, 4), replica_id=Utils.rank
          ),
          "sd_keyB": ShardedTensor.from_rank_offsets(
              "keyB", torch.ones(3, 5, 7), replica_id=Utils.rank
          ),
      }
  
      with TempNamedDir(
          tmp_path_dist_ckpt / "test_single_process_save_load"
      ) as ckpt_dir:
          save(sharded_state_dict, ckpt_dir)
  
          assert (ckpt_dir / "keyA").is_dir()
          assert (ckpt_dir / "keyB").is_dir()
          assert not (ckpt_dir / "keyC").exists()
          assert not (ckpt_dir / "sd_keyA").is_dir()
  
          load_ssd = {
              "load_sd_keyA": ShardedTensor.from_rank_offsets(
                  "keyA", torch.ones(2, 4), replica_id=Utils.rank
              ),
          }
          loaded_state_dict = load(load_ssd, ckpt_dir)
  
          assert set(loaded_state_dict.keys()) == {"load_sd_keyA"}
          assert isinstance(
              loaded_state_dict["load_sd_keyA"], torch.Tensor
          )
          assert loaded_state_dict["load_sd_keyA"].shape == (2, 4)
  
      Utils.destroy_model_parallel()
  
file_path: megatron-lm/tests/unit_tests/dist_checkpointing/test_serialization.py
imports: import numpy
  import pytest
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing import ShardedTensor, save, load
  from megatron.core.dist_checkpointing.core import CheckpointingException
  from megatron.core.dist_checkpointing.dict_utils import diff
  from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory, ShardedObject
  from megatron.core.dist_checkpointing.serialization import load_tensors_metadata
  from tests.unit_tests.dist_checkpointing import TempNamedDir
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 1 for calling the TestSerialization obj's underlying member method code with example usage: obj.test_single_process_save_load(tmp_path_dist_ckpt) without return types