code: def test_matches_pytorch(
      self,
      rtol: Optional[float] = None,
      atol: Optional[float] = None,
      num_layers: int = 11,
      layer_size: int = 7,
      batch_size: int = 3,
      num_steps: int = 3,
      micro_batch_steps: int = 3,
      adam_w_mode: bool = True,
      overlap_communication: bool = True,
      use_nosync: 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",
      contiguous_buffers: bool = False,
      store_params: bool = False,
      store_param_remainders: bool = False,
      bucket_cap_mb: float = 71 / (4 * 1024 * 1024),
      init_optim_func: Optional[
          Callable[[DistributedFusedAdam], None]
      ] = None,
  ):
      torch.manual_seed(self.seed + self.rank)
  
      # Identical models with data-parallel and ZeRO
      ref_model, ref_optim, dist_model, dist_optim = make_models(
          num_layers,
          layer_size,
          adam_w_mode=adam_w_mode,
          model_dtype=model_dtype,
          optim_dtype=optim_dtype,
          grad_sync_dtype=grad_sync_dtype,
          param_sync_dtype=param_sync_dtype,
          device=device,
          overlap_communication=overlap_communication,
          contiguous_buffers=contiguous_buffers,
          store_params=store_params,
          store_param_remainders=store_param_remainders,
          bucket_cap_mb=bucket_cap_mb,
      )
  
      # Initialize distributed optimizer
      if init_optim_func is not None:
          init_optim_func(dist_optim)
  
      # Training loop
      for step in range(num_steps):
          # Reset gradients
          ref_optim.zero_grad()
          dist_optim.zero_grad()
  
          # Forward and backward passes
          for micro_step in range(micro_batch_steps):
              # Synthetic data
              x = torch.rand(batch_size, layer_size) - 0.5
              dy = torch.rand_like(x) - 0.5
              x = x.to(dtype=model_dtype, device=device)
              dy = dy.to(dtype=model_dtype, device=device)
  
              # Reference implementation
              x_ref = x.detach().clone().requires_grad_(True)
              y_ref = ref_model(x_ref)
              y_ref.backward(dy)
  
              # Distributed implementation
              x_dist = x.detach().clone().requires_grad_(True)
              y_dist = dist_model(x_dist)
              backward_context = dummy_context
              if use_nosync and micro_step < micro_batch_steps - 1:
                  backward_context = dist_optim.no_sync
              with backward_context():
                  y_dist.backward(dy)
  
              # Check that data tensors match
              torch.testing.assert_close(
                  y_dist, y_ref, rtol=rtol, atol=atol
              )
              torch.testing.assert_close(
                  x_dist.grad, x_ref.grad, rtol=rtol, atol=atol
              )
  
          # Optimization step
          ref_optim.step()
          dist_optim.step()
  
          # Check that parameters match
          for ref_param, dist_param in zip(
              ref_model.parameters(), dist_model.parameters()
          ):
              torch.testing.assert_close(
                  dist_param, ref_param, rtol=rtol, atol=atol
              )
  
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 the TestDistributedFusedAdam obj's underlying member method code with example usage: obj.test_matches_pytorch(rtol, atol, num_layers, layer_size, batch_size, num_steps, micro_batch_steps, adam_w_mode, overlap_communication, use_nosync, model_dtype, optim_dtype, grad_sync_dtype, param_sync_dtype, device, contiguous_buffers, store_params, store_param_remainders, bucket_cap_mb, init_optim_func) without return types