code: def test_fsdp_memory(fsdp, ckpt):
      expected = {
          ("ddp", "no_ckpt"): {
              "iter 0: start": 9,
              "iter 0: after fwd": 346,
              "iter 0: after loss": 346,
              "iter 0: after bwd": 14,
              "iter 0: after step": 17,
              "iter 0: done": 13,
              "iter 1: start": 13,
              "iter 1: after fwd": 350,
              "iter 1: after loss": 350,
              "iter 1: after bwd": 17,
              "iter 1: after step": 17,
              "iter 1: done": 13,
              "iter 2: start": 13,
              "iter 2: after fwd": 350,
              "iter 2: after loss": 350,
              "iter 2: after bwd": 17,
              "iter 2: after step": 17,
              "iter 2: done": 13,
              "iter 3: start": 13,
              "iter 3: after fwd": 350,
              "iter 3: after loss": 350,
              "iter 3: after bwd": 17,
              "iter 3: after step": 17,
              "iter 3: done": 13,
          },
          ("fsdp", "no_ckpt"): {
              "iter 0: start": 3,
              "iter 0: after fwd": 340,
              "iter 0: after loss": 340,
              "iter 0: after bwd": 16,
              "iter 0: after step": 18,
              "iter 0: done": 5,
              "iter 1: start": 5,
              "iter 1: after fwd": 342,
              "iter 1: after loss": 342,
              "iter 1: after bwd": 18,
              "iter 1: after step": 18,
              "iter 1: done": 5,
              "iter 2: start": 5,
              "iter 2: after fwd": 342,
              "iter 2: after loss": 342,
              "iter 2: after bwd": 18,
              "iter 2: after step": 18,
              "iter 2: done": 5,
              "iter 3: start": 5,
              "iter 3: after fwd": 342,
              "iter 3: after loss": 342,
              "iter 3: after bwd": 18,
              "iter 3: after step": 18,
              "iter 3: done": 5,
          },
          ("fsdp_amp_default", "no_ckpt"): {
              "iter 0: start": 28,
              "iter 0: after fwd": 630,
              "iter 0: after loss": 630,
              "iter 0: after bwd": 67,
              "iter 0: after step": 93,
              "iter 0: done": 54,
              "iter 1: start": 54,
              "iter 1: after fwd": 657,
              "iter 1: after loss": 657,
              "iter 1: after bwd": 93,
              "iter 1: after step": 93,
              "iter 1: done": 54,
              "iter 2: start": 54,
              "iter 2: after fwd": 657,
              "iter 2: after loss": 657,
              "iter 2: after bwd": 93,
              "iter 2: after step": 93,
              "iter 2: done": 54,
              "iter 3: start": 54,
              "iter 3: after fwd": 657,
              "iter 3: after loss": 657,
              "iter 3: after bwd": 93,
              "iter 3: after step": 93,
              "iter 3: done": 54,
          },
          ("fsdp_amp_compute_dtype32", "no_ckpt"): {
              "iter 0: start": 28,
              "iter 0: after fwd": 657,
              "iter 0: after loss": 657,
              "iter 0: after bwd": 67,
              "iter 0: after step": 93,
              "iter 0: done": 54,
              "iter 1: start": 54,
              "iter 1: after fwd": 684,
              "iter 1: after loss": 684,
              "iter 1: after bwd": 93,
              "iter 1: after step": 93,
              "iter 1: done": 54,
              "iter 2: start": 54,
              "iter 2: after fwd": 684,
              "iter 2: after loss": 684,
              "iter 2: after bwd": 93,
              "iter 2: after step": 93,
              "iter 2: done": 54,
              "iter 3: start": 54,
              "iter 3: after fwd": 684,
              "iter 3: after loss": 684,
              "iter 3: after bwd": 93,
              "iter 3: after step": 93,
              "iter 3: done": 54,
          },
          ("ddp", "ckpt"): {
              "iter 0: start": 9,
              "iter 0: after fwd": 57,
              "iter 0: after loss": 57,
              "iter 0: after bwd": 14,
              "iter 0: after step": 17,
              "iter 0: done": 13,
              "iter 1: start": 13,
              "iter 1: after fwd": 61,
              "iter 1: after loss": 61,
              "iter 1: after bwd": 17,
              "iter 1: after step": 17,
              "iter 1: done": 13,
              "iter 2: start": 13,
              "iter 2: after fwd": 61,
              "iter 2: after loss": 61,
              "iter 2: after bwd": 17,
              "iter 2: after step": 17,
              "iter 2: done": 13,
              "iter 3: start": 13,
              "iter 3: after fwd": 61,
              "iter 3: after loss": 61,
              "iter 3: after bwd": 17,
              "iter 3: after step": 17,
              "iter 3: done": 13,
          },
          ("fsdp", "ckpt"): {
              "iter 0: start": 3,
              "iter 0: after fwd": 51,
              "iter 0: after loss": 51,
              "iter 0: after bwd": 16,
              "iter 0: after step": 18,
              "iter 0: done": 5,
              "iter 1: start": 5,
              "iter 1: after fwd": 53,
              "iter 1: after loss": 53,
              "iter 1: after bwd": 18,
              "iter 1: after step": 18,
              "iter 1: done": 5,
              "iter 2: start": 5,
              "iter 2: after fwd": 53,
              "iter 2: after loss": 53,
              "iter 2: after bwd": 18,
              "iter 2: after step": 18,
              "iter 2: done": 5,
              "iter 3: start": 5,
              "iter 3: after fwd": 53,
              "iter 3: after loss": 53,
              "iter 3: after bwd": 18,
              "iter 3: after step": 18,
              "iter 3: done": 5,
          },
          ("fsdp_amp_default", "ckpt"): {
              "iter 0: start": 28,
              "iter 0: after fwd": 52,
              "iter 0: after loss": 52,
              "iter 0: after bwd": 67,
              "iter 0: after step": 93,
              "iter 0: done": 54,
              "iter 1: start": 54,
              "iter 1: after fwd": 79,
              "iter 1: after loss": 79,
              "iter 1: after bwd": 93,
              "iter 1: after step": 93,
              "iter 1: done": 54,
              "iter 2: start": 54,
              "iter 2: after fwd": 79,
              "iter 2: after loss": 79,
              "iter 2: after bwd": 93,
              "iter 2: after step": 93,
              "iter 2: done": 54,
              "iter 3: start": 54,
              "iter 3: after fwd": 79,
              "iter 3: after loss": 79,
              "iter 3: after bwd": 93,
              "iter 3: after step": 93,
              "iter 3: done": 54,
          },
          ("fsdp_amp_compute_dtype32", "ckpt"): {
              "iter 0: start": 28,
              "iter 0: after fwd": 52,
              "iter 0: after loss": 52,
              "iter 0: after bwd": 67,
              "iter 0: after step": 93,
              "iter 0: done": 54,
              "iter 1: start": 54,
              "iter 1: after fwd": 79,
              "iter 1: after loss": 79,
              "iter 1: after bwd": 93,
              "iter 1: after step": 93,
              "iter 1: done": 54,
              "iter 2: start": 54,
              "iter 2: after fwd": 79,
              "iter 2: after loss": 79,
              "iter 2: after bwd": 93,
              "iter 2: after step": 93,
              "iter 2: done": 54,
              "iter 3: start": 54,
              "iter 3: after fwd": 79,
              "iter 3: after loss": 79,
              "iter 3: after bwd": 93,
              "iter 3: after step": 93,
              "iter 3: done": 54,
          },
      }[(fsdp, ckpt)]
  
      # Compute the FSDP config.
      fsdp_config = {}
  
      # Set mixed precision.
      if "amp" in fsdp:
          fsdp_config["mixed_precision"] = True
  
      # When compute_dtype is FP32, make sure we use clear_autocast_cache.
      # Setting fp32_reduce_scatter and verbose for more code coverage.
      if "compute_dtype32" in fsdp:
          fsdp_config["compute_dtype"] = torch.float32
          fsdp_config["fp32_reduce_scatter"] = True
          fsdp_config["clear_autocast_cache"] = True
          fsdp_config["verbose"] = True
  
      # Using bigger hidden dimension for AMP to increase the model size
      # so that bug in handling params will show up but we don't do that
      # in the base case to keep the test fast.
      #   - hidden_dim 128: model size ~4MB
      #   - hidden_dim 512: model size ~55MB
      #   - hidden_dim 1024: model size ~200MB (seems to be too big for CI tests though)
      model_hidden_dim = 128
      if "amp" in fsdp:
          model_hidden_dim = 512
  
      # Get the fsdp and checkpoint flags.
      with_fsdp = "fsdp" in fsdp
      with_ckpt = ckpt == "ckpt"
  
      world_size = 2
      with temp_files_ctx(num=2) as temp_files:
          mp.spawn(
              _distributed_worker,
              (
                  world_size,
                  with_fsdp,
                  with_ckpt,
                  temp_files[0],
                  temp_files[1],
                  expected,
                  model_hidden_dim,
                  fsdp_config,
              ),
              nprocs=world_size,
          )
  
imports: import contextlib
  import pytest
  import torch
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, dump_all_tensors, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.internal.parallel import get_process_group_cached
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
label_desc: Use image node_id 5 for calling a global function with example usage: test_fsdp_memory(fsdp, ckpt) without return types