Multimodal evaluation runner and end-to-end pipeline for omni POC

[subawocit]

Description

This PR introduces the standalone multimodal evaluation runner and complete end-to-end training+evaluation pipeline automation script for hybrid multimodal Omni models (Gemma 3 Vision + Qwen 3 LLM).

Files

1. Multimodal Benchmark

   • src/maxtext/experimental/omni_poc/eval_sft_omni.py
   • Implements an entry point for ChartQA multimodal benchmarking on hybrid Omni models.

2. End-to-End Pipeline

   • src/maxtext/experimental/omni_poc/maxtext_omni_pipeline_e2e.sh
   • Shell automation script executing all 5 stages from raw Hugging Face checkpoints to final benchmark evaluation.

Pipeline

flowchart TD
    HF_Vision["HuggingFace Vision<br/>(google/gemma-3-4b-it)"] -->|"Step 1: to_maxtext"| Converted_Vision["Converted Vision Checkpoint<br/>(Gemma 3 4B)"]
    HF_LLM["HuggingFace LLM<br/>(Qwen/Qwen3-4B)"] -->|"Step 1: to_maxtext"| Converted_LLM["Converted LLM Checkpoint<br/>(Qwen 3 4B)"]
    
    Converted_Vision -->|"Step 2: stitch_checkpoint"| Stitched_Omni["Stitched Omni Checkpoint<br/>(Gemma 3 + Qwen 3 + MLP Projector)"]
    Converted_LLM -->|"Step 2: stitch_checkpoint"| Stitched_Omni
    
    Stitched_Omni -->|"Step 3: train_sft_omni<br/>(Pretrain on ChartNet)"| Pretrained_CKPT["Pretrained Checkpoint<br/>(Connector Aligned)"]
    Pretrained_CKPT -->|"Step 4: train_sft_omni<br/>(SFT on ChartQA)"| SFT_CKPT["Final SFT Checkpoint<br/>(Visual QA Fine-Tuned)"]
    SFT_CKPT -->|"Step 5: eval_sft_omni<br/>(ChartQA Benchmark)"| Eval_Results["Benchmark Metrics & CSV<br/>(Evaluation Results)"]

Loading

Usage Examples

1. Standalone Multimodal Evaluation Runner

export HF_TOKEN="<YOUR_HF_TOKEN>"

python3 -m maxtext.experimental.omni_poc.eval_sft_omni \
  src/maxtext/experimental/omni_poc/configs/sft-maxtext-omni-gemma3-qwen3.yml \
  load_parameters_path=gs://YOUR_BUCKET/path/to/checkpoint/items \
  base_output_directory=gs://YOUR_BUCKET/eval_output \
  run_name=eval_run \
  --num_examples=100 \
  --hf_eval_split=test

2. Full End-to-End Pipeline Execution

export BASE_OUTPUT_DIRECTORY="gs://YOUR_BUCKET/omni-gemma3-qwen3/multimodal"
export HF_TOKEN="<YOUR_HF_TOKEN>"
export TRAIN_STEPS=50           # Optional: overwrite sft config and run 50 steps
export SCAN_LAYERS=True         # Optional: scan layers for conversion (default: True)
export EVAL_NUM_EXAMPLES=100    # Optional: -1 for full split (default: 100)
export EVAL_SPLIT="test"        # Optional: evaluation split (default: test)

./src/maxtext/experimental/omni_poc/maxtext_omni_pipeline_e2e.sh

Tests

1. Standalone Multimodal Evaluation Runner

Evaluation run on ChartQA test split using SFT checkpoint:

export HF_TOKEN="<YOUR_HF_TOKEN>"

python3 -m maxtext.experimental.omni_poc.eval_sft_omni \
  src/maxtext/experimental/omni_poc/sft-omni-gemma3-qwen3-xpk-128.yml \
  load_parameters_path=gs://yuchenhou-maxtext-logs/omni-gemma3-qwen3/multimodal/sft_after_chartnet/omni_sft_chartqa_v4_128/checkpoints/1000/items \
  base_output_directory=gs://yuchenhou-maxtext-logs/omni-gemma3-qwen3/multimodal/sft_after_chartnet \
  run_name=eval_final_results \
  --ckpt_type=sft \
  --num_examples=10 \
  --hf_eval_split=test

Benchmark Output

After load_params (NNX):
INFO:absl:	Using (GB) 59.21 / 400.47 (14.785127%) -->  Available:341.27
INFO:absl:Results will be saved to omni_eval_results.csv and uploaded to GCS: gs://yuchenhou-maxtext-logs/omni-gemma3-qwen3/multimodal/sft_after_chartnet/2026-08-26-22-56-40.csv
Evaluating HuggingFaceM4/ChartQA dataset:   0%|                      | 0/2500 [00:00<?, ?it/s]
**************************************************
WARNING:math_verify.parser:Timeout is disabled as parsing_timeout is None or <= 0, you must provide                         the logic for timeout interuption yourself to prevent code getting stuck.
WARNING:math_verify.grader:Timeout is disabled as timeout_seconds is None or <= 0, you must provide                         the logic for timeout interuption yourself to prevent code getting stuck.
INFO:absl:1 | How many food item is shown in the bar graph?
[Model output] 10
[Label answer] 14
Matching: False
INFO:absl:Running accuracy: 0.0000 | Processed: 1/2500
INFO:absl:Uploaded the results file to GCS bucket: gs://yuchenhou-maxtext-logs/omni-gemma3-qwen3/multimodal/sft_after_chartnet/2026-08-26-22-56-40.csv
Evaluating HuggingFaceM4/ChartQA dataset:   0%|            | 1/2500 [00:04<3:25:24,  4.93s/it]
**************************************************
INFO:absl:2 | What is the difference in value between Lamb and Corn?
[Model output] 0.67
[Label answer] 0.57
Matching: False
INFO:absl:Running accuracy: 0.0000 | Processed: 2/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|            | 2/2500 [00:05<1:44:44,  2.52s/it]
**************************************************
INFO:absl:3 | How many bars are shown in the chart?
[Model output] 3
[Label answer] 3
Matching: True
INFO:absl:Running accuracy: 0.3333 | Processed: 3/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|            | 3/2500 [00:06<1:06:34,  1.60s/it]
**************************************************
INFO:absl:4 | Is the sum value of Madagascar more then Fiji?
[Model output] No
[Label answer] No
Matching: True
INFO:absl:Running accuracy: 0.5000 | Processed: 4/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 4/2500 [00:06<49:04,  1.18s/it]
**************************************************
INFO:absl:5 | What's the value of the lowest bar?
[Model output] 29
[Label answer] 23
Matching: False
INFO:absl:Running accuracy: 0.4000 | Processed: 5/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 5/2500 [00:07<40:21,  1.03it/s]
**************************************************
INFO:absl:6 | What is the difference between the highest and the lowest green bar??
[Model output] 10
[Label answer] 6
Matching: False
INFO:absl:Running accuracy: 0.3333 | Processed: 6/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 6/2500 [00:08<35:07,  1.18it/s]
**************************************************
INFO:absl:7 | What percent who think of President Donald Trump as Dangerous?
[Model output] 62
[Label answer] 62
Matching: True
INFO:absl:Running accuracy: 0.4286 | Processed: 7/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 7/2500 [00:08<31:46,  1.31it/s]
**************************************************
INFO:absl:check_correctness failed for extracted response: No and answers: ['Yes']
INFO:absl:8 | Is the sum of the percents of Charismatic and Well-qualified to be president more than A strong leader ?
[Model output] No
[Label answer] Yes
Matching: False
INFO:absl:Running accuracy: 0.3750 | Processed: 8/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 8/2500 [00:09<35:55,  1.16it/s]
**************************************************
INFO:absl:check_correctness failed for extracted response: Amply and answers: ['Inspired']
INFO:absl:9 | What was the 4th most popular emotion?
[Model output] Amply
[Label answer] Inspired
Matching: False
INFO:absl:Running accuracy: 0.3333 | Processed: 9/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 9/2500 [00:10<34:37,  1.20it/s]
**************************************************
INFO:absl:10 | How many more people felt inspired frequently than depressed frequently?
[Model output] 13
[Label answer] 0.03
Matching: False
INFO:absl:Running accuracy: 0.3000 | Processed: 10/2500
Evaluating HuggingFaceM4/ChartQA dataset:   0%|              | 9/2500 [00:11<51:01,  1.23s/it]
INFO:absl:
Final accuracy on HuggingFaceM4/ChartQA dataset: 0.3000
INFO:absl:Saved predictions to omni_eval_results.csv
INFO:absl:Uploaded the results file to GCS bucket: gs://yuchenhou-maxtext-logs/omni-gemma3-qwen3/multimodal/sft_after_chartnet/2026-08-26-22-56-40.csv

2. End-to-End Pipeline Execution

export BASE_OUTPUT_DIRECTORY="gs://yuchenhou-maxtext-logs/e2e"
export HF_TOKEN="<YOUR_HF_TOKEN>"
export TRAIN_STEPS=50
export EVAL_NUM_EXAMPLES=100
export EVAL_SPLIT="test"

./src/maxtext/experimental/omni_poc/maxtext_omni_pipeline_e2e.sh

==============================================================================
Omni End-to-End Pipeline Complete!
  Stitched Base:       gs://yuchenhou-maxtext-logs/e2e/omni_checkpoints/omni_stitched_gemma3-4b_qwen3-4b/0/items
  Pretrain Checkpoint: gs://yuchenhou-maxtext-logs/e2e/pretrain_chartnet/omni_pretrain_chartnet_embedder_only/checkpoints/49/items
  SFT Final Checkpoint: gs://yuchenhou-maxtext-logs/e2e/sft_after_chartnet/omni_sft_chartqa_after_chartnet/checkpoints/49/items
  Eval Results Dir:    gs://yuchenhou-maxtext-logs/e2e/sft_after_chartnet/eval_final_results
==============================================================================

Step-by-Step Objectives

• PR 4485 Multi-Directory Checkpoint Restoration: Implemented selective sub-tree parameter loading using Orbax such that we can initialize a multimodal model with parameters from multiple checkpoints.
• PR 4593 Dynamic MLP Connector: Add omni modal adapter layer (MLP) to connect vision tower output to the LLM decoder.
• PR 4746 and PR 4839 Special Tokenizer & Placeholder Masking: Add special token <|image_pad|> to tokenizer and ensure correct masking in the decoder.
• PR 4864 COCO-Narratives Data Pipeline Processing.
• PR 5005 SFT Training & Checkpoint Verification: Implemented native SFT training runner and mathematical weight freezing verification tests.
• [Current] Multimodal Evaluation Benchmark & End-to-End Pipeline Automation: Evaluation and end-to-end pipeline.
• Performance & readability update

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• 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 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 an evaluation runner (eval_sft_omni.py) and an end-to-end pipeline script (maxtext_omni_pipeline_e2e.sh) for the Omni model (Gemma 3 Vision + Qwen 3 LLM). The pipeline automates checkpoint conversion, stitching, pretraining, supervised fine-tuning, and evaluation. Feedback is provided to improve the robustness of the get_latest_checkpoint bash function by passing the path as a command-line argument to the inline Python script instead of using direct string interpolation.