Merge pull request AI-Hypercomputer#4485 from AI-Hypercomputer:multi-dir-ckpt-load

PiperOrigin-RevId: 953422205

Author: Google-ML-Automation

src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml

@@ -0,0 +1,39 @@
+# model config for omni-gemma3-qwen3
+# Combines LLM backbone from Qwen 3 4B with Vision features from Gemma 3 4B
+
+# NOTE: model_name is set to "gemma3-4b" so encoders.py loads Gemma 3's Vision Tower without modification.
+# Meanwhile, decoders.py can build Qwen 3 LLM decoder directly from reading `decoder_block: "qwen3"` below.
+model_name: "gemma3-4b"
+use_multimodal: true
+
+# Multimodal config for vision model (from gemma3-4b.yml)
+image_size_for_vit: 896
+num_channels_for_vit: 3
+patch_size_for_vit: 14
+conv_stride_for_vit: 14
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_hidden_layers_for_vit: 27
+num_attention_heads_for_vit: 16
+image_placeholder: <start_of_image>
+
+# LLM backbone config (from qwen3-4b.yml)
+base_emb_dim: 2560
+base_num_query_heads: 32
+base_num_kv_heads: 8
+base_mlp_dim: 9728
+base_num_decoder_layers: 36
+head_dim: 128
+mlp_activations: ["silu", "linear"]
+vocab_size: 151936
+decoder_block: "qwen3"
+normalization_layer_epsilon: 1.0e-6
+rope_max_timescale: 1000000
+use_qk_norm: true
+logits_via_embedding: true
+normalize_embedding_logits: false
+enable_dropout: false
+tokenizer_type: "huggingface"
+
+# Ensure no pretrained weights or checkpoints are loaded initially
+load_parameters_path: ""

src/maxtext/experimental/omni_poc/prepare_checkpoint.sh

@@ -0,0 +1,115 @@
+#!/bin/bash
+# Step 1 of 5 maxtext multimodal alignment proof of concept project.
+#
+# This file:
+#   1. downloads and converts a vision-language model (e.g. gemma3-4b) from Hugging Face to MaxText format
+#   2. downloads and converts an text-only model (e.g. qwen3-4b) from Hugging Face to MaxText format
+#   3. stitches the vision component and LLM checkpoints into a single omni checkpoint
+#   4. saves the stitched checkpoints to an output directory
+#
+# Example usage:
+# HF_TOKEN="your_token" BASE_OUTPUT_DIRECTORY="gs://your_bucket/omni_checkpoints" ./prepare_checkpoint.sh
+
+set -e
+
+# Hugging Face Token & Login
+HF_TOKEN="${HF_TOKEN:-}"
+
+if [ -z "${BASE_OUTPUT_DIRECTORY}" ]; then
+  echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it as an environment variable."
+  echo "Example: export BASE_OUTPUT_DIRECTORY=\"gs://your_bucket/omni_checkpoints\""
+  exit 1
+fi
+BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY%/}"
+
+if [ -n "$HF_TOKEN" ]; then
+  if command -v hf &> /dev/null; then
+    echo "Logging into Hugging Face using hf..."
+    hf auth login --token "$HF_TOKEN"
+  elif command -v huggingface-cli &> /dev/null; then
+    echo "Logging into Hugging Face using huggingface-cli..."
+    huggingface-cli login --token "$HF_TOKEN"
+  else
+    echo "Neither hf nor huggingface-cli found. Skipping Hugging Face login."
+  fi
+fi
+
+# Configuration & Paths
+VISION_MAXTEXT_MODEL="gemma3-4b"
+VISION_HF_REPO="google/gemma-3-4b-it"
+
+LLM_MAXTEXT_MODEL="qwen3-4b"
+LLM_HF_REPO="Qwen/Qwen3-4B"
+
+# Automatically find maxtext package directory
+MAXTEXT_PKG_DIR=$(python3 -c "import os, maxtext; print(os.path.dirname(maxtext.__file__))")
+OMNI_CONFIG_PATH="${MAXTEXT_PKG_DIR}/experimental/omni_poc/omni-gemma3-qwen3.yml"
+
+VISION_CKPT_DIR="${BASE_OUTPUT_DIRECTORY}/${VISION_MAXTEXT_MODEL}_converted"
+LLM_CKPT_DIR="${BASE_OUTPUT_DIRECTORY}/${LLM_MAXTEXT_MODEL}_converted"
+STITCHED_CKPT_DIR="${BASE_OUTPUT_DIRECTORY}/omni_stitched_${VISION_MAXTEXT_MODEL}_${LLM_MAXTEXT_MODEL}"
+
+VISION_ITEMS_PATH="${VISION_CKPT_DIR}/0/items"
+LLM_ITEMS_PATH="${LLM_CKPT_DIR}/0/items"
+STITCHED_ITEMS_PATH="${STITCHED_CKPT_DIR}/0/items"
+
+echo "Base Output Directory:  ${BASE_OUTPUT_DIRECTORY}"
+echo "Vision Converted Path:  ${VISION_ITEMS_PATH}"
+echo "LLM Converted Path:     ${LLM_ITEMS_PATH}"
+echo "Stitched Target Path:   ${STITCHED_ITEMS_PATH}"
+echo ""
+
+export JAX_PLATFORMS=cpu
+
+# Helper to check if local/GCS paths exist using python etils (same as python script)
+path_exists() {
+  python3 -c "from etils import epath; import sys; sys.exit(0 if epath.Path(sys.argv[1]).exists() else 1)" "$1"
+}
+
+# Step 1: Download & Convert Vision Model from Hugging Face -> MaxText
+echo "============================================================"
+if ! path_exists "$VISION_ITEMS_PATH"; then
+  echo "Converting Vision Model (${VISION_MAXTEXT_MODEL}) from Hugging Face (${VISION_HF_REPO})..."
+  python3 -m maxtext.checkpoint_conversion.to_maxtext \
+    "${MAXTEXT_PKG_DIR}/configs/base.yml" \
+    "model_name=${VISION_MAXTEXT_MODEL}" \
+    "base_output_directory=${VISION_CKPT_DIR}" \
+    "use_multimodal=True" \
+    "scan_layers=True" \
+    "skip_jax_distributed_system=True" \
+    "--eager_load_method=transformers" \
+    "--lazy_load_tensors=False" \
+    "log_config=False"
+  echo "Vision checkpoint conversion successful!"
+  echo ""
+else
+  echo "Step 1: Vision checkpoint already exists at ${VISION_ITEMS_PATH}"
+fi
+
+# Step 2: Download & Convert Language Model from Hugging Face -> MaxText
+echo "============================================================"
+if ! path_exists "$LLM_ITEMS_PATH"; then
+  echo "Converting Language Model (${LLM_MAXTEXT_MODEL}) from Hugging Face (${LLM_HF_REPO})..."
+  python3 -m maxtext.checkpoint_conversion.to_maxtext \
+    "${MAXTEXT_PKG_DIR}/configs/base.yml" \
+    "model_name=${LLM_MAXTEXT_MODEL}" \
+    "base_output_directory=${LLM_CKPT_DIR}" \
+    "scan_layers=True" \
+    "skip_jax_distributed_system=True" \
+    "--eager_load_method=transformers" \
+    "--lazy_load_tensors=False" \
+    "log_config=False"
+  echo "LLM checkpoint conversion successful!"
+  echo ""
+else
+  echo "Step 2: LLM checkpoint already exists at ${LLM_ITEMS_PATH}"
+fi
+
+# Step 3: Checkpoint Stitching (Vision Tower + LLM Decoder + Fresh Projector)
+echo "============================================================"
+echo "Stitching Vision and LLM subtrees into unified Omni checkpoint..."
+python3 -m maxtext.experimental.omni_poc.utils.stitch_checkpoint \
+  "$OMNI_CONFIG_PATH" \
+  "vision_load_path=${VISION_ITEMS_PATH}" \
+  "llm_load_path=${LLM_ITEMS_PATH}" \
+  "stitched_output_path=${STITCHED_ITEMS_PATH}"