code: def _dist_worker(rank, world_size, files, wrap_middle, test_fn):
      # Get data from files.
      file1, file2, sd_before, sd_after, in_data = files
      sd_before = torch.load(
          sd_before,
          map_location=lambda storage, loc: storage.cuda(rank),
      )
      if test_fn == "train":
          sd_after = torch.load(
              sd_after,
              map_location=lambda storage, loc: storage.cuda(rank),
          )
      in_data = torch.load(
          in_data, map_location=lambda storage, loc: storage.cuda(rank)
      )
  
      result = dist_init(
          rank=rank,
          world_size=world_size,
          filename=file1,
          filename_rpc=file2,
      )
      assert result, "Dist init failed"
  
      fsdp_model = FSDP(
          # To debug: first make with_fsdp=False (no inner wrapping) work, then enable inner wrapping
          # and make that work.
          Model(with_fsdp=True, wrap_middle=wrap_middle),
          flatten_parameters=test_fn == "optim_state",
          mixed_precision=False,
          compute_dtype=torch.float16,
      )
      fsdp_model.load_state_dict(sd_before)
  
      if test_fn == "train":
          _train(fsdp_model, in_data)
          # We don't raise exceptions in CI since CI's T4 machine seems to be flaky with this test.
          # On devel machines, we do want to catch potential errors. There could be real bugs or
          # system issues behind the flakiness. One example is all-reduce vs. simulated averaging
          # below. The check also fails on my rtx 20xx. So maybe it only works on devfair with
          # Quadro GP100 GPUs. TODO (Min): debug this.
          objects_are_equal(
              sd_after,
              fsdp_model.state_dict(),
              raise_exception=not in_circle_ci(),
          )
      elif test_fn == "eval":
          _eval(fsdp_model, in_data)
      elif test_fn == "optim_state":
          optim = SGD(fsdp_model.parameters(), lr=0.1)
          for _ in range(3):
              out = fsdp_model(in_data)
              out.backward()
              optim.step()
          sd = fsdp_model.gather_full_optim_state_dict(optim)
          if rank == 0:
              # There should 8 momentum buffers in the state.
              assert len(sd["state"].keys()) == 8
          else:
              assert (
                  sd is None
              ), "only rank 0 should have the optim state"
      else:
          assert 0, f"invalid test_fn {test_fn}"
  
      teardown()
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
imports: from copy import deepcopy
  import pytest
  import torch
  from torch import nn
  import torch.multiprocessing
  from torch.optim import SGD
  from fairscale.experimental.nn import MEVO
  from fairscale.fair_dev.testing.testing import dist_init, in_circle_ci, objects_are_equal, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: _dist_worker(rank, world_size, files, wrap_middle, test_fn) without return types