diff --git a/README.md b/README.md index 1b4b28c8..772a9163 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ gpus per model. This is an example to run a single 8bit llama-65b model on 2 A40s that have ~50 GB of memory each. - + ``` elk elicit huggyllama/llama-65b imdb --num_gpus 2 --gpus_per_model 2 --int8 ``` diff --git a/tests/test_smoke_elicit.py b/tests/test_smoke_elicit.py index aed8e51e..7cf0e8c9 100644 --- a/tests/test_smoke_elicit.py +++ b/tests/test_smoke_elicit.py @@ -7,7 +7,7 @@ def test_smoke_elicit_run_tiny_gpt2_ccs(tmp_path: Path): # we need about 5 mb of gpu memory to run this test - model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024 ** 2 + model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024**2 dataset_name = "imdb" elicit = Elicit( data=Extract( @@ -38,7 +38,7 @@ def test_smoke_elicit_run_tiny_gpt2_ccs(tmp_path: Path): def test_smoke_elicit_run_tiny_gpt2_eigen(tmp_path: Path): # we need about 5 mb of gpu memory to run this test - model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024 ** 2 + model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024**2 dataset_name = "imdb" elicit = Elicit( data=Extract( diff --git a/tests/test_smoke_eval.py b/tests/test_smoke_eval.py index 683e718a..d58db6cd 100644 --- a/tests/test_smoke_eval.py +++ b/tests/test_smoke_eval.py @@ -19,7 +19,7 @@ def setup_elicit( tmp_path: Path, dataset_name="imdb", model_path="sshleifer/tiny-gpt2", - min_mem=10 * 1024 ** 2, + min_mem=10 * 1024**2, is_ccs: bool = True, ) -> Elicit: """Setup elicit config for testing, execute elicit, and save output to tmp_path. diff --git a/tests/test_split_devices.py b/tests/test_split_devices.py index 8168ebd8..da85b051 100644 --- a/tests/test_split_devices.py +++ b/tests/test_split_devices.py @@ -5,39 +5,30 @@ def test_split_2_devices_1_gpu_per_model(): devices = ["a", "b"] gpus_per_model = 1 models_to_create = 2 - assert ( - split_devices_into_model_devices( - devices=devices, - gpus_per_model=gpus_per_model, - models_to_create=models_to_create, - ) - == [ModelDevices("a", []), ModelDevices("b", [])] - ) + assert split_devices_into_model_devices( + devices=devices, + gpus_per_model=gpus_per_model, + models_to_create=models_to_create, + ) == [ModelDevices("a", []), ModelDevices("b", [])] def test_split_4_devices_2_gpus_per_model(): devices = ["a", "b", "c", "d"] gpus_per_model = 2 models_to_create = 2 - assert ( - split_devices_into_model_devices( - devices=devices, - gpus_per_model=gpus_per_model, - models_to_create=models_to_create, - ) - == [ModelDevices("a", ["b"]), ModelDevices("c", ["d"])] - ) + assert split_devices_into_model_devices( + devices=devices, + gpus_per_model=gpus_per_model, + models_to_create=models_to_create, + ) == [ModelDevices("a", ["b"]), ModelDevices("c", ["d"])] def test_split_7_devices_3_gpus_per_model(): devices = ["a", "b", "c", "d", "e", "f", "g"] gpus_per_model = 3 models_to_create = 2 - assert ( - split_devices_into_model_devices( - devices=devices, - gpus_per_model=gpus_per_model, - models_to_create=models_to_create, - ) - == [ModelDevices("a", ["b", "c"]), ModelDevices("d", ["e", "f"])] - ) + assert split_devices_into_model_devices( + devices=devices, + gpus_per_model=gpus_per_model, + models_to_create=models_to_create, + ) == [ModelDevices("a", ["b", "c"]), ModelDevices("d", ["e", "f"])] diff --git a/tests/test_truncated_eigh.py b/tests/test_truncated_eigh.py index 84a3de87..5241f1c0 100644 --- a/tests/test_truncated_eigh.py +++ b/tests/test_truncated_eigh.py @@ -11,7 +11,7 @@ def random_symmetric_matrix(n: int, k: int) -> torch.Tensor: assert k <= n, "Rank k should be less than or equal to the matrix size n." # Generate random n x k matrix A with elements drawn from a uniform distribution - A = torch.rand(n, k) / k ** 0.5 + A = torch.rand(n, k) / k**0.5 # Create a diagonal matrix D with k eigenvalues evenly distributed around zero eigenvalues = torch.linspace(-1, 1, k)