Enable SFT for multimodal llama4

[aireenmei]

Description

Vanilla SFT support for multimodal llama4

• Using "HuggingFaceM4/ChartQA" dataset and hf data pipeline, some preprocessing steps are specific to this dataset, serves as demo for other custom preprocessing.
• Update multimodal_utils to have preprocessing functions using np instead of jax, so that it won't try to access TPU. Because we want these preprocessing functions to run on CPU while TPU is doing model computation.

Tests

• Tested pre-train with reduced-layer model on v5p, with this command. training output

TODOs for future PRs:

• unit tests
• doc
• script for testing on TPU cluseters

Checklist

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

• I have performed a self-review of my code.
• 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.

[hengtaoguo]

Awesome! Thanks for the quick implementation!

[RissyRan]

LGTM! One question, have a you checked the ckpt from SFT (i.e. 50-100 steps), and do some decoding and see if it makes sense in the output?

[aireenmei]

LGTM! One question, have a you checked the ckpt from SFT (i.e. 50-100 steps), and do some decoding and see if it makes sense in the output?

We haven't checked yet, after the functionality added with this PR. We plan to run more SFT tests, ideally with a bigger dataset since llama4 models are quite big.

[gagika]

Thanks!

[RissyRan]

Discussed offline. The decoding test will be verified as follow up.