Enable SFT for multimodal llama4

Author: aireenmei

MaxText/configs/base.yml

@@ -409,7 +409,9 @@ expansion_factor_real_data: -1 # if -1 then all hosts will load real data, else
 eval_per_device_batch_size: 0.0
 max_corpus_chars: 10_000_000
 train_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected"
+train_image_column: 'image'
 eval_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected"
+eval_image_column: 'image'
 packing: True
 num_epoch: 1  # only grain and tfds pipeline supports num_epoch > 1
 
@@ -732,7 +734,7 @@ dtype_mm: "float32"  # Data type for multimodal model's vision encoder
 remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options.
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding
-
+image_placeholder: "<|image|>"
 
 ### llama4 multi modal configs
 hidden_size_for_vit: 1408 

MaxText/configs/sft-vision-chartqa.yml

@@ -0,0 +1,31 @@
+# Copyright 2025 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.
+
+base_config: "base.yml"
+
+use_sft: True
+use_multimodal: True
+# For vision, the prompt contrains image, we only train on completion tokens
+sft_train_on_completion_only: True
+packing: False  # packing is not supported yet
+freeze_vision_encoder_params: True
+learning_rate: 2.e-5
+
+# -------------- HF pipeline --------------
+dataset_type: hf
+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

MaxText/decode.py

@@ -100,7 +100,7 @@ def main(argv: Sequence[str]) -> None:
   prefill_length = config.max_prefill_predict_length
   processor_output = multimodal_utils.PreprocessorOutput()
   if config.use_multimodal:
-    text = multimodal_utils.reformat_prompt(text, config.model_name)
+    text = multimodal_utils.reformat_prompt(text, image_placeholder=config.image_placeholder, model_name=config.model_name)
     # TODO(hengtaoguo): Support multiple images as input.
     images = multimodal_utils.load_image_from_path(config.image_path)
     processor_output = multimodal_utils.pre_process_image(images, model_name=config.model_name)

MaxText/input_pipeline/_hf_data_processing.py

@@ -32,6 +32,116 @@
 from MaxText import multihost_dataloading
 
 
+def vision_sft_preprocessing_pipeline(
+    dataset,
+    config,
+    dataloading_host_index,
+    dataloading_host_count,
+    global_mesh,
+    text_columns,
+    image_column,
+    global_batch_size,
+):
+  """pipeline for multimodal SFT with HF dataset"""
+
+  assert len(text_columns) == 2, f"Need two text_columns for query and response, received {text_columns=}"
+
+  if config.enable_data_shuffling:
+    dataset = dataset.shuffle(seed=config.data_shuffle_seed)
+
+  dataset = dataset.select_columns(text_columns + [image_column])
+  dataset = dataset.map(
+      _input_pipeline_utils.reformat_prompt,
+      fn_kwargs={"column": text_columns[0], "image_placeholder": config.image_placeholder, "model_name": config.model_name},
+  )
+  dataset = dataset.map(
+      _input_pipeline_utils.reformat_response,
+      fn_kwargs={"column": text_columns[1], "model_name": config.model_name},
+  )
+  if image_column != "images":
+    dataset = dataset.rename_column(image_column, "images")
+
+  dataset = dataset.map(
+      _input_pipeline_utils.pre_process_image_sft,
+      fn_kwargs={"image_column": "images", "model_name": config.model_name},
+  )
+
+  tokenizer = transformers.AutoTokenizer.from_pretrained(
+      config.tokenizer_path,
+      add_bos_token=False,
+      add_eos_token=False,
+      legacy=False,
+      token=config.hf_access_token,
+  )
+  if tokenizer.pad_token_id is not None:
+    pad_id = tokenizer.pad_token_id
+  elif tokenizer.unk_token_id is not None:
+    pad_id = tokenizer.unk_token_id
+  else:
+    pad_id = -1
+
+  dataset = dataset.map(
+      _input_pipeline_utils.tokenization,
+      batched=True,
+      fn_kwargs={
+          "hf_tokenizer": tokenizer,
+          "truncation": False,
+          "max_length": config.max_target_length,
+          "column_names": text_columns,
+      },
+  )
+  dataset = dataset.map(
+      _input_pipeline_utils.prepare_text_for_image_fusion,
+      fn_kwargs={"column_name": text_columns[0], "model_name": config.model_name},
+  )
+
+  dataset = _input_pipeline_utils.HFDataSource(
+      dataset=dataset,
+      dataloading_host_index=dataloading_host_index,
+      dataloading_host_count=dataloading_host_count,
+      num_threads=1,
+      generate_padding_example=True,
+      max_target_length=config.max_target_length,
+      data_column_names=text_columns,
+  )
+  operations = []
+  operations.append(
+      _input_pipeline_utils.SFTPromptMaskingVision(
+          query_column=text_columns[0],
+          response_column=text_columns[1],
+          max_target_length=config.max_target_length,
+          unk_id=pad_id,
+      )
+  )
+  # TODO(aireenmei, hengtaoguo): support packing
+  operations.append(_input_pipeline_utils.PadToMaxLength(config.max_target_length, pad_id))
+  operations.append(grain.Batch(batch_size=global_batch_size // jax.process_count(), drop_remainder=True))
+  operations.append(_input_pipeline_utils.ShiftData(ignored_ids=[pad_id], axis=1))
+  dummy_index_sampler = grain.IndexSampler(
+      num_records=len(dataset),
+      num_epochs=1,
+      shard_options=grain.ShardOptions(
+          shard_index=dataloading_host_index, shard_count=dataloading_host_count, drop_remainder=False
+      ),
+      shuffle=False,
+      seed=0,
+  )
+
+  dataloader = grain.DataLoader(
+      data_source=dataset,
+      operations=operations,
+      sampler=dummy_index_sampler,
+      worker_count=1,  # only supports <=1 for now, more workers results in duplicated data
+      worker_buffer_size=1,
+      read_options=grain.ReadOptions(num_threads=1, prefetch_buffer_size=128),
+  )
+
+  multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(dataloader, global_mesh)
+
+  # Return multi-host jax.Array prep iterator
+  return multihost_gen
+
+
 def preprocessing_pipeline(
     dataloading_host_index,
     dataloading_host_count,
@@ -212,27 +322,39 @@ def make_hf_train_iterator(
       streaming=True,
       token=config.hf_access_token,
   )
-  train_iter = preprocessing_pipeline(
-      dataloading_host_index=process_indices_train.index(jax.process_index()),
-      dataloading_host_count=len(process_indices_train),
-      global_mesh=global_mesh,
-      dataset=train_ds,
-      data_column_names=config.train_data_columns,
-      tokenize=config.tokenize_train_data,
-      tokenizer_path=config.tokenizer_path,
-      hf_access_token=config.hf_access_token,
-      global_batch_size=config.global_batch_size_to_load,
-      max_target_length=config.max_target_length,
-      shuffle=config.enable_data_shuffling,
-      data_shuffle_seed=config.data_shuffle_seed,
-      add_bos=config.add_bos,
-      add_eos=config.add_eos,
-      packing=config.packing,
-      generate_padding_example=False,
-      use_dpo=config.use_dpo,
-      use_sft=config.use_sft,
-      sft_train_on_completion_only=config.sft_train_on_completion_only,
-  )
+  if config.use_sft and config.use_multimodal:
+    train_iter = vision_sft_preprocessing_pipeline(
+        dataset=train_ds,
+        config=config,
+        dataloading_host_index=process_indices_train.index(jax.process_index()),
+        dataloading_host_count=len(process_indices_train),
+        global_mesh=global_mesh,
+        text_columns=config.train_data_columns,
+        image_column=config.train_image_column,
+        global_batch_size=config.global_batch_size_to_load,
+    )
+  else:
+    train_iter = preprocessing_pipeline(
+        dataloading_host_index=process_indices_train.index(jax.process_index()),
+        dataloading_host_count=len(process_indices_train),
+        global_mesh=global_mesh,
+        dataset=train_ds,
+        data_column_names=config.train_data_columns,
+        tokenize=config.tokenize_train_data,
+        tokenizer_path=config.tokenizer_path,
+        hf_access_token=config.hf_access_token,
+        global_batch_size=config.global_batch_size_to_load,
+        max_target_length=config.max_target_length,
+        shuffle=config.enable_data_shuffling,
+        data_shuffle_seed=config.data_shuffle_seed,
+        add_bos=config.add_bos,
+        add_eos=config.add_eos,
+        packing=config.packing,
+        generate_padding_example=False,
+        use_dpo=config.use_dpo,
+        use_sft=config.use_sft,
+        sft_train_on_completion_only=config.sft_train_on_completion_only,
+    )
   return train_iter
 
 
@@ -252,25 +374,37 @@ def make_hf_eval_iterator(
   )
 
   eval_generate_padding_example = config.eval_steps > 0
-  eval_iter = preprocessing_pipeline(
-      dataloading_host_index=process_indices_eval.index(jax.process_index()),
-      dataloading_host_count=len(process_indices_eval),
-      global_mesh=global_mesh,
-      dataset=eval_ds,
-      data_column_names=config.eval_data_columns,
-      tokenize=config.tokenize_eval_data,
-      tokenizer_path=config.tokenizer_path,
-      hf_access_token=config.hf_access_token,
-      global_batch_size=config.global_batch_size_to_load_eval,
-      max_target_length=config.max_target_length,
-      shuffle=False,
-      data_shuffle_seed=config.data_shuffle_seed,
-      add_bos=config.add_bos,
-      add_eos=config.add_eos,
-      packing=config.packing,
-      generate_padding_example=eval_generate_padding_example,
-      use_dpo=config.use_dpo,
-      use_sft=config.use_sft,
-      sft_train_on_completion_only=config.sft_train_on_completion_only,
-  )
+  if config.use_sft and config.use_multimodal:
+    eval_iter = vision_sft_preprocessing_pipeline(
+        dataset=eval_ds,
+        config=config,
+        dataloading_host_index=process_indices_eval.index(jax.process_index()),
+        dataloading_host_count=len(process_indices_eval),
+        global_mesh=global_mesh,
+        text_columns=config.eval_data_columns,
+        image_column=config.eval_image_column,
+        global_batch_size=config.global_batch_size_to_load_eval,
+    )
+  else:
+    eval_iter = preprocessing_pipeline(
+        dataloading_host_index=process_indices_eval.index(jax.process_index()),
+        dataloading_host_count=len(process_indices_eval),
+        global_mesh=global_mesh,
+        dataset=eval_ds,
+        data_column_names=config.eval_data_columns,
+        tokenize=config.tokenize_eval_data,
+        tokenizer_path=config.tokenizer_path,
+        hf_access_token=config.hf_access_token,
+        global_batch_size=config.global_batch_size_to_load_eval,
+        max_target_length=config.max_target_length,
+        shuffle=False,
+        data_shuffle_seed=config.data_shuffle_seed,
+        add_bos=config.add_bos,
+        add_eos=config.add_eos,
+        packing=config.packing,
+        generate_padding_example=eval_generate_padding_example,
+        use_dpo=config.use_dpo,
+        use_sft=config.use_sft,
+        sft_train_on_completion_only=config.sft_train_on_completion_only,
+    )
   return eval_iter

MaxText/input_pipeline/_input_pipeline_utils.py

@@ -27,6 +27,7 @@
 import tensorflow as tf
 from MaxText import max_logging
 from MaxText import tokenizer
+from MaxText import multimodal_utils
 
 Features = Dict[str, tf.Tensor]
 AUTOTUNE = tf.data.experimental.AUTOTUNE
@@ -68,6 +69,37 @@ def add_segmentation_and_position(x, data_columns, padding_token=0):
 ########## Functions used by HF pipeline
 
 
+def reformat_prompt(example, column, image_placeholder, model_name):
+  """reformat prompt for multimodal SFT"""
+  example[column] = multimodal_utils.reformat_prompt(example[column], image_placeholder, model_name)
+  return example
+
+
+def reformat_response(example, column, model_name):
+  """reformat response for multimodal SFT"""
+  example[column] = multimodal_utils.reformat_response(example[column], model_name)
+  return example
+
+
+def pre_process_image_sft(example, image_column, model_name):
+  """pre-process image for multimodal SFT"""
+  image = multimodal_utils.convert_to_RGB(example[image_column])
+  # TODO(aireenmei, hengtaoguo): add support for different image sizes
+  image = multimodal_utils.resize_image(image, model_name)
+  image = np.array(image)
+  example[image_column] = multimodal_utils.pre_process_image(image, model_name)
+  return example
+
+
+def prepare_text_for_image_fusion(example, column_name, model_name):
+  """prepare text for image fusion for multimodal SFT"""
+  example[column_name] = multimodal_utils.prepare_text_for_image_fusion(
+      example[column_name], model_name, processor_output=example["images"]
+  )
+  example["images"] = example["images"].pixel_values
+  return example
+
+
 def combine_columns(example, columns, data_column):
   """Combine columns such as 'prompt' and 'completion' for sft training"""
   assert len(columns) > 1
@@ -192,6 +224,26 @@ def map(self, element):
     }
 
 
+@dataclasses.dataclass
+class SFTPromptMaskingVision(grain.MapTransform):
+  """SFT prompt masking for multimodal"""
+
+  def __init__(self, query_column, response_column, max_target_length, unk_id):
+    self.query_column = query_column
+    self.response_column = response_column
+    self.max_target_length = max_target_length
+    self.unk_id = unk_id
+
+  def map(self, element):
+    inputs = np.concatenate((element[self.query_column], element[self.response_column]))
+    targets = np.concatenate((np.asarray([self.unk_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),
+        "images": element["images"],
+    }
+
+
 @dataclasses.dataclass
 class HFNormalizeFeatures(grain.MapTransform):
   """Normalize feature keys for HuggingFace input"""
@@ -413,10 +465,12 @@ def _pad(x, max_length, pad_id):
 
     data_columns = list(element.keys())
     for data_column in data_columns:
-      element[f"{data_column}_segmentation"] = (element[data_column] != self.pad_id).astype(np.int32)
-      element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32)
+      if data_column != "images":
+        element[f"{data_column}_segmentation"] = (element[data_column] != self.pad_id).astype(np.int32)
+        element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32)
     for key, _ in element.items():
-      element[key] = _pad(element[key], self.max_length, self.pad_id)
+      if key != "images":
+        element[key] = _pad(element[key], self.max_length, self.pad_id)
     return element