diff --git a/tests/assets/logits_generation/generate_hf_golden_logits.py b/tests/assets/logits_generation/generate_hf_golden_logits.py index c57d58c380..5b59214d62 100644 --- a/tests/assets/logits_generation/generate_hf_golden_logits.py +++ b/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}") diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh new file mode 100644 index 0000000000..56a8e12073 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh @@ -0,0 +1,124 @@ +#!/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. Run a forward pass check to compare the logits and KL divergence between +# the MaxText checkpoint and the Hugging Face checkpoint. +# 3. (Optional) Convert the resulting MaxText checkpoint back to Hugging Face format. + +# Pre-requisites: +# 1. Set HF_TOKEN environment variable to your Hugging Face access token. +# export HF_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= + +export MODEL_BUCKET= + +# Path to Maxtext converted to HF checkpoint +export 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 + + +# 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=${HF_MODEL} \ + --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 \ + 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 \ + --max_kl_div=0.1 \ + --run_hf_model=true \ + --hf_model_path=${HF_MODEL} \ + override_model_config=true +fi + +# Optional: Convert checkpoint back to HF format +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 + diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 3236486f74..1be6934357 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -34,7 +34,39 @@ # 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. Comparing MaxText against a pre-generated HF golden logits file +(supported for both text-only and multimodal models; recommended for large models and repeated tests): +For multimodal models, you must pre-generate the HuggingFace golden logits file first using generate_hf_golden_logits.py: + +python3 -m tests.assets.logits_generation.generate_hf_golden_logits \ + --model-id=Qwen/Qwen3-VL-2B-Instruct --hf-model-path=path/to/qwen3-vl-2b \ + --prompts="Describe this image" --image-paths=tests/assets/test_image.jpg \ + --output-path=golden_qwen3-vl-2b_vision.jsonl \ + --apply-chat-template --output-format=json + +Then run the following command to check if the MaxText checkpoint matches the golden logits: + +python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \ + tokenizer_path=path/to/qwen3-vl-2b \ + load_parameters_path=/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_qwen3-vl-2b_vision.jsonl + +Note: Pre-generating golden logits is also supported and recommended for large text-only models and repeated tests. + +2. Running the HuggingFace model on-the-fly to compare against MaxText (supported for text-only models): + +python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \ + tokenizer_path=path/to/qwen3-0.6b \ + load_parameters_path=/qwen3-0.6b/maxtext_ckpt/0/items model_name=qwen3-0.6b \ + use_multimodal=false per_device_batch_size=1 dtype=float32 \ + --max_kl_div=0.1 --run_hf_model=true --hf_model_path=path/to/qwen3-0.6b +""" import argparse import functools @@ -169,7 +201,7 @@ def check_kl_divergence(model_logits, golden_logits, atol=0.02): log_target=False, ) - max_logging.log(f"\nAverage KL divergence per token (D_KL(P_golden || Q_model)): {kl_div_value.item():.4e}") + max_logging.log(f"\nAverage KL divergence per token (D_KL(P_golden || Q_model)): {kl_div_value.item():.2e}") # To find the max KL divergence for any single token in the set # use reduction='none'. @@ -180,13 +212,13 @@ def check_kl_divergence(model_logits, golden_logits, atol=0.02): ) # Sum over the vocab dim to get a single KL value per token # Log per-token KL divergences - formatted_list = [f"{x:.4e}" for x in kl_divs_per_token.tolist()] + formatted_list = [f"{x:.2e}" for x in kl_divs_per_token.tolist()] max_logging.log(f"Per-token KL Divergences: \n{formatted_list}") max_kl_div = kl_divs_per_token.max() - max_logging.log(f"\nMax KL divergence for a single token in the set: {max_kl_div.item():.4e}") + max_logging.log(f"\nMax KL divergence for a single token in the set: {max_kl_div.item():.2e}") - assert max_kl_div < atol, f"KL divergence values {max_kl_div.item():.4e} exceed the threshold {atol}" + assert max_kl_div < atol, f"KL divergence values {max_kl_div.item():.2e} exceed the threshold {atol}" def get_data(golden_data_point, config): @@ -211,6 +243,16 @@ 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))) pixel_values = np.stack([pixel_values for _ in range(config.global_batch_size_to_train_on)]) else: pixel_values = None @@ -234,13 +276,35 @@ def get_data(golden_data_point, config): prompt = golden_data_point["formatted_prompt"] else: prompt = golden_data_point["prompt"] - max_logging.log(f' prompt="{prompt}" raw ids={original_ids}, logits.shape = {logits.shape}') + max_logging.log(f' prompt="{prompt}" raw ids (first 20)={original_ids[:20]}, logits.shape = {logits.shape}') 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 @@ -251,6 +315,30 @@ def main(config, test_args): # pylint: disable=W0621 devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) + # Load tokenizer so it is available for: + # 1. Pre-loaded golden logits comparison (multimodal input) + # 2. On-the-fly HuggingFace model comparison (text only input) + hf_token = config.hf_access_token + try: + if test_args.hf_model_path: + max_logging.log(f"Loading tokenizer from {test_args.hf_model_path}.") + tokenizer = AutoTokenizer.from_pretrained( + test_args.hf_model_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + else: + max_logging.log(f"Loading tokenizer from {config.tokenizer_path}.") + tokenizer = AutoTokenizer.from_pretrained( + config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Tokenizer loading error: {e}.\nLoading tokenizer from {config.tokenizer_path}.") + tokenizer = AutoTokenizer.from_pretrained( + config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + + if config.model_name.startswith(("llama3.1", "mixtral")): + tokenizer.pad_token = tokenizer.eos_token + if not test_args.run_hf_model: """Comparing maxtext/huggingface model with pre-loaded golden logitis""" max_logging.log("Initializing MaxText model") @@ -350,13 +438,29 @@ def main(config, test_args): # pylint: disable=W0621 max_rel_diff_val = rel_diff[max_rel_diff_idx] msg = ( "\n[numerical difference]\n" - f"Max absolute difference: {max_abs_diff_val:.4e} at index {max_abs_diff_idx}\n" - f" (Train: {train_logits_slice[max_abs_diff_idx]:.4e}, Golden: {golden_logits_slice[max_abs_diff_idx]:.4e})\n" - f"Max relative difference: {max_rel_diff_val:.4e} at index {max_rel_diff_idx}\n" - f" (Train: {train_logits_slice[max_rel_diff_idx]:.4e}, Golden: {golden_logits_slice[max_rel_diff_idx]:.4e})" + f"Max absolute difference: {max_abs_diff_val:.2e} at index {max_abs_diff_idx}\n" + f" (Train: {train_logits_slice[max_abs_diff_idx]:.2e}, Golden: {golden_logits_slice[max_abs_diff_idx]:.2e})\n" + f"Max relative difference: {max_rel_diff_val:.2e} at index {max_rel_diff_idx}\n" + f" (Train: {train_logits_slice[max_rel_diff_idx]:.2e}, Golden: {golden_logits_slice[max_rel_diff_idx]:.2e})" ) max_logging.log(msg) + # --- Compare logits for top-k token comparison table --- + mt_last_token_logits = ( + convert_jax_weight_to_torch(train_logits_slice[-1:, :]) + if isinstance(train_logits_slice, jax.Array) + else torch.tensor(train_logits_slice[-1:, :]) + ) + hf_last_token_logits = ( + convert_jax_weight_to_torch(golden_logits_slice[-1:, :]) + if isinstance(golden_logits_slice, jax.Array) + else torch.tensor(golden_logits_slice[-1:, :]) + ) + + tokens_maxtext = get_top_k_tokens_scores(mt_last_token_logits, tokenizer, k=10, description="MaxText model") + tokens_hf = get_top_k_tokens_scores(hf_last_token_logits, tokenizer, k=10, description="HF model") + compare_top_tokens(converted_tokens=tokens_maxtext, golden_tokens=tokens_hf) + if test_args.clip_logits_epsilon is not None: model_probabilities = jnp.clip(jax.nn.softmax(train_logits_slice, axis=-1), min=test_args.clip_logits_epsilon) golden_probabilities = jnp.clip(jax.nn.softmax(golden_logits_slice, axis=-1), min=test_args.clip_logits_epsilon) @@ -370,17 +474,41 @@ 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) + + # print per token kl for debugging, omit for multimodal as the seq length is usually long + per_token_kl = "" if config.use_multimodal else f"KL divergence = {kl_div}, " max_logging.log( f"\n[KL divergence]\n" - f"KL divergence = {kl_div}, max KL divergence = {max_kl_div_val} at index {max_kl_div_idx}, " + f"{per_token_kl} max KL divergence = {max_kl_div_val} at index {max_kl_div_idx}, " f"the corresponding token id is {ids[0, max_kl_div_idx + start_index]}" ) 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(), } @@ -430,7 +558,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, torch_dtype=torch_dtype, token=hf_token, trust_remote_code=test_args.trust_remote_code ) hf_lora_path = config.hf_lora_adapter_path @@ -442,25 +577,6 @@ def main(config, test_args): # pylint: disable=W0621 raise ImportError("peft library is required to load HF LoRA adapter. Run `pip install peft`.") from exc hf_model = PeftModel.from_pretrained(hf_model, hf_lora_path) - # Load tokenizer: `test_args.hf_model_path` or fallback to `config.tokenizer_path` - try: - # Try loading from `test_args.hf_model_path` - max_logging.log(f"Loading tokenizer from {test_args.hf_model_path}.") - tokenizer = AutoTokenizer.from_pretrained( - test_args.hf_model_path, token=hf_token, trust_remote_code=test_args.trust_remote_code - ) - except Exception as e: # pylint: disable=broad-except - # Fallback to `config.tokenizer_path`. local hf directory may not contain tokenizer, read from remote tokenizer - max_logging.log(f"Tokenizer loading error: {e}.\nLoading tokenizer from {config.tokenizer_path}.") - tokenizer = AutoTokenizer.from_pretrained( - config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code - ) - - # maxtext model prefix, use eos token as pad token - pad_token_prefixes = ["llama3.1", "mixtral"] - if any(config.model_name.startswith(prefix) for prefix in pad_token_prefixes): - tokenizer.pad_token = tokenizer.eos_token - quant = quantizations.configure_quantization(config) if config.pure_nnx_decoder and config.enable_nnx: maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)