Add Grain support for multimodal SFT data processing pipeline

[snehalv2002]

Description

This PR migrates the Multimodal Supervised Fine-Tuning (SFT) data pipeline from dataset_type=hf to dataset_type=grain, with support for single-image (ChartQA), multi-image (SlideVQA), and MRoPE-based models (Qwen3-VL).

Summary of Changes:

1. Grain Multimodal SFT Pipeline:
   • Implemented vision_sft_preprocessing_pipeline in src/maxtext/input_pipeline/grain_data_processing.py to handle prompt formatting, image decoding, Grain-native tokenization (grain_tokenizer.TokenizeAndTrim), prompt masking, batching, image-folding, MRoPE 3D position computation (ComputeQwen3OmniPositions), and sequence shifting in Grain.
   • Handled both single-image (train_image_column: 'image') and multi-image columns (train_image_column: ['page_1', ..., 'page_20']) via explicit image_column parameter.
   • Updated get_tokenizer_and_pad_id to allow overriding add_bos and add_eos to False for multimodal SFT pipelines to prevent duplicate control tokens.
   • Updated _get_pipeline_fn in grain_data_processing.py to route to vision_sft_preprocessing_pipeline when config.use_sft and config.use_multimodal.
2. Multimodal Utilities:
   • Enhanced convert_to_RGB in src/maxtext/multimodal/utils.py to support dictionary-wrapped bytes ({'bytes': ...} as formatted in Parquet/Grain datasets), raw bytes, and file paths.
3. Configs, Tests & Documentation:
   • Updated src/maxtext/configs/post_train/sft-vision-chartqa.yml and src/maxtext/configs/post_train/sft-vision-slidevqa.yml to default to dataset_type: grain and grain_file_type: parquet.
   • Updated docs/tutorials/posttraining/multimodal.md and tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh to use dataset_type=grain.
   • Consolidated unit tests in tests/unit/grain_sft_test.py marked with @pytest.mark.cpu_only covering ChartQA, SlideVQA multi-image merging, elastic iterator validation, BOS/EOS exclusion, Qwen3-VL MRoPE 3D positions, and exact output batch array equality matching between Hugging Face and Grain pipelines.

Tests

• Unit Tests:
  • PYTHONPATH=src pytest -m cpu_only tests/unit/grain_sft_test.py (Passed 6/6 on CPU).
  • PYTHONPATH=src pytest tests/unit/hf_data_processing_test.py (Passed 3/3).
  • PYTHONPATH=src pytest tests/unit/maxtext_utils_test.py -k test_qwen_and_gemma_omit_image_masks (Passed 1/1).
• Remote TPU Verification:
  • Gemma3-4B SFT: Executed a 5-step native SFT training run on ChartQA via Grain on TPU v4 (snehalv-tpu-v4).
  • Qwen3-VL-2B SFT: Executed a native SFT training run with MRoPE positions on ChartQA via Grain on TPU v4 (snehalv-tpu-v4) with model_name=qwen3-vl-2b, completing training steps and saving checkpoints.
• Pre-commit Checks:
  • pylint (10.00/10), pyink clean.

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a new multimodal SFT input pipeline using Grain, including the vision_sft_preprocessing_pipeline for handling text and image preprocessing, updated configurations, and unit tests. The reviewer feedback highlights several important improvements for robustness: explicitly raising an error for the unsupported ElasticIterator in multimodal SFT (and simplifying the corresponding batching logic), reordering tokenizer checks in tokenization to prevent HuggingFace tokenizers from bypassing truncation and max length limits, and ensuring robust handling of both list and tuple types when parsing data and image columns.

[gemini-code-assist]

ElasticIterator is not supported yet for multimodal SFT because post-batch transformations (like folding images) cannot be easily applied to the iterator. We should explicitly raise a ValueError if config.grain_use_elastic_iterator is enabled to prevent silent failures or cryptic runtime errors.

Suggested change

	assert len(data_columns) == 2, f"Need two data_columns for query and response, received {data_columns=}"
	text_columns = list(data_columns)
	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)

[gemini-code-assist]

Since grain_use_elastic_iterator is not supported for multimodal SFT (and we raise an error at the start of the function), we can simplify this block and remove the dead/broken conditional branch.

  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))

[gemini-code-assist]

Checking hasattr(hf_tokenizer, "encode") first will always evaluate to True for HuggingFace tokenizers, completely bypassing the callable check. This means HuggingFace tokenizers will be called via encode without truncation and max_length arguments, silently ignoring these limits. We should reorder the checks to prioritize the callable/tokenizer wrapper paths and only fall back to encode if they are not applicable.

  def _encode(text):
    if hasattr(hf_tokenizer, "tokenizer"):
      return hf_tokenizer.tokenizer(text, truncation=truncation, max_length=max_length)["input_ids"]
    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"]
        return res
      except TypeError:
        pass
    if hasattr(hf_tokenizer, "encode"):
      return hf_tokenizer.encode(text)
    raise ValueError(f"Unsupported tokenizer: {hf_tokenizer}")

[gemini-code-assist]

Comparing data_columns directly to config.eval_data_columns can be fragile if one is parsed as a list and the other as a tuple. Converting both to lists before comparison ensures robust matching.

Suggested change

	if data_columns == getattr(config, "eval_data_columns", None):
	image_column = getattr(config, "eval_image_column", "image")
	else:
	image_column = getattr(config, "train_image_column", "image")
	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")

[gemini-code-assist]

If image_column is parsed as a tuple, checking isinstance(image_column, list) will evaluate to False, causing it to be treated as a single string and leading to malformed columns. We should check for both list and tuple types.

Suggested change

	if isinstance(image_column, list):
	columns_to_parse = text_columns + image_column
	else:
	columns_to_parse = text_columns + [image_column]
	if isinstance(image_column, (list, tuple)):
	columns_to_parse = text_columns + list(image_column)
	else:
	columns_to_parse = text_columns + [image_column]

[gemini-code-assist]

Similarly, we should support tuple types when merging multiple image columns.

Suggested change

	if isinstance(image_column, list):
	dataset = dataset.map(
	functools.partial(
	input_pipeline_utils.merge_image_columns,
	image_columns=image_column,
	max_num_images_per_example=config.max_num_images_per_example,
	)
	)
	image_column = "images"
	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"

[codecov]

Codecov Report

❌ Patch coverage is 67.74194% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/input_pipeline/grain_data_processing.py	72.91%	8 Missing and 5 partials ⚠️
src/maxtext/multimodal/utils.py	12.50%	3 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

[aireenmei]

Could you see if we can also migrate https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/post_train/sft-vision-slidevqa.yml with current change? Feel free to leave it to future PR if additional changes needed

[aireenmei]

input_pipeline_utils.tokenization is specific to hf input pipeline, grain_tokenizer.TokenizeAndTrim is the counterpart for grain

[aireenmei]

Pass image_column as a param similar to data_columns.

[aireenmei]

The hf vision_sft_pipeline disables bos and eos here regardless of the add_bos and add_eos flags in base.yml. This is because the reformat prompt and response steps earlier already handles bos/eos. Update get_tokenizer_and_pad_id to allow overriding add_bos and add_eos to false. Also add comments to explain and a unit test make sure no bos, eos in the tokenization output.

[aireenmei]

1. can we add some tests to make sure the data batch of the two pipelines are matching, not just shape.
2. test_slidevqa_config_initialization seems redundant as sft-vision-slidevqa.yml is already included in BASE_CONFIGS in tests/unit/configs_test.py.
3. multimodal_sft_grain_test seems very specific to this feature, can we group the test into tests/unit/grain_data_processing_test.py? Or we can make a separate grain_sft_test.py, and we expand it later as we migrate other sft pipelines. Up to you, but try to group the tests properly and leverage existing files and avoid having too many specific unit test files.

[aireenmei]

for 1. tests/post_training/unit/sft_data_processing_test.py is a good example.

[aireenmei]

Thanks for the work! Could you check if the newly added unit tests are running on CPU or TPU? I think they should be fine on CPU since our TPU resource is limited

[aireenmei]

From the PR description, 3. Tokenizer & Worktree Interoperability and 4. IterDataset Conversion seems to be beyond the scope of this PR? Please double check if they are needed to this feature

[aireenmei]

From the PR description, 3. Tokenizer & Worktree Interoperability and 4. IterDataset Conversion seems to be beyond the scope of this PR? Please double check if they are needed to this feature