diff --git a/README.md b/README.md index 87524d732c..fb37366da3 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies ## 🔥 Latest news 🔥 +- [October 15, 2026] Qwen3.8 2.4T is now [supported](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md). - [October 8, 2026] MaxText now supports JEV-style decision scoring on TPU via vLLM; see the [Decision Models tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/decision_models.html) for a Qwen3.5-9B example with Bespoke-Nimble LoRA weights. - [August 19, 2026] Distributed low-communication training with **Streaming DiLoCo** is now supported in MaxText! Scale LLMs and MoEs across multi-slice TPU clusters and low-bandwidth DCN/WAN networks. See the [DiLoCo training tutorial](https://maxtext.readthedocs.io/en/maxtext-v0.2.5/tutorials/diloco_pretraining.html) and [DiLoCo theory reference](https://maxtext.readthedocs.io/en/maxtext-v0.2.5/reference/core_concepts/diloco.html) for details. - [August 5, 2026] DeepSeek V4 Flash (284B) is now supported. See the [user guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md) for checkpoint conversion and evaluation details. @@ -135,6 +136,7 @@ MaxText aims to provide you with the best OSS models, whether as a reference imp - Gemma 2 (2B, 9B, 27B) - Gemma 1 (2B, 7B) - Alibaba + - Qwen 3.8 MoE (2.4T) - Qwen 3.5 MoE (35B, 397B) - Qwen 3 Next (80B) - Qwen 3 MoE (30B, 235B), Qwen 3 MoE 2507 (235B, 480B) diff --git a/docs/reference/models/supported_models_and_architectures.md b/docs/reference/models/supported_models_and_architectures.md index 5213dc6dcc..8c6038377f 100644 --- a/docs/reference/models/supported_models_and_architectures.md +++ b/docs/reference/models/supported_models_and_architectures.md @@ -56,6 +56,11 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX. - **Variants** MoE with Hybrid Attention (35B-A3B, 397B-A17B) - **Notes**: GatedDeltaNet, GQA, Multimodal, MoE. +### Qwen3.8 + +- **Variants**: MoE with Hybrid Attention (2.4T-A95B) +- **Notes**: GatedDeltaNet, GQA, MoE. + ### GPT-OSS - **Variants**: 20B, 120B @@ -101,7 +106,7 @@ The following summarizes observed runtime efficiency and scaling behaviors of Ma - **Qwen**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md) | [Qwen3-Next Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) | [Qwen3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py) | [Qwen3-Next Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py) | - [Qwen3.5 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3_5.py) + [Qwen3.5 & Qwen3.8 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3_5.py) - **GPT-OSS**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/gpt_oss/run_gpt_oss.md) | [GPT-OSS Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gpt_oss.py) - **Kimi**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/kimi/Run_Kimi.md) | [K2 reuses DeepSeek Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/deepseek.py) diff --git a/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/1_test_qwen3.8_2.4t_a95b.sh b/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/1_test_qwen3.8_2.4t_a95b.sh new file mode 100644 index 0000000000..7e03681967 --- /dev/null +++ b/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/1_test_qwen3.8_2.4t_a95b.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +# This file is documentation for how to get started with Qwen3.8-2.4T. + +# This file runs Step 1 on CPU. +# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16): +# Scanned format is better for training; unscanned format is better for decoding. +# 2. Run logit check, pre-training, fine-tuning, and decoding. +# --- +# Example Usage: +# +# bash tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/1_test_qwen3.8_2.4t_a95b.sh +# --- + +set -ex + +export MODEL_NAME="${MODEL_NAME:-qwen3.8-2.4t-a95b}" +export TOKENIZER_PATH="${TOKENIZER_PATH:-Qwen/Qwen3.8-2.4T-A95B}" + +# (Optional) Path to your local Hugging Face checkpoint +export HF_MODEL_PATH="${HF_MODEL_PATH:-}" + +# Base output path for MaxText generations. +export BASE_OUTPUT_PATH="${BASE_OUTPUT_PATH:-gs://your-gcs-bucket/qwen3.8-2.4t-a95b_maxtext_ckpt}" + +if [ -z "${HF_TOKEN}" ] && [ -z "${HF_MODEL_PATH}" ]; then + echo "Error: HF_TOKEN environment variable is not set and HF_MODEL_PATH is not provided. Please export your Hugging Face token." + echo "Example: export HF_TOKEN=hf_..." + exit 1 +fi + +# Strip trailing slash from base path to avoid malformed URIs +BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/} +echo "Using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH}" + +# Install torch for checkpoint conversion and forward_pass_logit_checker.py +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +# Setup local HF path argument if one was provided +HF_LOCAL_ARG="" +if [ -n "${HF_MODEL_PATH}" ]; then + HF_LOCAL_ARG="hf_model_path=${HF_MODEL_PATH}" +fi + +# 1.1 Convert checkpoint to `scanned` format, more suitable for training +JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ + model_name=${MODEL_NAME} \ + base_output_directory=${BASE_OUTPUT_PATH}/scanned \ + hf_access_token=${HF_TOKEN} \ + scan_layers=true \ + use_multimodal=false \ + ${HF_LOCAL_ARG} + +# 1.2 Convert checkpoint to `unscanned` format, more suitable for decoding +JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ + model_name=${MODEL_NAME} \ + base_output_directory=${BASE_OUTPUT_PATH}/unscanned \ + hf_access_token=${HF_TOKEN} \ + scan_layers=false \ + use_multimodal=false \ + ${HF_LOCAL_ARG} diff --git a/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/2_test_qwen3.8_2.4t_a95b.sh b/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/2_test_qwen3.8_2.4t_a95b.sh new file mode 100644 index 0000000000..6221f72544 --- /dev/null +++ b/tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/2_test_qwen3.8_2.4t_a95b.sh @@ -0,0 +1,71 @@ +#!/bin/bash + +# This file is documentation for how to get started with Qwen3.8-2.4T. + +# This file runs Step 2 on v5p-512 on a daily basis. +# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16): +# Scanned format is better for training; unscanned format is better for decoding. +# 2. Run logit check, pretraining, finetuning, and decoding. + +# The golden logit can be generated by: +# python3 -m tests.assets.logits_generation.generate_hf_golden_logits --model-id=Qwen/Qwen3.8-2.4T-A95B --output-path=golden_data_qwen3.8-2.4t-a95b.jsonl --prompts='I love to' --hf-model-path=$local_bf16_path --trust-remote-code=False --hf-load-dtype=bfloat16 + +set -ex + +# Ensure JetStream and dependencies are installed for inference.decode +if ! python3 -c "import jetstream" &>/dev/null; then + python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf +fi + +export LIBTPU_INIT_ARGS="--xla_tpu_scoped_vmem_limit_kib=65536" +export MODEL_NAME='qwen3.8-2.4t-a95b' +export TOKENIZER_PATH='Qwen/Qwen3.8-2.4T-A95B' + +# Installing torch for checkpoint conversion and forward_pass_logit_checker.py +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +if [ -z "${HF_TOKEN}" ]; then + echo "Error: HF_TOKEN environment variable is not set. Please export your Hugging Face token." + echo "Example: export HF_TOKEN=hf_..." + exit 1 +fi + +if [ -z "${BASE_OUTPUT_PATH}" ]; then + # Non-Googlers please remember to point `BASE_OUTPUT_PATH` to GCS buckets that you own, this script uses internal buckets for testing. + # this bucket will store all the files generated by MaxText during a run + export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/$(date +%Y-%m-%d-%H-%M) + echo "BASE_OUTPUT_PATH is not set" +fi +BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/} +echo using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH} + +# Step 2: +# We define the checkpoint paths. This way it is easier to use these paths in the `train.py` and `decode.py` commands +# export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/0/items +# export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items +# Use a hard-coded golden checkpoint, rather than checkpoints generated by Step 1 as it is not in daily test. +SCANNED_CKPT_PATH=gs://maxtext-model-checkpoints/qwen3.8-2.4t-a95b/scanned/0/items +UNSCANNED_CKPT_PATH=gs://maxtext-model-checkpoints/qwen3.8-2.4t-a95b/unscanned/0/items +# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data +export DATASET_PATH=gs://maxtext-dataset + +# Test whether the forward pass logits match the golden logits +# default golden_logits_path=/deps/tests/assets/golden_logits/golden_data_{MODEL_NAME}.jsonl, copied from gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl +GOLDEN_LOGITS_DISK_LOCATION="/deps/tests/assets/golden_logits/golden_data_${MODEL_NAME}.jsonl" +if [ ! -f "${GOLDEN_LOGITS_DISK_LOCATION}" ]; then + GOLDEN_LOGITS_PATH="gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl" + GOLDEN_LOGITS_DISK_LOCATION=/tmp/golden_data.jsonl + gcloud storage cp ${GOLDEN_LOGITS_PATH} ${GOLDEN_LOGITS_DISK_LOCATION} +fi + +python3 -m tests.utils.forward_pass_logit_checker ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=forward_logits_check load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=true attention=dot_product per_device_batch_size=1 model_name=${MODEL_NAME} max_prefill_predict_length=4 max_target_length=4 async_checkpointing=false sparse_matmul=True megablox=True ici_fsdp_parallelism=64 ici_expert_parallelism=-1 weight_dtype=float32 dtype=float32 activations_in_float32=true matmul_precision=highest float32_logits=true float32_qk_product=true --golden_logits_path=${GOLDEN_LOGITS_DISK_LOCATION} --atol=1.5 --rtol=1.5 --max_kl_div=0.2 --skip_first_token + +# Run pre-training +python3 -m maxtext.trainers.pre_train.train ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=q38_pre_training model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=grain grain_file_type=tfrecord dataset_path=${DATASET_PATH} enable_checkpointing=false attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=1 steps=5 max_target_length=256 ici_fsdp_parallelism=64 ici_expert_parallelism=-1 + +# Run fine-tuning +python3 -m maxtext.trainers.pre_train.train ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=q38_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=grain grain_file_type=tfrecord dataset_path=${DATASET_PATH} enable_checkpointing=false load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=True attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=1 steps=5 max_target_length=256 ici_fsdp_parallelism=64 ici_expert_parallelism=-1 + +# Run decoding +# Note decode requires the access token for huggingface tokenizer even if the model is not gated +python3 -m maxtext.inference.decode ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=decode model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} hf_access_token=${HF_TOKEN} load_parameters_path=${UNSCANNED_CKPT_PATH} scan_layers=False attention=dot_product sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=1 max_prefill_predict_length=64 max_target_length=128 ici_fsdp_parallelism=1 ici_tensor_parallelism=1 ici_expert_parallelism=-1 prompt="An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and outputs are all vectors. The output is " diff --git a/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md b/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md index 3d95f94cfb..5672bf8d04 100644 --- a/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md +++ b/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md @@ -15,6 +15,8 @@ Qwen3 is a family of open-source large language models from the Qwen team at Ali - **Qwen3.5-35B-A3B** +- **Qwen3.8-2.4T-A95B** + For more details on Qwen3 architecture, see the [Qwen3 Technical Report](https://arxiv.org/abs/2505.09388). For more details on Qwen3-Omni architecture, see the [Qwen3-Omni Technical Report](https://arxiv.org/abs/2509.17765). @@ -32,38 +34,57 @@ To get started, you first need a MaxText-compatible checkpoint. 1. **Download the Model**: Download the official model from Hugging Face. You can use a tool like `hf_transfer` for a fast download. + ```bash + # Example for Qwen3.8-2.4T-A95B + hf download Qwen/Qwen3.8-2.4T-A95B --local-dir /path/to/qwen3.8_2_4t_hf_checkpoint + ``` + ```bash # Example for Qwen3.5-35B-A3B - hf_transfer download Qwen/Qwen3.5-35B-A3B --local-dir /path/to/qwen3.5_35b_hf_checkpoint + hf download Qwen/Qwen3.5-35B-A3B --local-dir /path/to/qwen3.5_35b_hf_checkpoint ``` ```bash # Example for Qwen3-235B-A22B - hf_transfer download Qwen/Qwen3-235B-A22B-Thinking-2507 --local-dir /path/to/qwen3_hf_checkpoint + hf download Qwen/Qwen3-235B-A22B-Thinking-2507 --local-dir /path/to/qwen3_hf_checkpoint ``` 2. **Convert the Checkpoint**: **Checkpointing Util**: Supports bidirection conversion from MaxText <-> HF. For more info look at [convert_checkpoint](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/checkpointing_solutions/convert_checkpoint.md) doc. + - Ex. Qwen3.8-2.4T-A95B + + ```bash + # Set scan_layers=false for unscanned checkpoint + JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ + model_name=qwen3.8-2.4t-a95b \ + base_output_directory=gs://your-gcs-bucket/qwen3.8_2_4t_maxtext_ckpt \ + hf_access_token=${HF_TOKEN} \ + scan_layers=true \ + use_multimodal=false + ``` + - Ex. Qwen3.5-35B ```bash + # Set scan_layers=false for unscanned checkpoint JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ model_name=qwen3.5-35b-a3b \ base_output_directory=gs://your-gcs-bucket/qwen3.5_35b_maxtext_ckpt \ hf_access_token=${HF_TOKEN} \ - scan_layers=true \ # Set to false for unscanned checkpoint + scan_layers=true \ use_multimodal=false ``` - Ex. Qwen3-MoE - + ```bash + # Set scan_layers=false for unscanned checkpoint JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ model_name= \ base_output_directory=gs://your-gcs-bucket/qwen3_maxtext_ckpt \ hf_access_token=${HF_TOKEN} \ - scan_layers=true \ # Set to false for unscanned checkpoint + scan_layers=true \ hf_model_path=${HF_MODEL_PATH} # Optional ``` @@ -207,4 +228,23 @@ export BASE_OUTPUT_PATH=gs:///qwen3.5-397b-a17b_maxtext_ckpt # 3. Execute the conversion and validation script bash tests/end_to_end/tpu/qwen/moe/qwen3.5-397b-a17b/1_test_qwen3.5_397b_a17b.sh -``` \ No newline at end of file +``` + +### Qwen3.8-2.4T-A95B + +```bash +# 1. Export your Hugging Face token +export HF_TOKEN="your_hf_token_here" + +# 2. Set the base path for conversion outputs +export BASE_OUTPUT_PATH=gs:///qwen3.8-2.4t-a95b_maxtext_ckpt + +# (Optional) Set the path to your local Hugging Face checkpoint +# export HF_MODEL_PATH=/path/to/local/qwen3.8-2.4t-a95b_hf_checkpoint + +# 3. Execute the checkpoint conversion script (CPU) +bash tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/1_test_qwen3.8_2.4t_a95b.sh + +# 4. (Optional) Execute the end-to-end TPU validation script (v5p-512) +# bash tests/end_to_end/tpu/qwen/moe/qwen3.8-2.4t-a95b/2_test_qwen3.8_2.4t_a95b.sh +``` diff --git a/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md b/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md index 0a5f542e92..4c391b8546 100644 --- a/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md +++ b/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md @@ -43,7 +43,7 @@ To get started, you first need a MaxText-compatible checkpoint. ``` # Example for Qwen3-Next-80B-A3B-Instruct - hf_transfer download Qwen/Qwen3-Next-80B-A3B-Instruct --local-dir /path/to/qwen3_next_hf_checkpoint + hf download Qwen/Qwen3-Next-80B-A3B-Instruct --local-dir /path/to/qwen3_next_hf_checkpoint ``` 2. **Convert the Checkpoint**: Run the `convert_qwen3_next_scanned.py` script to convert the downloaded Hugging Face weights into the Orbax format required by MaxText.