code: def exam_state_dict(
      placement_config,
      shard: bool,
      model_name: str,
      size_per_shard: int,
      tp_size: int,
      zero_size: int,
  ):
      (model_fn, data_gen_fn, output_transform_fn, _, _) = next(
          iter(model_zoo.get_sub_registry(model_name).values())
      )
      criterion = lambda x: x.mean()
      enable_all_optimization = True if tp_size > 1 else False
      extra_dp_size = dist.get_world_size() // (zero_size * tp_size)
      plugin = GeminiPlugin(
          **placement_config,
          precision="fp16",
          initial_scale=(2**14),
          tp_size=tp_size,
          extra_dp_size=extra_dp_size,
          enable_all_optimization=enable_all_optimization,
      )
      booster = Booster(plugin=plugin)
  
      model = model_fn()
      new_model = model_fn()
      optimizer = HybridAdam(model.parameters(), lr=0.001)
      model, optimizer, criterion, _, _ = booster.boost(
          model, optimizer, criterion
      )
      new_optimizer = HybridAdam(new_model.parameters(), lr=0.001)
      new_model, new_optimizer, criterion, _, _ = booster.boost(
          new_model, new_optimizer, criterion
      )
  
      data = data_gen_fn()
      data = {
          k: v.to("cuda")
          if torch.is_tensor(v) or "Tensor" in v.__class__.__name__
          else v
          for k, v in data.items()
      }
      output = model(**data)
      output = output_transform_fn(output)
      output_key = list(output.keys())[0]
      loss = criterion(output[output_key])
  
      booster.backward(loss, optimizer)
      optimizer.step()
  
      with shared_tempdir() as tempdir:
          model_ckpt_path = f"{tempdir}/model"
          optimizer_ckpt_path = f"{tempdir}/optimizer"
          booster.save_model(
              model,
              model_ckpt_path,
              shard=shard,
              size_per_shard=size_per_shard,
          )
  
          booster.save_optimizer(
              optimizer,
              optimizer_ckpt_path,
              shard=shard,
              size_per_shard=size_per_shard,
          )
          dist.barrier()
  
          booster.load_model(new_model, model_ckpt_path)
          check_state_dict_equal(
              model.state_dict(only_rank_0=False),
              new_model.state_dict(only_rank_0=False),
              False,
              ignore_dtype=True,
          )
  
          booster.load_optimizer(new_optimizer, optimizer_ckpt_path)
          check_state_dict_equal(
              optimizer.state_dict(only_rank_0=False),
              new_optimizer.state_dict(only_rank_0=False),
              False,
          )
  
          # Check the new model/optimizer can successfully run.
          data = data_gen_fn()
          data = {
              k: v.to("cuda")
              if torch.is_tensor(v) or "Tensor" in v.__class__.__name__
              else v
              for k, v in data.items()
          }
          output = new_model(**data)
          output = output_transform_fn(output)
          output_key = list(output.keys())[0]
          loss = criterion(output[output_key])
          booster.backward(loss, new_optimizer)
          new_optimizer.step()
          booster.save_model(new_model, model_ckpt_path, shard=shard)
          booster.save_optimizer(
              new_optimizer, optimizer_ckpt_path, shard=shard
          )
  
file_path: ColossalAI/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
imports: import os
  import pytest
  import torch
  import torch.distributed
  from transformers import LlamaForCausalLM
  from utils import shared_tempdir
  import colossalai
  from colossalai.testing import skip_if_not_enough_gpus
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin
  from colossalai.lazy import LazyInitContext
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.testing import check_state_dict_equal, clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 2 for calling a global function with example usage: exam_state_dict(placement_config, shard, model_name, size_per_shard, tp_size, zero_size) without return types