code: def test_pp_basic(self, inputs, tmpdir):
      # basic test case, mp_size=2, pp_size=2, verify ckpt saving/loading.
      args_defaults = {
          "num_layers": 8,
          "hidden_size": 128,
          "num_attention_heads": 8,
          "max_position_embeddings": 128,
      }
      mp_size = self.mp_size
      pp_size = self.pp_size
      world_size = self.world_size
  
      topo = get_topology(mp_size, pp_size, world_size)
      gpt2_pipe_model = GPT2ModelPipe(
          num_layers=8,
          num_stages=pp_size,
          mp_size=mp_size,
          args_others=args_defaults,
          topo=topo,
      )
      model = get_deepspeed_model(gpt2_pipe_model)
  
      tag = "pp_basic"
      state_dict = {}
      state_dict["checkpoint_version"] = get_megatron_version()
      model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
  
      if model.is_first_stage() or model.is_last_stage():
          loader = RepeatingLoader([(inputs[0], 0)])
          data_iter = iter(loader)
      else:
          data_iter = None
  
      baseline = model.eval_batch(
          data_iter=data_iter, compute_loss=False, reduce_output=None
      )
  
      dist.barrier()
      model.load_checkpoint(
          tmpdir,
          tag=tag,
          load_optimizer_states=False,
          load_lr_scheduler_states=False,
      )
      dist.barrier()
  
      test = model.eval_batch(
          data_iter=data_iter, compute_loss=False, reduce_output=None
      )
  
      if test is not None:
          assert len(baseline) == len(test)
          # Compare outputs of each microbatch
          for mb in range(len(baseline)):
              for b, t in zip(baseline[mb], test[mb]):
                  if b.is_floating_point():  # don't compare masks
                      assert torch.allclose(
                          b, t, atol=1e-07
                      ), f"Baseline output {baseline} is not equal to save-then-load output {test}"
  
imports: import os
  import torch
  import deepspeed
  import pytest
  import random
  import numpy
  import deepspeed.comm
  from unit.common import DistributedTest, DistributedFixture
  from unit.megatron_model import get_megatron_version
  from unit.megatron_model import MockGPT2ModelPipe
  from deepspeed.utils import RepeatingLoader
label_desc: Use image node_id 1 for calling the TestConfigurablePP obj's underlying member method code with example usage: obj.test_pp_basic(inputs, tmpdir) without return types