Add images in SFT loss_func for multimodal post training

[dgouju]

Description

MaxText's post-training wrapper function (use_maxtext_loss_function in train_sft.py) has a hardcoded function signature accepting only text keyword arguments (inputs, targets, etc.), but Tunix's batch iterator unpacks all batch dictionary keys—including images and image_masks—into the loss function during multimodal training.

I was working on post-train Gemma3-4B on Robotic dataset, previously converted from HF LeRobot to Parquet format.

Tests

Tested by post-training Gemma3-4B with the following command:

python3 -m maxtext.trainers.post_train.sft.train_sft \
    src/maxtext/configs/post_train/sft-vision-chartqa.yml \
    run_name=${RUN_NAME} \
    model_name="gemma3-4b" \
    tokenizer_path="google/gemma-3-4b-it" \
    tokenizer_type="huggingface" \
    train_split="train" \
    train_data_columns="['prompt','completion']" \
    train_image_column="image" \
    sharding_tolerance=0.05 \
    load_parameters_path=${MAXTEXT_CKPT_PATH} \
    base_output_directory=${BASE_OUTPUT_DIRECTORY} \
    dataset_type="hf" \
    hf_path="parquet" \
    hf_train_files="test/dataset.parquet" \
    hf_eval_files="test/dataset.parquet" \
    train_split="train" \
    hf_eval_split="train" \
    train_data_columns="['prompt','completion']" \
    eval_data_columns="['prompt','completion']" \
    train_image_column="image" \
    eval_image_column="image" \
    max_num_images_per_example=4 \
    use_sft=true \
    use_multimodal=true \
    sft_train_on_completion_only=true \
    freeze_vision_encoder_params=true \
    learning_rate=2e-5 \
    per_device_batch_size=1 \
    max_prefill_predict_length=1024 \
    max_target_length=2048 \
    steps=1000 \
    dtype="bfloat16" \
    weight_dtype="bfloat16" \
    async_checkpointing=false \
    attention="dot_product"

Checklist

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

• [ X ] 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.
• [ X ] 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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.