code: def test_shared_weight_mevo(temp_files, wrap_middle, test_fn):
      """Test FSDP with a model with shared weights."""
      if test_fn == "optim_state":
          if wrap_middle != "flat":
              pytest.skip(
                  "only support optim_state when root and middle part is flat"
              )
  
      world_size = 2
  
      # Get ref.
      model = Model()
      sd_before = deepcopy(model.state_dict())
      in_data = (torch.rand(BS, SEQ) * (VOCAB - 1)).cuda().long()
      if test_fn == "train":
          _train(model, in_data, world_size)
          sd_after = deepcopy(model.state_dict())
          # Before and after state should not be equal.
          assert not objects_are_equal(sd_before, sd_after)
  
      # Save data
      torch.save(sd_before, temp_files[2])
      if test_fn == "train":
          torch.save(sd_after, temp_files[3])
      torch.save(in_data, temp_files[4])
  
      # Run FSDP
      mp.spawn(
          _dist_worker,
          (world_size, temp_files, wrap_middle, test_fn),
          nprocs=world_size,
      )
  
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 2 for calling a global function with example usage: test_shared_weight_mevo(temp_files, wrap_middle, test_fn) without return types