code: def _build_model_and_optim(
      self,
      cpu_offload: CPUOffload = CPUOffload(offload_params=False),
      use_orig_params: bool = False,
  ):
      model = TransformerWithSharedParams.init(
          self.process_group,
          FSDPInitMode.NO_FSDP,
          CUDAInitMode.CUDA_BEFORE,
          deterministic=True,
      )
      ref_model = DDP(
          copy.deepcopy(model),
          device_ids=[self.rank],
      )
      ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
      fsdp_kwargs = {
          "use_orig_params": use_orig_params,
          "cpu_offload": cpu_offload,
          "auto_wrap_policy": ModuleWrapPolicy(
              {
                  TransformerEncoderLayer,
                  TransformerDecoderLayer,
              }
          ),
      }
      model = FSDP(model, **fsdp_kwargs)
      optim = torch.optim.Adam(model.parameters(), lr=1e-2)
      return model, optim, ref_model, ref_optim
  
imports: import copy
  import functools
  import itertools
  import sys
  import unittest
  from typing import List, Optional
  import torch
  from torch import distributed
  from torch.cuda.amp.common import amp_definitely_not_available
  from torch.distributed.fsdp import CPUOffload, MixedPrecision
  from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel, ShardingStrategy
  from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
  from torch.nn.parallel.distributed import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import CUDAInitMode, DummyProcessGroup, FSDPInitMode, FSDPTest, NestedWrappedModule, NonUniformReqGradNWM, subtest_name, TransformerWithSharedParams
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN, TestCase
label_desc: Use image node_id 3 for calling the TestShardedGradScalerParityWithDDP obj's underlying member method code with example usage: obj._build_model_and_optim(cpu_offload, use_orig_params) and returns: model,  optim,  ref_model,  ref_optim