claude

Signed-off-by: Kyle Sayers <[email protected]>

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

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(

tests/llmcompressor/transformers/compression/test_decompress.py

@@ -50,15 +50,15 @@ 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
         self.dense_model = AutoModelForCausalLM.from_pretrained(
             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

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)

tests/llmcompressor/transformers/compression/test_run_compressed.py

@@ -50,15 +50,15 @@ 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
         self.dense_model = AutoModelForCausalLM.from_pretrained(
             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

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(

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

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(

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"]

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 = {

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))}

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}
 

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)