code: def exam_from_pretrained(
      plugin_type: str, model_name: str, shard=True, size_per_shard=32
  ):
      (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) = next(
          iter(model_zoo.get_sub_registry(model_name).values())
      )
      criterion = loss_fn
  
      if plugin_type == "ddp":
          plugin = TorchDDPPlugin()
      elif plugin_type == "zero":
          plugin = LowLevelZeroPlugin(
              stage=2, max_norm=1.0, initial_scale=32
          )
      elif plugin_type == "gemini":
          plugin = GeminiPlugin(precision="fp16", initial_scale=32)
      else:
          raise ValueError(
              f"Plugin with type {plugin_type} is invalid, please check your argument."
          )
  
      booster = Booster(plugin=plugin)
  
      model = model_fn().cuda()
      model_huggingface_cls = model.__class__
      optimizer = HybridAdam(model.parameters(), lr=0.001)
      model, optimizer, criterion, _, _ = booster.boost(
          model, 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)
      loss = criterion(output)
  
      booster.backward(loss, optimizer)
      optimizer.step()
  
      with shared_tempdir() as tempdir:
          model_ckpt_path = f"{tempdir}/model"
          booster.save_model(
              model,
              model_ckpt_path,
              shard=shard,
              size_per_shard=size_per_shard,
          )
          dist.barrier()
  
          new_model = model_huggingface_cls.from_pretrained(
              model_ckpt_path
          )
          new_model = new_model.cuda()
          new_optimizer = HybridAdam(new_model.parameters(), lr=0.001)
          new_model, new_optimizer, criterion, _, _ = booster.boost(
              new_model, new_optimizer, criterion
          )
  
          if plugin_type == "gemini":
              check_state_dict_equal(
                  model.state_dict(only_rank_0=False),
                  new_model.state_dict(only_rank_0=False),
                  False,
              )
          else:
              check_state_dict_equal(
                  model.unwrap().state_dict(),
                  new_model.unwrap().state_dict(),
                  False,
              )
          dist.barrier()
  
file_path: colossalai/tests/test_checkpoint_io/test_plugins_huggingface_compatibility.py
imports: import pytest
  import torch
  import torch.distributed
  from utils import shared_tempdir
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin
  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 1 for calling a global function with example usage: exam_from_pretrained(plugin_type, model_name, shard, size_per_shard) without return types