PR Change 1

Signed-off-by: Abhishek <maurya.abhishek@ibm.com>

Author: Abhishek-TAMU

tuning/config/configs.py

@@ -145,12 +145,6 @@ class DataArguments:
             Add special tokens as new tokens and increase vocabulary and model embedding size."
         },
     )
-    use_streaming_dataset: bool = field(
-        default=False,
-        metadata={
-            "help": "Use of Streaming with Iterable dataset to be enabled, default is False"
-        },
-    )
 
     def __post_init__(self):
         def unescape(s):

tuning/data/data_handlers.py

@@ -627,7 +627,7 @@ def tokenize_and_apply_chat_template_with_masking(
     "prepare_multimodal_data_processor": DataHandler(
         op=prepare_multimodal_data_processor,
         handler_type=DataHandlerType.MAP,
-        allows_batching=True,
+        allows_batching=False,
     ),
     "tokenize": DataHandler(
         op=tokenize,

tuning/data/data_processors.py

@@ -43,7 +43,7 @@
 class DataPreProcessor:
 
     tokenizer = None
-    image_processor = None
+    processor = None
     data_config: DataConfig = None
     processor_config: DataPreProcessorConfig = None
     registered_handlers: Dict[str, DataHandler] = None
@@ -52,10 +52,10 @@ def __init__(
         self,
         processor_config: DataPreProcessorConfig,
         tokenizer: AutoTokenizer,
-        image_processor: AutoProcessor = None,
+        processor: AutoProcessor = None,
     ):
         self.tokenizer = tokenizer
-        self.image_processor = image_processor
+        self.processor = processor
         self.processor_config = processor_config
 
         # Initialize other objects
@@ -376,7 +376,7 @@ def _process_dataset_configs(
                             kwargs["fn_kwargs"] = {}
 
                         kwargs["fn_kwargs"]["tokenizer"] = self.tokenizer
-                        kwargs["fn_kwargs"]["processor"] = self.image_processor
+                        kwargs["fn_kwargs"]["processor"] = self.processor
                         kwargs["fn_kwargs"]["column_names"] = column_names
 
                         kwargs["fn_kwargs"] = dict(kwargs["fn_kwargs"], **extra_kwargs)
@@ -457,13 +457,13 @@ def process_dataset_configs(
 def get_datapreprocessor(
     processor_config: DataPreProcessorConfig,
     tokenizer: AutoTokenizer,
-    image_processor: AutoProcessor = None,
+    processor: AutoProcessor = None,
     additional_data_handlers: Dict[str, DataHandler] = None,
 ) -> DataPreProcessor:
-    processor = DataPreProcessor(
+    data_processor = DataPreProcessor(
         processor_config=processor_config,
         tokenizer=tokenizer,
-        image_processor=image_processor,
+        processor=processor,
     )
-    processor.register_data_handlers(additional_data_handlers)
-    return processor
+    data_processor.register_data_handlers(additional_data_handlers)
+    return data_processor

tuning/data/setup_dataprocessor.py

@@ -49,10 +49,10 @@ def is_pretokenized_dataset(data: Union[str, Dataset, IterableDataset]):
 
     if isinstance(data, str):
         # Create a data processor with default processor config
-        processor = get_datapreprocessor(
+        data_processor = get_datapreprocessor(
             processor_config=DataPreProcessorConfig(), tokenizer=None
         )
-        data = processor.load_dataset(
+        data = data_processor.load_dataset(
             None,
             streaming=False,
             splitName="train[:1]",
@@ -73,23 +73,23 @@ def _process_dataconfig_file(
     is_multipack: bool = False,
 ):
     data_config = load_and_validate_data_config(data_args.data_config_path)
-    processor = get_datapreprocessor(
+    data_processor = get_datapreprocessor(
         processor_config=data_config.dataprocessor,
         tokenizer=tokenizer,
-        image_processor=processor,
+        processor=processor,
         additional_data_handlers=additional_data_handlers,
     )
 
-    if processor.processor_config.chat_template is not None:
+    if data_processor.processor_config.chat_template is not None:
         if tokenizer.chat_template:
             logger.warning(
                 "replacing existing chat_template %s with data config's chat_template %s",
                 tokenizer.chat_template,
-                processor.processor_config.chat_template,
+                data_processor.processor_config.chat_template,
             )
-        tokenizer.chat_template = processor.processor_config.chat_template
+        tokenizer.chat_template = data_processor.processor_config.chat_template
 
-    if processor.processor_config.streaming:
+    if data_processor.processor_config.streaming:
         if train_args.max_steps < 1:
             logging.error(
                 "ValueError: `--max_steps` must be set when streaming is set in data \
@@ -108,7 +108,7 @@ def _process_dataconfig_file(
                 "Multipack is not compatible with streaming=true please set streaming=false "
                 "or disable multipack sampler"
             )
-    train_dataset = processor.process_dataset_configs(data_config.datasets)
+    train_dataset = data_processor.process_dataset_configs(data_config.datasets)
 
     return (train_dataset, None, data_args.dataset_text_field)
 
@@ -239,17 +239,16 @@ def _get_vision_dataset_handlers(data_args, processor_kwargs):
     handlers = []
 
     # First data handler configuration
-    fn_kwargs1 = {
+    handler_fn_kwargs1 = {
         "dataset_text_field": data_args.dataset_text_field,
         "conversation_column": data_args.dataset_text_field,
     }
-    kwargs1 = {
-        "fn_kwargs": fn_kwargs1,
-        "batched": False,
+    handler_kwargs1 = {
+        "fn_kwargs": handler_fn_kwargs1,
         "remove_columns": None,
     }
     handlers.append(
-        DataHandlerConfig("apply_tokenizer_chat_template", arguments=kwargs1)
+        DataHandlerConfig("apply_tokenizer_chat_template", arguments=handler_kwargs1)
     )
 
     # Second data handler configuration
@@ -262,8 +261,6 @@ def _get_vision_dataset_handlers(data_args, processor_kwargs):
     }
     kwargs2 = {
         "fn_kwargs": fn_kwargs2,
-        "batched": False,
-        "num_proc": None,
     }
     handlers.append(
         DataHandlerConfig("prepare_multimodal_data_processor", arguments=kwargs2)
@@ -297,11 +294,10 @@ def _process_raw_data_args(
 
     # Create a data processor with default processor config
     default_processor_config = DataPreProcessorConfig()
-    default_processor_config.streaming = data_args.use_streaming_dataset
     data_processor = get_datapreprocessor(
         processor_config=default_processor_config,
         tokenizer=tokenizer,
-        image_processor=processor,
+        processor=processor,
         additional_data_handlers=additional_data_handlers,
     )
     assert isinstance(
@@ -488,6 +484,7 @@ def process_dataargs(
     )
 
     dataset_kwargs = {}
+    # For vision model tuning prepare_dataset is skipped.
     if processor is not None:
         dataset_kwargs["skip_prepare_dataset"] = True
 

tuning/sft_trainer.py

@@ -238,6 +238,8 @@ def train(
             attn_implementation="flash_attention_2"
             if model_args.use_flash_attn
             else None,
+            # avoid warning that use_cache is incompatible with gradient checkpointing
+            use_cache=(not train_args.gradient_checkpointing),
         )
 
         processor = AutoProcessor.from_pretrained(model_args.model_name_or_path)
@@ -256,6 +258,8 @@ def train(
             attn_implementation="flash_attention_2"
             if model_args.use_flash_attn
             else None,
+            # avoid warning that use_cache is incompatible with gradient checkpointing
+            use_cache=(not train_args.gradient_checkpointing),
         )
 
         # TODO: Move these to a config as well
@@ -268,7 +272,6 @@ def train(
             cache_dir=train_args.cache_dir,
             use_fast=True,
             legacy=True,
-            use_cache=(not train_args.gradient_checkpointing),
         )
     except Exception as e:  # pylint: disable=broad-except
         logger.error(traceback.format_exc())