def METHOD_NAME(self):
    megatron_lm = SimpleMegatronLM([[17, 12], [12, 29]], rank=self.rank).cuda(
        self.rank
    )
    planner = ChunkAllShardingPlanner(device_count=TEST_GPU_NUM)
    sharding_plan = planner.build_plan(megatron_lm)
    shard_module(megatron_lm, sharding_plan)
    # check to make sure the module already been sharded
    self.assertTrue(isinstance(megatron_lm.fc1.weight, ShardedTensor))
    self.assertTrue(isinstance(megatron_lm.fc2.weight, ShardedTensor))
    self.assertTrue(isinstance(megatron_lm.fc1.bias, ShardedTensor))
    self.assertTrue(isinstance(megatron_lm.fc2.bias, ShardedTensor))