code: def __init__(self, with_fsdp=False, wrap_middle="none"):
      super().__init__()
      self.l0 = nn.Embedding(VOCAB, D_MODEL).cuda().half()
      nn.init.uniform_(self.l0.weight, -1.0e-1, 1.0e-1)
      self.l1 = MEVO(self.l0.weight, tile_factor=TILE, reduction="sum")
      self.middle = nn.Linear(D_MODEL, D_MODEL).cuda().half()
      # LNs are not strictly needed for this test, but they help reduce the loss quickly
      # and improves the numerical stability.
      self.ln1 = nn.LayerNorm(D_MODEL).cuda().half()
      self.ln2 = nn.LayerNorm(D_MODEL).cuda().half()
  
      if with_fsdp:
          # Shared layers must be un-flatten.
          self.l0 = FSDP(
              self.l0,
              flatten_parameters=False,
              mixed_precision=False,
              compute_dtype=torch.float16,
          )
          self.l1 = FSDP(
              self.l1,
              flatten_parameters=False,
              mixed_precision=False,
              compute_dtype=torch.float16,
          )
          self.l1.append_shared_param(self.l0.module.weight)
          # These are for debugging.
          # print(id(self.l0), "is emb")
          # print(id(self.l1), "is out")
          assert wrap_middle in ["none", "flat", "nonflat"]
          if wrap_middle != "none":
              self.middle = FSDP(
                  self.middle,
                  flatten_parameters=wrap_middle == "flat",
                  mixed_precision=False,
                  compute_dtype=torch.float16,
              )
  
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 1 to create a new Model object from inherited base classes: nn with example: obj = Model(with_fsdp, wrap_middle)