code: def exam_gpt_fwd_bwd(
      placement_config,
      keep_gather,
      model_name: str,
      use_grad_checkpoint: bool = False,
      master_weights: bool = True,
  ):
      init_device = get_current_device()
      (
          model_builder,
          data_gen_fn,
          output_transform_fn,
          loss_fn,
          *_,
      ) = next(iter(model_zoo.get_sub_registry(model_name).values()))
  
      set_seed(42)
      model = model_builder()
  
      set_seed(42)
      torch_model = model_builder().cuda()
      for torch_p, p in zip(
          torch_model.parameters(), model.parameters()
      ):
          torch_p.data.copy_(p.data)
  
      if use_grad_checkpoint:
          model.gradient_checkpointing_enable()
          torch_model.gradient_checkpointing_enable()
  
      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_gather
      model = GeminiDDP(
          model,
          config_dict,
          init_device,
          pin_memory=True,
          **placement_config,
          master_weights=master_weights,
      )
      optimizer = HybridAdam(model.parameters(), lr=1e-3)
      zero_optim = GeminiOptimizer(optimizer, model, initial_scale=1)
  
      rank = dist.get_rank()
      amp_config = dict(
          opt_level="O2",
          keep_batchnorm_fp32=False,
          loss_scale=1,
          master_weights=master_weights,
      )
      torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3)
      torch_model, torch_optim = convert_to_apex_amp(
          torch_model, torch_optim, amp_config
      )
      torch_model = DDP(torch_model, device_ids=[rank])
  
      set_seed(rank)
  
      data = data_gen_fn()
      data = {
          k: v.cuda() if isinstance(v, torch.Tensor) else v
          for k, v in data.items()
      }
  
      torch_optim.zero_grad()
      zero_optim.zero_grad()
  
      # set random seed is same as torch_model.eval()
      set_seed(42)
      torch_loss = run_fwd_bwd(
          torch_model,
          data,
          output_transform_fn,
          loss_fn,
          optimizer=torch_optim,
      )
      set_seed(42)
      loss = run_fwd_bwd(
          model,
          data,
          output_transform_fn,
          loss_fn,
          optimizer=zero_optim,
      )
  
      assert_close(torch_loss.float(), loss.float())
  
      check_grad(model, torch_model)
  
file_path: ColossalAI/tests/test_zero/test_gemini/test_fwd_bwd.py
imports: import pytest
  import torch
  import torch.distributed
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing import assert_close
  import colossalai
  from colossalai.legacy.amp import convert_to_apex_amp
  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.utils.device import get_current_device
  from colossalai.zero import GeminiDDP, GeminiOptimizer
  from colossalai.zero.gemini.chunk import search_chunk_configuration
  from tests.kit.model_zoo import model_zoo, run_fwd_bwd
label_desc: Use image node_id 2 for calling a global function with example usage: exam_gpt_fwd_bwd(placement_config, keep_gather, model_name, use_grad_checkpoint, master_weights) without return types