code: def exam_zero_optim_state_dict(placement_config, keep_gathered):
      set_seed(431)
      model_builder, data_gen_fn, output_transform_fn, *_ = next(
          iter(
              model_zoo.get_sub_registry("transformers_gpt_lm").values()
          )
      )
  
      model = model_builder()
  
      set_seed(451)
  
      world_size = torch.distributed.get_world_size()
      config_dict, *_ = search_chunk_configuration(
          model, search_range_m=1, search_interval=100
      )
      config_dict[world_size]["chunk_size"] = 5000
      config_dict[world_size]["keep_gathered"] = keep_gathered
  
      model = GeminiDDP(
          model, config_dict, **placement_config, pin_memory=True
      )
  
      optimizer = HybridAdam(model.parameters())
      optim = GeminiOptimizer(
          optimizer, model, initial_scale=32
      )  # initialize the link between chunk16 and chunk32
  
      set_seed(dist.get_rank() * 3 + 128)
      model.train()
      data = data_gen_fn()
      data = {
          k: v.cuda() if isinstance(v, torch.Tensor) else v
          for k, v in data.items()
      }
  
      optim.zero_grad()
      outputs = model(**data)
      outputs = output_transform_fn(outputs)
      loss = next(iter(outputs.values())).sum()
      optim.backward(loss)
      optim.step()
  
      optim_state_dict = optim.state_dict()
      optim.load_state_dict(optim_state_dict)
      new_state = optim.state_dict()["state"]
      org_state = optim_state_dict["state"]
  
      for k, v in org_state.items():
          w = new_state[k]
          for n, m in v.items():
              if isinstance(m, torch.Tensor):
                  o = w[n]
                  assert torch.equal(m, o)
              else:
                  assert m == w[n]
  
file_path: ColossalAI/tests/test_zero/test_gemini/test_zerooptim_state_dict.py
imports: import pytest
  import torch
  import torch.distributed
  import colossalai
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
  from colossalai.utils import set_seed
  from colossalai.zero import GeminiDDP, GeminiOptimizer
  from colossalai.zero.gemini.chunk import search_chunk_configuration
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 1 for calling a global function with example usage: exam_zero_optim_state_dict(placement_config, keep_gathered) without return types