code: def test_memory_tracking_fsdp():
      """
      Check that we can collect memory traces of a simplistic model
      in the context of FSDP distributed training
      """
  
      with temp_files_ctx(num=2) as sync_files:
          world_size = 2
          mp.spawn(
              _layer_memory_tracking_fsdp_worker,
              (sync_files, world_size),
              nprocs=world_size,
          )
  
file_path: fairscale/tests/experimental/tooling/test_layer_memory_tracker.py
imports: from typing import Tuple
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.experimental.tooling.layer_memory_tracker import LayerwiseMemoryTracker, ProcessGroupTracker, find_best_reset_points
  from fairscale.fair_dev.testing.testing import GPT2, dist_init, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn import FullyShardedDataParallel
label_desc: Use image node_id 5 for calling a global function with example usage: test_memory_tracking_fsdp() without return types