add Qwen3 VL tests, update logit checker, and golden logit generator

tests/assets/logits_generation/generate_hf_golden_logits.py

@@ -80,6 +80,10 @@ def save_golden_logits(
     from transformers import Llama4ForConditionalGeneration  # pylint: disable=import-outside-toplevel
 
     model_class = Llama4ForConditionalGeneration
+  elif "qwen3-vl" in model_id.lower():
+    from transformers import Qwen3VLForConditionalGeneration  # pylint: disable=import-outside-toplevel
+
+    model_class = Qwen3VLForConditionalGeneration
   else:
     from transformers import AutoModelForCausalLM  # pylint: disable=import-outside-toplevel
 
@@ -151,7 +155,10 @@ def save_golden_logits(
     for key, value in inputs.items():
       new_key = "tokens" if key == "input_ids" else key
       val_np = value.cpu().numpy()
-      data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np
+      if key == "pixel_values" and "qwen3-vl" in model_id.lower():
+        data_to_save[new_key] = val_np
+      else:
+        data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np
     data_to_save["logits"] = logits[0]
 
     print(f"Token length is {len(data_to_save['tokens'])} for prompt: {prompt_text}")

tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh

@@ -0,0 +1,129 @@
+#!/bin/bash
+
+# This script is both an end-to-end test and documentation for converting a
+# Qwen3-VL-2B MaxText checkpoint to Hugging Face format. Can be run on a v4-8.
+
+# The flow of this script is as follows:
+# 1. Convert an original Hugging Face model checkpoint to MaxText format.
+# 2. Convert the resulting MaxText checkpoint back to Hugging Face format.
+# 3. Run a forward pass check to compare the logits and KL divergence between
+#    the MaxText checkpoint and the Hugging Face checkpoint.
+
+# Pre-requisites:
+# 1. Set HF_TOKEN environment variable to your Hugging Face access token.
+#    export HF_TOKEN=<Hugging Face access token>
+# 2. Configure USE_MULTIMODAL (true for multimodal, false for text-only).
+# 3. Configure USE_SCAN_LAYERS (true if checkpoint was trained with scanned layers, false otherwise).
+
+set -ex
+
+
+MODEL_NAME='qwen3-vl-2b'
+export MODEL_VARIATION='vl_2b'
+export HF_MODEL=Qwen/Qwen3-VL-2B-Instruct
+
+idx=$(date +%Y-%m-%d-%H-%M)
+
+# Set USE_SCAN_LAYERS=true if the checkpoint was trained with scanned layers
+USE_SCAN_LAYERS=false
+if ${USE_SCAN_LAYERS}; then export CHECKPOINT_TYPE=scanned; else export CHECKPOINT_TYPE=unscanned; fi
+
+USE_MULTIMODAL=true
+
+
+export HF_TOKEN=<hf_token>
+
+export MODEL_BUCKET=<your_gcs_bucket_path>
+
+# Path to Maxtext converted to HF checkpoint
+export LOCAL_PATH=<your_local_path>/hf/${MODEL_NAME}/${idx}
+
+
+# Installing torch for deps in forward_pass_logit_checker.py
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+python3 -m pip install decord
+
+# Check point conversion
+python3 -m maxtext.checkpoint_conversion.to_maxtext \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    model_name=${MODEL_NAME} \
+    base_output_directory=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx} \
+    scan_layers=false \
+    hf_access_token=${HF_TOKEN} \
+    weight_dtype=bfloat16 \
+    hardware=cpu \
+    skip_jax_distributed_system=True \
+    checkpoint_storage_use_ocdbt=False \
+    checkpoint_storage_use_zarr3=False \
+    --eager_load_method=safetensors \
+    --lazy_load_tensors=False
+
+# Path to MaxText checkpoint
+export CKPT_PATH=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx}/0/items
+
+python3 -m maxtext.checkpoint_conversion.to_huggingface \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    model_name=${MODEL_NAME} \
+    hf_access_token=${HF_TOKEN} \
+    load_parameters_path=${CKPT_PATH} \
+    base_output_directory=${LOCAL_PATH} \
+    use_multimodal=${USE_MULTIMODAL} \
+    scan_layers=${USE_SCAN_LAYERS} \
+    override_model_config=true
+
+# Run forward pass logit checker to validate the converted checkpoint.
+if [ "${USE_MULTIMODAL}" == true ]; then
+    TEST_PROMPT='Describe this image'
+    TEST_IMAGE='tests/assets/test_image.jpg'
+    export GOLDEN_LOGITS_PATH=/tmp/golden_qwen3_vl_2b_vision.jsonl
+
+    python3 -m tests.assets.logits_generation.generate_hf_golden_logits \
+        --model-id=${HF_MODEL} \
+        --output-path=${GOLDEN_LOGITS_PATH} \
+        --prompts="${TEST_PROMPT}" \
+        --image-paths=${TEST_IMAGE} \
+        --hf-model-path=${LOCAL_PATH} \
+        --apply-chat-template \
+        --output-format=json
+
+    echo "=== Running MaxText Forward Pass Logit Checker ==="
+    python3 -m tests.utils.forward_pass_logit_checker \
+        "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+        tokenizer_path=${HF_MODEL} \
+        load_parameters_path=${CKPT_PATH} \
+        model_name=${MODEL_NAME} \
+        use_multimodal=${USE_MULTIMODAL} \
+        scan_layers=${USE_SCAN_LAYERS} \
+        dtype=float32 \
+        wi_tile_fwd_embed_dim=512 \
+        wi_tile_fwd_mlp_dim=512 \
+        wo_tile_fwd_embed_dim=512 \
+        wo_tile_fwd_mlp_dim=512 \
+        matmul_precision=highest \
+        per_device_batch_size=1 \
+        attention=dot_product \
+        prompt="${TEST_PROMPT}" \
+        image_path=${TEST_IMAGE} \
+        --max_kl_div=0.1 \
+        --golden_logits_path=${GOLDEN_LOGITS_PATH} \
+        override_model_config=true
+else
+    echo "=== Running MaxText Forward Pass Logit Checker ==="
+    python3 -m tests.utils.forward_pass_logit_checker \
+        "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+        tokenizer_path=${HF_MODEL} \
+        load_parameters_path=${CKPT_PATH} \
+        model_name=${MODEL_NAME} \
+        use_multimodal=${USE_MULTIMODAL} \
+        scan_layers=${USE_SCAN_LAYERS} \
+        per_device_batch_size=1 \
+        dtype=float32 \
+        wi_tile_fwd_embed_dim=512 \
+        wi_tile_fwd_mlp_dim=512 \
+        wo_tile_fwd_embed_dim=512 \
+        wo_tile_fwd_mlp_dim=512 \
+        --max_kl_div=0.1 \
+        --run_hf_model=true \
+        --hf_model_path=${LOCAL_PATH} \
+        override_model_config=true
+fi

tests/utils/forward_pass_logit_checker.py

@@ -34,7 +34,46 @@
 # For example:
 #   tests/assets/logits_generation/golden_llama2-7b_export.ipynb
 
-"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs"""
+"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs.
+
+Usage:
+
+1. For multimodal models (comparing MaxText against a pre-generated HF golden logits file):
+
+python3 -m tests.utils.forward_pass_logit_checker \
+    src/maxtext/configs/base.yml \
+    tokenizer_path=<your-bucket>/qwen3-vl-2b \
+    load_parameters_path=<your-bucket>/qwen3-vl-2b/maxtext_ckpt/0/items \
+    model_name=qwen3-vl-2b \
+    use_multimodal=true \
+    scan_layers=false \
+    dtype=float32 \
+    per_device_batch_size=1 \
+    prompt="Describe this image" \
+    image_path=tests/assets/test_image.jpg \
+    --max_kl_div=0.1 \
+    --golden_logits_path=golden_qwen2-vl-7b_vision.jsonl \
+    override_model_config=true
+
+Note: For multimodal models, running the HuggingFace model on-the-fly inside this script is not supported.
+You must pre-generate the HuggingFace golden logits file first using
+tests/assets/logits_generation/generate_hf_golden_logits.py.
+
+2. For text-only models (running HuggingFace model on-the-fly to compare against MaxText):
+
+python3 -m tests.utils.forward_pass_logit_checker \
+    src/maxtext/configs/base.yml \
+    tokenizer_path=<your-bucket>/gemma-2-2b \
+    load_parameters_path=<your-bucket>/gemma-2-2b/maxtext_ckpt/0/items \
+    model_name=gemma-2-2b \
+    use_multimodal=false \
+    per_device_batch_size=1 \
+    dtype=float32 \
+    --max_kl_div=0.1 \
+    --run_hf_model=true \
+    --hf_model_path=/path/to/local/hf/gemma-2-2b \
+    override_model_config=true
+"""
 
 import argparse
 import functools
@@ -211,6 +250,18 @@ def get_data(golden_data_point, config):
         pixel_values = np.transpose(pixel_values, (1, 2, 0))
     elif model_prefix in ["llama4"]:
       pixel_values = pixel_values[None, :]
+    elif model_prefix in ["qwen3"]:
+      grid_thw = np.asarray(golden_data_point["image_grid_thw"])
+      if grid_thw.ndim == 2:
+        grid_t, grid_h, grid_w = grid_thw[0]
+      else:
+        grid_t, grid_h, grid_w = grid_thw
+      tps = config.temporal_patch_size_for_vit
+      p = config.patch_size_for_vit
+      c = config.num_channels_for_vit
+      pixel_values = np.reshape(pixel_values, (c, int(grid_t * tps), int(grid_h * p), int(grid_w * p)))
+    else:
+      pixel_values = np.transpose(pixel_values, (1, 2, 0))
     pixel_values = np.stack([pixel_values for _ in range(config.global_batch_size_to_train_on)])
   else:
     pixel_values = None
@@ -238,9 +289,31 @@ def get_data(golden_data_point, config):
 
   decoder_segment_ids = np.zeros(s, dtype=np.int32)
   decoder_segment_ids[:, :seq_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR
-  decoder_positions = np.stack(
-      [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)]
-  )
+
+  # For Qwen3-VL model, compute RoPE embeddings to generate 3D position IDs
+  if model_prefix in ["qwen3"] and config.use_mrope:
+    from maxtext.multimodal import processor_qwen3_omni  # pylint: disable=import-outside-toplevel
+
+    image_grid_thw = np.atleast_2d(golden_data_point["image_grid_thw"]) if "image_grid_thw" in golden_data_point else None
+    video_grid_thw = np.atleast_2d(golden_data_point["video_grid_thw"]) if "video_grid_thw" in golden_data_point else None
+
+    attention_mask = np.zeros(s, dtype=np.int32)
+    attention_mask[:, :seq_len] = 1
+
+    position_ids, _ = processor_qwen3_omni.get_rope_index(
+        input_ids=ids,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=video_grid_thw,
+        attention_mask=attention_mask,
+        spatial_merge_size=config.spatial_merge_size_for_vit,
+        position_id_per_seconds=config.position_id_per_seconds,
+        config=config,
+    )
+    decoder_positions = position_ids.astype(np.int32)
+  else:  # For non-Qwen3-VL models, keep 1D position IDs
+    decoder_positions = np.stack(
+        [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)]
+    )
   return ids, decoder_segment_ids, decoder_positions, logits, seq_len, pixel_values
 
 
@@ -371,6 +444,27 @@ def main(config, test_args):  # pylint: disable=W0621
         max_logging.log(f"{model_probabilities[1]=}")
 
       kl_div = jax.numpy.sum(jax.scipy.special.kl_div(golden_probabilities, model_probabilities), axis=-1)
+
+      # Mask out vision placeholder tokens for KL calculation
+      ignore_token_ids = []
+      if "qwen3" in config.model_name.lower():
+        from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
+
+        qwen_tokens = QwenTokens(config)
+        ignore_token_ids = [
+            qwen_tokens.vision_start,
+            qwen_tokens.vision_end,
+            qwen_tokens.image_pad,
+            qwen_tokens.video_pad,
+        ]
+
+      if ignore_token_ids:
+        slice_ids = ids[0, start_index:token_size]
+        mask = jnp.ones_like(slice_ids, dtype=jnp.bool_)
+        for ignore_id in ignore_token_ids:
+          mask = mask & (slice_ids != ignore_id)
+        kl_div = jnp.where(mask, kl_div, 0.0)
+
       max_kl_div_val = jax.numpy.max(kl_div)
       max_kl_div_idx = jax.numpy.argmax(kl_div)
       max_logging.log(
@@ -381,7 +475,7 @@ def main(config, test_args):  # pylint: disable=W0621
 
       if jax.process_index() == 0 and test_args.output_logits_path:
         data_to_save = {
-            "prompt": golden_data[golden_data_index]["prompt"],
+            "prompt": golden_data_point["prompt"],
             "tokens": ids[0, :seq_len].tolist(),
             "logits": full_train_logits[0].tolist(),
         }
@@ -431,7 +525,14 @@ def main(config, test_args):  # pylint: disable=W0621
     torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16)
     max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})")
 
-    hf_model = AutoModelForCausalLM.from_pretrained(
+    if "qwen3-vl" in config.model_name.lower():
+      from transformers import Qwen3VLForConditionalGeneration  # pylint: disable=import-outside-toplevel
+
+      model_class = Qwen3VLForConditionalGeneration
+    else:
+      model_class = AutoModelForCausalLM
+
+    hf_model = model_class.from_pretrained(
         test_args.hf_model_path, dtype=torch_dtype, token=hf_token, trust_remote_code=test_args.trust_remote_code
     )
     hf_lora_path = config.hf_lora_adapter_path