code: def run_grad_acc_test(test_args):
      model_fn, *_ = next(
          iter(
              model_zoo.get_sub_registry("transformers_gpt_lm").values()
          )
      )
      model = model_fn()
      optimizer = HybridAdam(model.parameters())
      origin_model = copy.deepcopy(model).cuda()
      origin_optimizer = HybridAdam(origin_model.parameters())
  
      plugin = HybridParallelPlugin(
          tp_size=test_args["tp"],
          pp_size=test_args["pp"],
          pp_style=test_args["pp_style"],
          zero_stage=test_args["zero"],
          num_model_chunks=test_args["num_model_chunks"],
          enable_fused_normalization=True,
          num_microbatches=test_args["num_microbatches"],
          precision=test_args["precision"],
      )
      booster = Booster(plugin=plugin)
  
      dataset = RandomDataset(
          num_samples=test_args["batch_size"]
          * test_args["num_steps"]
          * plugin.dp_size,
          max_length=test_args["max_length"],
          vocab_size=model.config.vocab_size,
      )
      dataloader = plugin.prepare_dataloader(
          dataset,
          batch_size=test_args["batch_size"],
          shuffle=True,
          drop_last=True,
      )
  
      model, optimizer, _, dataloader, _ = booster.boost(
          model, optimizer, dataloader=dataloader
      )
  
      grad_accu_step = test_args["gradient_accumulation_step"]
      for step, batch in enumerate(dataloader):
          batch = move_to_cuda(batch)
          # train origin model
          origin_output = origin_model(**batch)
          origin_loss = origin_output[0] / grad_accu_step
          origin_loss.backward()
  
          if (step + 1) % grad_accu_step != 0 and test_args[
              "zero"
          ] != 2:
              ctx = booster.no_sync(model, optimizer)
          else:
              ctx = nullcontext()
  
          with ctx:
              if plugin.stage_manager is not None:
                  batch = iter([batch])
                  booster.execute_pipeline(
                      batch,
                      model,
                      criterion=lambda outputs, inputs: outputs[0]
                      / grad_accu_step,
                      optimizer=optimizer,
                      return_loss=False,
                  )
              else:
                  outputs = model(**batch)
                  loss = outputs[0] / grad_accu_step
                  booster.backward(loss, optimizer)
  
          if (step + 1) % grad_accu_step == 0:
              # update origin model weight
              origin_optimizer.step()
              origin_optimizer.zero_grad()
  
              # update sharded model
              optimizer.step()
              optimizer.zero_grad()
  
      # tricky code here, shard the origin model inorder to check the parameters in the same stage.
      origin_model, origin_optimizer, _, dataloader, _ = booster.boost(
          origin_model, origin_optimizer, dataloader=dataloader
      )
      for p1, p2 in zip(
          model.unwrap().parameters(),
          origin_model.unwrap().parameters(),
      ):
          assert_close(p1.to(p2.dtype), p2, atol=1e-2, rtol=1e-2)
  
imports: import copy
  from contextlib import nullcontext
  from typing import Optional
  import torch
  import torch.distributed
  from torch.testing import assert_close
  from torch.utils.data import Dataset
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import HybridParallelPlugin
  from colossalai.fx import is_compatible_with_meta
  from colossalai.lazy.lazy_init import LazyInitContext
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
  from colossalai.utils import get_current_device, set_seed
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 4 for calling a global function with example usage: run_grad_acc_test(test_args) without return types