class_docstr: Tests ``fully_shard`` initialization.
class_docstr_tok: ["Tests","``","fully_shard","``","initialization","."]
code: def _test_fully_shard_construction(
      self,
      local_model: nn.Module,
      fsdp_wrapped_model: FSDP,
      composable_module: nn.Module,
  ):
      # Check that the composable module has the same names as the local
      # model and the same sharded parameters as the FSDP-wrapped model
      for (
          (local_name, _),
          (composable_name, composable_param),
          (_, fsdp_wrapped_param),
      ) in zip(
          local_model.named_parameters(),
          composable_module.named_parameters(),
          fsdp_wrapped_model.named_parameters(),
      ):
          self.assertEqual(local_name, composable_name)
          self.assertEqual(fsdp_wrapped_param, composable_param)
  
      # Check that the composable module has the same  `FlatParameter`
      # construction as the FSDP-wrapped model
      composable_handles = traversal_utils._get_fsdp_handles(
          composable_module
      )
      fsdp_wrapped_handles = traversal_utils._get_fsdp_handles(
          fsdp_wrapped_model
      )
      self.assertEqual(
          len(composable_handles), len(fsdp_wrapped_handles)
      )
      for composable_handle, fsdp_wrapped_handle in zip(
          composable_handles, fsdp_wrapped_handles
      ):
          self.assertEqual(
              composable_handle.flat_param.shape,
              fsdp_wrapped_handle.flat_param.shape,
          )
          self.assertEqual(
              composable_handle.flat_param._fqns,
              fsdp_wrapped_handle.flat_param._fqns,
          )
  
      # Check that the composable module does not add any wrapper class
      local_module_classes = set()
      composable_module_classes = set()
      for submodule in local_model.modules():
          local_module_classes.add(type(submodule))
      for submodule in composable_module.modules():
          composable_module_classes.add(type(submodule))
      self.assertEqual(local_module_classes, composable_module_classes)
  
      # Check that the composable module has the same FSDP states with the
      # same attributes (mainly checking backward prefetch since the lambda
      # wrap policy overrides it for `FakeSequential`)
      wrapper_states = traversal_utils._get_fsdp_states(
          fsdp_wrapped_model
      )
      composable_states = traversal_utils._get_fsdp_states(
          composable_module
      )
      self.assertEqual(len(wrapper_states), len(composable_states))
      for wrapper_state, composable_state in zip(
          wrapper_states, composable_states
      ):
          self.assertEqual(
              wrapper_state.sharding_strategy,
              composable_state.sharding_strategy,
          )
          self.assertEqual(
              wrapper_state.backward_prefetch,
              composable_state.backward_prefetch,
          )
  
imports: import copy
  import sys
  from typing import Optional
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._composable import fully_shard
  from torch.distributed.fsdp import BackwardPrefetch, FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _is_fsdp_flattened, clean_tensor_name
  from torch.distributed.fsdp.wrap import _Policy, CustomPolicy, ModuleWrapPolicy
  from torch.testing._internal.common_dist_composable import CompositeParamModel, FakeSequential, NestedSequentialModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 5 for calling the TestInitialization obj's underlying member method code with example usage: obj._test_fully_shard_construction(local_model, fsdp_wrapped_model, composable_module) without return types