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.

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, tokenization, prompt masking, batching, image-folding, and padding in Grain.
   • 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. Tokenizer Interoperability:
   • Updated tokenization in src/maxtext/input_pipeline/input_pipeline_utils.py to support SentencePiece, TikToken, and Hugging Face tokenizers.
4. IterDataset Conversion:
   • Updated apply_multiprocessing_and_prefetch in src/maxtext/input_pipeline/data_processing_utils.py to convert MapDataset to IterDataset via .to_iter_dataset() when needed.
5. Configs, End-to-End Tests & Documentation:
   • Updated src/maxtext/configs/post_train/sft-vision-chartqa.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.
   • Added unit tests in tests/unit/input_pipeline/multimodal_sft_grain_test.py.

Tests

• Unit Tests:
  • PYTHONPATH=src pytest tests/unit/input_pipeline/multimodal_sft_grain_test.py (Passed 2/2).
• Remote TPU Verification:
  • Executed a 5-step native SFT training run with Gemma3-4B on ChartQA via Grain on TPU v4 (snehalv-tpu-v4):

  python -m maxtext.trainers.post_train.sft.train_sft_native \
    src/maxtext/configs/post_train/sft-vision-chartqa.yml \
    model_name=gemma3-4b per_device_batch_size=1 steps=5 \
    max_prefill_predict_length=1024 max_target_length=2048 \
    scan_layers=false async_checkpointing=False attention=dot_product \
    tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3 \
    base_output_directory=gs://runner-maxtext-logs/gemma3-4b/multimodal/sft \
    dataset_type=grain grain_file_type=parquet \
    grain_train_files=gs://aireenmei-multipod/dataset/hf/chartqa/train-*

  • Result: Completed steps 0–4 with checkpoints saved and expected throughput.
• Pre-commit Checks:
  • pre-commit run --files ... passed (codespell, pylint, pyink, mdformat, yamllint).

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 52.56410% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/input_pipeline_utils.py	37.03%	13 Missing and 4 partials ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	68.29%	7 Missing and 6 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