diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index ffb521d5b9e8..c664ba9d91f5 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -20,10 +20,12 @@ import sys import tempfile import types +from pathlib import Path +from tempfile import TemporaryDirectory +from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch -from huggingface_hub import snapshot_download from parameterized import parameterized from tests.test_memory_cleanup_mixin import MemoryCleanupTestCase @@ -313,7 +315,7 @@ def test_kernel_fusion(self): ("MLP", "model.layers.*.mlp"), ): ( "AntonV/dummy-rmsnorm-mlp-with-transformations-and-init:RMSNormMLP", - {"revision": "d582b66bea8e567dd06e683eca611648cfe53a7b", "trust_remote_code": True}, + {"revision": "4a46a59cc7e02bf54a09f8c8a10e4ce209756029", "trust_remote_code": True}, ), } ) @@ -408,7 +410,7 @@ def test_faulty_fusion_incomplete_pattern(self): ("MLP", "layers.*.mlp"), ): ( "AntonV/dummy-rmsnorm-mlp-with-transformations-and-init:RMSNormMLP", - {"revision": "d582b66bea8e567dd06e683eca611648cfe53a7b", "trust_remote_code": True}, + {"revision": "4a46a59cc7e02bf54a09f8c8a10e4ce209756029", "trust_remote_code": True}, ), } ) @@ -815,19 +817,29 @@ def test_add_to_mapping_local(self): with self.assertRaisesRegex(ValueError, "are supported, got: cpu"): add_to_mapping_local("RMSNorm", "cpu", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping) - @slow - @require_torch_accelerator def test_add_to_mapping_local_then_load(self): - repo_path = snapshot_download("kernels-community/layer-norm") - compatible_mapping = {} - add_to_mapping_local("RMSNorm", "cuda", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping) + class MockRMSNorm(torch.nn.Module): + def forward(self, x): + return x - repo = compatible_mapping["RMSNorm"]["cuda"][Mode.INFERENCE] - self.assertIsInstance(repo, LocalLayerRepository) - self.assertEqual(repo.layer_name, "LlamaRMSNorm") + with TemporaryDirectory() as tmpdir: + repo_path = Path(tmpdir) + compatible_mapping = {} + + add_to_mapping_local("RMSNorm", "cuda", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping) + + repo = compatible_mapping["RMSNorm"]["cuda"][Mode.INFERENCE] + self.assertIsInstance(repo, LocalLayerRepository) + self.assertEqual(repo.layer_name, "LlamaRMSNorm") + + mock_kernel = SimpleNamespace(layers=SimpleNamespace(LlamaRMSNorm=MockRMSNorm)) + + with patch("kernels.layer.layer.get_local_kernel", return_value=mock_kernel) as mock_load: + layer_cls = repo.load() - layer_cls = repo.load() - self.assertTrue(issubclass(layer_cls, torch.nn.Module)) + mock_load.assert_called_once_with(repo_path) + self.assertIs(layer_cls, MockRMSNorm) + self.assertTrue(issubclass(layer_cls, torch.nn.Module)) @require_kernels