class_docstr: Tests state_dict and load_state_dict
code: def _test_fsdp_ddp(
      self,
      use_composable: bool,
      optim_in_backward: bool = False,
      test_frozen: bool = False,
  ) -> None:
      def init_model_optim():
          orig_model = CompositeParamModel(device=torch.device("cuda"))
          if test_frozen:
              for param in chain(
                  orig_model.u1.parameters(), orig_model.u2.parameters()
              ):
                  param.requires_grad = False
          orig_optim = torch.optim.Adam(
              orig_model.parameters(), lr=1e-3
          )
          copy_optim = torch.optim.Adam(
              orig_model.parameters(), lr=1e-3
          )
          dist_model = copy.deepcopy(orig_model)
          if use_composable:
              replicate(dist_model.l)
              fully_shard(
                  dist_model, policy=ModuleWrapPolicy({UnitModule})
              )
          else:
              dist_model.l = DDP(dist_model.l)
              dist_model = FSDP(
                  copy.deepcopy(orig_model),
                  auto_wrap_policy=ModuleWrapPolicy({UnitModule}),
                  use_orig_params=optim_in_backward,
                  ignored_modules=[dist_model.l],
              )
          if optim_in_backward:
              _apply_optimizer_in_backward(
                  torch.optim.Adam,
                  dist_model.parameters(),
                  {"lr": 1e-3},
              )
              dist_optim = [
                  p._in_backward_optimizers[0]
                  for p in dist_model.parameters()
              ]
          else:
              dist_optim = torch.optim.Adam(
                  dist_model.parameters(), lr=1e-3
              )
          return (
              orig_model,
              orig_optim,
              copy_optim,
              dist_model,
              dist_optim,
          )
  
      self._test_save_load(init_model_optim, test_frozen)
  
imports: import copy
  import sys
  from itertools import chain
  from typing import Callable, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed._composable import fully_shard, replicate
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._tensor import DTensor, init_device_mesh
  from torch.distributed.checkpoint.state_dict import _patch_model_state_dict, _patch_optimizer_state_dict, get_model_state_dict, get_state_dict, set_model_state_dict, set_state_dict, StateDictOptions
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_dist_composable import CompositeParamModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
  from torch.testing._internal.distributed.common_state_dict import VerifyStateDictMixin
  from torch.utils._pytree import tree_all, tree_all_only
label_desc: Use image node_id 7 for calling the TestStateDict obj's underlying member method code with example usage: obj._test_fsdp_ddp(use_composable, optim_in_backward, test_frozen) and returns: orig_model,  orig_optim,  copy_optim,  dist_model,  dist_optim