Refactor train_sft_native.py and update multimodal documentation

Author: SurbhiJainUSC

docs/tutorials/posttraining/multimodal.md

@@ -27,46 +27,85 @@ Multimodal Large Language Models (LLMs) extend traditional text-only models by i
 ![Illustration of multimodal MaxText.](../../_static/multimodal_overview.png)
 *Figure 1: Overview of multimodal dataflow in MaxText.*
 
-## Checkpoint Conversion
+## Install MaxText
+
+For instructions on installing MaxText on your VM, please refer to the [official documentation](../../install_maxtext.md#from-source) and use the `maxtext[tpu]` installation path to include all necessary dependencies.
 
-Recently we have onboarded a new centralized tool for bidirectional checkpoint conversion between MaxText and HuggingFace ([README](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/checkpoint_conversion/README.md)).
+> **Note:** If you have previously installed MaxText with a different option, we strongly recommend using a fresh virtual environment for `maxtext[tpu]` to avoid potential library version conflicts.
 
-Install pytorch:
+## Setup environment variables
 
+Login to Hugging Face. Provide your access token when prompted:
+
+```bash
+hf auth login
 ```
-python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+Set up the following environment variables to configure your training run. Replace
+placeholders with your actual values.
+
+```bash
+# -- Model configuration --
+# The MaxText model name. See `src/maxtext/configs/types.py` for `ModelName` for a
+# full list of supported models.
+export MODEL=<MODEL_NAME> # e.g., 'gemma3-4b'
+
+# -- MaxText configuration --
+# Use a GCS bucket you own to store logs and checkpoints. Ideally in the same
+# region as your TPUs to minimize latency and costs.
+# You can list your buckets and their locations in the
+# [Cloud Console](https://console.cloud.google.com/storage/browser).
+export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # e.g., gs://my-bucket/maxtext-runs
+
+# An arbitrary string to identify this specific run.
+# We recommend to include the model, user, and timestamp.
+export RUN_NAME=<RUN_NAME>
+
+export STEPS=<STEPS> # e.g., 1000
 ```
 
-Then use this command to convert an unscanned checkpoint from HuggingFace to MaxText, and save it to `MAXTEXT_CKPT_PATH`:
+## Checkpoint Conversion
+
+This section explains how to prepare your model checkpoint for use with MaxText.
+
+### Option 1: Using an existing MaxText checkpoint
+
+If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.
+
+```sh
+export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
+```
+
+### Option 2: Converting a Hugging Face checkpoint
+
+Refer to [Hugging Face to MaxText](hf-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have correct checkpoint files converted and saved. Use this command to convert an unscanned checkpoint from HuggingFace to MaxText:
+
+```sh
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
-```shell
-# Your Hugging Face access token. Required to download gated models like Llama.
-# You can generate one at https://huggingface.co/settings/tokens.
 # We explicitly set lazy_load_tensors to False here as lazy loading of tensors
 # is not supported when use_multimodal is True.
-export HF_TOKEN=<Hugging Face access token>
-export MAXTEXT_CKPT_PATH=<Checkpoint GCS path> # gs://my-bucket/path
-python -m maxtext.checkpoint_conversion.to_maxtext \
-    model_name=gemma3-4b \
-    hf_access_token=${HF_TOKEN?} \
-    base_output_directory=${MAXTEXT_CKPT_PATH?} \
+python3 -m maxtext.checkpoint_conversion.to_maxtext \
+    model_name=${MODEL?} \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
     use_multimodal=true \
     scan_layers=false \
-    --lazy_load_tensors=False
+    --lazy_load_tensors=False \
+    --eager_load_method="transformers"
 ```
 
-For the Llama4 model family, we are using a separate checkpoint conversion script (of note, we will gradually migrate all checkpoint conversion scripts to the above consolidated tool soon):
+After conversion finishes, set `MAXTEXT_CKPT_PATH` to the converted MaxText checkpoint path.
 
-```shell
-export LOCAL_HF_MODEL_PATH=...  # Need to pre-download the safetensors from HuggingFace
-export MAXTEXT_CKPT_PATH=<Checkpoint GCS path> # gs://my-bucket/path
-python -m maxtext.checkpoint_conversion.standalone_scripts.llama4_ckpt_unscanned \
-    --model-size=llama4-17b-16e \
-    --huggingface-checkpoint=True \
-    --base-model-path=${LOCAL_HF_MODEL_PATH?} \
-    --maxtext-model-path=${MAXTEXT_CKPT_PATH?}
+```sh
+export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
 ```
 
+> [!IMPORTANT]
+> **Matching the `scan_layers` Parameter:**
+> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+>
+> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
+
 ## Multimodal Decode
 
 MaxText supports multimodal decoding, allowing you to input text with multiple images to get a text output. To use this feature, you need three main settings:
@@ -77,23 +116,18 @@ MaxText supports multimodal decoding, allowing you to input text with multiple i
 
 Since each model uses a unique native chatting template from its pretraining, we've implemented these specific templates within `multimodal_utils.py` and apply them directly to your prompt.
 
-### Decode with text+image
+### Decode with Text + Image
 
 To run a forward pass and verify the model's output, use the following command:
 
-```shell
-# Gemma3 decode
-python -m maxtext.inference.decode \
-    model_name=gemma3-4b \
-    hf_access_token=${HF_TOKEN?} \
+```sh
+python3 -m maxtext.inference.decode \
+    model_name=${MODEL?} \
     tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3 \
-    load_parameters_path=${MAXTEXT_CKPT_PATH?}/0/items \
+    load_parameters_path=${MAXTEXT_CKPT_PATH?} \
     per_device_batch_size=1 \
-    run_name=ht_test \
     max_prefill_predict_length=272 \
     max_target_length=300 \
-    steps=1 \
-    async_checkpointing=false \
     scan_layers=false \
     use_multimodal=true \
     prompt='Describe image <start_of_image>' \
@@ -109,7 +143,7 @@ Describe image <start_of_image><end_of_turn>
 <start_of_turn>model
 ` -> `Here's a description of the image:
 
-**Overall Impression:** The image is a bright, expansive cityscape view of Seattle, Washington, with`
+**Overall Impression:** The image is a bright, expansive cityscape view of Seattle, Washington,`
 ```
 
 To decode with multiple images at once, you can provide multiple image paths like this:
@@ -118,8 +152,8 @@ To decode with multiple images at once, you can provide multiple image paths lik
 export TARGET_LENGTH=...  # Adjust to fit expected output length
 export PREDICT_LENGTH=...  # Adjust to fit image tokens + text prompt
 
-python -m maxtext.inference.decode \
-    model_name=gemma3-4b \
+python3 -m maxtext.inference.decode \
+    model_name=${MODEL?} \
     ... \
     max_prefill_predict_length=${PREDICT_LENGTH?}  # Adjust to fit image tokens + text prompt \
     max_target_length=${TARGET_LENGTH?} \
@@ -130,22 +164,19 @@ python -m maxtext.inference.decode \
 
 For larger models such as Llama4-Scout/Maverick, we suggest to run the decoding on a TPU cluster such as v5p-16.
 
-### Decode with text+video+audio
+### Decode with Text + Video + Audio
 
 For models that support video input (e.g., Qwen3-Omni and Qwen3.5), pass a video file via `video_path`. For Qwen3-Omni, which also supports audio, set `use_audio_in_video=true` to additionally process the embedded audio track. Since the required token budget scales with video length and resolution, set `max_prefill_predict_length` accordingly.
 
-```shell
+```sh
 # Qwen3-Omni decode with video + audio
-export MAXTEXT_CKPT_PATH=<Checkpoint GCS path>  # gs://my-bucket/path for Qwen3-Omni
-python -m maxtext.inference.decode \
+export MAXTEXT_CKPT_PATH=<Checkpoint GCS path>  # gs://my-bucket/path/0/items for Qwen3-Omni
+python3 -m maxtext.inference.decode \
     model_name=qwen3-omni-30b-a3b \
     tokenizer_path=Qwen/Qwen3-Omni-30B-A3B-Instruct \
     tokenizer_type=huggingface \
-    load_parameters_path=${MAXTEXT_CKPT_PATH?}/0/items \
+    load_parameters_path=${MAXTEXT_CKPT_PATH?} \
     per_device_batch_size=1 \
-    run_name=ht_test \
-    steps=1 \
-    async_checkpointing=false \
     scan_layers=false \
     use_multimodal=true \
     use_audio_in_video=true \
@@ -173,15 +204,10 @@ Supervised Fine-Tuning (SFT) of multimodal LLMs in MaxText focuses specifically
 **Text+image SFT is supported for all models listed above.** The following example uses Gemma3-4B with the [ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA) dataset:
 
 ```shell
-export MAXTEXT_CKPT_PATH=<your-checkpoints-path>  # either set to an already available MaxText ckpt or to the one we just converted in the previous step
-export BASE_OUTPUT_DIRECTORY=gs://<GCS_BUCKET>
-export STEPS=1000
-python -m maxtext.trainers.post_train.sft.train_sft_native \
+python3 -m maxtext.trainers.post_train.sft.train_sft_native \
     src/maxtext/configs/post_train/sft-vision-chartqa.yml \
-    run_name="chartqa-sft" \
-    model_name=gemma3-4b \
-    tokenizer_path="google/gemma-3-4b-it" \
-    hf_access_token=${HF_TOKEN?} \
+    run_name=${RUN_NAME?} \
+    model_name=${MODEL?} \
     load_parameters_path=${MAXTEXT_CKPT_PATH?} \
     base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
     per_device_batch_size=1 \

src/maxtext/trainers/post_train/sft/train_sft_native.py

@@ -14,182 +14,27 @@
 
 "Training loop for Supervised Fine-Tuning (SFT)."
 
-import datetime
-import os
 from typing import Sequence
 
 from absl import app
 
-import numpy as np
-
-import tensorflow as tf
-import jax
-
-from flax.linen import partitioning as nn_partitioning
-
-from maxtext.configs import pyconfig
-from maxtext.trainers.pre_train.train import (
-    eval_step,
-    get_first_step,
-    train_step,
-)
-from maxtext.common import checkpointing, profiler
-from maxtext.common.data_loader import DataLoader
 from maxtext.common.goodput import (
-    GoodputEvent,
-    RECORD_JOB_END_TIME,
     RECORD_JOB_START_TIME,
-    create_goodput_recorder,
     maybe_monitor_goodput,
-    maybe_record_goodput,
     record_goodput,
 )
-from maxtext.common.metric_logger import MetricLogger
-from maxtext.utils import exceptions
-from maxtext.utils import gcs_utils
-from maxtext.utils import max_utils
-from maxtext.utils import max_logging
-from maxtext.utils import maxtext_utils
-from maxtext.utils import sharding
-from maxtext.utils import train_utils
-
-
-def train_loop(config, recorder, state=None):
-  """Main training loop for SFT."""
-  if not config.use_sft:
-    raise TypeError("Set use_sft to True to run Supervised Fine Tuning.")
-
-  (
-      init_rng,
-      checkpoint_manager,
-      state_mesh_shardings,
-      model,
-      mesh,
-      learning_rate_schedule,
-      data_iterator,
-      _,
-      _,
-      eval_data_iterator,
-      state,
-  ) = train_utils.setup_train_loop(config, recorder)
-
-  params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
-
-  p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
-      config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator, params_shardings
-  )
-
-  with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
-    data_sharding = sharding.get_input_data_sharding(config, mesh)
-    shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
-    compiled = p_train_step.lower(state, shaped_batch, init_rng).compile()
-    compiled_stats = compiled.memory_analysis()
-    max_utils.print_compiled_memory_stats(compiled_stats)
-
-  start_step = get_first_step(model, state)  # this is the start_step for training
-  prof = profiler.Profiler(config, offset_step=start_step)
-  data_loader = DataLoader(config, mesh, data_iterator, recorder)
-  metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule)
-
-  # Write train config params, num model params, and XLA flags to tensorboard
-  metric_logger.write_setup_info_to_tensorboard(state.params)
-
-  _job_completed_gracefully = False
-  try:
-    last_step_completion = datetime.datetime.now()
-    for step in np.arange(start_step, config.steps):
-      prof.maybe_activate_profiler(step, state)
-
-      with jax.profiler.StepTraceAnnotation("train", step_num=step):
-        example_batch = data_loader.load_next_batch()
-        # pylint: disable=not-callable
-        nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
-        with maybe_record_goodput(recorder, GoodputEvent.STEP, step):
-          with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
-            state, metrics = p_train_step(state, example_batch, nextrng)
-
-      step_time_delta = datetime.datetime.now() - last_step_completion
-
-      checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step)
-
-      if config.dump_hlo and step == start_step:
-        jax.block_until_ready(state)  # Ensure compilation has finished.
-        gcs_utils.upload_dump(
-            config.dump_hlo_local_dir,
-            config.dump_hlo_gcs_dir,
-            module_name=config.dump_hlo_module_name,
-            delete_local_after=config.dump_hlo_delete_local_after,
-            all_host_upload=config.dump_hlo_upload_all,
-        )
-
-      if config.eval_interval > 0 and step > start_step and (step + 1) % config.eval_interval == 0:
-        assert eval_data_iterator
-        # Explicitly reset the eval iterator and counters before starting the eval loop
-        eval_data_iterator.reset()
-        metric_logger.reset_eval_metrics()
-        max_logging.log(f"Starting eval after train step {step}")
-        eval_step_count = 0
-        last_eval_step_completion = datetime.datetime.now()
-        # pylint: disable=not-callable
-        for eval_batch in eval_data_iterator:
-          if config.eval_steps > 0 and eval_step_count >= config.eval_steps:
-            break
-          with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
-            eval_metrics = p_eval_step(state, eval_batch, nextrng)
-          eval_step_time_delta = datetime.datetime.now() - last_eval_step_completion
-          last_eval_step_completion = datetime.datetime.now()
-          metric_logger.buffer_and_write_metrics(
-              eval_metrics, eval_step_count, step_time_delta=eval_step_time_delta, is_training=False
-          )
-          max_logging.log(f"Completed eval step {eval_step_count}")
-          eval_step_count += 1
-
-      prof.maybe_deactivate_profiler(step, state)
-
-      if step == start_step:
-        max_utils.print_mem_stats("After params initialized")
-
-      last_step_completion = datetime.datetime.now()
-      metric_logger.buffer_and_write_metrics(metrics, step, step_time_delta)
-
-    if config.save_checkpoint_on_completion:
-      checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator)
-    elif checkpoint_manager is not None:
-      # in case the last checkpoint_period checkpoint is still in progress
-      checkpoint_manager.wait_until_finished()
-    _job_completed_gracefully = True
-  except exceptions.StopTraining as e:
-    prof.deactivate()
-    max_logging.log(f"Training stopped: {str(e)}")
-    _job_completed_gracefully = True
-  finally:
-    if _job_completed_gracefully:
-      record_goodput(recorder, RECORD_JOB_END_TIME)
-    metric_logger.flush_metrics_and_cleanup()
-
-  return state
+from maxtext.trainers.pre_train.train import get_train_func, initialize
 
 
 def main(argv: Sequence[str]) -> None:
-  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
-  # TF allocates extraneous GPU memory when using TFDS data
-  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
-  tf.config.set_visible_devices([], "GPU")
-  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
-  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
-    os.environ["LIBTPU_INIT_ARGS"] = (
-        os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
-    )
-  config = pyconfig.initialize(argv, use_tunix_gradient_accumulation=False)
-  jax.config.update("jax_use_shardy_partitioner", config.shardy)
-  max_utils.print_system_information()
-  train_utils.validate_train_config(config)
-  os.environ["TFDS_DATA_DIR"] = config.dataset_path
-
-  recorder = create_goodput_recorder(config)
+  argv = list(argv)
+  argv.append("use_sft=True")
+  argv.append("use_tunix_gradient_accumulation=False")
+  config, recorder = initialize(argv)
   record_goodput(recorder, RECORD_JOB_START_TIME)
+  train_func = get_train_func(config, recorder, argv)
   with maybe_monitor_goodput(config):
-    train_loop(config, recorder)
+    train_func()
 
 
 if __name__ == "__main__":