docstr: Test state_dict and load_state_dict functions
  
  Two models are constructed, possibly on different process
  groups. One of the models is trained for a few steps, a
  checkpoint is saved, and the checkpoint is loaded on the other
  model. Both models are then trained for a few steps and
  checked to make sure that they produce identical results.
  
  Arguments:
      rtol (float): Relative tolerance for numerical checks (see
          torch.allclose).
      atol (float): Absolute tolerance for numerical checks (see
          torch.allclose).
      num_layers (int): Number of layers in test model.
      layer_size (int): Number of features in model layers.
      num_steps (int): Number of training steps to perform
          before and after checkpointing.
      save_group_size (int): Process group size for model that
          saves the checkpoint. Uses the default process group
          by default.
      load_group_size (int): Process group size for model that
          loads the checkpoint. Uses the default process group
          by default.
      save_model_kwargs (dict): keyword arguments passed to
          make_models when constructing the model that saves the
          checkpoint.
      load_model_kwargs (dict): keyword arguments passed to
          make_models when constructing the model that loads the
          checkpoint.
code: def test_checkpoint(
      self,
      rtol: Optional[float] = None,
      atol: Optional[float] = None,
      num_layers: int = 2,
      layer_size: int = 2,
      num_steps: int = 3,
      save_group_size: Optional[int] = None,
      load_group_size: Optional[int] = None,
      save_model_kwargs: Optional[dict] = None,
      load_model_kwargs: Optional[dict] = None,
  ):
      """Test state_dict and load_state_dict functions
  
      Two models are constructed, possibly on different process
      groups. One of the models is trained for a few steps, a
      checkpoint is saved, and the checkpoint is loaded on the other
      model. Both models are then trained for a few steps and
      checked to make sure that they produce identical results.
  
      Arguments:
          rtol (float): Relative tolerance for numerical checks (see
              torch.allclose).
          atol (float): Absolute tolerance for numerical checks (see
              torch.allclose).
          num_layers (int): Number of layers in test model.
          layer_size (int): Number of features in model layers.
          num_steps (int): Number of training steps to perform
              before and after checkpointing.
          save_group_size (int): Process group size for model that
              saves the checkpoint. Uses the default process group
              by default.
          load_group_size (int): Process group size for model that
              loads the checkpoint. Uses the default process group
              by default.
          save_model_kwargs (dict): keyword arguments passed to
              make_models when constructing the model that saves the
              checkpoint.
          load_model_kwargs (dict): keyword arguments passed to
              make_models when constructing the model that loads the
              checkpoint.
  
      """
  
      # Initialize process groups
      world_size = torch.distributed.get_world_size()
      if save_group_size is None:
          save_group_size = world_size
          save_group = None
      else:
          if save_group_size > world_size:
              self.skipTest(
                  f"Requires {save_group_size} ranks, found {world_size}"
              )
          save_ranks = list(range(save_group_size))
          save_group = torch.distributed.new_group(ranks=save_ranks)
      if load_group_size is None:
          load_group_size = world_size
          load_group = None
      else:
          if load_group_size > world_size:
              self.skipTest(
                  f"Requires {load_group_size} ranks, found {world_size}"
              )
          load_ranks = list(range(load_group_size))
          load_group = torch.distributed.new_group(ranks=load_ranks)
  
      # Construct two models with same config and different params
      torch.manual_seed(self.seed)
      if self.rank < save_group_size:
          if not save_model_kwargs:
              save_model_kwargs = {}
          _, _, model_save, optim_save = make_models(
              num_layers,
              layer_size,
              process_group=save_group,
              average_grad_sync=False,
              **save_model_kwargs,
          )
          optim_save.init_params(
              reversed(list(model_save.parameters()))
          )
      torch.manual_seed(self.seed + 1)
      if self.rank < load_group_size:
          if not load_model_kwargs:
              load_model_kwargs = {}
          _, _, model_load, optim_load = make_models(
              num_layers,
              layer_size,
              process_group=load_group,
              average_grad_sync=False,
              **load_model_kwargs,
          )
          optim_load.init_params(list(model_load.parameters()))
  
      batch_size = 2 * save_group_size * load_group_size
  
      def make_global_batch() -> torch.Tensor:
          """Generate random tensor on root rank and broadcast"""
          x = torch.empty(batch_size, layer_size, device="cuda")
          if self.rank == 0:
              torch.rand(x.size(), out=x)
              x -= 0.5
          torch.distributed.broadcast(x, src=0)
          return x
  
      def to_local_batch(
          global_batch: torch.Tensor,
          group: Optional[torch.distributed.ProcessGroup],
      ) -> Optional[torch.Tensor]:
          """Get local portion of tensor that is replicated across all ranks"""
          group_size = torch.distributed.get_world_size(group)
          if group_size < 0:
              return None
          local_batch_size = batch_size // group_size
          batch_start = self.rank * local_batch_size
          batch_end = (self.rank + 1) * local_batch_size
          return global_batch[batch_start:batch_end, ...]
  
      def to_global_batch(
          local_batch: torch.Tensor,
          group: Optional[torch.distributed.ProcessGroup],
      ) -> torch.Tensor:
          """Gather distributed tensor and broadcast to all ranks"""
  
          # Allocate buffer
          global_batch = torch.empty(
              batch_size, layer_size, device="cuda"
          )
  
          # Gather data on root rank
          group_size = torch.distributed.get_world_size(group)
          if group_size > 0:
              local_batches = None
              if self.rank == 0:
                  local_batch_size = batch_size // group_size
                  local_batches = [
                      global_batch[
                          rank
                          * local_batch_size : (rank + 1)
                          * local_batch_size,
                          ...,
                      ]
                      for rank in range(group_size)
                  ]
              torch.distributed.gather(
                  local_batch,
                  local_batches,
                  dst=0,
                  group=group,
              )
  
          # Broadcast data to all ranks
          torch.distributed.broadcast(global_batch, src=0)
          return global_batch
  
      # Train one of the models
      torch.manual_seed(self.seed + 2)
      for step in range(num_steps):
          if self.rank < save_group_size:
              optim_save.zero_grad()
          x = make_global_batch()
          dy = make_global_batch()
          if self.rank < save_group_size:
              x = to_local_batch(x, save_group)
              dy = to_local_batch(dy, save_group)
              y = model_save(x)
              y.backward(dy)
              optim_save.step()
  
      # Make sure models are different
      if self.rank < min(save_group_size, load_group_size):
          for param_save, param_load in zip(
              model_save.parameters(), model_load.parameters()
          ):
              self.assertRaises(
                  AssertionError,
                  torch.testing.assert_close,
                  param_load,
                  param_save,
                  rtol=rtol,
                  atol=atol,
              )
  
      # Save state
      state_bytes = None
      if self.rank < save_group_size:
          state_dict = {
              "model": model_save.state_dict(),
              "optim": optim_save.state_dict(),
          }
          byte_stream = io.BytesIO()
          torch.save(state_dict, byte_stream)
          state_bytes = byte_stream.getvalue()
  
      # Broadcast state from root rank and load
      if self.rank < load_group_size:
          if load_group_size != save_group_size:
              if self.rank != 0:
                  state_bytes = None
              state_bytes = [state_bytes]
              torch.distributed.broadcast_object_list(
                  state_bytes,
                  src=0,
                  group=load_group,
              )
              state_bytes = state_bytes[0]
          state_dict = torch.load(io.BytesIO(state_bytes))
          model_load.load_state_dict(state_dict["model"])
          optim_load.load_state_dict(state_dict["optim"])
  
      # Make sure models are identical
      if self.rank < min(save_group_size, load_group_size):
          for param_save, param_load in zip(
              model_save.parameters(), model_load.parameters()
          ):
              torch.testing.assert_close(
                  param_load, param_save, rtol=rtol, atol=atol
              )
  
      # Train both models
      torch.manual_seed(self.seed + 3)
      for step in range(num_steps):
          # Reset grads
          if self.rank < save_group_size:
              optim_save.zero_grad()
          if self.rank < load_group_size:
              optim_load.zero_grad()
  
          # Synthetic data
          x = make_global_batch()
          dy = make_global_batch()
  
          # Training step for model that saved checkpoint
          y_save = None
          dx_save = None
          if self.rank < save_group_size:
              x_save = to_local_batch(x, save_group)
              x_save = x_save.detach().clone().requires_grad_(True)
              dy_save = to_local_batch(dy, save_group)
              y_save = model_save(x_save)
              y_save.backward(dy_save)
              dx_save = x_save.grad
          y_save = to_global_batch(y_save, save_group)
          dx_save = to_global_batch(dx_save, save_group)
  
          # Training step for model that loaded checkpoint
          y_load = None
          dx_load = None
          if self.rank < load_group_size:
              x_load = to_local_batch(x, load_group)
              x_load = x_load.detach().clone().requires_grad_(True)
              dy_load = to_local_batch(dy, load_group)
              y_load = model_load(x_load)
              y_load.backward(dy_load)
              dx_load = x_load.grad
          y_load = to_global_batch(y_load, load_group)
          dx_load = to_global_batch(dx_load, load_group)
  
          # Check that data tensors match
          torch.testing.assert_close(
              y_load, y_save, rtol=rtol, atol=atol
          )
          torch.testing.assert_close(
              dx_load, dx_save, rtol=rtol, atol=atol
          )
  
          # Optimizer step
          if self.rank < save_group_size:
              optim_save.step()
          if self.rank < load_group_size:
              optim_load.step()
  
          # Check that parameters match
          if self.rank < min(save_group_size, load_group_size):
              for param_save, param_load in zip(
                  model_save.parameters(), model_load.parameters()
              ):
                  torch.testing.assert_close(
                      param_load,
                      param_save,
                      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 19 for calling the TestDistributedFusedAdam obj's underlying member method code with example usage: obj.test_checkpoint(rtol, atol, num_layers, layer_size, num_steps, save_group_size, load_group_size, save_model_kwargs, load_model_kwargs) and returns: x, global_batch, global_batch, None