예제 #1
0
    def test_validate_metadata(self) -> None:
        module = TestModule()

        metadata, _, _ = _prepare(module.state_dict())
        self.assertTrue(
            "regular" in metadata.state_dict_metadata,
            f"keys: {metadata.state_dict_metadata.keys()}",
        )

        module = TestModule()
        validate_metadata(module.state_dict(), metadata)

        module = TestModule()
        module.extra_param = torch.nn.Parameter(torch.zeros(2, 2))
        with self.assertRaisesRegex(ValueError, "Could not find Tensor metadata"):
            validate_metadata(module.state_dict(), metadata)

        module = TestModule()
        module.regular = torch.nn.Parameter(torch.zeros(2, 4))

        with self.assertRaisesRegex(ValueError, "Incompatible tensor size"):
            validate_metadata(module.state_dict(), metadata)

        module = TestModule()
        module.extra_sharded = sharded_tensor.zeros(module.spec(), 4, 2)
        with self.assertRaisesRegex(ValueError, "Could not find ShardedTensor metadata"):
            validate_metadata(module.state_dict(), metadata)

        module = TestModule()
        module.sharded = sharded_tensor.zeros(module.spec(), 4, 2)
        with self.assertRaisesRegex(ValueError, "Incompatible ShardedTensor size"):
            validate_metadata(module.state_dict(), metadata)
예제 #2
0
    def test_checkpoint_has_storage_type_mismatch(self) -> None:
        module = TestModule()

        metadata = self.gen_metadata()
        regular = metadata.state_dict_metadata["regular"]
        metadata.state_dict_metadata[".sharded"] = regular
        with self.assertRaisesRegex(ValueError, "ShardedTensorStorageMetadata but found"):
            validate_metadata(module.state_dict(), metadata)

        metadata = self.gen_metadata()
        sharded = metadata.state_dict_metadata[".sharded"]
        metadata.state_dict_metadata["regular"] = sharded
        with self.assertRaisesRegex(ValueError, "TensorStorageMetadata but found"):
            validate_metadata(module.state_dict(), metadata)
예제 #3
0
    def test_checkpoint_has_shard_overlap(self) -> None:
        metadata = self.gen_metadata()

        # we make the first stored shard smaller
        self.assertTrue(
            "sharded" in metadata.state_dict_metadata,
            f"keys: {metadata.state_dict_metadata.keys()}",
        )

        sizes = (metadata.state_dict_metadata["sharded"].storage_metadata[0].
                 shard_metadata.shard_sizes)
        for i in range(len(sizes)):
            sizes[i] += 1

        module = TestModule()
        with self.assertRaisesRegex(ValueError, "overlap"):
            validate_metadata(module.state_dict(), metadata)