code: def test_checkpoint_pipe_engine(self, zero_stage, tmpdir):
      config_dict = {
          "train_batch_size": 2,
          "train_micro_batch_size_per_gpu": 1,
          "steps_per_print": 1,
          "optimizer": {"type": "Adam", "params": {"lr": 1e-5}},
          "zero_optimization": {"stage": zero_stage},
          "fp16": {"enabled": zero_stage > 0},
          "scheduler": {
              "type": "OneCycle",
              "params": {
                  "cycle_first_step_size": 1000,
                  "cycle_first_stair_count": 500,
                  "cycle_second_step_size": 1000,
                  "cycle_second_stair_count": 500,
                  "decay_step_size": 1000,
                  "cycle_min_lr": 0.0001,
                  "cycle_max_lr": 0.0010,
                  "decay_lr_rate": 0.001,
                  "cycle_min_mom": 0.85,
                  "cycle_max_mom": 0.99,
                  "decay_mom_rate": 0.0,
              },
          },
      }
  
      models = [LinearStackPipe(num_stages=2) for _ in range(2)]
      checkpoint_correctness_verification(
          config_dict=config_dict,
          models=models,
          hidden_dim=models[0].hidden_dim,
          tmpdir=tmpdir,
          fp16=config_dict["fp16"]["enabled"],
          load_optimizer_states=True,
          load_lr_scheduler_states=True,
          train_batch=True,
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/checkpoint/test_pipeline.py
imports: from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from unit.common import DistributedTest
  from unit.simple_model import *
  from unit.checkpoint.common import checkpoint_correctness_verification
  import pytest
label_desc: Use image node_id 1 for calling the TestPipelineCheckpoint obj's underlying member method code with example usage: obj.test_checkpoint_pipe_engine(zero_stage, tmpdir) without return types