code: def test_state_dict_comparison(self, tmp_path_dist_ckpt):
      Utils.initialize_model_parallel(2, 4)
      with TempNamedDir(
          tmp_path_dist_ckpt / "test_state_dict_comparison_A"
      ) as ckpt_dir_A, TempNamedDir(
          tmp_path_dist_ckpt / "test_state_dict_comparison_B"
      ) as ckpt_dir_B:
          gpt_model_A = initialize_gpt_model(1)
          save(gpt_model_A.sharded_state_dict(), ckpt_dir_A)
          gpt_model_B = initialize_gpt_model(2)
          save(gpt_model_B.sharded_state_dict(), ckpt_dir_B)
  
          state_dict_A = load_plain_tensors(ckpt_dir_A)
          state_dict_A_dup = load_plain_tensors(ckpt_dir_A)
          state_dict_B = load_plain_tensors(ckpt_dir_B)
  
          # Test that A matches A
          diffs = diff(state_dict_A, state_dict_A_dup)
          assert not any(map(bool, diffs)), diffs
  
          # Test that A *keys* match B *keys*, but the tensors content is different
          only_left, only_right, mismatch = diff(
              state_dict_A, state_dict_B
          )
          assert not only_left and not only_right, (
              only_left,
              only_right,
          )
          assert len(mismatch) == len(state_dict_A), (
              len(mismatch),
              (len(state_dict_A)),
          )
  
file_path: megatron-lm/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py
imports: import pytest
  import torch
  from torch.distributed._tensor import DeviceMesh
  from megatron.core.dist_checkpointing import save, load, load_plain_tensors
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.dict_utils import diff
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_model import GPTModel
  from tests.unit_tests.dist_checkpointing import TempNamedDir
  from tests.unit_tests.test_utilities import Utils
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec
label_desc: Use image node_id 2 for calling the TestGPTModelReconfiguration obj's underlying member method code with example usage: obj.test_state_dict_comparison(tmp_path_dist_ckpt) without return types