code: def exam_zero_1_torch_ddp(
      world_size, dtype: torch.dtype, master_weights: bool
  ):
      """
      In this test, two pairs of model and optimizers are created.
      1. zero: use sharded optimizer and fp16 parameters
      2. torch: use torch DDP and fp32 parameters
  
      We feed these two sets of models with the same input and check if the
      differences in model output and updated parameters are within tolerance.
      """
      local_rank = torch.distributed.get_rank()
      seed_all(1453)
  
      # create models
      torch_model = MlpModel().cuda()
      zero_model = copy.deepcopy(torch_model).to(dtype)
  
      torch_model = DDP(torch_model.cuda(), static_graph=True).cuda()
  
      # create optimizer
      zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
  
      # we only test stage 1 here
      # in `check_sharded_param_consistency.py`, we will test whether
      # level 1 and 2 will produce exactly the same results
      zero_optimizer = LowLevelZeroOptimizer(
          zero_optimizer,
          overlap_communication=True,
          initial_scale=1,
          reduce_bucket_size=1024 * 1024,
          master_weights=master_weights,
      )
  
      torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
  
      seed_all(1453 + local_rank)
      # create
      input_data = torch.rand(32, 123).cuda()
  
      # zero-dp forward
      zero_output = zero_model(input_data.to(dtype))
  
      # torch-ddp forward
      torch_output = torch_model(input_data)
      loose_close(zero_output, torch_output, dtype=dtype)
  
      # zero-dp backward
      zero_optimizer.backward(zero_output.mean().float())
  
      # torch-ddp backward
      torch_output.mean().backward()
  
      # check grad
      for (n, p), z1p in zip(
          torch_model.named_parameters(), zero_model.parameters()
      ):
          if p.grad is not None:
              zero_grad_list = zero_optimizer._grad_store.get_partitioned_gradients_by_param_id(
                  0, id(z1p)
              )
              torch_grad_list = split_ddp_grad(p.grad, world_size)
              for zero_grad, torch_grad in zip(
                  zero_grad_list, torch_grad_list
              ):
                  loose_close(zero_grad, torch_grad, dtype=dtype)
  
      # zero-dp step
      zero_optimizer.step()
  
      # torch ddp step
      torch_optimizer.step()
  
      # check updated param
      for (n, p), z1p in zip(
          torch_model.named_parameters(), zero_model.parameters()
      ):
          loose_close(p.data, z1p.data, dtype=dtype)
  
file_path: colossalai/tests/test_zero/test_low_level/test_zero1_2.py
imports: import copy
  import pytest
  import torch
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing import assert_close
  import colossalai
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
  from colossalai.testing.random import seed_all
  from colossalai.zero import LowLevelZeroOptimizer
label_desc: Use image node_id 4 for calling a global function with example usage: exam_zero_1_torch_ddp(world_size, dtype, master_weights) without return types