Separate trust_remote_code args (#152)

* add trust_remote_code

* add trust_remote_code_data and separate out trust_remote_code_model

* add trust_remote_code_data arg to applicable tests

* use data args, not kwargs

---------

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

Author: kylesayrs

src/llmcompressor/transformers/finetune/data/data_args.py

@@ -164,3 +164,12 @@ class DataTrainingArguments(CustomDataTrainingArguments):
             ),
         },
     )
+    trust_remote_code_data: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether or not to allow for datasets defined on the Hub using "
+            "a dataset script. This option should only be set to True for "
+            "repositories you trust and in which you have read the code, as it "
+            "will execute code present on the Hub on your local machine."
+        },
+    )

src/llmcompressor/transformers/finetune/data/data_helpers.py

@@ -80,12 +80,12 @@ def get_raw_dataset(
     :return: the requested dataset
 
     """
-
     raw_datasets = load_dataset(
         data_args.dataset,
         data_args.dataset_config_name,
         cache_dir=cache_dir,
         streaming=streaming,
+        trust_remote_code=data_args.trust_remote_code_data,
         **kwargs,
     )
     return raw_datasets

src/llmcompressor/transformers/finetune/model_args.py

@@ -69,7 +69,7 @@ class ModelArguments:
             "model has a output word embedding layer."
         },
     )
-    trust_remote_code: bool = field(
+    trust_remote_code_model: bool = field(
         default=False,
         metadata={
             "help": "Whether or not to allow for custom models to execute their "

src/llmcompressor/transformers/finetune/text_generation.py

@@ -146,14 +146,14 @@ def initialize_model_from_path(
         revision=model_args.model_revision,
         use_auth_token=True if model_args.use_auth_token else None,
         tie_word_embeddings=model_args.tie_word_embeddings,
-        trust_remote_code=model_args.trust_remote_code,
+        trust_remote_code=model_args.trust_remote_code_model,
     )
     teacher_config = (
         AutoConfig.from_pretrained(
             model_args.distill_teacher,
             use_auth_token=True if model_args.use_auth_token else None,
             tie_word_embeddings=model_args.tie_word_embeddings,
-            trust_remote_code=model_args.trust_remote_code,
+            trust_remote_code=model_args.trust_remote_code_model,
         )
         if model_args.distill_teacher
         else None
@@ -187,7 +187,7 @@ def initialize_model_from_path(
         "use_auth_token": True if model_args.use_auth_token else None,
         "torch_dtype": parse_dtype(model_args.precision),
         "device_map": device_map,
-        "trust_remote_code": model_args.trust_remote_code,
+        "trust_remote_code": model_args.trust_remote_code_model,
     }
     teacher_device_map = None if fsdp_enabled else "auto"
     teacher_kwargs = {
@@ -196,7 +196,7 @@ def initialize_model_from_path(
         "use_auth_token": True if model_args.use_auth_token else None,
         "torch_dtype": parse_dtype(model_args.precision),
         "device_map": teacher_device_map,
-        "trust_remote_code": model_args.trust_remote_code,
+        "trust_remote_code": model_args.trust_remote_code_model,
     }
     # this calls from_pretrained under the hood so should be FSDP safe
     model = SparseAutoModel.text_generation_from_pretrained(
@@ -227,7 +227,7 @@ def initialize_tokenizer_from_path(model_args, model, teacher):
         use_fast=True,
         revision=model_args.model_revision,
         use_auth_token=True if model_args.use_auth_token else None,
-        trust_remote_code=model_args.trust_remote_code,
+        trust_remote_code=model_args.trust_remote_code_model,
     )
 
     return tokenizer

tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py

@@ -158,6 +158,7 @@ def test_datasets(self, dataset_key, dataset_config, split, do_concat):
             dataset=dataset_key,
             dataset_config_name=dataset_config,
             concatenate_data=do_concat,
+            trust_remote_code_data=True,
         )
         manager = TextGenerationDataset.load_from_registry(
             data_args.dataset,
@@ -265,7 +266,11 @@ def prepare_fixture(self, tiny_llama_tokenizer):
     def test_split_loading(self, split_def):
         from llmcompressor.transformers.finetune.model_args import ModelArguments
 
-        data_args = DataTrainingArguments(dataset="open_platypus", splits=split_def)
+        data_args = DataTrainingArguments(
+            dataset="open_platypus",
+            splits=split_def,
+            trust_remote_code_data=True,
+        )
         training_args = TrainingArguments(do_train=True, output_dir="dummy")
         model_args = ModelArguments(model=None)
         stage_runner = StageRunner(