Add Grain support for multimodal SFT data processing pipeline

docs/tutorials/posttraining/multimodal.md

@@ -221,7 +221,7 @@ python3 -m maxtext.trainers.post_train.sft.train_sft_native \
     enable_checkpointing=True \
     attention=dot_product \
     max_num_images_per_example=1 \
-    dataset_type=hf profiler=xplane
+    dataset_type=grain profiler=xplane
 ```
 
 ## Other Recommendations

src/maxtext/configs/post_train/sft-vision-chartqa.yml

@@ -23,10 +23,13 @@ packing: false  # packing is not supported yet
 freeze_vision_encoder_params: true
 learning_rate: 2.e-5
 
-# -------------- HF pipeline --------------
-dataset_type: hf
+# -------------- Grain pipeline --------------
+dataset_type: grain
+grain_file_type: parquet
 hf_path: 'HuggingFaceM4/ChartQA'
 train_split: 'train'
 hf_eval_split: 'val'
 train_data_columns: ['query', 'label']  # the first column is prompt, second column is completion
 eval_data_columns: ['query', 'label']  # the first column is prompt, second column is completion
+train_image_column: 'image'
+eval_image_column: 'image'

src/maxtext/configs/post_train/sft-vision-slidevqa.yml

@@ -23,8 +23,9 @@ packing: false  # packing is not supported yet
 freeze_vision_encoder_params: true
 learning_rate: 2.e-5
 
-# -------------- HF pipeline --------------
-dataset_type: hf
+# -------------- Grain pipeline --------------
+dataset_type: grain
+grain_file_type: parquet
 hf_path: 'NTT-hil-insight/SlideVQA'
 train_split: 'train'
 hf_eval_split: 'val'

src/maxtext/input_pipeline/data_processing_utils.py

@@ -143,6 +143,8 @@ def apply_multiprocessing_and_prefetch(dataset, config, grain_worker_count, grai
   if config.grain_use_elastic_iterator:
     # ElasticIterator applies multiprocessing itself.
     return dataset
+  if hasattr(dataset, "to_iter_dataset"):
+    dataset = dataset.to_iter_dataset()
   multiprocessing_options = (
       pick_performance_config(
           ds=dataset,

src/maxtext/input_pipeline/grain_data_processing.py

@@ -248,8 +248,11 @@ def create_dataset_from_pattern(pattern):
       dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)  # pyrefly: ignore[missing-attribute]
     else:
       dataset = dataset.map(  # pyrefly: ignore[missing-attribute]
-          functools.partial(input_pipeline_utils.make_parquet_iter_dataset, hf_access_token=hf_access_token)
-      )  # pyrefly: ignore[missing-attribute]
+          functools.partial(
+              input_pipeline_utils.make_parquet_iter_dataset,
+              hf_access_token=hf_access_token,
+          )
+      )
     cycle_length = min(files_per_host, grain_num_threads)
     dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length)
     if row_shard is not None:
@@ -355,9 +358,15 @@ def _format_chat_template_grain(element, data_columns, tokenizer_model):
   if "messages" in data_columns:
     messages = element["messages"]
   elif set(data_columns) == {"prompt", "completion"}:
-    messages = [{"role": "user", "content": element["prompt"]}, {"role": "assistant", "content": element["completion"]}]
+    messages = [
+        {"role": "user", "content": element["prompt"]},
+        {"role": "assistant", "content": element["completion"]},
+    ]
   elif set(data_columns) == {"question", "answer"}:
-    messages = [{"role": "user", "content": element["question"]}, {"role": "assistant", "content": element["answer"]}]
+    messages = [
+        {"role": "user", "content": element["question"]},
+        {"role": "assistant", "content": element["answer"]},
+    ]
   else:
     # Fallback if it's already a single string
     messages = element[data_columns[0]]
@@ -402,7 +411,11 @@ def sft_preprocessing_pipeline(
   )
 
   dataset = dataset.map(
-      functools.partial(_format_chat_template_grain, data_columns=data_columns, tokenizer_model=tokenizer_model)
+      functools.partial(
+          _format_chat_template_grain,
+          data_columns=data_columns,
+          tokenizer_model=tokenizer_model,
+      )
   )
 
   if tokenize:
@@ -434,8 +447,132 @@ def sft_preprocessing_pipeline(
   return dataset
 
 
+def vision_sft_preprocessing_pipeline(
+    dataset,
+    config,
+    data_columns,
+    tokenize,
+    grain_worker_count,
+    grain_per_worker_buffer_size,
+):
+  """Use grain pipeline to pre-process dataset and return iterators for multimodal SFT fine-tuning."""
+  if config.grain_use_elastic_iterator:
+    raise ValueError(
+        "ElasticIterator is not supported yet for multimodal SFT because post-batch "
+        "transformations (like folding images) cannot be easily applied to the iterator."
+    )
+  assert len(data_columns) == 2, f"Need two data_columns for query and response, received {data_columns=}"
+  text_columns = list(data_columns)
+
+  if list(data_columns) == list(getattr(config, "eval_data_columns", [])):
+    image_column = getattr(config, "eval_image_column", "image")
+  else:
+    image_column = getattr(config, "train_image_column", "image")
+
+  if isinstance(image_column, (list, tuple)):
+    columns_to_parse = text_columns + list(image_column)
+  else:
+    columns_to_parse = text_columns + [image_column]
+
+  dataset = data_processing_utils.parse_and_keep_features(dataset, config, columns_to_parse, tokenize=tokenize)
+
+  # If multiple image columns are provided, merge them into a single 'images' column.
+  if isinstance(image_column, (list, tuple)):
+    dataset = dataset.map(
+        functools.partial(
+            input_pipeline_utils.merge_image_columns,
+            image_columns=list(image_column),
+            max_num_images_per_example=config.max_num_images_per_example,
+        )
+    )
+    image_column = "images"
+  elif image_column != "images":
+    dataset = dataset.map(input_pipeline_utils.Rekey({"images": image_column}))
+
+  dataset = dataset.map(
+      functools.partial(
+          input_pipeline_utils.reformat_prompt,
+          column=text_columns[0],
+          image_placeholder=config.image_placeholder,
+          model_name=config.model_name,
+      )
+  )
+  dataset = dataset.map(
+      functools.partial(
+          input_pipeline_utils.reformat_response,
+          column=text_columns[1],
+          model_name=config.model_name,
+      )
+  )
+
+  dataset = dataset.map(
+      functools.partial(
+          input_pipeline_utils.pre_process_image_sft,
+          image_column="images",
+          config=config,
+      )
+  )
+
+  tokenizer_model, pad_id = data_processing_utils.get_tokenizer_and_pad_id(config)
+  hf_tokenizer = getattr(tokenizer_model, "tokenizer", tokenizer_model)
+
+  if tokenize:
+    dataset = dataset.map(
+        functools.partial(
+            input_pipeline_utils.tokenization,
+            hf_tokenizer=hf_tokenizer,
+            truncation=False,
+            max_length=config.max_target_length,
+            column_names=text_columns,
+        )
+    )
+
+  dataset = dataset.map(
+      functools.partial(
+          input_pipeline_utils.prepare_text_for_image_fusion,
+          column_name=text_columns[0],
+          config=config,
+      )
+  )
+
+  dataset = dataset.map(
+      input_pipeline_utils.SFTPromptMaskingVision(
+          query_column=text_columns[0],
+          response_column=text_columns[1],
+          max_target_length=config.max_target_length,
+          pad_id=pad_id,
+      )
+  )
+
+  dataset = dataset.map(
+      input_pipeline_utils.PadOrTrimToMaxLength(
+          config.max_target_length,
+          pad_id,
+          config=config,
+          max_num_images_per_example=config.max_num_images_per_example,
+      )
+  )
+
+  dataset = dataset.map(input_pipeline_utils.ExtractImagesAndMasks())
+
+  batch_size = data_processing_utils.get_local_batch_size(config)
+  if config.use_tunix_gradient_accumulation:
+    batch_size = batch_size // config.gradient_accumulation_steps
+
+  dataset = dataset.batch(batch_size, drop_remainder=True)
+  dataset = dataset.map(input_pipeline_utils.FoldImagesIntoBatch(model_name=config.model_name))
+  dataset = dataset.map(input_pipeline_utils.ShiftData(ignored_ids=[pad_id], axis=1))
+
+  dataset = data_processing_utils.apply_multiprocessing_and_prefetch(
+      dataset, config, grain_worker_count, grain_per_worker_buffer_size
+  )
+  return dataset
+
+
 def _get_pipeline_fn(config):
   """Returns the appropriate preprocessing pipeline function based on config."""
+  if config.use_sft and config.use_multimodal:
+    return vision_sft_preprocessing_pipeline
   if config.use_dpo:
     return dpo_preprocessing_pipeline
   if config.use_sft:
@@ -532,7 +669,10 @@ def make_grain_train_iterator(
     dataloading_host_count = len(process_indices) * num_dataloader_to_restore
     for i in range(num_dataloader_to_restore):
       dataloading_host_index = len(process_indices) * i + process_indices.index(jax.process_index())
-      train_ds = get_ds_fn(dataloading_host_index=dataloading_host_index, dataloading_host_count=dataloading_host_count)
+      train_ds = get_ds_fn(
+          dataloading_host_index=dataloading_host_index,
+          dataloading_host_count=dataloading_host_count,
+      )
       train_dataloader = preprocessing_fn(dataset=train_ds)
       train_dataloader_list.append(train_dataloader)
     return [
@@ -557,7 +697,12 @@ def make_grain_train_iterator(
         else None
     )
     train_dataloader = _make_elastic_iterator(
-        train_ds, config, preprocessing_fn, shard_index=shard_index, shard_count=shard_count, mp_opts=mp_options
+        train_ds,
+        config,
+        preprocessing_fn,
+        shard_index=shard_index,
+        shard_count=shard_count,
+        mp_opts=mp_options,
     )
   else:
     train_dataloader = preprocessing_fn(dataset=train_ds)

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -66,12 +66,17 @@ def add_segmentation_and_position(x, data_columns, padding_token=0):
   for data_column in data_columns:
     x[f"{data_column}_segmentation"] = tf.cast(x[data_column] != padding_token, tf.int32)
     x[f"{data_column}_position"] = tf.broadcast_to(
-        tf.range(x[data_column].shape[-1], dtype=np.int32)[None, :], x[data_column].shape
+        tf.range(x[data_column].shape[-1], dtype=np.int32)[None, :],
+        x[data_column].shape,
     )
   return x
 
 
-def TokenizeOp(tokenizer_model, features: Features, data_keys: Iterable[str] = ("inputs", "targets")) -> Features:
+def TokenizeOp(
+    tokenizer_model,
+    features: Features,
+    data_keys: Iterable[str] = ("inputs", "targets"),
+) -> Features:
   """Op for tokenization"""
   import tensorflow as tf  # pylint: disable=import-outside-toplevel
 
@@ -314,14 +319,39 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
 
 
 def tokenization(example, hf_tokenizer, truncation, max_length, column_names):
-  """Tokenize a HuggingFace dataset"""
+  """Tokenize text columns using HuggingFace or SentencePiece/Tiktoken tokenizer."""
+
+  def _encode(text):
+    if callable(hf_tokenizer):
+      try:
+        res = hf_tokenizer(text, truncation=truncation, max_length=max_length)
+        if isinstance(res, dict) and "input_ids" in res:
+          return res["input_ids"]
+        if isinstance(res, list):
+          return res
+      except TypeError:
+        pass
+
+    if hasattr(hf_tokenizer, "tokenizer") and callable(hf_tokenizer.tokenizer):
+      try:
+        res = hf_tokenizer.tokenizer(text, truncation=truncation, max_length=max_length)
+        if isinstance(res, dict) and "input_ids" in res:
+          return res["input_ids"]
+        if isinstance(res, list):
+          return res
+      except TypeError:
+        pass
+
+    if hasattr(hf_tokenizer, "encode"):
+      return hf_tokenizer.encode(text)
+
+    raise ValueError(f"Unsupported tokenizer: {hf_tokenizer}")
+
   for column_name in column_names:
     if isinstance(example[column_name], list):
-      example[column_name] = [
-          hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name]
-      ]
+      example[column_name] = [_encode(x) for x in example[column_name]]
     elif isinstance(example[column_name], str):
-      example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)["input_ids"]
+      example[column_name] = _encode(example[column_name])
   return example
 
 
@@ -369,7 +399,12 @@ def __init__(self, query_column, response_column, max_target_length, pad_id):
 
   def map(self, element):
     inputs = np.concatenate((element[self.query_column], element[self.response_column]))
-    targets = np.concatenate((np.asarray([self.pad_id] * len(element[self.query_column])), element[self.response_column]))
+    targets = np.concatenate(
+        (
+            np.asarray([self.pad_id] * len(element[self.query_column])),
+            element[self.response_column],
+        )
+    )
     return {
         "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32),
         "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32),
@@ -527,7 +562,11 @@ def compute_file_sharding(file_count, host_index, host_count):
       and the file's group has >1 reader; otherwise None.
   """
   if file_count >= host_count:
-    return slice(host_index, None, host_count), max(file_count // host_count, 1), None
+    return (
+        slice(host_index, None, host_count),
+        max(file_count // host_count, 1),
+        None,
+    )
   file_idx = host_index % file_count
   row_shard_idx = host_index // file_count
   row_shard_count = (host_count // file_count) + (1 if file_idx < (host_count % file_count) else 0)
@@ -864,11 +903,13 @@ def map(
             np.int32
         )
         element[f"{data_column}_position"] = np.arange(
-            element[data_column].shape[0], dtype=np.int32  # pyrefly: ignore[missing-attribute]
+            element[data_column].shape[0],
+            dtype=np.int32,  # pyrefly: ignore[missing-attribute]
         )  # pyrefly: ignore[missing-attribute]
         if self.add_true_length:
           element[f"{data_column}_true_length"] = np.array(
-              [element[data_column].shape[0]], dtype=np.int32  # pyrefly: ignore[missing-attribute]
+              [element[data_column].shape[0]],
+              dtype=np.int32,  # pyrefly: ignore[missing-attribute]
           )  # pyrefly: ignore[missing-attribute]
 
     for key, _ in element.items():

src/maxtext/multimodal/utils.py

@@ -14,6 +14,7 @@
 
 """General utility functions for multimodal processing."""
 
+import io
 import os
 from dataclasses import dataclass
 from typing import Optional, Union
@@ -62,7 +63,13 @@ class PreprocessorOutput:
 
 def convert_to_RGB(image):
   """Convert image to RGB format."""
-  if image.mode != "RGB":
+  if isinstance(image, dict) and "bytes" in image and image["bytes"] is not None:
+    image = Image.open(io.BytesIO(image["bytes"]))
+  elif isinstance(image, bytes):
+    image = Image.open(io.BytesIO(image))
+  elif isinstance(image, str):
+    image = Image.open(image)
+  if hasattr(image, "mode") and image.mode != "RGB":
     image = image.convert("RGB")
   return image
 

src/maxtext/utils/globals.py

@@ -24,7 +24,7 @@
 MAXTEXT_REPO_ROOT = os.environ.get(
     "MAXTEXT_REPO_ROOT",
     r
-    if os.path.isdir(
+    if os.path.exists(
         os.path.join(r := os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), ".git")
     )
     else MAXTEXT_PKG_DIR,