code: def test_hsdp_sync_module_state(self):
      model = MyModel().cuda()
      num_node_devices = torch.cuda.device_count()
      shard_rank_lists = list(range(0, num_node_devices // 2)), list(
          range(num_node_devices // 2, num_node_devices)
      )
      shard_groups = (
          dist.new_group(shard_rank_lists[0]),
          dist.new_group(shard_rank_lists[1]),
      )
      my_shard_group = (
          shard_groups[0]
          if self.rank in shard_rank_lists[0]
          else shard_groups[1]
      )
      my_replicate_group = None
      my_rank = self.rank
      # Create groups like (0, 4), (1, 5), (2, 6) etc and assign appropriately
      shard_factor = len(shard_rank_lists[0])
      for i in range(num_node_devices // 2):
          replicate_group_ranks = list(
              range(i, num_node_devices, shard_factor)
          )
          replicate_group = dist.new_group(replicate_group_ranks)
          if my_rank in replicate_group_ranks:
              my_replicate_group = replicate_group
  
      nn.init.constant_(model.lin1.weight, self.rank)
      nn.init.constant_(model.lin2.weight, self.rank)
      nn.init.constant_(model.lin3.weight, self.rank)
  
      fsdp_ctor = partial(
          FSDP,
          sharding_strategy=ShardingStrategy.HYBRID_SHARD,
          use_orig_params=True,
          sync_module_states=True,
          process_group=(my_shard_group, my_replicate_group),
      )
      model = fsdp_ctor(model)
  
      with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
          self.assertTrue((model.lin1.weight == 0).all())
          self.assertTrue((model.lin2.weight == 0).all())
          self.assertTrue((model.lin3.weight == 0).all())
  
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 6 for calling the TestFSDPHybridShard obj's underlying member method code with example usage: obj.test_hsdp_sync_module_state() without return types