code: def exam_lazy_from_pretrained():
      llama_path = os.environ["LLAMA_PATH"]
      plugin = GeminiPlugin()
      booster = Booster(plugin=plugin)
      orig_model = LlamaForCausalLM.from_pretrained(llama_path)
      orig_state_dict = {
          k: v.half() for k, v in orig_model.state_dict().items()
      }
      with LazyInitContext():
          model = LlamaForCausalLM.from_pretrained(llama_path)
      model, *_ = booster.boost(model)
      with shared_tempdir() as tempdir:
          save_path = os.path.join(tempdir, "model.pt")
          booster.save_model(model, save_path, shard=False)
          dist.barrier()
          state_dict = torch.load(save_path, map_location="cpu")
          check_state_dict_equal(
              state_dict, orig_state_dict, False, ignore_dtype=True
          )
  
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 3 for calling a global function with example usage: exam_lazy_from_pretrained() without return types