From 94095d29a0e1bb74c83b42fba9b3a870a38f89a3 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 6 Aug 2026 01:04:02 +0000 Subject: [PATCH] Add Grain support for multimodal SFT data processing pipeline --- docs/tutorials/posttraining/multimodal.md | 2 +- .../configs/post_train/sft-vision-chartqa.yml | 7 +- .../post_train/sft-vision-slidevqa.yml | 5 +- .../input_pipeline/data_processing_utils.py | 2 + .../input_pipeline/grain_data_processing.py | 159 +++++++++++++++++- .../input_pipeline/input_pipeline_utils.py | 63 +++++-- src/maxtext/multimodal/utils.py | 9 +- src/maxtext/utils/globals.py | 2 +- .../gemma3/4b/test_gemma3_multimodal_sft.sh | 9 +- .../multimodal_sft_grain_test.py | 134 +++++++++++++++ 10 files changed, 362 insertions(+), 30 deletions(-) create mode 100644 tests/unit/input_pipeline/multimodal_sft_grain_test.py diff --git a/docs/tutorials/posttraining/multimodal.md b/docs/tutorials/posttraining/multimodal.md index 90ceff00bc..97b2da3db0 100644 --- a/docs/tutorials/posttraining/multimodal.md +++ b/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 diff --git a/src/maxtext/configs/post_train/sft-vision-chartqa.yml b/src/maxtext/configs/post_train/sft-vision-chartqa.yml index e4e32eb539..aba2b78a50 100644 --- a/src/maxtext/configs/post_train/sft-vision-chartqa.yml +++ b/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' diff --git a/src/maxtext/configs/post_train/sft-vision-slidevqa.yml b/src/maxtext/configs/post_train/sft-vision-slidevqa.yml index f77d5910e6..4c6b766e68 100644 --- a/src/maxtext/configs/post_train/sft-vision-slidevqa.yml +++ b/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' diff --git a/src/maxtext/input_pipeline/data_processing_utils.py b/src/maxtext/input_pipeline/data_processing_utils.py index 3336b1fed3..ae6b7f19e6 100644 --- a/src/maxtext/input_pipeline/data_processing_utils.py +++ b/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, diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 382df3fd16..0557710e4d 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/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) diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 45b13b24a5..7091f71741 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/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(): diff --git a/src/maxtext/multimodal/utils.py b/src/maxtext/multimodal/utils.py index 5b7298bee3..4e193b76cd 100644 --- a/src/maxtext/multimodal/utils.py +++ b/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 diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index c06f4f4f10..a575216328 100644 --- a/src/maxtext/utils/globals.py +++ b/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, diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh index c916a0a3ac..0cd4a37265 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh @@ -54,11 +54,10 @@ python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DI max_prefill_predict_length=1024 \ max_target_length=2048 \ steps=5 \ - scan_layers=true \ - async_checkpointing=False \ - attention=\'dot_product\' \ - dataset_type=hf hf_path=parquet \ - hf_train_files=gs://aireenmei-multipod/dataset/hf/chartqa/train-* \ + scan_layers=false async_checkpointing=False \ + attention=dot_product \ + dataset_type=grain grain_file_type=parquet \ + grain_train_files=gs://aireenmei-multipod/dataset/hf/chartqa/train-* \ base_output_directory=${BASE_OUTPUT_DIRECTORY}/multimodal/sft \ load_parameters_path=${MULTIMODAL_SCANNED_CKPT_PATH} \ dtype=bfloat16 \ diff --git a/tests/unit/input_pipeline/multimodal_sft_grain_test.py b/tests/unit/input_pipeline/multimodal_sft_grain_test.py new file mode 100644 index 0000000000..21ac466f19 --- /dev/null +++ b/tests/unit/input_pipeline/multimodal_sft_grain_test.py @@ -0,0 +1,134 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for multimodal SFT grain input pipeline.""" + +import unittest +from PIL import Image +import grain.python as grain + +from maxtext.configs import pyconfig +from maxtext.input_pipeline.grain_data_processing import ( + vision_sft_preprocessing_pipeline, + _get_pipeline_fn, +) + + +class MultimodalSftGrainTest(unittest.TestCase): + """Tests for the Grain multimodal SFT pipeline.""" + + def setUp(self): + super().setUp() + self.config = pyconfig.initialize( + [ + "", + "src/maxtext/configs/post_train/sft-vision-chartqa.yml", + "model_name=gemma3-4b", + "per_device_batch_size=1", + "max_target_length=1024", + "max_prefill_predict_length=512", + "tokenize_train_data=True", + "tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3", + "dataset_type=grain", + "grain_worker_count=0", + ] + ) + + def test_get_pipeline_fn_routes_to_vision_sft(self): + """Verifies that _get_pipeline_fn routes to vision_sft_preprocessing_pipeline.""" + fn = _get_pipeline_fn(self.config) + self.assertEqual(fn, vision_sft_preprocessing_pipeline) + + def test_vision_sft_pipeline_execution(self): + """Tests executing the Grain vision SFT pipeline with dummy samples.""" + dummy_image = Image.new("RGB", (100, 100), color="red") + raw_data = [ + { + "query": f"What is the chart {i} about?", + "label": [f"Sales in 202{i}"], + "image": dummy_image, + } + for i in range(8) + ] + ds = grain.MapDataset.source(raw_data) + processed_ds = vision_sft_preprocessing_pipeline( + dataset=ds, + config=self.config, + data_columns=["query", "label"], + tokenize=True, + grain_worker_count=0, + grain_per_worker_buffer_size=1, + ) + + batches = list(processed_ds) + self.assertGreater(len(batches), 0) + first_batch = batches[0] + + for expected_key in [ + "inputs", + "targets", + "images", + "inputs_position", + "targets_position", + "inputs_segmentation", + "targets_segmentation", + ]: + self.assertIn(expected_key, first_batch) + + self.assertEqual(first_batch["inputs"].shape[-1], self.config.max_target_length) + self.assertEqual(first_batch["targets"].shape[-1], self.config.max_target_length) + + def test_elastic_iterator_unsupported_error(self): + """Verifies that enabling grain_use_elastic_iterator raises a ValueError.""" + config = pyconfig.initialize( + [ + "", + "src/maxtext/configs/post_train/sft-vision-chartqa.yml", + "model_name=gemma3-4b", + "grain_file_type=arrayrecord", + "grain_train_files=dummy", + "grain_use_elastic_iterator=True", + "tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3", + ] + ) + raw_data = [{"query": "q", "label": ["l"], "image": Image.new("RGB", (10, 10))}] + ds = grain.MapDataset.source(raw_data) + with self.assertRaises(ValueError): + vision_sft_preprocessing_pipeline( + dataset=ds, + config=config, + data_columns=["query", "label"], + tokenize=True, + grain_worker_count=0, + grain_per_worker_buffer_size=1, + ) + + def test_slidevqa_config_initialization(self): + """Verifies that SlideVQA config initializes cleanly with dataset_type=grain.""" + config = pyconfig.initialize( + [ + "", + "src/maxtext/configs/post_train/sft-vision-slidevqa.yml", + "model_name=gemma3-4b", + "tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3", + ] + ) + self.assertEqual(config.dataset_type, "grain") + self.assertEqual(config.grain_file_type, "parquet") + fn = _get_pipeline_fn(config) + self.assertEqual(fn, vision_sft_preprocessing_pipeline) + + +if __name__ == "__main__": + unittest.main()