code: def test(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)
      param_optimizer = list(model.named_parameters())
      mask1 = torch.zeros_like(param_optimizer[0][1].data)
      mask2 = torch.zeros_like(param_optimizer[0][1].data)
      for col in range(mask1.size()[1]):
          mask1[0][col] += 1
          mask2[1][col] += 1
      mask1 = torch.flatten(mask1)
      mask2 = torch.flatten(mask2)
  
      optimizer_grouped_parameters_1 = [
          {
              "params": [param_optimizer[0][1]],
              "weight_decay": 0.01,
              "exp_avg_mask": mask1,
          },
          {"params": [param_optimizer[1][1]], "weight_decay": 0.01},
      ]
  
      optimizer_grouped_parameters_2 = [
          {
              "params": [param_optimizer[0][1]],
              "weight_decay": 0.01,
              "exp_avg_mask": mask2,
          },
          {"params": [param_optimizer[1][1]], "weight_decay": 0.01},
      ]
  
      optimizer_grouped_parameters_3 = [
          {"params": [param_optimizer[0][1]], "weight_decay": 0.01},
          {"params": [param_optimizer[1][1]], "weight_decay": 0.01},
      ]
  
      model_1, optimizer_1, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=optimizer_grouped_parameters_1,
      )
      data_loader = random_dataloader(
          model=model_1,
          total_samples=10,
          hidden_dim=hidden_dim,
          device=model_1.device,
      )
      for n, batch in enumerate(data_loader):
          loss = model_1(batch[0], batch[1])
          model_1.backward(loss)
          model_1.step()
      # Test whether momentum mask still exist after saving checkpoint
      mask1 = mask1.to(
          device=optimizer_1.param_groups[0]["exp_avg_mask"].device
      )
      assert torch.allclose(
          optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07
      ), f"Incorrect momentum mask"
      save_folder = os.path.join(tmpdir, "saved_checkpoint")
      model_1.save_checkpoint(save_folder, tag=None)
      assert torch.allclose(
          optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07
      ), f"Momentum mask should not change after saving checkpoint"
  
      model_2, optimizer_2, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=optimizer_grouped_parameters_2,
      )
      # Test whether momentum mask stays the same after loading checkpoint
      mask2 = mask2.to(
          device=optimizer_2.param_groups[0]["exp_avg_mask"].device
      )
      assert torch.allclose(
          optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07
      ), f"Incorrect momentum mask"
      model_2.load_checkpoint(
          save_folder,
          tag=None,
          load_optimizer_states=True,
          load_lr_scheduler_states=True,
      )
      assert torch.allclose(
          optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07
      ), f"Momentum mask should not change after loading checkpoint"
      # Test whether worker&server error is reset
      for v in optimizer_2.state.values():
          assert "worker_error" not in v, f"Incorrect worker error"
          assert "server_error" not in v, f"Incorrect server error"
  
      model_3, optimizer_3, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=optimizer_grouped_parameters_3,
      )
      optimizer_3.optimizer.freeze_step = 20
      data_loader = random_dataloader(
          model=model_3,
          total_samples=50,
          hidden_dim=hidden_dim,
          device=model_3.device,
      )
      for n, batch in enumerate(data_loader):
          loss = model_3(batch[0], batch[1])
          model_3.backward(loss)
          model_3.step()
      # Test whether momentum mask stays the same after loading checkpoint
      assert (
          "exp_avg_mask" not in optimizer_3.param_groups[0]
      ), f"Incorrect momentum mask"
      model_3.load_checkpoint(
          save_folder,
          tag=None,
          load_optimizer_states=True,
          load_lr_scheduler_states=True,
      )
      assert (
          "exp_avg_mask" not in optimizer_3.param_groups[0]
      ), f"Momentum mask should not change after loading checkpoint"
      # Test whether worker&server error is reset
      for v in optimizer_3.state.values():
          assert "worker_error" not in v, f"Incorrect worker error"
          assert "server_error" not in v, f"Incorrect server error"
  
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 1 for calling the TestZeroOneAdamCheckpointing obj's underlying member method code with example usage: obj.test(tmpdir) without return types