code: def test_checkpoint_pipe_module(self, base_topo, test_topo, tmpdir):
      checkpoint_engine = TorchCheckpointEngine()
      base_model = LinearStackPipe(topology=base_topo)
      base_model.save_state_dict(
          tmpdir, checkpoint_engine=checkpoint_engine
      )
  
      dist.barrier()
  
      test_model = LinearStackPipe(topology=test_topo)
      test_model.load_state_dir(
          tmpdir, checkpoint_engine=checkpoint_engine
      )
  
      # Base and test can have different lengths, so make sure we map from the
      # smaller to larger model
      if len(base_model.forward_funcs) < len(test_model.forward_funcs):
          A = base_model
          B = test_model
      else:
          A = test_model
          B = base_model
  
      # Compare layers individually since partitions are different
      for idx, A_layer in enumerate(A.forward_funcs):
          if not hasattr(A_layer, "parameters"):
              # Skip functionals, etc.
              continue
  
          # Find the corresponding layer in B
          global_idx = idx + A._local_start
          B_local_idx = global_idx - B._local_start
          B_layer = B.forward_funcs[B_local_idx]
  
          # Compare layer parameters
          for p0, p1 in zip(A_layer.parameters(), B_layer.parameters()):
              assert torch.allclose(
                  p0, p1, atol=1e-07
              ), f"Model state {p0} is not equal to {p1}"
  
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 2 for calling the TestPipelineCheckpoint obj's underlying member method code with example usage: obj.test_checkpoint_pipe_module(base_topo, test_topo, tmpdir) without return types