docstr: Tests that using checkpoint_wrapper or the functional
  torch.utils.checkpoint (with the same reentrant config)
  results in the same maximum memory usage, i.e. they are
  equivalent memory usage wise.
docstr_tok: ["Tests","that","using","checkpoint_wrapper","or","the","functional","torch.utils.checkpoint","(","with","the","same","reentrant","config",")","results","in","the","same","maximum","memory","usage",",","i.e",".","they","are","equivalent","memory","usage","wise","."]
code: def test_checkpoint_wrapper_parity(self):
      """
      Tests that using checkpoint_wrapper or the functional
      torch.utils.checkpoint (with the same reentrant config)
      results in the same maximum memory usage, i.e. they are
      equivalent memory usage wise.
      """
  
      class Model(nn.Module):
          def __init__(
              self,
              n: int,
              use_cp: bool,
              use_wrapper: bool = False,
              use_reentrant: bool = True,
          ):
              super().__init__()
              self.layers = nn.ModuleList()
              self.n = n
              self.use_cp = use_cp
              self.use_wrapper = use_wrapper
              self.use_reentrant = use_reentrant
              wrp = partial(
                  checkpoint_wrapper,
                  checkpoint_impl=CheckpointImpl.REENTRANT
                  if use_reentrant
                  else CheckpointImpl.NO_REENTRANT,
              )
              for i in range(self.n):
                  l = nn.Sequential(
                      nn.Linear(256, 256),
                      nn.Linear(256, 256),
                      nn.Linear(256, 256),
                  )
                  use_checkpoint_wrapper = self.use_wrapper
                  if use_checkpoint_wrapper:
                      l = wrp(l)
                  self.layers.append(l)
  
          def forward(self, x):
              for i in range(self.n):
                  if self.use_wrapper or not self.use_cp:
                      x = self.layers[i](x)
                  else:
                      x = checkpoint(
                          self.layers[i],
                          x,
                          use_reentrant=self.use_reentrant,
                      )
              return x
  
      def test(use_checkpointing, use_wrapper, use_reentrant):
          a = Model(
              8,
              use_checkpointing,
              use_wrapper=use_wrapper,
              use_reentrant=use_reentrant,
          ).cuda()
          x = torch.randn(10000, 256, requires_grad=True).cuda()
          torch.cuda.reset_peak_memory_stats()
          loss = a(x).sum()
          loss.backward()
          return torch.cuda.max_memory_allocated()
  
      functional_no_reentrant = test(
          use_checkpointing=True, use_wrapper=False, use_reentrant=False
      )
      wrapper_no_reentrant = test(
          use_checkpointing=False, use_wrapper=True, use_reentrant=False
      )
      self.assertEqual(functional_no_reentrant, wrapper_no_reentrant)
  
      functional_reentrant = test(
          use_checkpointing=True, use_wrapper=False, use_reentrant=True
      )
      wrapper_reentrant = test(
          use_checkpointing=False, use_wrapper=True, use_reentrant=True
      )
      self.assertEqual(functional_reentrant, wrapper_reentrant)
  
file_path: pytorch/test/distributed/fsdp/test_checkpoint_wrapper.py
imports: import contextlib
  import unittest
  from copy import deepcopy
  from functools import partial
  import torch
  import torch.nn
  from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import apply_activation_checkpointing, checkpoint_wrapper, CheckpointImpl, CheckpointWrapper, offload_wrapper, OffloadWrapper
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.testing._internal.common_utils import run_tests, TestCase
  from torch.utils.checkpoint import checkpoint
label_desc: Use image node_id 4 for calling the CheckpointWrapperTest obj's underlying member method code with example usage: obj.test_checkpoint_wrapper_parity() and returns: torch, x