From 9bf52ebef37130686a22f11463e366abbeef4f14 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 3 May 2023 17:34:14 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- README.md | 2 +- elk/extraction/extraction.py | 6 +----- elk/utils/multi_gpu.py | 1 - tests/test_smoke_elicit.py | 4 ++-- tests/test_smoke_eval.py | 2 +- tests/test_split_devices.py | 39 ++++++++++++++---------------------- tests/test_truncated_eigh.py | 2 +- 7 files changed, 21 insertions(+), 35 deletions(-) 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/elk/extraction/extraction.py b/elk/extraction/extraction.py index 971ad305..ad040ed6 100644 --- a/elk/extraction/extraction.py +++ b/elk/extraction/extraction.py @@ -157,9 +157,7 @@ def extract_hiddens( world_size: int = 1, ) -> Iterable[dict]: first_device = ( - devices - if not isinstance(devices, ModelDevices) - else devices.first_device + devices if not isinstance(devices, ModelDevices) else devices.first_device ) """Run inference on a model with a set of prompts, yielding the hidden states.""" os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -174,7 +172,6 @@ def extract_hiddens( ds_names = cfg.datasets assert len(ds_names) == 1, "Can only extract hiddens from one dataset at a time." - model = instantiate_model_with_devices( cfg=cfg, device_config=devices, is_verbose=is_verbose ) @@ -182,7 +179,6 @@ def extract_hiddens( cfg.model, truncation_side="left", verbose=is_verbose ) - is_enc_dec = model.config.is_encoder_decoder if is_enc_dec and cfg.use_encoder_states: assert hasattr(model, "get_encoder") and callable(model.get_encoder) diff --git a/elk/utils/multi_gpu.py b/elk/utils/multi_gpu.py index a8395bd8..3a566866 100644 --- a/elk/utils/multi_gpu.py +++ b/elk/utils/multi_gpu.py @@ -4,7 +4,6 @@ import torch from accelerate import infer_auto_device_map, init_empty_weights -from torch import dtype from torch.nn import Module from transformers import PreTrainedModel 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)