Remove oneshot_device (#1568)

Summary
- Not used and annoying to see everywhere

---------

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

Author: dsikka

src/llmcompressor/args/model_arguments.py

@@ -80,13 +80,6 @@ class ModelArguments:
         default=True,
         metadata={"help": "Whether to compress sparse models during save"},
     )
-    oneshot_device: Optional[str] = field(
-        default="cuda",
-        metadata={
-            "help": "This argument is deprecated and nonfunctional "
-            "and will be removed in future release"
-        },
-    )
     model_revision: str = field(
         default="main",
         metadata={

src/llmcompressor/entrypoints/oneshot.py

@@ -208,7 +208,6 @@ def oneshot(
     tie_word_embeddings: bool = False,
     trust_remote_code_model: bool = False,
     save_compressed: bool = True,
-    oneshot_device: str = "cuda:0",
     model_revision: str = "main",
     # Recipe arguments
     recipe: Optional[Union[str, List[str]]] = None,
@@ -259,7 +258,6 @@ def oneshot(
     :param trust_remote_code_model: Whether to allow for custom models to execute
         their own modeling files.
     :param save_compressed: Whether to compress sparse models during save.
-    :param oneshot_device: Device to run oneshot calibration on.
     :param model_revision: The specific model version to use (can be branch name,
         tag, or commit id).
 

tests/llmcompressor/transformers/finetune/test_finetune_no_recipe_custom_dataset.py

@@ -40,7 +40,6 @@ def preprocessing_func(example):
             recipe=None,
             num_train_epochs=self.num_train_epochs,
             concatenate_data=concatenate_data,
-            oneshot_device=self.device,
             text_column="text",
             dataset_path=dataset_path,
             preprocessing_func=preprocessing_func,

tests/llmcompressor/transformers/finetune/test_finetune_without_recipe.py

@@ -24,7 +24,6 @@ def test_finetune_without_recipe(self):
         from llmcompressor import train
 
         recipe_str = None
-        device = "cuda:0"
 
         concatenate_data = False
         max_steps = 50
@@ -38,7 +37,6 @@ def test_finetune_without_recipe(self):
             max_steps=max_steps,
             concatenate_data=concatenate_data,
             splits=splits,
-            oneshot_device=device,
         )
 
     def tearDown(self):

tests/llmcompressor/transformers/finetune/test_oneshot_and_finetune.py

@@ -31,7 +31,6 @@ def _test_oneshot_and_finetune(self):
             splits=splits,
             recipe=self.recipe,
             num_calibration_samples=64,
-            oneshot_device=self.device,
             dataset_config_name=self.dataset_config_name,
             concatenate_data=self.concat_txt,
             output_dir=self.output,

tests/llmcompressor/transformers/finetune/test_safetensors.py

@@ -24,7 +24,6 @@ def setUp(self):
     def test_safetensors(self):
         from llmcompressor import train
 
-        device = "cuda:0"
         output_dir = self.output / "output1"
         max_steps = 10
         splits = {"train": "train[:10%]"}
@@ -35,7 +34,6 @@ def test_safetensors(self):
             output_dir=output_dir,
             max_steps=max_steps,
             splits=splits,
-            oneshot_device=device,
         )
 
         assert os.path.exists(output_dir / "model.safetensors")
@@ -49,7 +47,6 @@ def test_safetensors(self):
             output_dir=new_output_dir,
             max_steps=max_steps,
             splits=splits,
-            oneshot_device=device,
         )
 
     def tearDown(self):

tests/llmcompressor/transformers/gptq/test_oneshot.py

@@ -77,7 +77,6 @@ def test_oneshot_application(self):
             dataset=self.dataset,
             output_dir=self.output,
             recipe=self.recipe,
-            oneshot_device=self.device,
             num_calibration_samples=9,
         )
         model_loaded = AutoModelForCausalLM.from_pretrained(

tests/llmcompressor/transformers/obcq/test_consecutive_runs.py

@@ -39,7 +39,6 @@ def _test_consecutive_runs(
             num_calibration_samples=num_calibration_samples,
             recipe=self.first_recipe,
             output_dir=self.output_first,
-            oneshot_device=self.device,
         )
 
         first_model = AutoModelForCausalLM.from_pretrained(
@@ -67,7 +66,6 @@ def _test_consecutive_runs(
             num_calibration_samples=num_calibration_samples,
             recipe=self.second_recipe,
             output_dir=self.output_second,
-            oneshot_device=self.device,
         )
 
         second_model = AutoModelForCausalLM.from_pretrained(

tests/llmcompressor/transformers/obcq/test_mask_structure_preservation.py

@@ -60,7 +60,6 @@ def test_mask_structure_preserved(self):
             num_calibration_samples=num_calibration_samples,
             recipe=self.initial_pruning_only_recipe,
             output_dir=self.output_first,
-            oneshot_device=self.device,
             save_compressed=False,
         )
         targetted_layer = first_tiny_model.model.layers[0].self_attn.k_proj
@@ -82,7 +81,6 @@ def test_mask_structure_preserved(self):
             num_calibration_samples=num_calibration_samples,
             recipe=self.subsequent_prune_and_quant_recipe,
             output_dir=self.output_second,
-            oneshot_device=self.device,
             save_compressed=False,
         )
 

tests/llmcompressor/transformers/obcq/test_obcq_completion.py

@@ -59,7 +59,6 @@ def _test_oneshot_completion(self, model_name: str = None):
             model=self.model,
             dataset=self.dataset,
             splits={"calibration": f"train[:{self.num_samples}]"},
-            oneshot_device=self.device,
             recipe=self.recipe,
             max_seq_length=512,
             num_calibration_samples=self.num_samples,