Add end-to-end XLML test scripts and documentation for Qwen3.8-2.4T-A95B

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)

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)
 

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}

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 "

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=<qwen3-30b-a3b|qwen3-235b-a22b|qwen3-480b-a35b> \
         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://<YOUR-GCS-BUCKET>/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
-```
+```
+
+### 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://<YOUR-GCS-BUCKET>/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
+```

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.