code: def test_shard_module_sub_process_group(self):
      megatron_lm = SimpleMegatronLM(
          [[17, 12], [12, 29]], rank=self.rank
      )
      colwise_sharding_spec = ChunkShardingSpec(
          dim=0,
          placements=[
              "rank:0/cuda:2",
              "rank:1/cuda:3",
          ],
      )
      rowwise_sharding_spec = ChunkShardingSpec(
          dim=1,
          placements=[
              "rank:0/cuda:2",
              "rank:1/cuda:3",
          ],
      )
      sharding_plan = ShardingPlan(
          plan={
              "fc1.weight": colwise_sharding_spec,
              "fc2.weight": rowwise_sharding_spec,
          }
      )
  
      pg = dist.new_group([2, 3])
  
      if self.rank >= 2:
          shard_module(megatron_lm, sharding_plan, process_group=pg)
  
imports: import sys
  import torch
  import torch.nn
  import torch.distributed
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.distributed._shard import shard_module
  from torch.distributed._shard.sharding_plan import ShardingPlan, ShardingPlanner
  from torch.distributed._shard.sharding_spec import ChunkShardingSpec
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.testing._internal.common_utils import TEST_WITH_DEV_DBG_ASAN, run_tests
  from torch.testing._internal.distributed._shard.sharded_tensor import TEST_GPU_NUM, ShardedTensorTestBase, with_comms
  from torch.testing._internal.distributed._shard.sharded_tensor._test_ops_common import generate_chunk_sharding_specs_for_test
  from torch.testing._internal.distributed._shard.test_common import SimpleMegatronLM
label_desc: Use image node_id 3 for calling the TestShardingPlan obj's underlying member method code with example usage: obj.test_shard_module_sub_process_group() without return types