code: def make_models(
      num_layers: int,
      size: int,
      adam_w_mode: bool = True,
      model_dtype: torch.dtype = torch.float32,
      optim_dtype: Optional[torch.dtype] = None,
      grad_sync_dtype: Optional[torch.dtype] = None,
      param_sync_dtype: Optional[torch.dtype] = None,
      device: torch.device = "cuda",
      process_group: Optional[torch.distributed.ProcessGroup] = None,
      average_grad_sync: bool = True,
      overlap_communication: bool = True,
      contiguous_buffers: bool = False,
      store_params: bool = False,
      store_param_remainders: bool = False,
      bucket_cap_mb: float = 71 / (4 * 1024 * 1024),
  ):
      # Construct models with same parameters
      ref_model = SimpleModel(num_layers, size).to(
          dtype=model_dtype, device=device
      )
      dist_model = SimpleModel(num_layers, size).to(
          dtype=model_dtype, device=device
      )
      with torch.no_grad():
          for ref_param, dist_param in zip(
              dist_model.parameters(), ref_model.parameters()
          ):
              dist_param.copy_(ref_param)
  
      # Initialize reference model with data-parallelism
      rank = torch.distributed.get_rank()
      ref_model = torch.nn.parallel.DistributedDataParallel(
          ref_model,
          device_ids=[rank] if device == "cuda" else None,
          output_device=rank if device == "cuda" else None,
          process_group=process_group,
      )
  
      # Construct optimizers with same hyperparameters
      if optim_dtype is None:
          optim_dtype = model_dtype
      optim_args = dict(
          lr=0.1, betas=(0.1, 0.2), eps=0.25, weight_decay=0.1
      )
      ref_optim_class = (
          torch.optim.AdamW if adam_w_mode else torch.optim.Adam
      )
      ref_optim = ref_optim_class(
          [
              {"params": list(ref_model.parameters())[1::2], "lr": 0.2},
              {"params": list(ref_model.parameters())[0::2]},
          ],
          **optim_args,
      )
      dist_optim = DistributedFusedAdam(
          [
              {
                  "params": list(dist_model.parameters())[1::2],
                  "lr": 0.2,
              },
              {"params": list(dist_model.parameters())[0::2]},
          ],
          adam_w_mode=adam_w_mode,
          overlap_grad_sync=overlap_communication,
          overlap_param_sync=overlap_communication,
          bucket_cap_mb=bucket_cap_mb,
          dtype=optim_dtype,
          grad_sync_dtype=grad_sync_dtype,
          param_sync_dtype=param_sync_dtype,
          process_group=process_group,
          average_grad_sync=average_grad_sync,
          contiguous_param_buffer=contiguous_buffers,
          contiguous_grad_buffer=contiguous_buffers,
          store_params=store_params,
          store_param_remainders=store_param_remainders,
          **optim_args,
      )
  
      return ref_model, ref_optim, dist_model, dist_optim
  
imports: from contextlib import contextmanager
  import io
  from typing import Callable, Optional, Tuple
  import unittest
  import warnings
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling a global function with example usage: make_models(num_layers, size, adam_w_mode, model_dtype, optim_dtype, grad_sync_dtype, param_sync_dtype, device, process_group, average_grad_sync, overlap_communication, contiguous_buffers, store_params, store_param_remainders, bucket_cap_mb) and returns: ref_model,  ref_optim,  dist_model,  dist_optim