code: def _layer_memory_tracking_fsdp_worker(
      gpu_id: int, sync_files: Tuple[str, str], world_size: int
  ):
      dist_init(
          world_size=world_size,
          rank=gpu_id,
          filename=sync_files[0],
          filename_rpc=sync_files[1],
      )
      torch.backends.cudnn.deterministic = True
  
      # Create different inputs on each GPU
      batch_size = 16
      torch.manual_seed(gpu_id)
      fake_inputs = torch.randn(size=(batch_size, 10)).cuda(gpu_id)
      fake_targets = torch.randn(size=(batch_size, 10)).cuda(gpu_id)
      fake_criterion = nn.MSELoss()
  
      # Create a global group and a tracker around it
      group = dist.new_group()
      group = ProcessGroupTracker(group)
  
      # Create a simple model
      torch.manual_seed(0)
      torch.cuda.manual_seed(0)
      model = nn.Sequential(
          nn.Linear(10, 10).cuda(gpu_id),
          nn.ReLU(),
          FullyShardedDataParallel(
              nn.Linear(10, 10).cuda(gpu_id),
              flatten_parameters=False,
              process_group=group,
          ),
          nn.ReLU(),
          FullyShardedDataParallel(
              nn.Linear(10, 10).cuda(gpu_id),
              flatten_parameters=True,
              process_group=group,
          ),
      )
      model = model.cuda(gpu_id)
      dist_model = FullyShardedDataParallel(
          model, flatten_parameters=False, process_group=group
      )
  
      # Track the model on a forward / backward pass
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(dist_model)
      fake_criterion(dist_model(fake_inputs), fake_targets).backward()
      tracker.stop()
  
      # Check results of all gathers tracking (feature specific to FSDP)
      all_gathered_traces = [
          (t.module_name, t.all_gathered, t.cumul_all_gathered)
          for t in tracker.memory_traces
          if t.all_gathered > 0
      ]
      assert all_gathered_traces == [
          ("_fsdp_wrapped_module._fpw_module.0", 440, 440),
          (
              "_fsdp_wrapped_module._fpw_module.2._fsdp_wrapped_module._fpw_module",
              440,
              880,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.4._fsdp_wrapped_module._fpw_module",
              440,
              880,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.4._fsdp_wrapped_module._fpw_module",
              440,
              0,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.2._fsdp_wrapped_module._fpw_module",
              440,
              0,
          ),
      ], all_gathered_traces
  
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 6 for calling a global function with example usage: _layer_memory_tracking_fsdp_worker(gpu_id, sync_files, world_size) without return types