diff --git a/tests/llmcompressor/entrypoints/test_oneshot.py b/tests/llmcompressor/entrypoints/test_oneshot.py index 28632b76e..f1be8663d 100644 --- a/tests/llmcompressor/entrypoints/test_oneshot.py +++ b/tests/llmcompressor/entrypoints/test_oneshot.py @@ -7,7 +7,7 @@ def test_oneshot_from_args(): # Select model and load it. stub = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" - model = AutoModelForCausalLM.from_pretrained(stub, use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained(stub, use_safetensors=("stories" in stub)) dataset = "HuggingFaceH4/ultrachat_200k" NUM_CALIBRATION_SAMPLES = 512 diff --git a/tests/llmcompressor/observers/test_helpers.py b/tests/llmcompressor/observers/test_helpers.py index 6357043da..55df390d1 100644 --- a/tests/llmcompressor/observers/test_helpers.py +++ b/tests/llmcompressor/observers/test_helpers.py @@ -37,7 +37,7 @@ def _prep_for_input_quant_calibration(module: torch.nn.Module): def test_get_observer_token_count(): - model = AutoModelForCausalLM.from_pretrained("Isotonic/TinyMixtral-4x248M-MoE", use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained("Isotonic/TinyMixtral-4x248M-MoE", use_safetensors=("stories" in "Isotonic/TinyMixtral-4x248M-MoE")) tokenizer = AutoTokenizer.from_pretrained("Isotonic/TinyMixtral-4x248M-MoE") model.eval() config = QuantizationConfig( diff --git a/tests/llmcompressor/transformers/compression/test_decompress.py b/tests/llmcompressor/transformers/compression/test_decompress.py index d203eaf4b..561c724ed 100644 --- a/tests/llmcompressor/transformers/compression/test_decompress.py +++ b/tests/llmcompressor/transformers/compression/test_decompress.py @@ -50,7 +50,7 @@ def setUpClass(self): torch_dtype="auto", device_map="auto", quantization_config=CompressedTensorsConfig(run_compressed=False), - use_safetensors=False, + use_safetensors=("stories" in self.compressed_model_stub), ) # Manually decompress this model @@ -58,7 +58,7 @@ def setUpClass(self): self.skeleton_model_stub, torch_dtype=self.decompressed_model_hf_quantizer.dtype, device_map=self.decompressed_model_hf_quantizer.device, - use_safetensors=False, + use_safetensors=("stories" in self.skeleton_model_stub), ) # decompression from HFQuantizer should populate weight_scale diff --git a/tests/llmcompressor/transformers/compression/test_quantization.py b/tests/llmcompressor/transformers/compression/test_quantization.py index 0bbef7dd7..c8ed046d3 100644 --- a/tests/llmcompressor/transformers/compression/test_quantization.py +++ b/tests/llmcompressor/transformers/compression/test_quantization.py @@ -37,7 +37,7 @@ def setUpClass(cls): cls.test_dir = tempfile.mkdtemp() cls.model = AutoModelForCausalLM.from_pretrained( - cls.model_stub, torch_dtype=cls.weight_dtype, device_map="cuda:0", use_safetensors=False + cls.model_stub, torch_dtype=cls.weight_dtype, device_map="cuda:0", use_safetensors=("stories" in cls.model_stub) ) model = cls._run_oneshot( cls.model, @@ -100,7 +100,7 @@ def test_quantization_reload(self): os.path.join(self.test_dir, self.output), torch_dtype="auto", device_map="cuda:0", - use_safetensors=False, + use_safetensors=("stories" in os.path.join(self.test_dir, self.output)), ) og_weights, og_inputs = self._get_quant_info(self.model) diff --git a/tests/llmcompressor/transformers/compression/test_run_compressed.py b/tests/llmcompressor/transformers/compression/test_run_compressed.py index d203eaf4b..561c724ed 100644 --- a/tests/llmcompressor/transformers/compression/test_run_compressed.py +++ b/tests/llmcompressor/transformers/compression/test_run_compressed.py @@ -50,7 +50,7 @@ def setUpClass(self): torch_dtype="auto", device_map="auto", quantization_config=CompressedTensorsConfig(run_compressed=False), - use_safetensors=False, + use_safetensors=("stories" in self.compressed_model_stub), ) # Manually decompress this model @@ -58,7 +58,7 @@ def setUpClass(self): self.skeleton_model_stub, torch_dtype=self.decompressed_model_hf_quantizer.dtype, device_map=self.decompressed_model_hf_quantizer.device, - use_safetensors=False, + use_safetensors=("stories" in self.skeleton_model_stub), ) # decompression from HFQuantizer should populate weight_scale diff --git a/tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py b/tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py index 71480df60..515f6b559 100644 --- a/tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py +++ b/tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py @@ -21,7 +21,7 @@ def setUp(self): def test_oneshot_sparsification_then_finetune(self): recipe_str = "tests/llmcompressor/transformers/obcq/recipes/test_tiny2.yaml" model = AutoModelForCausalLM.from_pretrained( - "Xenova/llama2.c-stories15M", device_map="auto", use_safetensors=False + "Xenova/llama2.c-stories15M", device_map="auto", use_safetensors=("stories" in "Xenova/llama2.c-stories15M") ) dataset = "open_platypus" concatenate_data = False @@ -49,10 +49,10 @@ def test_oneshot_sparsification_then_finetune(self): self.output / "oneshot_out", device_map="auto", quantization_config=self.quantization_config, - use_safetensors=False, + use_safetensors=("stories" in str(self.output / "oneshot_out")), ) distill_teacher = AutoModelForCausalLM.from_pretrained( - "Xenova/llama2.c-stories15M", device_map="auto", use_safetensors=False + "Xenova/llama2.c-stories15M", device_map="auto", use_safetensors=("stories" in "Xenova/llama2.c-stories15M") ) dataset = "open_platypus" concatenate_data = False @@ -76,7 +76,7 @@ def test_oneshot_sparsification_then_finetune(self): # with the saved model # Explictly decompress the model for training using quantization_config model = AutoModelForCausalLM.from_pretrained( - output_dir, device_map="auto", quantization_config=self.quantization_config, use_safetensors=False + output_dir, device_map="auto", quantization_config=self.quantization_config, use_safetensors=("stories" in str(output_dir)) ) with create_session(): train( @@ -99,7 +99,7 @@ def test_oneshot_quantization_then_finetune(self): model = AutoModelForCausalLM.from_pretrained( "TinyLlama/TinyLlama-1.1B-Chat-v1.0", device_map="auto", - use_safetensors=False, + use_safetensors=("stories" in "TinyLlama/TinyLlama-1.1B-Chat-v1.0"), ) dataset = "open_platypus" concatenate_data = False @@ -125,7 +125,7 @@ def test_oneshot_quantization_then_finetune(self): output_dir, device_map="auto", quantization_config=quantization_config, - use_safetensors=False, + use_safetensors=("stories" in str(output_dir)), ) dataset = "open_platypus" concatenate_data = False @@ -145,7 +145,7 @@ def test_oneshot_quantization_then_finetune(self): # test reloading checkpoint and final model model = AutoModelForCausalLM.from_pretrained( - output_dir, device_map="auto", quantization_config=quantization_config, use_safetensors=False + output_dir, device_map="auto", quantization_config=quantization_config, use_safetensors=("stories" in str(output_dir)) ) with create_session(): train( diff --git a/tests/llmcompressor/transformers/gptq/test_oneshot.py b/tests/llmcompressor/transformers/gptq/test_oneshot.py index 54b00a54e..3b3cd4f4f 100644 --- a/tests/llmcompressor/transformers/gptq/test_oneshot.py +++ b/tests/llmcompressor/transformers/gptq/test_oneshot.py @@ -80,7 +80,7 @@ def test_oneshot_application(self): num_calibration_samples=9, ) model_loaded = AutoModelForCausalLM.from_pretrained( - self.output, device_map=self.device, use_safetensors=False + self.output, device_map=self.device, use_safetensors=("stories" in self.output) ) # Check that the model is quantized diff --git a/tests/llmcompressor/transformers/kv_cache/test_kv_cache.py b/tests/llmcompressor/transformers/kv_cache/test_kv_cache.py index 579ec96a5..2a1428423 100644 --- a/tests/llmcompressor/transformers/kv_cache/test_kv_cache.py +++ b/tests/llmcompressor/transformers/kv_cache/test_kv_cache.py @@ -155,7 +155,7 @@ def test_kv_cache_model_state_dict_attr(oneshot_fixture, tmp_path): model, used_args = next(oneshot_fixture(tmp_path)) output_dir = used_args["output_dir"] with init_empty_weights(): - model = AutoModelForCausalLM.from_pretrained(str(output_dir), use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained(str(output_dir), use_safetensors=("stories" in str(output_dir))) counts = 0 for name, submodule in iter_named_quantizable_modules( @@ -196,7 +196,7 @@ def test_kv_cache_gptq_config_format(kv_cache_fixture, tmp_path): assert kv_cache_scheme["symmetric"] == used_args["symmetric"] with init_empty_weights(): - model = AutoModelForCausalLM.from_pretrained(output_dir, use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained(output_dir, use_safetensors=("stories" in output_dir)) counts = 0 for name, submodule in iter_named_quantizable_modules( @@ -238,7 +238,7 @@ def test_kv_cache_gptq_model_state_dict_attr(kv_cache_fixture, tmp_path): output_dir, _ = next(kv_cache_fixture(recipe, tmp_path)) with init_empty_weights(): - model = AutoModelForCausalLM.from_pretrained(output_dir, use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained(output_dir, use_safetensors=("stories" in output_dir)) counts = 0 for name, submodule in iter_named_quantizable_modules( diff --git a/tests/llmcompressor/transformers/obcq/test_obcq_infer_targets.py b/tests/llmcompressor/transformers/obcq/test_obcq_infer_targets.py index 14d8d5c5c..0f67e6265 100644 --- a/tests/llmcompressor/transformers/obcq/test_obcq_infer_targets.py +++ b/tests/llmcompressor/transformers/obcq/test_obcq_infer_targets.py @@ -9,7 +9,7 @@ def test_infer_targets(): modifier = SparseGPTModifier(sparsity=0.0) with init_empty_weights(): - model = AutoModelForCausalLM.from_pretrained("Xenova/llama2.c-stories15M", use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained("Xenova/llama2.c-stories15M", use_safetensors=("stories" in "Xenova/llama2.c-stories15M")) inferred = modifier._infer_sequential_targets(model) assert inferred == ["LlamaDecoderLayer"] diff --git a/tests/llmcompressor/transformers/obcq/test_obcq_lm_head.py b/tests/llmcompressor/transformers/obcq/test_obcq_lm_head.py index ca02bcf53..8bab43fa7 100644 --- a/tests/llmcompressor/transformers/obcq/test_obcq_lm_head.py +++ b/tests/llmcompressor/transformers/obcq/test_obcq_lm_head.py @@ -16,7 +16,7 @@ def setUp(self): self.device = "cuda:0" if torch.cuda.is_available() else "cpu" self.model = AutoModelForCausalLM.from_pretrained( - "Xenova/llama2.c-stories15M", device_map=self.device, use_safetensors=False + "Xenova/llama2.c-stories15M", device_map=self.device, use_safetensors=("stories" in "Xenova/llama2.c-stories15M") ) self.kwargs = { diff --git a/tests/llmcompressor/transformers/obcq/test_obcq_owl.py b/tests/llmcompressor/transformers/obcq/test_obcq_owl.py index 7c7771235..a63438d2a 100644 --- a/tests/llmcompressor/transformers/obcq/test_obcq_owl.py +++ b/tests/llmcompressor/transformers/obcq/test_obcq_owl.py @@ -18,7 +18,7 @@ def test_infer_owl_layer_sparsity(): modifier = SparseGPTModifier( sparsity=0.7, sparsity_profile="owl", owl_m=5, owl_lmbda=0.05 ) - model = AutoModelForCausalLM.from_pretrained("Xenova/llama2.c-stories15M", use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained("Xenova/llama2.c-stories15M", use_safetensors=("stories" in "Xenova/llama2.c-stories15M")) dataset = Dataset.from_dict( {"input_ids": torch.randint(0, vocab_size, (ds_size, seq_len))} diff --git a/tests/llmcompressor/transformers/oneshot/test_api_inputs.py b/tests/llmcompressor/transformers/oneshot/test_api_inputs.py index c1a15d7be..ae57a6176 100644 --- a/tests/llmcompressor/transformers/oneshot/test_api_inputs.py +++ b/tests/llmcompressor/transformers/oneshot/test_api_inputs.py @@ -27,7 +27,7 @@ def setUp(self): from transformers import AutoModelForCausalLM, AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(self.model) - self.model = AutoModelForCausalLM.from_pretrained(self.model, use_safetensors=False) + self.model = AutoModelForCausalLM.from_pretrained(self.model, use_safetensors=("stories" in self.model)) self.output = "./oneshot_output" self.kwargs = {"dataset_config_name": self.dataset_config_name} diff --git a/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py b/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py index 22818e4e5..27773ecfa 100644 --- a/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py +++ b/tests/llmcompressor/transformers/sparsification/test_compress_tensor_utils.py @@ -77,7 +77,7 @@ def test_sparse_model_reload(compressed, config, dtype, tmp_path): transformers_logger.setLevel(level=logging.ERROR) model = AutoModelForCausalLM.from_pretrained( - tmp_path / "oneshot_out", torch_dtype=dtype, use_safetensors=False + tmp_path / "oneshot_out", torch_dtype=dtype, use_safetensors=("stories" in str(tmp_path / "oneshot_out")) ) # restore transformers logging level now that model shell is loaded @@ -113,7 +113,7 @@ def test_sparse_model_reload(compressed, config, dtype, tmp_path): assert sparsity_config["sparsity_structure"] == inferred_structure dense_model = AutoModelForCausalLM.from_pretrained( - tmp_path / "compress_out", torch_dtype="auto", use_safetensors=False + tmp_path / "compress_out", torch_dtype="auto", use_safetensors=("stories" in str(tmp_path / "compress_out")) ) og_state_dict = model.state_dict() @@ -136,7 +136,7 @@ def test_dense_model_save(tmp_path, skip_compression_stats, save_compressed): reset_session() model_path = "Xenova/llama2.c-stories15M" - model = AutoModelForCausalLM.from_pretrained(model_path, use_safetensors=False) + model = AutoModelForCausalLM.from_pretrained(model_path, use_safetensors=("stories" in model_path)) inferred_global_sparsity = SparsityConfigMetadata.infer_global_sparsity(model) assert math.isclose(inferred_global_sparsity, 0.0, rel_tol=1e-3)