Add qwen2 implementation

[ChingTsai]

Description

Changes

• Added Qwen2 implementation.
  • Implemented Qwen2 layers, these are mostly identical to Qwen3 but include the ability to apply attention bias.
  • Enabled attention bias specifically for QKV (excluding O) when using Qwen2. ref
  • Renamed Qwen3 to Qwen in hf_shape and param_mapping, and adding conversion for attention bias weights
  • Added end-to-end testing scripts.

b/471703114

Tests

Logit Verification

JAX_PLATFORMS=cpu python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml run_name=forward_pass_test_unscanned model_name=qwen2.5-7b tokenizer_path=Qwen/Qwen2.5-7B-Instruct load_parameters_path=${CHECKPOINT_PATH} max_prefill_predict_length=4 max_target_length=4 dataset_type=synthetic scan_layers=true per_device_batch_size=1 skip_jax_distributed_system=True --max_kl_div=0.017 --run_hf_model=True weight_dtype=bfloat16 --hf_model_path=Qwen/Qwen2.5-7B-Instruct

• qwen2.5-14B scanned
• qwen2.5-14B unscanned
• qwen2.5-7B scanned
• qwen2.5-7B unscanned

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 27.27273% with 72 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/layers/qwen2.py	41.81%	32 Missing ⚠️
.../MaxText/integration/tunix/weight_mapping/qwen2.py	0.00%	15 Missing ⚠️
...xText/utils/ckpt_conversion/utils/param_mapping.py	0.00%	13 Missing ⚠️
...rc/MaxText/utils/ckpt_conversion/utils/hf_shape.py	0.00%	4 Missing ⚠️
...xText/integration/tunix/weight_mapping/__init__.py	0.00%	3 Missing ⚠️
src/MaxText/layers/decoders.py	0.00%	1 Missing and 1 partial ⚠️
...xt/utils/ckpt_conversion/utils/hf_model_configs.py	0.00%	2 Missing ⚠️
src/MaxText/pyconfig.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @ChingTsai, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

📋 Review Summary

This Pull Request introduces the implementation of Qwen2 models, including new decoder layers, weight mappings, and configuration updates. The changes integrate Qwen2 into the existing MaxText framework, extending its model compatibility.

🔍 General Feedback

• The generalization of Qwen3 mappings and hook functions to a unified Qwen approach in hf_shape.py and param_mapping.py is a good practice, improving code reusability and maintainability.
• New configuration files for Qwen2.5 models are well-structured and consistent with existing model configurations.
• Ensure consistent handling of attention biases across the model definition and weight mapping to prevent potential runtime issues.

[RissyRan]

Thanks for bringing up new models! We usually verify implementation using this script against the HF version. Please let us know if you meet any issues.

[RissyRan]

cc @parambole who is working on Qwen3 for helping review PRs

[ChingTsai]

Thanks for bringing up new models! We usually verify implementation using this script against the HF version. Please let us know if you meet any issues.

Hi @RissyRan, I noticed that the 7b scanned checkpoint has a higher max KL divergence of 0.016245 (see logs). I've updated the threshold (0.015 -> 0.017) to allow this to pass, but please let me know if this level of divergence is a concern.