code: def test_multiple_forward_checkpoint(
      precision, flatten, wrap_bn, model_type, bn_type
  ):
      mixed_precision = precision == "mixed"
      flatten = flatten == "flatten"
      wrap_bn = wrap_bn == "auto_wrap_bn"
      fp32_reduce_scatter = True if mixed_precision else None
      with_model2 = model_type == "model2"
      with_sync_bn = bn_type == "sync_bn"
  
      if (
          torch_version() >= (1, 7, 0)
          and torch_version() < (1, 8, 0)
          and with_sync_bn
      ):
          # SyncBN is buggy in 1.7, errors like:
          # E         File "/home/circleci/venv/lib/python3.8/site-packages/torch/nn/modules/_functions.py", line 13, in forward
          # E           dtype=running_mean.dtype,
          # E       AttributeError: 'NoneType' object has no attribute 'dtype'
          pytest.skip("SyncBatchNorm in 1.7 is buggy")
  
      if with_sync_bn and not wrap_bn:
          pytest.skip("SyncBatchNorm requires auto_wrap_bn")
  
      if torch_version() < (1, 8, 0) and flatten:
          # 1.6 and 1.7 throws this error:
          #   RuntimeError: Trying to backward through the graph a second time, but the saved
          #   intermediate results have already been freed. Specify retain_graph=True when calling
          #   backward the first time.
          pytest.skip("older pytorch throws error when flatten is used")
  
      world_size = 2
      expected_losses = None
  
      # Ensure ddp == fsdp when modules are called multiple times per forward pass with/without checkpointing, forward
      # counters and reducer bucketing.
      #
      # The bucketing check exists because the asynchronous gradient reduction it induces can interact with multiple
      # forward passes in complex ways. For example, in the midst of a sharded backward pass, `parameter.grad` may only be
      # `None` or an unsharded gradient tensor. The sharded tensor is then set at the end of the backwards pass. But a
      # unit test with bucketing enabled might not catch violations of this invariant. For very small models, like the
      # kind used in this unit test, bucketing will delay gradient reduction until after all the gradient computation is
      # done. If the reduction incorrectly sets `.grad` to the _sharded_ variant, the test might not fail, since the
      # gradient computations have already happened. Toggling bucketing helps verify that gradient reduction and
      # computation interact correctly.
      combinations = []
      for with_fsdp in [False, True]:
          for with_checkpoint in [False, True]:
              if not with_fsdp and with_checkpoint:
                  continue
              for with_bucketing in [False, True]:
                  if not with_fsdp and with_bucketing:
                      continue
                  combinations.append(
                      (with_fsdp, with_checkpoint, with_bucketing)
                  )
      print("")
      print("Testing the following configurations:")
      for with_fsdp, with_checkpoint, with_bucketing in combinations:
          print(
              f"  fsdp {with_fsdp} ckpt {with_checkpoint} bucketing {with_bucketing}"
          )
  
      for with_fsdp, with_checkpoint, with_bucketing in combinations:
          if with_bucketing:
              bucket_cap_mb = 25
          else:
              bucket_cap_mb = 0
          final_losses = _get_cached_results(
              world_size,
              with_model2,
              with_sync_bn,
              with_fsdp,
              with_checkpoint,
              mixed_precision,
              flatten,
              wrap_bn,
              fp32_reduce_scatter,
              bucket_cap_mb,
          )
          if expected_losses is None:
              expected_losses = final_losses
          else:
              print(
                  f"checking: fsdp {with_fsdp} ckpt {with_checkpoint} bucketing {with_bucketing} with ddp+no_ckpt"
              )
  
              def check(exp, res):
                  assert list(exp.keys()) == list(
                      res.keys()
                  ), f"{list(exp.keys())} vs. {list(res.keys())}"
                  rtol = 1e-4
                  atol = 1e-5
                  if (
                      with_model2
                      and mixed_precision
                      and torch_version() >= (1, 9, 0)
                  ):
                      # On CI, with longer model2, mixed precsion and 1.9, even ddp vs. ddp+ckpt has
                      # larger errors.
                      rtol = 1e-3
                      atol = 1e-4
                  for key in exp.keys():
                      exp_loss = exp[key]
                      res_loss = res[key]
                      torch.testing.assert_allclose(
                          exp_loss, res_loss, rtol=rtol, atol=atol
                      )
  
              check(expected_losses, final_losses)
  
imports: import contextlib
  import pickle
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
  from fairscale.nn.wrap import enable_wrap, wrap
label_desc: Use image node_id 4 for calling a global function with example usage: test_multiple_forward_checkpoint(precision, flatten, wrap_bn, model_type, bn_type) without return types