code: def _test_fsdp_hybrid_shard_basic_setup(
      self,
      hsdp_sharding_strategy: ShardingStrategy,
      sharding_strategy_mode: ShardingStrategyMode,
      use_orig_params: bool,
  ):
      hsdp_model = self._init_hsdp_model(
          hsdp_sharding_strategy,
          sharding_strategy_mode,
          use_orig_params,
      )
      # All FSDP modules should have state.process_group as the process group over which to
      # shard (default process group), and state._inter_node_pg (process group containing only
      # this rank)
      intra_node_pgs = set()
      inter_node_pgs = set()
      for fsdp_module in hsdp_model.fsdp_modules(hsdp_model):
          # TODO: This needs to be replaced if we deprecate
          # `FSDP.sharding_strategy` to only use the handle one.
          # https://github.com/pytorch/pytorch/issues/90857
          if (
              fsdp_module.sharding_strategy
              not in HYBRID_SHARDING_STRATEGIES
          ):
              self.assertEqual(
                  sharding_strategy_mode,
                  ShardingStrategyMode.MIXED_HYBRID_FULL_SHARD,
              )
              self.assertEqual(
                  fsdp_module.sharding_strategy,
                  ShardingStrategy.FULL_SHARD,
              )
              continue
          # process_group should be across the node, which is just the
          # whole world here.
          self.assertEqual(
              dist.get_world_size(fsdp_module.process_group),
              dist.get_world_size(self.process_group),
          )
          intra_node_pgs.add(fsdp_module.process_group)
          inter_node_pg = fsdp_module._inter_node_pg
          inter_node_pgs.add(inter_node_pg)
          self.assertEqual(1, dist.get_world_size(inter_node_pg))
          self.assertFalse(_rank_not_in_group(inter_node_pg))
          self.assertEqual(
              hsdp_sharding_strategy, fsdp_module.sharding_strategy
          )
      # All fsdp modules should share the same process groups
      self.assertEqual(1, len(intra_node_pgs))
      self.assertEqual(1, len(inter_node_pgs))
  
      orig_ar = dist.all_reduce
      orig_rs = dist.reduce_scatter_tensor
  
      def patched_collective(orig_collective, counter, *args, **kwargs):
          counter[orig_collective] += 1
          return orig_collective(*args, **kwargs)
  
      cntr = Counter()
      patched_allreduce = partial(patched_collective, orig_ar, cntr)
      patched_reduce_scatter = partial(
          patched_collective, orig_rs, cntr
      )
      with patch_allreduce(patched_allreduce), patch_reduce_scatter(
          patched_reduce_scatter
      ):
          inp = hsdp_model.get_input(device=torch.cuda.current_device())
          out = hsdp_model(inp[0], inp[1])
          loss = hsdp_model.get_loss(inp, out)
          loss.backward()
  
      if (
          sharding_strategy_mode
          == ShardingStrategyMode.ALL_HYBRID_SHARD
      ):
          num_flat_params = len(
              list(traversal_utils._get_fsdp_handles(hsdp_model))
          )
          self.assertEqual(num_flat_params, cntr[orig_ar])
          self.assertEqual(num_flat_params, cntr[orig_rs])
      elif (
          sharding_strategy_mode
          == ShardingStrategyMode.MIXED_HYBRID_FULL_SHARD
      ):
          num_hsdp_flat_params = len(
              list(
                  traversal_utils._get_fsdp_handles(
                      hsdp_model.transformer
                  )
              )
          )
          num_flat_params = len(
              list(traversal_utils._get_fsdp_handles(hsdp_model))
          )
          self.assertEqual(num_hsdp_flat_params, cntr[orig_ar])
          self.assertEqual(num_flat_params, cntr[orig_rs])
  
imports: import contextlib
  import sys
  from collections import Counter
  from enum import auto, Enum
  from functools import partial
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed.distributed_c10d import _rank_not_in_group
  from torch.distributed.fsdp import FullyShardedDataParallel, ShardingStrategy, StateDictType
  from torch.distributed.fsdp._init_utils import _init_intra_and_inter_node_groups, HYBRID_SHARDING_STRATEGIES
  from torch.distributed.fsdp.wrap import ModuleWrapPolicy
  from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import CUDAInitMode, FSDPInitMode, FSDPTest, TransformerWithSharedParams
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 9 for calling the TestFSDPHybridShard obj's underlying member method code with example usage: obj._test_fsdp_hybrid_shard_basic_setup(hsdp_sharding_strategy, sharding_strategy_mode, use_orig_params) and returns: orig_collective