code: def test_overflow(self, tmpdir):
      config_dict = {
          "train_batch_size": 2,
          "steps_per_print": 1,
          "optimizer": {
              "type": "ZeroOneAdam",
              "params": {
                  "lr": 0.00015,
                  "weight_decay": 0.01,
                  "var_freeze_step": 4,
                  "var_update_scaler": 1,
                  "local_step_scaler": 1,
                  "local_step_clipper": 2,
                  "cuda_aware": False,
                  "comm_backend_name": "nccl",
              },
          },
          "gradient_clipping": 1.0,
          "fp16": {
              "enabled": True,
              "loss_scale": 0,
              "initial_scale_power": 16,
          },
      }
      hidden_dim = 10
  
      model = SimpleModel(hidden_dim)
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=model.parameters(),
      )
      data_loader = random_dataloader(
          model=model,
          total_samples=100,
          hidden_dim=hidden_dim,
          device=model.device,
      )
      save_folder = os.path.join(tmpdir, "saved_checkpoint")
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          if dist.get_rank() == 0 and n >= 10:
              loss = loss * 1000000.0
          model.backward(loss)
          dist.barrier()
          model.step()
          dist.barrier()
          model.save_checkpoint(save_folder, tag=None)
  
imports: import torch
  import torch.nn
  import deepspeed.comm
  import deepspeed
  import pytest
  import copy
  import os
  import numpy
  from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
  from deepspeed.ops.op_builder import OpBuilder
  from deepspeed.runtime.pipe.module import PipelineModule
  from unit.common import DistributedTest
  from unit.simple_model import SimpleModel, random_dataloader
  from unit.alexnet_model import AlexNetPipe, train_cifar
label_desc: Use image node_id 2 for calling the TestZeroOneAdamCheckpointing obj's underlying member method code with example usage: obj.test_overflow(tmpdir) without return types