From 31951ab3616bef8fad914c9f32823aa69b9f2485 Mon Sep 17 00:00:00 2001 From: Nuojin Cheng Date: Thu, 13 Aug 2026 03:26:04 +0000 Subject: [PATCH 1/3] feat(qwen3-next): support Qwen3-Next 80B with GDNv3, fine-grained remat, MoE routing, and Muon sharding --- .dockerignore | 2 + run_qwen3_80b_aot.sh | 158 +++ run_qwen3_80b_aot_fsdp.sh | 176 ++++ run_qwen3_80b_xpk.sh | 211 ++++ run_qwen3_80b_xpk_fsdp.sh | 196 ++++ .../dockerfiles/maxtext_runner.Dockerfile | 3 + .../utils/param_mapping.py | 352 ++++--- src/maxtext/configs/base.yml | 6 + .../configs/models/qwen3-next-80b-a3b.yml | 62 +- src/maxtext/configs/types.py | 29 +- src/maxtext/kernels/megablox/backend.py | 23 +- src/maxtext/kernels/ragged/ragged_gather.py | 2 +- .../kernels/ragged/ragged_gather_reduce_v2.py | 2 +- src/maxtext/layers/attention_op.py | 5 +- src/maxtext/layers/decoders.py | 131 ++- src/maxtext/layers/moe.py | 60 +- src/maxtext/layers/nnx_decoders.py | 784 +++++++-------- src/maxtext/layers/nnx_scan.py | 19 +- src/maxtext/layers/quantizations.py | 5 +- src/maxtext/models/hybrid_gdn.py | 910 ++++++++++++++++++ src/maxtext/models/qwen3.py | 876 +++++++++++------ src/maxtext/models/qwen3_5.py | 1 + src/maxtext/optimizers/optimizers.py | 2 +- src/maxtext/trainers/pre_train/train.py | 38 +- src/maxtext/utils/maxtext_utils.py | 6 +- src/maxtext/utils/maxtext_utils_nnx.py | 142 ++- src/maxtext/utils/muon_utils.py | 17 +- src/maxtext/utils/sharding.py | 59 +- tests/unit/muon_utils_test.py | 28 +- tests/unit/nnx_decoders_test.py | 138 ++- tests/unit/param_mapping_test.py | 9 +- 31 files changed, 3500 insertions(+), 952 deletions(-) create mode 100755 run_qwen3_80b_aot.sh create mode 100755 run_qwen3_80b_aot_fsdp.sh create mode 100755 run_qwen3_80b_xpk.sh create mode 100755 run_qwen3_80b_xpk_fsdp.sh create mode 100644 src/maxtext/models/hybrid_gdn.py diff --git a/.dockerignore b/.dockerignore index e567ea2ff6..fe3d4214b5 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,2 +1,4 @@ .git maxtext_venv +.venv +venv13 diff --git a/run_qwen3_80b_aot.sh b/run_qwen3_80b_aot.sh new file mode 100755 index 0000000000..e9a77e5030 --- /dev/null +++ b/run_qwen3_80b_aot.sh @@ -0,0 +1,158 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment (~/.venv) +source /usr/local/google/home/chengnuojin/.venv/bin/activate + +# Execute in a subshell to ensure environment recovery automatically +( + # --- 1. Set XLA Flags --- + XLA_FLAGS_ARRAY=( + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_enable_concurrent_sparse_core_offloading=true" + "--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" + "--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" + "--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" + "--xla_tpu_sparse_core_all_gather_latency_multiplier=1" + "--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_use_tc_device_shape_on_sc=true" + "--xla_sc_enable_instruction_fusion=false" + "--xla_sc_disable_megacore_partitioning=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" + "--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" + "--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + "--xla_tpu_enable_all_experimental_scheduler_features=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=4" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=ENABLED" + "--xla_tpu_enable_ag_backward_pipelining=true" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" + "--xla_tpu_scheduler_percent_shared_memory_limit=150" + "--xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment=true" +) + export LIBTPU_INIT_ARGS="${XLA_FLAGS_ARRAY[*]}" + + # --- 2. Export Required Environment Variables --- + export PYTHONPATH=$PWD/src:$PWD/src/maxtext/src:$PYTHONPATH + export JAX_PLATFORMS='cpu' + export ENABLE_PJRT_COMPATIBILITY='true' + + # --- 3. Configuration --- + TIMESTAMP=$(date +%m%d%H%M%S) + export MODEL_NAME="qwen3-next-80b-a3b" + export BASE_OUTPUT_DIR="gs://chengnuojin-maxtext-logs/qwen3-next-80b-profiles/run-${TIMESTAMP}" + + # --- 4. Run Train Compile (AOT Compilation) --- + MAXTEXT_ARGS_ARRAY=( + "compile_xla_flags=${LIBTPU_INIT_ARGS}" + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "log_config=false" + "debug_sharding=false" + "ragged_gather_reduce_fallback=false" + "compile_topology=v6e-256" + "compile_topology_num_slices=1" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=True" + "ici_expert_parallelism=4" + "use_ring_of_experts=True" + "custom_mesh=hybrid_ring_64x4" + "use_ragged_sort=True" + "use_random_routing=True" + "num_moe_token_chunks=2" + "per_device_batch_size=9" + "opt_type=adamw" + "max_target_length=2048" + "ragged_buffer_factor=1.5" + "remat_policy=custom" + "reuse_example_batch=1" + "decoder_layer_input=device" + "ici_fsdp_parallelism=-1" + "steps=15" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=1024" + "sa_block_kv=1024" + "sa_block_kv_compute=512" + "sa_block_q_dkv=1024" + "sa_block_kv_dkv=1024" + "sa_block_kv_dkv_compute=1024" + "sa_fuse_reciprocal=false" + "use_splash_scheduler=true" + "sa_use_base2_exp=true" + "dq_reduction_steps=3" + "hardware=tpu" + "skip_jax_distributed_system=True" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=128" + "wi_tile_dlhs_batch_seq=128" + "wi_tile_drhs_batch_seq=128" + "wo_tile_fwd_batch_seq=128" + "wo_tile_dlhs_batch_seq=128" + "wo_tile_drhs_batch_seq=128" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=tiktoken" + "tokenizer_path=tokenizer_74B/" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "abort_on_nan_loss=False" + "profiler=xplane" + "profiler_steps=5" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" + ) + + echo "========================================================================" + echo "Running MaxText AOT train_compile for ${MODEL_NAME} on v6e-256" + echo "========================================================================" + + rm -f /tmp/libtpu_lockfile 2>/dev/null || true + python3 -m maxtext.trainers.pre_train.train_compile src/maxtext/configs/base.yml "${MAXTEXT_ARGS_ARRAY[@]}" +) diff --git a/run_qwen3_80b_aot_fsdp.sh b/run_qwen3_80b_aot_fsdp.sh new file mode 100755 index 0000000000..9a22c80fc1 --- /dev/null +++ b/run_qwen3_80b_aot_fsdp.sh @@ -0,0 +1,176 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment (~/.venv) +source /usr/local/google/home/chengnuojin/.venv/bin/activate + +# Execute in a subshell to ensure environment recovery automatically +( + # --- 1. Set XLA Flags --- + # XLA_FLAGS_ARRAY=( + # "--xla_msa_enable_sync_slice_replacement=false" + # "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + # "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + # "--xla_msa_enable_sync_copy_replacement=false" + # "--xla_tpu_scoped_vmem_limit_kib=81000" + # "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + # "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + # "--xla_tpu_offload_gather_to_sparsecore=true" + # "--xla_tpu_dvfs_p_state=7" + # "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + # "--xla_tpu_enable_async_collective_fusion=true" + # "--xla_tpu_overlap_compute_collective_tc=true" + # "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + # "--xla_tpu_enable_latency_hiding_scheduler=true" + # "--xla_latency_hiding_scheduler_rerun=10" + # "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + # "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + # "--xla_latency_hiding_scheduler_enable_selective_resources=true" + # "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + # "--xla_tpu_enable_all_experimental_scheduler_features=true" + # "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + # "--xla_tpu_host_transfer_overlap_limit=4" + # "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + # "--xla_lhs_prioritize_async_depth_over_stall=DISABLED" + # "--xla_tpu_enable_ag_backward_pipelining=true" + # "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + # "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + # "--xla_max_concurrent_host_send_recv=100" + # "--xla_tpu_scheduler_percent_shared_memory_limit=50" + # ) + XLA_FLAGS_ARRAY=( + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_enable_concurrent_sparse_core_offloading=true" + "--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" + "--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" + "--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" + "--xla_tpu_sparse_core_all_gather_latency_multiplier=1" + "--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_use_tc_device_shape_on_sc=true" + "--xla_sc_enable_instruction_fusion=false" + "--xla_sc_disable_megacore_partitioning=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" + "--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" + "--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + "--xla_tpu_enable_all_experimental_scheduler_features=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=4" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=ENABLED" + "--xla_tpu_enable_ag_backward_pipelining=true" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" + "--xla_tpu_scheduler_percent_shared_memory_limit=150" +) + export LIBTPU_INIT_ARGS="${XLA_FLAGS_ARRAY[*]}" + + # --- 2. Export Required Environment Variables --- + export PYTHONPATH=$PWD/src:$PWD/src/maxtext/src:$PYTHONPATH + export JAX_PLATFORMS='cpu' + export ENABLE_PJRT_COMPATIBILITY='true' + + # --- 3. Configuration --- + TIMESTAMP=$(date +%m%d%H%M%S) + export MODEL_NAME="qwen3-next-80b-a3b" + export BASE_OUTPUT_DIR="gs://chengnuojin-maxtext-logs/qwen3-next-80b-profiles/run-${TIMESTAMP}" + + # --- 4. Run Train Compile (AOT Compilation) --- + MAXTEXT_ARGS_ARRAY=( + "compile_xla_flags=${LIBTPU_INIT_ARGS}" + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "log_config=false" + "debug_sharding=false" + "compile_topology=v6e-256" + "compile_topology_num_slices=1" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=False" + "ici_expert_parallelism=1" + "per_device_batch_size=6" + "opt_type=adamw" + "max_target_length=2048" + "remat_policy=custom" + "reuse_example_batch=1" + "decoder_layer_input=device" + "ici_fsdp_parallelism=-1" + "steps=20" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=2048" + "sa_block_kv=2048" + "sa_block_kv_compute=1024" + "sa_block_q_dkv=2048" + "sa_block_kv_dkv=2048" + "sa_block_kv_dkv_compute=1024" + "hardware=tpu" + "skip_jax_distributed_system=True" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=64" + "wi_tile_dlhs_batch_seq=64" + "wi_tile_drhs_batch_seq=64" + "wo_tile_fwd_batch_seq=64" + "wo_tile_dlhs_batch_seq=64" + "wo_tile_drhs_batch_seq=64" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=tiktoken" + "tokenizer_path=tokenizer_74B/" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "profiler=xplane" + "profiler_steps=2" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" + ) + + echo "========================================================================" + echo "Running MaxText AOT train_compile for ${MODEL_NAME} on v6e-256" + echo "========================================================================" + + rm -f /tmp/libtpu_lockfile 2>/dev/null || true + python3 -m maxtext.trainers.pre_train.train_compile src/maxtext/configs/base.yml "${MAXTEXT_ARGS_ARRAY[@]}" +) diff --git a/run_qwen3_80b_xpk.sh b/run_qwen3_80b_xpk.sh new file mode 100755 index 0000000000..8a835caf35 --- /dev/null +++ b/run_qwen3_80b_xpk.sh @@ -0,0 +1,211 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment +source /usr/local/google/home/chengnuojin/.venv/bin/activate + +# --- Environment Variables --- +export PROJECT_ID="tpu-prod-env-one-vm" +export CLUSTER_NAME="bodaborg-v6e-256-lcscld-c" +export ZONE="southamerica-west1-a" + +# --- Configuration & Automated Image Build --- +TIMESTAMP=$(date +%m%d%H%M%S) +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:chengnuojin_${TIMESTAMP}" +export WORKLOAD_NAME="chengnuojin-qn80b-${TIMESTAMP}" +export DEVICE_TYPE="v6e-256" +export NUM_SLICES=1 +export PRIORITY="very-high" +export NUM_STEPS=15 +export MAX_RESTARTS=0 +export MODEL_NAME="qwen3-next-80b-a3b" +export BASE_OUTPUT_DIR="gs://chengnuojin-maxtext-logs/qwen3-next-80b-profiles/run-${TIMESTAMP}" + +echo "========================================================================" +echo "Building and uploading Docker runner image from /usr/local/google/home/chengnuojin/maxtext" +echo "Target Image: ${WORKLOAD_IMAGE}" +echo "========================================================================" + +( + cd /usr/local/google/home/chengnuojin/maxtext && \ + if ! docker image inspect maxtext_base_image &> /dev/null; then + echo "Local image 'maxtext_base_image' not found. Pulling gcr.io/tpu-prod-env-one-vm/param3_21jul:latest and tagging as maxtext_base_image..." + docker pull gcr.io/tpu-prod-env-one-vm/param3_21jul:latest + docker tag gcr.io/tpu-prod-env-one-vm/param3_21jul:latest maxtext_base_image + fi && \ + CLOUD_IMAGE_NAME="${WORKLOAD_IMAGE}" \ + BASE_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:latest" \ + bash src/dependencies/scripts/docker_upload_runner.sh +) + +echo "Docker image upload complete: ${WORKLOAD_IMAGE}" + +# --- XLA Flags --- +XLA_FLAGS_ARRAY=( + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_enable_concurrent_sparse_core_offloading=true" + "--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" + "--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" + "--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" + "--xla_tpu_sparse_core_all_gather_latency_multiplier=1" + "--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_use_tc_device_shape_on_sc=true" + "--xla_sc_enable_instruction_fusion=false" + "--xla_sc_disable_megacore_partitioning=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" + "--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" + "--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + "--xla_tpu_enable_all_experimental_scheduler_features=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=4" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=ENABLED" + "--xla_tpu_enable_ag_backward_pipelining=true" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" + "--xla_tpu_scheduler_percent_shared_memory_limit=150" + "--xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment=true" +) +export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}" + +# --- MaxText Workload Overrides --- +MAXTEXT_ARGS_ARRAY=( + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "log_config=false" + "debug_sharding=false" + "ragged_gather_reduce_fallback=false" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=True" + "ici_expert_parallelism=4" + "use_ring_of_experts=True" + "custom_mesh=hybrid_ring_64x4" + "use_ragged_sort=True" + "use_random_routing=True" + "num_moe_token_chunks=2" + "per_device_batch_size=9" + "opt_type=adamw" + "max_target_length=2048" + "ragged_buffer_factor=1.5" + "remat_policy=custom" + "reuse_example_batch=1" + "decoder_layer_input=device" + "ici_fsdp_parallelism=-1" + "steps=15" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=1024" + "sa_block_kv=1024" + "sa_block_kv_compute=512" + "sa_block_q_dkv=1024" + "sa_block_kv_dkv=1024" + "sa_block_kv_dkv_compute=1024" + "sa_fuse_reciprocal=false" + "use_splash_scheduler=true" + "sa_use_base2_exp=true" + "dq_reduction_steps=3" + "hardware=tpu" + "skip_jax_distributed_system=False" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=128" + "wi_tile_dlhs_batch_seq=128" + "wi_tile_drhs_batch_seq=128" + "wo_tile_fwd_batch_seq=128" + "wo_tile_dlhs_batch_seq=128" + "wo_tile_drhs_batch_seq=128" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=tiktoken" + "tokenizer_path=tokenizer_74B/" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "abort_on_nan_loss=False" + "profiler=xplane" + "profiler_steps=5" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" +) +MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}" + +# The command to run inside the container +RUN_COMMAND="set -e && \ +export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && \ +export JAX_PLATFORMS='tpu,cpu' && \ +export ENABLE_PJRT_COMPATIBILITY='true' && \ +export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && \ +export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && \ +python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS}" + +# --- XPK Workload Creation --- +echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" + +PYTHONPATH=/usr/local/google/home/chengnuojin/xpk/src python3 -m xpk.main workload create \ + --cluster="${CLUSTER_NAME}" \ + --project="${PROJECT_ID}" \ + --zone="${ZONE}" \ + --priority="${PRIORITY}" \ + --max-restarts="${MAX_RESTARTS}" \ + --device-type="${DEVICE_TYPE}" \ + --num-slices="${NUM_SLICES}" \ + --docker-image="${WORKLOAD_IMAGE}" \ + --workload="${WORKLOAD_NAME}" \ + --command="${RUN_COMMAND}" + +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" +TB_URL="https://tensorboard.corp.google.com/?logdir=${BASE_OUTPUT_DIR}/param-3/tensorboard" + +echo "========================================================================" +echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" +echo "${LOGS_URL}" +echo "" +echo "☸️ GKE Workload Details:" +echo "${GKE_URL}" +echo "" +echo "πŸ“Š GCS TensorBoard Link:" +echo "${TB_URL}" +echo "========================================================================" diff --git a/run_qwen3_80b_xpk_fsdp.sh b/run_qwen3_80b_xpk_fsdp.sh new file mode 100755 index 0000000000..9fc49351ca --- /dev/null +++ b/run_qwen3_80b_xpk_fsdp.sh @@ -0,0 +1,196 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment +source /usr/local/google/home/chengnuojin/.venv/bin/activate + +# --- Environment Variables --- +export PROJECT_ID="tpu-prod-env-one-vm" +export CLUSTER_NAME="bodaborg-v6e-256-lcscld-c" +export ZONE="southamerica-west1-a" + +# --- Configuration & Automated Image Build --- +TIMESTAMP=$(date +%m%d%H%M%S) +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:chengnuojin_${TIMESTAMP}" +export WORKLOAD_NAME="chengnuojin-qn80b-fsdp-${TIMESTAMP}" +export DEVICE_TYPE="v6e-256" +export NUM_SLICES=1 +export PRIORITY="very-high" +export MAX_RESTARTS=1 +export NUM_STEPS=20 +export MODEL_NAME="qwen3-next-80b-a3b" +export BASE_OUTPUT_DIR="gs://chengnuojin-maxtext-logs/qwen3-next-80b-profiles/run-${TIMESTAMP}" + +echo "========================================================================" +echo "Building and uploading Docker runner image from /usr/local/google/home/chengnuojin/maxtext" +echo "Target Image: ${WORKLOAD_IMAGE}" +echo "========================================================================" + +( + cd /usr/local/google/home/chengnuojin/maxtext && \ + if ! docker image inspect maxtext_base_image &> /dev/null; then + echo "Local image 'maxtext_base_image' not found. Pulling gcr.io/tpu-prod-env-one-vm/param3_21jul:latest and tagging as maxtext_base_image..." + docker pull gcr.io/tpu-prod-env-one-vm/param3_21jul:latest + docker tag gcr.io/tpu-prod-env-one-vm/param3_21jul:latest maxtext_base_image + fi && \ + CLOUD_IMAGE_NAME="${WORKLOAD_IMAGE}" \ + BASE_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:latest" \ + bash src/dependencies/scripts/docker_upload_runner.sh +) + +echo "Docker image upload complete: ${WORKLOAD_IMAGE}" + +# --- XLA Flags --- +XLA_FLAGS_ARRAY=( +"--xla_msa_enable_sync_slice_replacement=false" +"--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" +"--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" +"--xla_msa_enable_sync_copy_replacement=false" +"--xla_tpu_scoped_vmem_limit_kib=81000" +"--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" +"--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" +"--xla_tpu_enable_concurrent_sparse_core_offloading=true" +"--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" +"--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" +"--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" +"--xla_tpu_sparse_core_all_gather_latency_multiplier=1" +"--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" +"--xla_tpu_offload_gather_to_sparsecore=true" +"--xla_tpu_dvfs_p_state=7" +"--xla_tpu_disable_sparse_core_collective_offload_remover=true" +"--xla_tpu_use_tc_device_shape_on_sc=true" +"--xla_sc_enable_instruction_fusion=false" +"--xla_sc_disable_megacore_partitioning=true" +"--xla_tpu_enable_async_collective_fusion=true" +"--xla_tpu_overlap_compute_collective_tc=true" +"--xla_tpu_enable_async_collective_fusion_multiple_steps=true" +"--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" +"--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" +"--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" +"--xla_tpu_enable_latency_hiding_scheduler=true" +"--xla_latency_hiding_scheduler_rerun=10" +"--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" +"--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" +"--xla_latency_hiding_scheduler_enable_selective_resources=true" +"--xla_tpu_enable_ilp_latency_hiding_scheduler=true" +"--xla_tpu_enable_all_experimental_scheduler_features=true" +"--xla_tpu_enable_scheduler_memory_pressure_tracking=true" +"--xla_tpu_host_transfer_overlap_limit=4" +"--xla_tpu_aggressive_opt_barrier_removal=ENABLED" +"--xla_lhs_prioritize_async_depth_over_stall=ENABLED" +"--xla_tpu_enable_ag_backward_pipelining=true" +"--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" +"--xla_should_add_loop_invariant_op_in_chain=ENABLED" +"--xla_max_concurrent_host_send_recv=100" +"--xla_tpu_scheduler_percent_shared_memory_limit=150" +) +export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}" + +# --- MaxText Workload Overrides --- +MAXTEXT_ARGS_ARRAY=( + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=False" + "ici_expert_parallelism=1" + "per_device_batch_size=6" + "opt_type=adamw" + "max_target_length=2048" + "remat_policy=custom" + "reuse_example_batch=1" + "decoder_layer_input=device" + "ici_fsdp_parallelism=-1" + "steps=20" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=2048" + "sa_block_kv=2048" + "sa_block_kv_compute=1024" + "sa_block_q_dkv=2048" + "sa_block_kv_dkv=2048" + "sa_block_kv_dkv_compute=1024" + "hardware=tpu" + "skip_jax_distributed_system=False" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=64" + "wi_tile_dlhs_batch_seq=64" + "wi_tile_drhs_batch_seq=64" + "wo_tile_fwd_batch_seq=64" + "wo_tile_dlhs_batch_seq=64" + "wo_tile_drhs_batch_seq=64" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=tiktoken" + "tokenizer_path=tokenizer_74B/" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "profiler=xplane" + "profiler_steps=2" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" +) +MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}" + +# The command to run inside the container +RUN_COMMAND="set -e && \ +export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && \ +export JAX_PLATFORMS='tpu,cpu' && \ +export ENABLE_PJRT_COMPATIBILITY='true' && \ +export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && \ +export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && \ +python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS}" + +# --- XPK Workload Creation --- +echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" + +PYTHONPATH=/usr/local/google/home/chengnuojin/xpk/src python3 -m xpk.main workload create \ + --cluster="${CLUSTER_NAME}" \ + --project="${PROJECT_ID}" \ + --zone="${ZONE}" \ + --priority="${PRIORITY}" \ + --max-restarts="${MAX_RESTARTS}" \ + --device-type="${DEVICE_TYPE}" \ + --num-slices="${NUM_SLICES}" \ + --docker-image="${WORKLOAD_IMAGE}" \ + --workload="${WORKLOAD_NAME}" \ + --command="${RUN_COMMAND}" + +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" +TB_URL="https://tensorboard.corp.google.com/?logdir=${BASE_OUTPUT_DIR}/param-3/tensorboard" + +echo "========================================================================" +echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" +echo "${LOGS_URL}" +echo "" +echo "☸️ GKE Workload Details:" +echo "${GKE_URL}" +echo "" +echo "πŸ“Š GCS TensorBoard Link:" +echo "${TB_URL}" +echo "========================================================================" diff --git a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile index d85511c848..02b138931d 100644 --- a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile @@ -14,6 +14,9 @@ ENV MAXTEXT_REPO_ROOT=/deps # Set the working directory in the container WORKDIR /deps +# Install GDN v3 Tokamax commit +RUN pip install --no-deps --no-cache-dir --force-reinstall git+https://github.com/openxla/tokamax.git@b626dd8b54d708047788cf2ec538cba63a4e3739 + # Copy assets separately COPY ${PACKAGE_DIR}/maxtext/assets/ "${MAXTEXT_ASSETS_ROOT}" diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 26359cdddc..2cf8503e7f 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -1268,7 +1268,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - # pyrefly: ignore[unsupported-operation] + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = ( process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] ) @@ -1386,100 +1386,225 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F } if scan_layers: - # 2. Scan over block cycles - for block_idx in range(layer_cycle_interval): - hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval)) - prefix = f"params-decoder-layers-layer_{block_idx}" + num_blocks = num_main_layers // layer_cycle_interval + num_scanned = num_blocks * layer_cycle_interval + num_remaining = num_main_layers % layer_cycle_interval - # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.input_layernorm.weight" for i in hf_indices - ] # pyrefly: ignore[bad-assignment] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices - ] + def hf_layer(idx, suffix): + return f"model.layers.{idx}.{suffix}" - # Handle Interleaved Attention (Linear vs Full) - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 + local_prefix = "params-decoder-scanned_blocks-local_layers" + local_positions = list(range(layer_cycle_interval - 1)) - if is_full_attention_layer: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-attention-attention-query-kernel": [ - f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key-kernel": [ - f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-value-kernel": [ - f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-out-kernel": [ - f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-query_norm-scale": [ - f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key_norm-scale": [ - f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices - ], - } + # Local / linear attention layers (nested [block][local]) + mapping.update( + { + f"{local_prefix}-input_layernorm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "input_layernorm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-post_attention_layernorm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "post_attention_layernorm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-in_proj_qkvz-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.in_proj_qkvz.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-in_proj_ba-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.in_proj_ba.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-conv1d-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.conv1d.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-A_log": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.A_log") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-dt_bias": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.dt_bias") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-norm-rms_norm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.norm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-out_proj-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.out_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-routed_experts-gate-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.gate.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wi_0-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.gate_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wi_1-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.up_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wo-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.down_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert_gate-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert_gate.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-routed_experts-wi_0": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.gate_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{local_prefix}-mlp-routed_experts-wi_1": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.up_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{local_prefix}-mlp-routed_experts-wo": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.down_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + } + ) + + global_prefix = "params-decoder-scanned_blocks-global_layer" + global_position = layer_cycle_interval - 1 + + # Global attention layer (flat over blocks) + mapping.update( + { + f"{global_prefix}-input_layernorm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "input_layernorm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-post_attention_layernorm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "post_attention_layernorm.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-query-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.q_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-key-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.k_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-value-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.v_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-out-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.o_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-query_norm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.q_norm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-key_norm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.k_norm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-mlp-routed_experts-gate-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.gate.weight") for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wi_0-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.gate_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wi_1-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.up_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wo-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.down_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert_gate-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert_gate.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-routed_experts-wi_0": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.gate_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{global_prefix}-mlp-routed_experts-wi_1": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.up_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{global_prefix}-mlp-routed_experts-wo": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.down_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + } + ) + + # Remainder layers if any + if num_remaining > 0: + for rem_idx in range(num_remaining): + hf_layer_idx = num_scanned + rem_idx + prefix = f"params-decoder-layers_{hf_layer_idx}" + layer_in_block = rem_idx % layer_cycle_interval + is_full_attention_layer = (layer_in_block + 1) % layer_cycle_interval == 0 + mapping[f"{prefix}-input_layernorm-scale"] = f"model.layers.{hf_layer_idx}.input_layernorm.weight" + mapping[f"{prefix}-post_attention_layernorm-scale"] = ( + f"model.layers.{hf_layer_idx}.post_attention_layernorm.weight" ) - else: - # Linear/Hybrid Attention Block - mapping.update( # pyrefly: ignore[no-matching-overload] + if is_full_attention_layer: + mapping.update( + { + f"{prefix}-attention-attention-query-kernel": f"model.layers.{hf_layer_idx}.self_attn.q_proj.weight", + f"{prefix}-attention-attention-key-kernel": f"model.layers.{hf_layer_idx}.self_attn.k_proj.weight", + f"{prefix}-attention-attention-value-kernel": f"model.layers.{hf_layer_idx}.self_attn.v_proj.weight", + f"{prefix}-attention-attention-out-kernel": f"model.layers.{hf_layer_idx}.self_attn.o_proj.weight", + f"{prefix}-attention-attention-query_norm-scale": f"model.layers.{hf_layer_idx}.self_attn.q_norm.weight", + f"{prefix}-attention-attention-key_norm-scale": f"model.layers.{hf_layer_idx}.self_attn.k_norm.weight", + } + ) + else: + mapping.update( + { + f"{prefix}-attention-in_proj_qkvz-kernel": f"model.layers.{hf_layer_idx}.linear_attn.in_proj_qkvz.weight", + f"{prefix}-attention-in_proj_ba-kernel": f"model.layers.{hf_layer_idx}.linear_attn.in_proj_ba.weight", + f"{prefix}-attention-conv1d-kernel": f"model.layers.{hf_layer_idx}.linear_attn.conv1d.weight", + f"{prefix}-attention-A_log": f"model.layers.{hf_layer_idx}.linear_attn.A_log", + f"{prefix}-attention-dt_bias": f"model.layers.{hf_layer_idx}.linear_attn.dt_bias", + f"{prefix}-attention-norm-rms_norm-scale": f"model.layers.{hf_layer_idx}.linear_attn.norm.weight", + f"{prefix}-attention-out_proj-kernel": f"model.layers.{hf_layer_idx}.linear_attn.out_proj.weight", + } + ) + mapping.update( { - f"{prefix}-attention-in_proj_qkvz-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-gate-kernel": f"model.layers.{hf_layer_idx}.mlp.gate.weight", + f"{prefix}-mlp-shared_expert-wi_0-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.gate_proj.weight", + f"{prefix}-mlp-shared_expert-wi_1-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.up_proj.weight", + f"{prefix}-mlp-shared_expert-wo-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.down_proj.weight", + f"{prefix}-mlp-shared_expert_gate-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert_gate.weight", + f"{prefix}-mlp-routed_experts-wi_0": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) ], - f"{prefix}-attention-in_proj_ba-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-wi_1": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts) ], - f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices], - f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices], - f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices], - f"{prefix}-attention-norm-rms_norm-scale": [ - f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices - ], - f"{prefix}-attention-out_proj-kernel": [ - f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-wo": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts) ], } ) - - # 3. Handle MLP: Gates and Shared Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices], - f"{prefix}-mlp-shared_expert-wi_0-kernel": [ - f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel": [ - f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel": [ - f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert_gate-kernel": [ - f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices - ], - } - ) - - # 4. Handle MoE Routed Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-wi_0": [ - [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1": [ - [f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo": [ - [f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - } - ) else: # Unscanned layer mapping for i in range(num_main_layers): @@ -1571,20 +1696,8 @@ def permute_conv(input_tensor, target_shape=None): "params-decoder-logits_dense-kernel": transpose, } - layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval - num_main_layers = config["num_hidden_layers"] - loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers) - - for i in loop_indices: - if scan_layers: - prefix = f"params-decoder-layers-layer_{i}" - block_idx = i - else: - prefix = f"params-decoder-layers_{i}" - block_idx = i % layer_cycle_interval - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 - - if is_full_attention_layer: + def _attach_block_hooks(prefix, is_global): + if is_global: for key in ["query", "key", "value", "out"]: hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: @@ -1604,6 +1717,16 @@ def permute_conv(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose + if scan_layers: + _attach_block_hooks("params-decoder-scanned_blocks-local_layers", is_global=False) + _attach_block_hooks("params-decoder-scanned_blocks-global_layer", is_global=True) + else: + for i in range(config.base_num_decoder_layers): + prefix = f"params-decoder-layers_{i}" + block_idx = i % config.inhomogeneous_layer_cycle_interval + is_full_attention_layer = (block_idx + 1) % config.inhomogeneous_layer_cycle_interval == 0 + _attach_block_hooks(prefix, is_global=is_full_attention_layer) + return hooks @@ -1953,7 +2076,6 @@ def interleave(input_tensor, target_shape=None): hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose # `composite_mt_key`: A hook for combining multiple MaxText params. hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave # pyrefly: ignore[unsupported-operation] - # pyrefly: ignore[unsupported-operation] hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = ( interleave # pyrefly: ignore[unsupported-operation] ) @@ -2875,10 +2997,10 @@ def _spec_active(gate): local_positions = list(range(attention_pattern_length - 1)) for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, suffix) for l in local_positions] for b in range(num_blocks) ] - mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, "self_attn.v_proj.weight") for l in local_positions] for b in range(num_blocks) ] @@ -2888,11 +3010,11 @@ def _spec_active(gate): global_position = attention_pattern_length - 1 for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, suffix) for b in range(num_blocks) ] if not share_kv_projections: - mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, "self_attn.v_proj.weight") for b in range(num_blocks) ] @@ -3754,17 +3876,15 @@ def QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=Fal scan_layers=scan_layers, saving_to_hf=saving_to_hf, ) + mapping.update(text_hooks) def process_wi_0_wi_1_fused(input_tensor, target_shape=None): if saving_to_hf: wi_0, wi_1 = input_tensor - gate_up = np.concatenate([wi_0, wi_1], axis=-1) - return gate_up + return np.concatenate([wi_0, wi_1], axis=-1) else: - gate, up = np.split(input_tensor, 2, axis=-1) - return np.stack([gate, up], axis=-1) - - mapping.update(text_hooks) + wi_0, wi_1 = np.split(input_tensor, 2, axis=-1) + return np.stack([wi_0, wi_1], axis=-1) # Special case for Qwen3-VL: MaxText model has separate wi_0 and wi_1 weights # for the MoE block, but the HF model expects a single fused weight. @@ -4015,7 +4135,7 @@ def get_hf_expert_keys(expert_subpath_template): } ) - mapping.update(layer_map) # pyrefly: ignore[no-matching-overload] + mapping.update(layer_map) if not scan_layers: for i in range(n_layers): @@ -4039,7 +4159,7 @@ def transpose(input_tensor, target_shape=None): return np.transpose(input_tensor) def ones_norm(input_tensor, target_shape=None): - return np.ones(target_shape, dtype=np.float32) # pyrefly: ignore[no-matching-overload] + return np.ones(target_shape, dtype=np.float32) def identity(input_tensor, target_shape=None): return input_tensor @@ -4070,9 +4190,9 @@ def reshape_transpose_o_a(input_tensor, target_shape=None): if saving_to_hf: tensor = np.transpose(input_tensor, (0, 2, 1)) return tensor.reshape(target_shape) - num_heads = target_shape[0] # pyrefly: ignore[unsupported-operation] - embed_dim = target_shape[1] # pyrefly: ignore[unsupported-operation] - kv_lora_rank = target_shape[2] # pyrefly: ignore[unsupported-operation] + num_heads = target_shape[0] + embed_dim = target_shape[1] + kv_lora_rank = target_shape[2] tensor = input_tensor.reshape((num_heads, kv_lora_rank, embed_dim)) return np.transpose(tensor, (0, 2, 1)) @@ -4228,7 +4348,6 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, - "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4282,7 +4401,6 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 227c7e36f1..4305751533 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -288,6 +288,7 @@ use_2d_fsdp_sharding: false # deepseek moe first_num_dense_layers: 0 # number of initial dense layers in the model shared_experts: 0 +moe_shared_expert_gate: false routed_scaling_factor: 1.0 # scaling factor for routing scores routed_score_func: "" # scoring function for routing routed_bias: false # a flag if a learnable bias is added for routing @@ -305,6 +306,7 @@ batch_split_factor: 1 # the factor by which to split the batch. Only used if use # inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope] # which can only be scanned together in one large block of inhomogeneous_layer_cycle_interval=4 layers. inhomogeneous_layer_cycle_interval: 1 +full_attention_layer_offset: 0 # pipeline parallelism # The number of decoder layers is equal to the product of num_stages, num_layers_per_pipeline_stage and num_pipeline_repeats. @@ -1283,6 +1285,10 @@ gdn_num_value_heads: 32 gdn_chunk_size: 64 # Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel. use_qk_norm_in_gdn: true +# Whether to use GDN Pallas kernel +use_gdn_kernel: false +# Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward +use_hybrid_gdn: false # The ratio of dimension to apply ROPE on partial_rotary_factor: 1.0 diff --git a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml index 765977f1b5..fa9565942a 100644 --- a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml +++ b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml @@ -18,35 +18,67 @@ decoder_block: "qwen3_next" # Core Architectural Parameters -base_emb_dim: 2048 -base_num_decoder_layers: 48 -base_num_query_heads: 16 -base_num_kv_heads: 2 -head_dim: 256 -vocab_size: 151936 +base_emb_dim: 3072 +base_num_decoder_layers: 40 +base_num_query_heads: 64 +base_num_kv_heads: 8 +head_dim: 64 +vocab_size: 128008 normalization_layer_epsilon: 1.0e-6 # MoE Specific Parameters -# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models. -base_mlp_dim: 512 -base_moe_mlp_dim: 512 -num_experts: 512 +# base_mlp_dim sizes the dense-prefix layer's MLP (see first_num_dense_layers below); +# base_moe_mlp_dim sizes every other (MoE) layer's routed + shared experts. +base_mlp_dim: 10240 +base_moe_mlp_dim: 1536 +num_experts: 128 shared_experts: 1 -num_experts_per_tok: 10 +num_experts_per_tok: 8 norm_topk_prob: true +# Router parity with reference: fp32 gate-logit matmul. +float32_gate_logits: true -# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net) -inhomogeneous_layer_cycle_interval: 4 +# DeepSeek-V3-style router: sigmoid scoring, aux-loss-free expert-bias balancing +# (primary mechanism) plus a small complementary aux loss (secondary safety net, +# matching DeepSeek-V3's own combined approach), and a top-k weight scaling factor. +routed_score_func: "sigmoid" +routed_bias: true +routed_bias_update_rate: 1.0e-3 +routed_scaling_factor: 2.5 +load_balance_loss_weight: 1.0e-3 + +# Explicit (diverges from the reference's False): keeps the shared-expert gate +# that real Qwen3-Next uses. +moe_shared_expert_gate: true + +# The first layer is a dense MLP (no MoE) and always uses full attention, +# mirroring DeepSeek V3's dense-prefix pattern. +first_num_dense_layers: 1 + +# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net). +# Attention schedule (2 GDN : 1 full-attention, full-attention first in each cycle): +# [0,1,1,0,1,1,0,1,1,...] where 0=full attention, 1=Gated Delta Net. +inhomogeneous_layer_cycle_interval: 3 +full_attention_layer_offset: 0 gdn_conv_kernel_dim: 4 gdn_key_head_dim: 128 gdn_value_head_dim: 128 gdn_num_key_heads: 16 gdn_num_value_heads: 32 gdn_chunk_size: 64 +# L2-norm on Q/K inside the Gated Delta Rule is standard practice for linear-attention +# architectures (stabilizes the recurrent state, since linear attention lacks softmax's +# implicit boundedness) and matches real Qwen3-Next's own Gated Delta Net. This is also +# MaxText's own default (base.yml); kept explicit here for clarity. +use_qk_norm_in_gdn: true # RoPE Settings -rope_max_timescale: 10000000 -partial_rotary_factor: 0.25 +rope_max_timescale: 10000 +partial_rotary_factor: 1 + +# Hyper-connections: mHC enabled +mhc_expansion_rate: 4 +sinkhorn_iterations: 20 # General Model Settings enable_dropout: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 4865557b5a..8f15d360f4 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1039,6 +1039,7 @@ class DeepSeekMoE(BaseModel): first_num_dense_layers: NonNegativeInt = Field(0, description="Number of initial dense layers in the model.") shared_experts: NonNegativeInt = Field(0, description="Number of shared experts.") + moe_shared_expert_gate: bool = Field(False, description="Whether to use a gate on shared experts.") routed_scaling_factor: float = Field(1.0, description="Scaling factor for routing scores.") routed_score_func: str = Field("", description="Scoring function for routing (e.g., 'softmax', 'sigmoid').") routed_bias: bool = Field(False, description="Whether to add a bias term for routing.") @@ -1081,6 +1082,14 @@ class Qwen3Next(BaseModel): True, description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.", ) + use_gdn_kernel: bool = Field( + False, + description="Whether to use GDN Pallas kernel.", + ) + use_hybrid_gdn: bool = Field( + False, + description="Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward.", + ) partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on") @@ -1107,6 +1116,7 @@ class HardwareAndMesh(BaseModel): ) shard_mode: ShardMode = Field("auto", description="can be either auto or explicit") inhomogeneous_layer_cycle_interval: int = Field(1, description="The interval of repeated inhomogeneous layer patterns.") + full_attention_layer_offset: int = Field(0, description="Offset for full attention layer.") scan_layers: bool = Field( True, description=( @@ -1827,6 +1837,10 @@ class Muon(BaseModel): None, description="If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).", ) + muon_ns_steps: int = Field( + 5, + description="Number of Newton-Schulz iterations for Muon optimizer.", + ) class PositionalEmbedding(BaseModel): @@ -3591,12 +3605,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if ( self.routed_bias and self.routed_bias_update_rate > 0.0 - and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4) + and self.decoder_block + not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, DecoderBlockType.QWEN3_NEXT) ): - raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.") - if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4: raise ValueError( - "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. " + "Loss-free load balancing is only supported for the DeepSeek, DeepSeek4, and Qwen3-Next decoder blocks." + ) + if not self.pure_nnx and self.routed_bias and self.decoder_block in (DecoderBlockType.DEEPSEEK4, DecoderBlockType.QWEN3_NEXT): + raise ValueError( + "Auxiliary-loss-free routed bias for DeepSeek V4 and Qwen3-Next is only supported in pure NNX mode. " "Please set pure_nnx=True or disable routed_bias." ) if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: @@ -3910,6 +3927,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, DecoderBlockType.QWEN3, + DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_MOE, + DecoderBlockType.QWEN3_5, + DecoderBlockType.QWEN3_CUSTOM_MOE, DecoderBlockType.GEMMA3, DecoderBlockType.LLAMA2, ]: diff --git a/src/maxtext/kernels/megablox/backend.py b/src/maxtext/kernels/megablox/backend.py index 618965c840..5ad8b8dd11 100644 --- a/src/maxtext/kernels/megablox/backend.py +++ b/src/maxtext/kernels/megablox/backend.py @@ -30,6 +30,19 @@ import qwix.pallas as qpl +def _make_shape_dtype_struct(shape, dtype, varying_axes=()): + try: + manual_axis_type = jax.sharding.ManualAxisType(varying=frozenset(varying_axes)) + return jax.ShapeDtypeStruct(shape, dtype, manual_axis_type=manual_axis_type) + except TypeError: + try: + manual_axis_type = jax.sharding.ManualAxisType(varying=frozenset(varying_axes)) + return jax.ShapeDtypeStruct(shape, dtype, manual_type=manual_axis_type) + except (TypeError, AttributeError): + return jax.ShapeDtypeStruct(shape, dtype) + + + def _validate_args( *, lhs: jnp.ndarray, @@ -524,9 +537,7 @@ def out_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): } call_gmm = qpl.pallas_call( kernel, - out_shape=jax.ShapeDtypeStruct( - (m, n), preferred_element_type, manual_axis_type=jax.sharding.ManualAxisType(varying=frozenset(varying_axes)) - ), + out_shape=_make_shape_dtype_struct((m, n), preferred_element_type, varying_axes), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, in_specs=[ @@ -783,11 +794,7 @@ def out_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): } call_gmm = qpl.pallas_call( kernel, - out_shape=jax.ShapeDtypeStruct( - (num_actual_groups, k, n), - preferred_element_type, - manual_axis_type=jax.sharding.ManualAxisType(varying=frozenset(varying_axes)), - ), + out_shape=_make_shape_dtype_struct((num_actual_groups, k, n), preferred_element_type, varying_axes), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, in_specs=[ diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py index 0c2aabb1c1..bf2a7c78fb 100644 --- a/src/maxtext/kernels/ragged/ragged_gather.py +++ b/src/maxtext/kernels/ragged/ragged_gather.py @@ -410,7 +410,7 @@ def ragged_gather( # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. - if enforce_fallback or jax.devices()[0].platform != "tpu": + if enforce_fallback: return _fallback_implementation(x, indices, weights, has_weights) sc_info = pltpu.get_tpu_info().sparse_core diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index 0594676b74..a1cfb09d91 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -660,7 +660,7 @@ def ragged_gather_reduce( # Step 1: Choose the implementation (TensorCore fallback or SparseCore). # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. - if enforce_fallback or jax.devices()[0].platform != "tpu": + if enforce_fallback: return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size) sc_info = pltpu.get_tpu_info().sparse_core diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 25dc782ab7..92657ef171 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -75,7 +75,10 @@ from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules import numpy as np from tokamax._src.ops.attention import base as tokamax_attention_base -from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton +try: + from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton +except ImportError: + tokamax_pallas_triton = None from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask as tokamax_splash_mask # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..63f17facb3 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -877,7 +877,7 @@ def __call__( ) mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) - if cfg.mhc_expansion_rate > 1: + if cfg.mhc_expansion_rate > 1 and cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4): # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) y = mhc_expand(y) @@ -1070,6 +1070,18 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, ) + elif cfg.decoder_block == DecoderBlockType.QWEN3_NEXT: + y = self._apply_qwen3_next_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: y = self._apply_deepseek4_scanned_blocks( y, @@ -1275,7 +1287,7 @@ def __call__( assert isinstance(y, jax.Array) # After the final transformer layer, `y` holds the raw, un-normalized hidden state. - if cfg.mhc_expansion_rate > 1: + if cfg.mhc_expansion_rate > 1 and cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4): if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: hidden_state = mhc.DeepSeek4HyperHeadToLinen( config=cfg, @@ -1422,6 +1434,121 @@ def _apply_gemma3_scanned_blocks( return y + def _apply_qwen3_next_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + block_pattern_len = cfg.inhomogeneous_layer_cycle_interval + num_full_blocks = cfg.num_decoder_layers // block_pattern_len + remainder_layers = cfg.num_decoder_layers % block_pattern_len + + if num_full_blocks > 0: + ScannableBlockToLinen = qwen3.Qwen3NextScannableBlockToLinen + policy = self.get_remat_policy() + + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, num_full_blocks, block_pattern_len, stack=True + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (None, nn.broadcast), # bidirectional_mask + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + # For a fully scanned block, apply it inside an nn.scan over the calculated number of full blocks + y, returned_kv_cache = nn.scan( + ScannableBlockToLinen, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + split_rngs={"params": True, "dropout": cfg.enable_dropout}, + in_axes=in_axes_tuple, + length=num_full_blocks, + unroll=1, + metadata_params={ + nn.PARTITION_NAME: "layers", + "abstract_init": False, + }, + )( + config=cfg, + mesh=mesh, + quant=self.quant, + model_mode=model_mode, + num_of_layers=block_pattern_len, + remat_policy_fn=policy, + apply_internal_remat=True, + name="scanned_blocks", + )( + y, *broadcast_args + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True + ) + + # Process any remaining layers that don't fit into a full scanned block + for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): + layer = qwen3.Qwen3NextDecoderLayerToLinen( + config=cfg, + mesh=mesh, + model_mode=model_mode, + quant=self.quant, + layer_idx=layer_id, + ) + kv_cache = kv_caches[layer_id] if kv_caches is not None else None + + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_cache, + attention_metadata, + ) + + y_and_kv = layer(y, *remainder_args) + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + kv_caches[layer_id] = new_kv + + return y + def _apply_gemma4_scanned_blocks( self, y, diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 16a2935d85..4d29b1f938 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -46,7 +46,10 @@ from maxtext.utils.sharding import get_logical_axis_rules, remove_expert_from_partition_spec, remove_mesh_axes_from_partition_spec import numpy as np import qwix -from qwix.contrib.sparsity import sparsity_module +try: + from qwix.contrib.sparsity import sparsity_module +except ImportError: + sparsity_module = None import qwix.pallas as qpl import tokamax @@ -369,14 +372,14 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. _initializing, out_sharding=output_sharding, ) - pre_bias_logits = None - if self.score_func: output = linears._convert_to_activation_function(self.score_func)(output) - # NOTE: deepseek2 has a different pattern - if self.model_name.startswith(("deepseek3", "deepseek4")): - pre_bias_logits = output + # Snapshot pre-bias logits unconditionally (cheap, already-computed array). Only + # consumed by callers that need it (DeepSeek V3/V4, or Qwen3-Next opting into + # bias/grouped routing via RoutedMoE.uses_grouped_or_bias_routing()); harmless + # no-op for every other decoder block. + pre_bias_logits = output if self.use_bias: bias = jnp.asarray(self.bias[...], self.dtype) @@ -705,6 +708,18 @@ def should_update_load_balance(self): """ return self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and not self.is_hash_routing + def uses_grouped_or_bias_routing(self) -> bool: + """Whether expert selection uses post-bias/grouped logits, weighted by pre-bias logits. + + True for DeepSeek V3/V4 (kept as a model_name check, not decoder_block, since DeepSeek V2 + and V3 share DecoderBlockType.DEEPSEEK), and additively for Qwen3-Next once it opts in via + routed_bias (aux-loss-free bias) or n_routing_groups (expert grouping). + """ + return self.config.model_name.startswith(("deepseek3", "deepseek4")) or ( + self.config.decoder_block == ctypes.DecoderBlockType.QWEN3_NEXT + and (self.config.routed_bias or self.config.n_routing_groups != -1) + ) + def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): """get topk.""" # shape of top_k_weights & top_k_indices: @@ -727,8 +742,7 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): # Cast input_ids to int32 to safely index the hash routing table top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)] top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) - # NOTE: deepseek2 has a different pattern - elif self.config.model_name.startswith(("deepseek3", "deepseek4")): + elif self.uses_grouped_or_bias_routing(): top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits) elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) @@ -737,11 +751,19 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): else: top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): + # Qwen3-Next takes the DeepSeek-style scaling path once it opts into a non-default + # routed_score_func / routed_scaling_factor / grouped-or-bias routing; a fully-default + # qwen3_next config falls through to the plain softmax branch below, identical to + # before this option existed. + qwen3_next_needs_deepseek_style_scaling = self.config.decoder_block == ctypes.DecoderBlockType.QWEN3_NEXT and ( + bool(self.config.routed_score_func) + or self.config.routed_scaling_factor != 1.0 + or self.uses_grouped_or_bias_routing() + ) + if self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK or qwen3_next_needs_deepseek_style_scaling: top_k_weights = self.deepseek_scale_weights(top_k_weights) - else: - if self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): - top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype) + elif self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): + top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype) # Normalization of router weights (e.g. used by Qwen3, Gemma4). if self.config.norm_topk_prob: @@ -1450,9 +1472,11 @@ def get_tokamax_group_sizes(group_sizes, inputs, _kernel): elif self.config.attention in ("vllm_rpa", "vllm_batched_rpa"): return group_sizes else: + num_groups = group_sizes.shape[0] + avg_size = inputs.shape[0] // num_groups return tokamax.RaggedDotGroupSizes( group_sizes, - inputs.shape[0], + (avg_size,) * num_groups, ) def get_quantization_dtypes(): @@ -1595,8 +1619,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): + if self.uses_grouped_or_bias_routing(): pre_bias_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) else: # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 @@ -2361,8 +2384,7 @@ def sparse_matmul_route_and_compute( input_axes = (batch_logical_axis, "activation_norm_length", None) gate_logits_axes = (batch_logical_axis, "activation_norm_length", None) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): + if self.uses_grouped_or_bias_routing(): pre_bias_logits_axes = (batch_logical_axis, "activation_norm_length", None) else: pre_bias_logits_axes = None @@ -2676,9 +2698,7 @@ def dense_matmul( """Dense matrix multiplication.""" # gate_logits: batch, length, expert gate_logits = self._maybe_shard_with_logical(gate_logits, ("activation_batch_moe", "activation_length_moe", None)) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): - # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 + if self.uses_grouped_or_bias_routing(): pre_bias_logits = self._maybe_shard_with_logical( pre_bias_logits, ("activation_batch_moe", "activation_length_moe", None) ) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 0963dc6403..da32232baa 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -37,7 +37,7 @@ ShardMode, ) from maxtext.layers import initializers, linears, mhc, normalizations, quantizations -from maxtext.layers import nnx_scan, nnx_wrappers +from maxtext.layers import nnx_wrappers from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding from maxtext.layers.normalizations import RMSNorm @@ -48,7 +48,6 @@ deepseek4, deepseek_batchsplit, deepseek_batchsplit_fp8, - envy, gemma, gemma2, gemma3, @@ -213,7 +212,7 @@ def __call__( layer_output = next_layer_addition_dropped_out + inputs layer_output = _maybe_shard_with_logical(layer_output, logical_axis_names) - if getattr(cfg, "record_internal_nn_metrics", False): + if cfg.record_internal_nn_metrics: self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output)) self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output)) self.sow( @@ -357,17 +356,13 @@ def layer_fn(carry, scanned_vars): **kwargs, ) new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out - # Avoid returning and stacking read-only parameters inside the scan body. - # This prevents huge unnecessary memory allocation. - _, _, updated_state = nnx.split(layer, nnx.Param, ...) - return new_carry, updated_state + return new_carry, nnx.state(layer) final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state)) if scan_axis != 0: scanned_params, scanned_other = scanned_state.split(nnx.Param, ...) - if scanned_params: - scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params) + scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params) scanned_state = nnx.State.merge(scanned_params, scanned_other) nnx.update(self.scanned_layers, scanned_state) @@ -433,17 +428,12 @@ def __init__( self.scanned_layers = None self.is_deepseek = self.config.decoder_block == DecoderBlockType.DEEPSEEK - self.is_deepseek4 = self.config.decoder_block == DecoderBlockType.DEEPSEEK4 self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL - - if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: - self.hc_head = mhc.DeepSeek4HyperHead( - config=config, - mesh=self.mesh, - rngs=self.rngs, - ) + self.is_qwen3_next_with_dense = ( + self.config.decoder_block == DecoderBlockType.QWEN3_NEXT and self.config.first_num_dense_layers > 0 + ) self._init_decoder_layers(decoder_block_classes, rngs, mesh) @@ -454,12 +444,12 @@ def _init_decoder_layers(self, decoder_block_classes, rngs, mesh): if self.is_gemma4_small: # Gemma4 E2B/E4B: per-layer-index KV-share donor threading and a distinct attention_type # per layer are not expressible inside nn.scan; pipeline parallelism is also unsupported. - if getattr(config, "using_pipeline_parallelism", False) or getattr(config, "scan_layers", False): + if config.using_pipeline_parallelism or config.scan_layers: raise ValueError("gemma4_small (Gemma4 E2B/E4B) does not support pipeline parallelism or scan_layers.") self._init_gemma4_small_layers(rngs) - elif getattr(config, "using_pipeline_parallelism", False): + elif config.using_pipeline_parallelism: self._init_pipeline_layers(decoder_block_classes, rngs, mesh) - elif getattr(config, "scan_layers", False): + elif config.scan_layers: self._init_scanned_layers(decoder_block_classes, rngs, mesh) else: self._init_sequential_layers(decoder_block_classes, rngs) @@ -511,13 +501,13 @@ def _init_pipeline_deepseek(self, decoder_block_classes, rngs): else: self.num_dense_layers = config.first_num_dense_layers for i in range(self.num_dense_layers): - self._create_and_register_layer(dense_cls, rngs, "dense_layers", i) + self._create_and_register_named_layer(dense_cls, rngs, "dense_layers", i) self.num_moe_outside_pipeline = ( config.num_decoder_layers - config.first_num_dense_layers ) - config.pipeline_parallel_layers if self.num_moe_outside_pipeline > 0: for i in range(self.num_moe_outside_pipeline): - self._create_and_register_layer(moe_cls, rngs, "moe_layers_outside_pipeline", i) + self._create_and_register_named_layer(moe_cls, rngs, "moe_layers_outside_pipeline", i) def _init_pipeline_generic(self, decoder_block_classes, rngs): """Initializes generic decoder layers outside pipeline.""" @@ -535,44 +525,21 @@ def _init_pipeline_generic(self, decoder_block_classes, rngs): else: self.num_layers_outside_pipeline = remaining_layers for i in range(self.num_layers_outside_pipeline): - self._create_and_register_layer(base_cls, rngs, "layers_outside_pipeline", i) + self._create_and_register_named_layer(base_cls, rngs, "layers_outside_pipeline", i) def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): """Initializes decoder layers with scanning (non-pipeline).""" if self.is_deepseek: self._init_scanned_deepseek(decoder_block_classes, rngs) - elif self.is_deepseek4: - self._init_scanned_deepseek4(rngs) elif self.is_gemma3: self._init_scanned_gemma3(decoder_block_classes, rngs, mesh) elif self.is_gemma4: self._init_scanned_gemma4(decoder_block_classes, rngs, mesh) + elif self.is_qwen3_next_with_dense: + self._init_scanned_qwen3_next_with_dense(rngs) else: self._init_scanned_generic(decoder_block_classes, rngs) - def _init_scanned_deepseek4(self, rngs): - """Initializes DeepSeek V4 scanned layers: unrolls prefix hash layers and scans remaining full blocks.""" - config = self.config - num_hash_layers = config.first_num_hash_layers - for layer_idx in range(num_hash_layers): - self._create_and_register_layer( - deepseek4.DeepSeek4DecoderLayer, - rngs, - "layers", - layer_idx, - layer_idx=layer_idx, - ) - - num_remaining_layers = config.num_decoder_layers - num_hash_layers - num_full_blocks = num_remaining_layers // 2 - if num_full_blocks > 0: - self.scanned_blocks = self._create_scanned_layers( - deepseek4.DeepSeek4ScannableBlock, - length=num_full_blocks, - metadata_axis_name="scanned_blocks", - rngs=rngs, - ) - def _init_scanned_deepseek(self, decoder_block_classes, rngs): """Initializes scanned DeepSeek layers with optional Engram support.""" config = self.config @@ -649,6 +616,70 @@ def _init_scanned_deepseek_standard(self, dense_cls, moe_cls, rngs): num_moe = config.num_decoder_layers - config.first_num_dense_layers self.moe_layers = self._create_scanned_layers(moe_cls, length=num_moe, metadata_axis_name="moe_layers", rngs=rngs) + def _init_scanned_qwen3_next_with_dense(self, rngs): + """Initializes scanned Qwen3-Next layers with a `first_num_dense_layers` dense prefix. + + Splits the stack into three pieces: an unscanned dense prefix (real global + `layer_idx`, forced full attention), a scanned middle of uniform + `Qwen3NextScannableBlock` repeats, and β€” since `num_decoder_layers - + first_num_dense_layers` need not be a multiple of `inhomogeneous_layer_cycle_interval` + β€” an unscanned remainder tail, mirroring `_init_scanned_gemma3`'s scan+remainder + split. `layer_idx_offset` keeps the attention-type cycle globally correct across + all three pieces despite the dense prefix breaking its natural period. + """ + config = self.config + n_dense = config.first_num_dense_layers + cycle = config.inhomogeneous_layer_cycle_interval + remaining = config.num_decoder_layers - n_dense + scan_length = remaining // cycle + num_remaining = remaining % cycle + + self.dense_layers = self._create_scanned_layers( + qwen3.Qwen3NextDecoderLayer, + length=n_dense, + metadata_axis_name="dense_layers", + rngs=rngs, + layer_idx=0, + is_dense_layer=True, + ) + + policy = self.get_remat_policy() + layer_kwargs = { + "num_of_layers": cycle, + "layer_idx_offset": n_dense, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + rem_layer_kwargs = { + "num_of_layers": num_remaining, + "layer_idx_offset": n_dense + scan_length * cycle, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + + if scan_length > 0: + self.layers = self._create_scanned_layers( + qwen3.Qwen3NextScannableBlock, + length=scan_length, + metadata_axis_name="layers", + rngs=rngs, + **layer_kwargs, + ) + else: + self.layers = nnx.List([]) + + if num_remaining > 0: + self.layers_remainder = qwen3.Qwen3NextScannableBlock( + config=config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + rngs=rngs, + **rem_layer_kwargs, + ) + else: + self.layers_remainder = None + def _init_scanned_gemma3(self, decoder_block_classes, rngs, mesh): """Initializes scanned Gemma3 layers.""" config = self.config @@ -684,24 +715,14 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN) scan_length = config.num_decoder_layers // attention_pattern_length num_remaining_layers = config.num_decoder_layers % attention_pattern_length - policy = self.get_remat_policy() - # The pure-NNX decoder skips block-level remat (skip_block_remat=True below), - # so the block rematerializes its own local/global layers instead. - layer_kwargs = { - "num_of_layers": attention_pattern_length, - "remat_policy_fn": policy, - "apply_internal_remat": True, - } - rem_layer_kwargs = { - "num_of_layers": num_remaining_layers, - "remat_policy_fn": policy, - "apply_internal_remat": True, - } + layer_kwargs = {"num_of_layers": attention_pattern_length} + + rem_layer_kwargs = {"num_of_layers": num_remaining_layers} RemattedGemma4Block = gemma4.Gemma4ScannableBlock if scan_length > 0: - self.scanned_blocks = self._create_scanned_layers( + self.layers = self._create_scanned_layers( RemattedGemma4Block, length=scan_length, metadata_axis_name="layers", @@ -728,10 +749,6 @@ def _init_scanned_generic(self, decoder_block_classes, rngs): "nope_layer_interval": self.config.nope_layer_interval, "interleave_moe_layer_step": self.config.interleave_moe_layer_step, } - if config.decoder_block == DecoderBlockType.ENVY: - layer_kwargs = { - "interleave_moe_layer_step": self.config.interleave_moe_layer_step, - } if num_layers > 0: self.layers = self._create_scanned_layers( @@ -741,9 +758,12 @@ def _init_scanned_generic(self, decoder_block_classes, rngs): rngs=rngs, **layer_kwargs, ) + else: + self.layers = nnx.List([]) def _init_sequential_layers(self, decoder_block_classes, rngs): """Initializes decoder layers sequentially (no scanning).""" + self.layers = nnx.List([]) if self.is_deepseek: self._init_sequential_deepseek(decoder_block_classes, rngs) @@ -755,9 +775,9 @@ def _init_sequential_deepseek(self, decoder_block_classes, rngs): config = self.config dense_cls, moe_cls = decoder_block_classes for i in range(config.first_num_dense_layers): - self._create_and_register_layer(dense_cls, rngs, "dense_layers", i) + self._create_and_register_layer(dense_cls, rngs, "dense_layer", i) for i in range(config.num_decoder_layers - config.first_num_dense_layers): - self._create_and_register_layer(moe_cls, rngs, "moe_layers", i) + self._create_and_register_layer(moe_cls, rngs, "moe_layer", i) def _init_sequential_generic(self, decoder_block_classes, rngs): """Initializes sequential generic decoder layers with per-architecture layer_kwargs.""" @@ -775,15 +795,9 @@ def _init_sequential_generic(self, decoder_block_classes, rngs): "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval), "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step), } - elif config.decoder_block == DecoderBlockType.ENVY: - layer_kwargs = { - "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0, - } - elif config.decoder_block in { DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, - DecoderBlockType.DEEPSEEK4, }: layer_kwargs = {"layer_idx": lyr} elif config.decoder_block == DecoderBlockType.GPT_OSS: @@ -804,6 +818,7 @@ def _init_gemma4_small_layers(self, rngs): ``_create_and_register_layer``. """ cfg = self.config + self.layers = nnx.List([]) # Only register the PLE submodule when it exists (mirrors the optional position_embedder # pattern); assigning None first would make nnx treat the attribute as static. if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0: @@ -821,6 +836,7 @@ def _init_gemma4_small_layers(self, rngs): rngs=rngs, ) setattr(self, f"layers_{lyr}", layer) + self.layers.append(layer) def _get_pipeline_stage_module(self, decoder_blocks, rngs): """Retrieves the wrapper module formatted for single pipeline stage execution.""" @@ -850,7 +866,14 @@ def _get_pipeline_stage_module(self, decoder_blocks, rngs): ) def _create_and_register_layer(self, layer_cls, rngs, base_name, i, **layer_kwargs): - """Creates a layer registered ONLY via named attribute.""" + attr_name = f"{base_name}_{i}" + layer = self._create_single_layer(layer_cls, rngs, **layer_kwargs) + setattr(self, attr_name, layer) + self.layers.append(layer) + + def _create_and_register_named_layer(self, layer_cls, rngs, base_name, i, **layer_kwargs): + """Creates a layer registered ONLY via named attribute. Used by pipeline-outside paths + to avoid double-registration when self.layers list is also tracked elsewhere.""" attr_name = f"{base_name}_{i}" layer = self._create_single_layer(layer_cls, rngs, **layer_kwargs) setattr(self, attr_name, layer) @@ -885,20 +908,94 @@ def _create_scanned_layers( **layer_kwargs, ): """Creates a scanned stack of layers using jax.lax.scan for memory-efficient initialization.""" - return nnx_scan.create_scanned_layers( - lambda layer_rngs: decoder_layer_class( - config=self.config, - mesh=self.mesh, - model_mode=self.model_mode, - quant=self.quant, - rngs=layer_rngs, - **layer_kwargs, - ), - length=length, - param_scan_axis=self.config.param_scan_axis, - metadata_axis_name=metadata_axis_name, - rngs=rngs, + if length == 0: + return None + scan_axis = self.config.param_scan_axis + + # Fork rngs to get per-layer RNG states for scanning + try: + forked_rngs = rngs.fork(split=length) + except: # pylint: disable=bare-except + pass + + rngs_graphdef, rngs_state = nnx.split(forked_rngs) + + first_rng_state = jax.tree.map(lambda x: x[0], rngs_state) + ref_rngs = nnx.merge(rngs_graphdef, first_rng_state) + ref_layer = decoder_layer_class( + config=self.config, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + rngs=ref_rngs, + **layer_kwargs, ) + layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...) + del ref_layer + + def scan_body(carry, rng_state_slice): + layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice) + layer = decoder_layer_class( + config=self.config, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + rngs=layer_rngs, + **layer_kwargs, + ) + _, params, rest = nnx.split(layer, nnx.Param, ...) + return carry, (params, rest) + + _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state) + + if scan_axis != 0: + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), stacked_params) + + def _add_scan_metadata(state, axis): + def _update_leaf(leaf): + if hasattr(leaf, "replace") and hasattr(leaf, "value"): + replace_kwargs = {} + if hasattr(leaf, "get_metadata"): + replace_kwargs.update(leaf.get_metadata()) + + replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name + replace_kwargs["param_scan_axis"] = axis + + for key in [ + "sharding", + "out_sharding", + "kernel_axes", + "sharding_names", + ]: + val = getattr(leaf, key, None) + if val is None and key in replace_kwargs: + val = replace_kwargs[key] + + if val is not None: + if isinstance(val, str): + val = (val,) + if isinstance(val, tuple): + l = list(val) + # Safely insert the scan axis into the logical axes string + if metadata_axis_name not in l: + insert_idx = min(axis, len(l)) + l.insert(insert_idx, metadata_axis_name) + replace_kwargs[key] = tuple(l) + + return leaf.replace(**replace_kwargs) + return leaf + + # We must use a custom is_leaf to catch the VariableState instances + return jax.tree.map( + _update_leaf, + state, + is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"), + ) + + stacked_params = _add_scan_metadata(stacked_params, scan_axis) + stacked_rest = _add_scan_metadata(stacked_rest, 0) + + return nnx.merge(layer_graphdef, stacked_params, stacked_rest) def _apply_layer_with_remat(self, layer: nnx.Module, y: jax.Array, policy: Any, prevent_cse: bool, **kwargs): """Helper to cleanly apply jax.checkpoint to a single unscanned layer or block.""" @@ -930,11 +1027,6 @@ def _apply_layers_sequentially( ): """Runs the layer stack using nnx.scan. - This is the fuller of the two NNX scan appliers: it also threads external - (vLLM) KV caches via a static unroll and re-applies scan-axis metadata. - ``nnx_scan.apply_scanned_layers`` is a leaner, model-agnostic alternative - (currently used only by Gemma4); unifying the two is a follow-up cleanup. - Args: layers: The stacked NNX module whose params are scanned over. x_in: The carry (hidden state) fed into the first layer. @@ -990,7 +1082,7 @@ def _extract_matching_state(template, full): def layer_fn(carry, scanned_vars): # Ensure metadata rank matches the sliced values - scanned_vars = maxtext_utils_nnx.nnx_remove_scan_axis(scanned_vars, "layers") + scanned_vars = maxtext_utils_nnx.nnx_remove_scan_axis(scanned_vars, metadata_axis_name) # Unpack the sliced variables for THIS layer if use_kv: @@ -1038,8 +1130,6 @@ def layer_fn(carry, scanned_vars): return new_carry, new_current_state if skip_block_remat: - # The scanned module applies its own remat internally; wrapping the whole - # body again would double-remat and recompute the entire block. layer_fn_wrapped = layer_fn else: layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse) @@ -1113,23 +1203,38 @@ def get_deepseek(): DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer], DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer], DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer], - DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock), + DecoderBlockType.GEMMA4: get_scannable( + gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock + ), DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer], DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer], DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer], DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer], DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer], - DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer], + DecoderBlockType.QWEN3_CUSTOM_MOE: [ + qwen3_custom.Qwen3CustomMoeDecoderLayer + ], DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: get_deepseek(), - DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), - DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), - DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), - DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), - DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), - DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), - DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), + DecoderBlockType.DEEPSEEK4: get_scannable( + deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock + ), + DecoderBlockType.GPT_OSS: get_scannable( + gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock + ), + DecoderBlockType.QWEN3_NEXT: get_scannable( + qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock + ), + DecoderBlockType.QWEN3_5: get_scannable( + qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock + ), + DecoderBlockType.LLAMA4: get_scannable( + llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock + ), + DecoderBlockType.OLMO3: get_scannable( + olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock + ), } if cfg.decoder_block not in layer_map: @@ -1143,7 +1248,6 @@ def minimal_policy(self, with_context=False, with_quantization=False): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", "out_proj", "mlpwi_0", @@ -1191,7 +1295,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", "context", "out_proj", @@ -1201,7 +1304,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", "out_proj", "mlpwo", @@ -1211,7 +1313,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", "out_proj", ) @@ -1220,7 +1321,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", ) elif cfg.remat_policy == "qkv_proj_offloaded": @@ -1230,7 +1330,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", ], offload_src="device", offload_dst="pinned_host", @@ -1242,7 +1341,6 @@ def get_remat_policy(self): "query_proj", "value_proj", "key_proj", - "kv_proj", "qkv_proj", "out_proj", "mlpwi_0", @@ -1275,7 +1373,6 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.MISTRAL, DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, - DecoderBlockType.DEEPSEEK4, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, @@ -1290,7 +1387,6 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.SIMPLE_MLP, DecoderBlockType.LLAMA4, DecoderBlockType.OLMO3, - DecoderBlockType.ENVY, }: return functools.partial( RMSNorm, @@ -1356,9 +1452,6 @@ def _apply_embedding( "llama4-17b-16e", "llama4-17b-128e", "qwen3-omni-30b-a3b", - "qwen3-vl-2b", - "qwen3-vl-4b", - "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", }: @@ -1372,14 +1465,7 @@ def _apply_embedding( raise ValueError(f"Unsupported model_name for multimodal: {cfg.model_name}") if video_embeddings is not None and cfg.use_multimodal: - if cfg.model_name in { - "qwen3-omni-30b-a3b", - "qwen3-vl-2b", - "qwen3-vl-4b", - "qwen3-vl-30b-a3b", - "qwen3.5-35b-a3b", - "qwen3.5-397b-a17b", - }: + if cfg.model_name in {"qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"}: y = mm_utils.merge_mm_embeddings( text_embeddings=y, multimodal_embeddings=video_embeddings, @@ -1492,11 +1578,6 @@ def _apply_single_engram_layer(self, y, layer_name, *args, **kwargs): decoder_input_tokens = kwargs.get("decoder_input_tokens") layer_kwargs = kwargs.get("layer_kwargs", {}) - # Create a copy of layer_kwargs and pop decoder_input_tokens if it exists - # to avoid passing it twice (once explicitly and once via **layer_kwargs). - layer_kwargs = dict(layer_kwargs) - layer_kwargs.pop("decoder_input_tokens", None) - out = layer(y, *args, decoder_input_tokens=decoder_input_tokens, **layer_kwargs) if isinstance(out, tuple): y = out[0] @@ -1606,27 +1687,16 @@ def __call__( multimodal_input=multimodal_input, ) - mhc_reduce = None - if hasattr(cfg, "mhc_expansion_rate"): - mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) - if cfg.mhc_expansion_rate > 1: - # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) - y = mhc_expand(y) + mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) + if cfg.mhc_expansion_rate > 1 and cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4): + # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) + y = mhc_expand(y) layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) - layer_kwargs = { - "slot": slot, - "previous_chunk": previous_chunk, - } + layer_kwargs = {} # Extract the bidirectional mask locally for layer configurations - bidirectional_mask = None - if multimodal_input is not None: - bidirectional_mask = ( - multimodal_input.bidirectional_mask - if multimodal_input.bidirectional_mask is not None - else multimodal_input.bidirectional_mask_video - ) + bidirectional_mask = multimodal_input.bidirectional_mask if multimodal_input is not None else None if cfg.decoder_block in {DecoderBlockType.GEMMA3, DecoderBlockType.GEMMA4}: layer_kwargs["bidirectional_mask"] = bidirectional_mask @@ -1634,10 +1704,7 @@ def __call__( if attention_metadata is not None: layer_kwargs["attention_metadata"] = attention_metadata - if cfg.engram_layers and decoder_input_tokens is not None: - layer_kwargs["decoder_input_tokens"] = decoder_input_tokens - - if getattr(cfg, "using_pipeline_parallelism", False): + if cfg.using_pipeline_parallelism: logical_partition_spec = ( self.pipeline_module.get_weight_sharding() if (cfg.pipeline_fsdp_ag_once or cfg.pipeline_fsdp_ag_per_repeat) @@ -1646,6 +1713,10 @@ def __call__( if self.is_deepseek: # Pre-pipeline: dense layers + outside-pipeline MoE layers under PP-as-DP axis rules. + ds_layer_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + } logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(cfg.logical_axis_rules) with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): if cfg.scan_layers: @@ -1655,7 +1726,7 @@ def __call__( y, *layer_args, length=cfg.first_num_dense_layers, - **layer_kwargs, + **ds_layer_kwargs, ) if hasattr(self, "moe_layers_outside_pipeline") and self.moe_layers_outside_pipeline is not None: num_moe_outside = (cfg.num_decoder_layers - cfg.first_num_dense_layers) - cfg.pipeline_parallel_layers @@ -1664,29 +1735,18 @@ def __call__( y, *layer_args, length=num_moe_outside, - **layer_kwargs, + **ds_layer_kwargs, ) else: # Unscanned: iterate registered layers by name. for i in range(getattr(self, "num_dense_layers", 0)): layer = getattr(self, f"dense_layers_{i}") - call_kwargs = dict(layer_kwargs) - if kv_caches is not None: - call_kwargs["kv_cache"] = kv_caches[i] - out = layer(y, *layer_args, **call_kwargs) + out = layer(y, *layer_args, **ds_layer_kwargs) y = out[0] if isinstance(out, tuple) else out - if kv_caches is not None and isinstance(out, tuple) and len(out) > 1 and out[1] is not None: - kv_caches[i] = out[1] for i in range(getattr(self, "num_moe_outside_pipeline", 0)): layer = getattr(self, f"moe_layers_outside_pipeline_{i}") - call_kwargs = dict(layer_kwargs) - kv_idx = getattr(self, "num_dense_layers", 0) + i - if kv_caches is not None: - call_kwargs["kv_cache"] = kv_caches[kv_idx] - out = layer(y, *layer_args, **call_kwargs) + out = layer(y, *layer_args, **ds_layer_kwargs) y = out[0] if isinstance(out, tuple) else out - if kv_caches is not None and isinstance(out, tuple) and len(out) > 1 and out[1] is not None: - kv_caches[kv_idx] = out[1] y = self.pipeline_module( y, @@ -1696,13 +1756,6 @@ def __call__( model_mode, logical_partition_spec=logical_partition_spec, ) - elif self.is_gemma4: - y = self._apply_gemma4_scanned_blocks( - y, - layer_args, - layer_kwargs, - kv_caches=kv_caches, - ) else: # Standard pipeline run (non-DeepSeek, incl. Gemma4 β€” matches Linen decoders.py). # Gemma4 routes through the pipeline here; _apply_gemma4_scanned_blocks is @@ -1736,14 +1789,8 @@ def __call__( elif (not cfg.scan_layers) and hasattr(self, "num_layers_outside_pipeline"): for i in range(self.num_layers_outside_pipeline): layer = getattr(self, f"layers_outside_pipeline_{i}") - call_kwargs = dict(layer_kwargs) - kv_idx = cfg.pipeline_parallel_layers + i - if kv_caches is not None: - call_kwargs["kv_cache"] = kv_caches[kv_idx] - out = layer(y, *layer_args, **call_kwargs) + out = layer(y, *layer_args, **layer_kwargs) y = out[0] if isinstance(out, tuple) else out - if kv_caches is not None and isinstance(out, tuple) and len(out) > 1 and out[1] is not None: - kv_caches[kv_idx] = out[1] else: if self.is_gemma4_small: @@ -1762,6 +1809,10 @@ def __call__( ) elif cfg.scan_layers: if self.is_deepseek: + layer_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + } if cfg.engram_layers: common_kwargs = { @@ -1803,7 +1854,7 @@ def __call__( policy = self.get_remat_policy() mock_params = self._build_linen_params(self.moe_layers) - if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: + if cfg.use_qwix_quantization and not cfg.use_manual_quantization: y = deepseek_batchsplit_fp8.scan_batch_split_layers( y, mock_params, @@ -1833,31 +1884,38 @@ def __call__( length=num_moe, **layer_kwargs, ) - - elif self.is_deepseek4: - y = self._apply_deepseek4_scanned_blocks( + elif self.is_gemma3: + y = self._apply_gemma3_scanned_blocks( y, decoder_segment_ids, decoder_positions, deterministic, model_mode, - slot, + bidirectional_mask, previous_chunk, - decoder_input_tokens, - ) - elif self.is_gemma3: - y = self._apply_gemma3_scanned_blocks( - y, - layer_args, - layer_kwargs, + slot, kv_caches=kv_caches, ) elif self.is_gemma4: y = self._apply_gemma4_scanned_blocks( y, - layer_args, - layer_kwargs, - kv_caches=kv_caches, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, + ) + elif self.is_qwen3_next_with_dense: + y = self._apply_qwen3_next_dense_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, ) else: scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) @@ -1884,48 +1942,27 @@ def __call__( ) else: prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - dynamic_graph_init = bool(getattr(self, "disable_quant_stats_update", False)) - def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): + # Hoisted function to preserve XLA cache ID + def pure_layer_fn(graphdef, state_in, y_in, kv_in): + if cfg.parameter_memory_host_offload: state_in = jax.tree.map( lambda x: jax.device_put(x, max_utils.device_space()), state_in, ) - merged_layer = nnx.merge(graphdef_in, state_in) - out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs) - state_out = nnx.state(merged_layer) - if dynamic_graph_init: - new_graphdef, _, _ = nnx.split(merged_layer, nnx.Param, ...) - return out_y, out_kv, state_out, new_graphdef - else: - return out_y, out_kv, state_out, graphdef_in + merged_layer = nnx.merge(graphdef, state_in) + out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs) + return out_y, out_kv, nnx.state(merged_layer) checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse) - for lyr in range(cfg.num_decoder_layers): - if self.is_deepseek: - if lyr < cfg.first_num_dense_layers: - layer = getattr(self, f"dense_layers_{lyr}", None) - else: - moe_idx = lyr - cfg.first_num_dense_layers - layer = getattr(self, f"moe_layers_{moe_idx}", None) - else: - layer = getattr(self, f"layers_{lyr}", None) - if layer is None and hasattr(self, "layers") and self.layers: - layer = self.layers[lyr] - - if layer is None: - raise AttributeError(f"Could not locate decoder layer at index {lyr} in {self.__class__.__name__}") - + for lyr, layer in enumerate(self.layers): graphdef, state = nnx.split(layer) if kv_caches is not None: - if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5) and cfg.attention not in ( - "vllm_rpa", - "vllm_batched_rpa", - ): - if (lyr + 1) % cfg.inhomogeneous_layer_cycle_interval == 0: + if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): + if layer.is_full_attention_layer: kv_cache = ( kv_caches["key_cache"][lyr], kv_caches["value_cache"][lyr], @@ -1937,36 +1974,16 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): else: kv_cache = None - input_tokens = ( - decoder_input_tokens if (cfg.engram_layers or cfg.decoder_block == DecoderBlockType.DEEPSEEK4) else None - ) + input_tokens = decoder_input_tokens if cfg.engram_layers else None if input_tokens is not None: layer_kwargs["decoder_input_tokens"] = input_tokens - if cfg.remat_policy != "none": - y, kv_cache, new_state, new_graphdef = checkpointed_fn(graphdef, state, y, kv_cache) - else: - y, kv_cache, new_state, new_graphdef = pure_layer_fn(graphdef, state, y, kv_cache) - - if dynamic_graph_init: - new_layer = nnx.merge(new_graphdef, new_state) - if self.is_deepseek: - if lyr < cfg.first_num_dense_layers: - setattr(self, f"dense_layers_{lyr}", new_layer) - else: - moe_idx = lyr - cfg.first_num_dense_layers - setattr(self, f"moe_layers_{moe_idx}", new_layer) - else: - setattr(self, f"layers_{lyr}", new_layer) - else: - nnx.update(layer, new_state) + y, kv_cache, new_state = checkpointed_fn(graphdef, state, y, kv_cache) + nnx.update(layer, new_state) if kv_caches is not None and kv_cache is not None: - if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5) and cfg.attention not in ( - "vllm_rpa", - "vllm_batched_rpa", - ): - if (lyr + 1) % cfg.inhomogeneous_layer_cycle_interval == 0: + if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): + if layer.is_full_attention_layer: kv_caches["key_cache"][lyr] = kv_cache[0] kv_caches["value_cache"][lyr] = kv_cache[1] else: @@ -1980,25 +1997,20 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): assert isinstance(y, jax.Array) # After the final transformer layer, `y` holds the raw, un-normalized hidden state. - if getattr(cfg, "mhc_expansion_rate", 1) > 1: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: - hidden_state = self.hc_head(y) - else: - # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) - hidden_state = mhc_reduce(y) + if getattr(cfg, "mhc_expansion_rate", 1) > 1 and cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4): + # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) + hidden_state = mhc_reduce(y) else: hidden_state = y # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage. - if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"): + if cfg.attention == "vllm_rpa": logits = None # When in the Indexer Dense Warm-up stage, skip the expensive output head projection # for efficiency, as the main model is frozen and the LM loss is not needed. elif ( - getattr(cfg, "use_indexer", False) - and getattr(cfg, "indexer_loss_scaling_factor", 0.0) > 0.0 - and not getattr(cfg, "indexer_sparse_training", False) + cfg.use_indexer and cfg.indexer_loss_scaling_factor > 0.0 and not cfg.indexer_sparse_training ) and model_mode == MODEL_MODE_TRAIN: logits = None @@ -2067,8 +2079,13 @@ def _apply_deepseek4_scanned_blocks( def _apply_gemma3_scanned_blocks( self, y, - layer_args, - layer_kwargs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask=None, + previous_chunk=None, + slot=None, kv_caches=None, ): """Applies Gemma3 scanned decoder blocks, handling main scan and remainders.""" @@ -2079,62 +2096,46 @@ def _apply_gemma3_scanned_blocks( attention_pattern_length = len(gemma3.GEMMA3_ATTENTION_PATTERN) scan_length = cfg.num_decoder_layers // attention_pattern_length + layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) + layer_kwargs = {"bidirectional_mask": bidirectional_mask} + # Apply the main scan over the full blocks if scan_length > 0: - grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( - kv_caches, scan_length, attention_pattern_length, stack=False - ) y, self.layers, _ = self._apply_layers_sequentially( - self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs - ) - maxtext_utils.update_kv_caches_after_scan( - kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False + self.layers, + y, + *layer_args, + length=scan_length, + kv_caches_stacked=kv_caches, + skip_block_remat=True, + unroll=1, + **layer_kwargs, ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length if num_remaining_layers > 0: - policy = self.get_remat_policy() - prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - - remainder_kv = None - if kv_caches is not None: - start_idx = scan_length * attention_pattern_length - remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) - - def pure_gemma_fn(graphdef, state_in, y_in, kv_in): - merged_layer = nnx.merge(graphdef, state_in) - call_kwargs = dict(layer_kwargs) - if kv_in is not None: - call_kwargs["kv_cache"] = kv_in - out_res = merged_layer(y_in, *layer_args, **call_kwargs) - if isinstance(out_res, tuple): - out_y = out_res[0] - out_kv = out_res[1] if len(out_res) > 1 else None - else: - out_y = out_res - out_kv = None - return out_y, out_kv, nnx.state(merged_layer) - - checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) - - graphdef, state = nnx.split(self.layers_remainder) - y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) - nnx.update(self.layers_remainder, new_state) - - if kv_caches is not None and updated_remainder_kv is not None: - start_idx = scan_length * attention_pattern_length - for offset, updated_item in enumerate(updated_remainder_kv): - kv_caches[start_idx + offset] = updated_item + out = self.layers_remainder( + y, + *layer_args, + previous_chunk=previous_chunk, + slot=slot, + **layer_kwargs, + ) + y = out[0] if isinstance(out, tuple) else out return y def _apply_gemma4_scanned_blocks( self, y, - layer_args, - layer_kwargs, - kv_caches=None, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, ): """Applies Gemma4 scanned decoder blocks, handling main scan and remainders.""" @@ -2144,79 +2145,79 @@ def _apply_gemma4_scanned_blocks( attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN) scan_length = cfg.num_decoder_layers // attention_pattern_length - # Apply the main scan over the full blocks. Gemma4ScannableBlock applies - # per-layer remat internally (local scan + global layer), so skip the - # block-level remat here to avoid double rematerialization. Unrolling the - # block loop (one iteration per repeated block) lets XLA pipeline/free block - # activations across iterations (memory + overlap knob). - block_unroll = max(1, scan_length) + layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) + layer_kwargs = {"bidirectional_mask": bidirectional_mask, "slot": slot, "previous_chunk": previous_chunk} + + # Apply the main scan over the full blocks if scan_length > 0: - grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( - kv_caches, scan_length, attention_pattern_length, stack=False - ) - y, self.scanned_blocks, _ = self._apply_layers_sequentially( - self.scanned_blocks, + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, y, *layer_args, length=scan_length, - kv_caches_stacked=grouped_kv_caches, skip_block_remat=True, - unroll=block_unroll, + unroll=1, **layer_kwargs, ) - maxtext_utils.update_kv_caches_after_scan( - kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False - ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length if num_remaining_layers > 0: - policy = self.get_remat_policy() - prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - - remainder_kv = None - if kv_caches is not None: - start_idx = scan_length * attention_pattern_length - remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) - - if cfg.use_qwix_quantization or cfg.lora.lora_weight_qtype: - call_kwargs = dict(layer_kwargs) - if remainder_kv is not None: - call_kwargs["kv_cache"] = remainder_kv - out_res = self.layers_remainder(y, *layer_args, **call_kwargs) - if isinstance(out_res, tuple): - y = out_res[0] - updated_remainder_kv = out_res[1] if len(out_res) > 1 else None - else: - y = out_res - updated_remainder_kv = None - else: + out = self.layers_remainder( + y, + *layer_args, + **layer_kwargs, + ) + y = out[0] if isinstance(out, tuple) else out - def pure_gemma_fn(graphdef, state_in, y_in, kv_in): - merged_layer = nnx.merge(graphdef, state_in) - call_kwargs = dict(layer_kwargs) - if kv_in is not None: - call_kwargs["kv_cache"] = kv_in - out_res = merged_layer(y_in, *layer_args, **call_kwargs) - if isinstance(out_res, tuple): - out_y = out_res[0] - out_kv = out_res[1] if len(out_res) > 1 else None - else: - out_y = out_res - out_kv = None - nnx.pop(merged_layer, (nnx.RngState, nnx.Intermediate)) - return out_y, out_kv, nnx.state(merged_layer) + return y - checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) + def _apply_qwen3_next_dense_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + ): + """Applies Qwen3-Next's dense prefix, scanned middle, and remainder tail (see + `_init_scanned_qwen3_next_with_dense`), mirroring `_apply_gemma3_scanned_blocks`.""" + cfg = self.config + layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) + layer_kwargs = {"previous_chunk": previous_chunk, "slot": slot} - graphdef, state = nnx.split(self.layers_remainder) - y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) - nnx.update(self.layers_remainder, new_state) + n_dense = cfg.first_num_dense_layers + if n_dense > 0: + y, self.dense_layers, _ = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=n_dense, + metadata_axis_name="dense_layers", + ) - if kv_caches is not None and updated_remainder_kv is not None: - start_idx = scan_length * attention_pattern_length - for offset, updated_item in enumerate(updated_remainder_kv): - kv_caches[start_idx + offset] = updated_item + remaining = cfg.num_decoder_layers - n_dense + cycle = cfg.inhomogeneous_layer_cycle_interval + scan_length = remaining // cycle + if scan_length > 0: + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, + y, + *layer_args, + length=scan_length, + **layer_kwargs, + ) + + num_remaining = remaining % cycle + if num_remaining > 0: + out = self.layers_remainder( + y, + *layer_args, + **layer_kwargs, + ) + y = out[0] if isinstance(out, tuple) else out return y @@ -2249,7 +2250,7 @@ def _apply_gemma4_small_layers( cache_index_of = gemma4_small.kv_cache_slot_map(layer_types, num_kv_shared) for lyr in range(cfg.num_decoder_layers): - layer = getattr(self, f"layers_{lyr}") + layer = self.layers[lyr] donor_idx = gemma4_small.kv_donor_layer_idx(lyr, layer_types, num_kv_shared) is_donor = gemma4_small.is_kv_donor_layer(lyr, layer_types, num_kv_shared) @@ -2294,63 +2295,6 @@ def _apply_gemma4_small_layers( return y, kv_caches - def get_layers(self) -> list[nnx.Module]: - """Returns all decoder layer modules/blocks in their forward-pass execution order.""" - layers = [] - seen = set() - - def _add(module): - if module is not None and id(module) not in seen: - seen.add(id(module)) - layers.append(module) - - def _append_unscanned(prefix): - i = 0 - while hasattr(self, f"{prefix}_{i}"): - _add(getattr(self, f"{prefix}_{i}")) - i += 1 - - def _append_scanned(name): - if hasattr(self, name): - val = getattr(self, name) - if name == "layers_remainder" and getattr(val, "num_of_layers", 0) == 0: - return - if isinstance(val, (nnx.Module, list)): - _add(val) - - if self.is_deepseek: - _append_scanned("dense_layers") - _append_unscanned("dense_layers") - - _append_scanned("moe_layers") - _append_unscanned("moe_layers") - - _append_scanned("moe_layers_outside_pipeline") - _append_unscanned("moe_layers_outside_pipeline") - - if hasattr(self, "pipeline_module"): - _add(getattr(self.pipeline_module, "layers", None)) - else: - if hasattr(self, "pipeline_module"): - _add(getattr(self.pipeline_module, "layers", None)) - - _append_scanned("scanned_blocks") # Gemma 4 - _append_scanned("layers") - _append_unscanned("layers") - - _append_scanned("layers_remainder") # Gemma 3/4 - - _append_scanned("layers_outside_pipeline") - _append_unscanned("layers_outside_pipeline") - - # Fallback for dynamic/chunked layer attributes (e.g. Engram dense_layers_0_3) - if not layers: - for k, m in vars(self).items(): - if k.startswith(("dense_layers", "moe_layers", "layers", "scanned_blocks")) and isinstance(m, (nnx.Module, list)): - _add(m) - - return layers - def decoder_as_linen( config: Config, diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index e43a4d47bf..765f82d1d9 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -122,23 +122,20 @@ def apply_scanned_layers( if length <= 0: return carry - layer_graphdef, params, state = nnx.split(layers, nnx.Param, ...) + layer_graphdef, params, rest = nnx.split(layers, nnx.Param, ...) if param_scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params) def scan_body(current_carry, scanned_state): - current_params, current_state = scanned_state - current_layer = nnx.merge(layer_graphdef, current_params, current_state) + current_params, current_rest = scanned_state + current_layer = nnx.merge(layer_graphdef, current_params, current_rest) next_carry = apply_fn(current_layer, current_carry) - return next_carry, nnx.state(current_layer) + # Avoid returning and stacking read-only parameters inside the scan body. + _, _, updated_rest = nnx.split(current_layer, nnx.Param, ...) + return next_carry, updated_rest scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body - final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), unroll=unroll) + final_carry, scanned_rest = jax.lax.scan(scan_fn, carry, (params, rest), unroll=unroll) - if param_scan_axis != 0: - scanned_params, scanned_other = scanned_state.split(nnx.Param, ...) - scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_params) - scanned_state = nnx.State.merge(scanned_params, scanned_other) - - nnx.update(layers, scanned_state) + nnx.update(layers, scanned_rest) return final_carry diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index a5dade3e19..4a2b8ead8a 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -30,7 +30,10 @@ import qwix from qwix._src.core import numerics from qwix._src.core import dot_general_qt -from qwix._src.core import sparsity +try: + from qwix._src.core import sparsity +except ImportError: + sparsity = None import jax import jax.numpy as jnp diff --git a/src/maxtext/models/hybrid_gdn.py b/src/maxtext/models/hybrid_gdn.py new file mode 100644 index 0000000000..8f2ca59d39 --- /dev/null +++ b/src/maxtext/models/hybrid_gdn.py @@ -0,0 +1,910 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hybrid Gated Delta Net (GDN) implementations for MaxText using Tokamax GDN v3 forward + Pallas Custom VJP backward.""" + +import functools +from typing import Any, Optional, Tuple + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + + +def _pallas_gdn_bwd_kernel( + padded_pre_conv_qkv_ref, + qkv_ref, + b_ref, + a_ref, + a_log_ref, + dt_bias_ref, + do_ref, + chunk_states_ref, + conv_weight_ref, + seq_lens_ref, + d_qkv_ref, + d_b_ref, + d_a_ref, + d_conv_weight_ref, + d_a_log_ref, + d_dt_bias_ref, + padded_pre_conv_qkv_vmem, + qkv_vmem, + b_vmem, + a_vmem, + a_log_vmem, + dt_bias_vmem, + do_vmem, + chunk_states_vmem, + d_qkv_vmem, + d_b_vmem, + d_a_vmem, + d_conv_weight_scratch, + d_a_log_scratch, + d_dt_bias_scratch, + sem_pre_conv_qkv, + sem_qkv, + sem_b, + sem_a, + sem_a_log, + sem_dt_bias, + sem_do, + sem_chunk_states, + sem_d_qkv, + sem_d_b, + sem_d_a, + sem_d_conv_weight, + sem_d_a_log, + sem_d_dt_bias, + *, + batch_size: int, + num_chunks: int, + chunk_size: int, + dim_size: int, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int, + pad_len: int = 8, + use_qk_norm_in_gdn: bool = False, +): + seq_idx = pl.program_id(0) + + d_conv_weight_scratch[...] = jnp.zeros_like(d_conv_weight_scratch) + + num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2) + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + v_size = num_v_heads * v_head_dim + repeats = num_v_heads // num_kq_heads + + def chunk_forward(q, k, v, b_val, a_val, a_log_val, dt_bias_val, state_prev): + q = q.astype(jnp.float32) + k = k.astype(jnp.float32) + v = v.astype(jnp.float32) + if use_qk_norm_in_gdn: + q = q / (jnp.linalg.norm(q, axis=-1, keepdims=True) + 1e-6) + k = k / (jnp.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + scale = 1.0 / jnp.sqrt(kq_head_dim) + q = q * scale + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.astype(jnp.float32) + q_rep = jnp.repeat(q, repeats, axis=1) + k_rep = jnp.repeat(k, repeats, axis=1) + + beta = jax.nn.sigmoid(b_val) + + # EXACT GDN v3 gating formula + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + v_beta = v * beta[:, :, None] + + # Fast MXU cumsum replacement + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + diff = cumsum_log_g[:, None, :] - cumsum_log_g[None, :, :] + mask = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype)) + safe_diff = jnp.where(mask[:, :, None] == 1.0, diff, -1e4) + G = jnp.exp(safe_diff) * mask[:, :, None] + G = jnp.transpose(G, (2, 0, 1)) + + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v_beta, (1, 0, 2)) + + # Dual Causal Attention + attn = jnp.einsum('hck,hdk->hcd', q_h, k_h) + attn_causal = attn * G + + out_intra = jnp.einsum('hcd,hdv->hcv', attn_causal, v_h) + out_intra = jnp.transpose(out_intra, (1, 0, 2)) + + cross_decay = jnp.exp(cumsum_log_g) + q_scaled = q_rep * cross_decay[:, :, None] + out_cross = jnp.einsum('chk,hkv->chv', q_scaled, state_prev) + + out = out_intra + out_cross + + state_decay_end = G[:, chunk_size - 1, :] + state_prev_decayed = state_prev * cross_decay[-1, :, None, None] + + k_scaled = k_h * state_decay_end[:, :, None] + state_new_intra = jnp.einsum('hck,hcv->hkv', k_scaled, v_h) + + state_new = state_prev_decayed + state_new_intra + + return out, state_new + + def fetch_inputs(chunk_idx, slot): + start_idx = chunk_idx * chunk_size + copy_len = chunk_size + pad_len + pltpu.make_async_copy( + padded_pre_conv_qkv_ref.at[seq_idx, pl.ds(start_idx, copy_len)], + padded_pre_conv_qkv_vmem.at[slot], + sem_pre_conv_qkv.at[slot], + ).start() + pltpu.make_async_copy(qkv_ref.at[seq_idx, chunk_idx], qkv_vmem.at[slot], sem_qkv.at[slot]).start() + pltpu.make_async_copy(b_ref.at[seq_idx, chunk_idx], b_vmem.at[slot], sem_b.at[slot]).start() + pltpu.make_async_copy(a_ref.at[seq_idx, chunk_idx], a_vmem.at[slot], sem_a.at[slot]).start() + pltpu.make_async_copy(do_ref.at[seq_idx, chunk_idx], do_vmem.at[slot], sem_do.at[slot]).start() + pltpu.make_async_copy( + chunk_states_ref.at[seq_idx, chunk_idx], chunk_states_vmem.at[slot], sem_chunk_states.at[slot] + ).start() + + def wait_inputs(slot): + pltpu.make_async_copy( + padded_pre_conv_qkv_vmem.at[slot], padded_pre_conv_qkv_vmem.at[slot], sem_pre_conv_qkv.at[slot] + ).wait() + pltpu.make_async_copy(qkv_vmem.at[slot], qkv_vmem.at[slot], sem_qkv.at[slot]).wait() + pltpu.make_async_copy(b_vmem.at[slot], b_vmem.at[slot], sem_b.at[slot]).wait() + pltpu.make_async_copy(a_vmem.at[slot], a_vmem.at[slot], sem_a.at[slot]).wait() + pltpu.make_async_copy(do_vmem.at[slot], do_vmem.at[slot], sem_do.at[slot]).wait() + pltpu.make_async_copy(chunk_states_vmem.at[slot], chunk_states_vmem.at[slot], sem_chunk_states.at[slot]).wait() + + def store_outputs(chunk_idx, slot): + pltpu.make_async_copy(d_qkv_vmem.at[slot], d_qkv_ref.at[seq_idx, chunk_idx], sem_d_qkv.at[slot]).start() + pltpu.make_async_copy(d_b_vmem.at[slot], d_b_ref.at[seq_idx, chunk_idx], sem_d_b.at[slot]).start() + pltpu.make_async_copy(d_a_vmem.at[slot], d_a_ref.at[seq_idx, chunk_idx], sem_d_a.at[slot]).start() + + def wait_outputs(slot): + pltpu.make_async_copy(d_qkv_vmem.at[slot], d_qkv_vmem.at[slot], sem_d_qkv.at[slot]).wait() + pltpu.make_async_copy(d_b_vmem.at[slot], d_b_vmem.at[slot], sem_d_b.at[slot]).wait() + pltpu.make_async_copy(d_a_vmem.at[slot], d_a_vmem.at[slot], sem_d_a.at[slot]).wait() + + # Prologue + fetch_inputs(num_chunks - 1, 0) + + # Fetch static inputs (a_log and dt_bias) + pltpu.make_async_copy(a_log_ref.at[seq_idx], a_log_vmem, sem_a_log.at[0]).start() + pltpu.make_async_copy(dt_bias_ref.at[seq_idx], dt_bias_vmem, sem_dt_bias.at[0]).start() + pltpu.make_async_copy(a_log_vmem, a_log_vmem, sem_a_log.at[0]).wait() + pltpu.make_async_copy(dt_bias_vmem, dt_bias_vmem, sem_dt_bias.at[0]).wait() + + d_state = jnp.zeros((num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32) + d_conv_weight_acc = jnp.zeros((kernel_size, dim_size), dtype=jnp.float32) + d_a_log_acc = jnp.zeros((num_v_heads,), dtype=jnp.float32) + d_dt_bias_acc = jnp.zeros((num_v_heads,), dtype=jnp.float32) + + def loop_body(i, carry): + d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc = carry + + chunk_idx = num_chunks - 1 - i + slot = i % 2 + next_slot = (i + 1) % 2 + + wait_inputs(slot) + + def fetch_next(): + fetch_inputs(chunk_idx - 1, next_slot) + + jax.lax.cond(i < num_chunks - 1, fetch_next, lambda: None) + + def wait_prev_out(): + wait_outputs(slot) + + jax.lax.cond(i > 1, wait_prev_out, lambda: None) + + padded_pre_conv_qkv_val = padded_pre_conv_qkv_vmem[slot, ...] + qkv_val = qkv_vmem[slot, ...] + b_val = b_vmem[slot, ...] + a_val = a_vmem[slot, ...] + do_val = do_vmem[slot, ...] + state_prev_val = chunk_states_vmem[slot, ...] + a_log_val = a_log_vmem[...] + dt_bias_val = dt_bias_vmem[...] + + q = qkv_val[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)) + k = qkv_val[:, q_size : q_size + k_size].reshape((chunk_size, num_kq_heads, kq_head_dim)) + v = qkv_val[:, q_size + k_size :].reshape((chunk_size, num_v_heads, v_head_dim)) + + _, vjp_fn = jax.vjp(chunk_forward, q, k, v, b_val, a_val, a_log_val, dt_bias_val, state_prev_val) + d_q, d_k, d_v, d_b_val, d_a_val, d_a_log_val, d_dt_bias_val, d_state_prev = vjp_fn((do_val.astype(jnp.float32), d_state)) + + d_qkv = jnp.concatenate([ + d_q.reshape(chunk_size, q_size), + d_k.reshape(chunk_size, k_size), + d_v.reshape(chunk_size, v_size), + ], axis=-1) + + d_cw_rows = [] + for k_idx in range(kernel_size): + start_slice = pad_len - k_idx + shifted = padded_pre_conv_qkv_val[start_slice : start_slice + chunk_size] + d_cw_rows.append(jnp.sum(d_qkv * shifted.astype(jnp.float32), axis=0)) + d_cw = jnp.stack(d_cw_rows, axis=0) + + d_conv_weight_acc += d_cw.astype(jnp.float32) + d_a_log_acc += d_a_log_val.astype(jnp.float32) + d_dt_bias_acc += d_dt_bias_val.astype(jnp.float32) + d_state = d_state_prev.astype(jnp.float32) + + d_qkv_vmem[slot, ...] = d_qkv.astype(d_qkv_vmem.dtype) + d_b_vmem[slot, ...] = d_b_val.astype(d_b_vmem.dtype) + d_a_vmem[slot, ...] = d_a_val.astype(d_a_vmem.dtype) + + store_outputs(chunk_idx, slot) + + return d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc + + d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc = jax.lax.fori_loop( + 0, num_chunks, loop_body, (d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc) + ) + + def wait_last_out(): + wait_outputs((num_chunks - 1) % 2) + + jax.lax.cond(num_chunks > 0, wait_last_out, lambda: None) + + def wait_prev_last_out(): + wait_outputs((num_chunks - 2) % 2) + + jax.lax.cond(num_chunks > 1, wait_prev_last_out, lambda: None) + + d_conv_weight_scratch[...] = d_conv_weight_acc.astype(d_conv_weight_scratch.dtype) + d_a_log_scratch[...] = d_a_log_acc.astype(d_a_log_scratch.dtype) + d_dt_bias_scratch[...] = d_dt_bias_acc.astype(d_dt_bias_scratch.dtype) + + pltpu.make_async_copy(d_conv_weight_scratch, d_conv_weight_ref.at[seq_idx, ...], sem_d_conv_weight.at[0]).start() + pltpu.make_async_copy(d_a_log_scratch, d_a_log_ref.at[seq_idx, ...], sem_d_a_log.at[0]).start() + pltpu.make_async_copy(d_dt_bias_scratch, d_dt_bias_ref.at[seq_idx, ...], sem_d_dt_bias.at[0]).start() + + pltpu.make_async_copy(d_conv_weight_scratch, d_conv_weight_scratch, sem_d_conv_weight.at[0]).wait() + pltpu.make_async_copy(d_a_log_scratch, d_a_log_scratch, sem_d_a_log.at[0]).wait() + pltpu.make_async_copy(d_dt_bias_scratch, d_dt_bias_scratch, sem_d_dt_bias.at[0]).wait() + + +def pallas_fused_conv1d_gdn_bwd_computation( + pre_conv_qkv: jax.Array, + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + conv_weight: jax.Array, + seq_lens: Optional[jax.Array] = None, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Executes the Pallas reverse-chunk GDNv3 backward kernel.""" + batch_size, seq_len, dim_size = pre_conv_qkv.shape + num_chunks = seq_len // chunk_size + + pre_conv_qkv_4d = pre_conv_qkv.reshape(batch_size, num_chunks, chunk_size, dim_size) + qkv_4d = qkv.reshape(batch_size, num_chunks, chunk_size, dim_size) + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + do_4d = do.reshape(batch_size, num_chunks, chunk_size, num_v_heads, v_head_dim) + + if a_log.ndim == 1: + a_log_2d = jnp.broadcast_to(a_log[None, :], (batch_size, num_v_heads)) + else: + a_log_2d = a_log + if dt_bias.ndim == 1: + dt_bias_2d = jnp.broadcast_to(dt_bias[None, :], (batch_size, num_v_heads)) + else: + dt_bias_2d = dt_bias + + if conv_weight.ndim == 3: + conv_weight_2d = conv_weight.squeeze(1) + else: + conv_weight_2d = conv_weight + + if seq_lens is None: + seq_lens = jnp.full((batch_size,), seq_len, dtype=jnp.int32) + + pad_len = ((kernel_size - 1 + 7) // 8) * 8 + pre_conv_pad = jnp.zeros((batch_size, pad_len, dim_size), dtype=pre_conv_qkv.dtype) + padded_pre_conv_qkv = jnp.concatenate([pre_conv_pad, pre_conv_qkv], axis=1) + + grid = (batch_size,) + hbm_spec = pl.BlockSpec(memory_space=pl.ANY) + + d_qkv_shape = jax.ShapeDtypeStruct(qkv_4d.shape, qkv_4d.dtype) + d_b_shape = jax.ShapeDtypeStruct(b_4d.shape, b_4d.dtype) + d_a_shape = jax.ShapeDtypeStruct(a_4d.shape, a_4d.dtype) + d_conv_weight_shape = jax.ShapeDtypeStruct((batch_size, kernel_size, dim_size), conv_weight_2d.dtype) + d_a_log_shape = jax.ShapeDtypeStruct((batch_size, num_v_heads), a_log_2d.dtype) + d_dt_bias_shape = jax.ShapeDtypeStruct((batch_size, num_v_heads), dt_bias_2d.dtype) + + out_shapes = (d_qkv_shape, d_b_shape, d_a_shape, d_conv_weight_shape, d_a_log_shape, d_dt_bias_shape) + + scratch_shapes = ( + pltpu.VMEM((2, chunk_size + pad_len, dim_size), padded_pre_conv_qkv.dtype), + pltpu.VMEM((2, chunk_size, dim_size), qkv_4d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), b_4d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), a_4d.dtype), + pltpu.VMEM((num_v_heads,), a_log_2d.dtype), + pltpu.VMEM((num_v_heads,), dt_bias_2d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads, v_head_dim), do_4d.dtype), + pltpu.VMEM((2, num_v_heads, kq_head_dim, v_head_dim), chunk_states.dtype), + pltpu.VMEM((2, chunk_size, dim_size), d_qkv_shape.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), d_b_shape.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), d_a_shape.dtype), + pltpu.VMEM((kernel_size, dim_size), d_conv_weight_shape.dtype), + pltpu.VMEM((num_v_heads,), d_a_log_shape.dtype), + pltpu.VMEM((num_v_heads,), d_dt_bias_shape.dtype), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + ) + + d_qkv, d_b, d_a, d_conv_weight, d_a_log, d_dt_bias = pl.pallas_call( + functools.partial( + _pallas_gdn_bwd_kernel, + batch_size=batch_size, + num_chunks=num_chunks, + chunk_size=chunk_size, + dim_size=dim_size, + num_v_heads=num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + kernel_size=kernel_size, + pad_len=pad_len, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ), + out_shape=out_shapes, + grid=grid, + in_specs=[hbm_spec] * 10, + out_specs=[hbm_spec] * 6, + scratch_shapes=scratch_shapes, + compiler_params=pltpu.CompilerParams( + disable_bounds_checks=True, + ), + )(padded_pre_conv_qkv, qkv_4d, b_4d, a_4d, a_log_2d, dt_bias_2d, do_4d, chunk_states, conv_weight_2d, seq_lens) + + d_conv_weight_reduced = jnp.sum(d_conv_weight, axis=0) + d_a_log_reduced = jnp.sum(d_a_log, axis=0) + d_dt_bias_reduced = jnp.sum(d_dt_bias, axis=0) + + d_qkv_flat = d_qkv.reshape(batch_size, seq_len, dim_size) + d_b_flat = d_b.reshape(batch_size, seq_len, num_v_heads) + d_a_flat = d_a.reshape(batch_size, seq_len, num_v_heads) + + if conv_weight.ndim == 3: + d_conv_weight_out = d_conv_weight_reduced[:, None, :] + else: + d_conv_weight_out = d_conv_weight_reduced + + return d_qkv_flat, d_b_flat, d_a_flat, d_conv_weight_out, d_a_log_reduced, d_dt_bias_reduced + + +def pure_jax_fused_conv1d_gdn( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Pure-JAX composite of Conv1D + GDN used during backward pass autodiff.""" + from maxtext.models.qwen3 import jax_chunk_gated_delta_rule + + batch, seq_len, _ = qkv.shape + key_dim = num_k_heads * head_k_dim + + # --- Step B: Pure JAX 1D Convolution --- + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + conv_weight_cast = conv_weight.astype(qkv.dtype) + conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight_cast, + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=qkv.shape[-1], + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(qkv.dtype) + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(compute_dtype) + + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [key_dim, 2 * key_dim], axis=-1) + + # Reshape for GDN + query = q_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + key = k_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + value = v_conv.reshape(batch, seq_len, num_v_heads, head_v_dim) + + A_log_cast = jnp.asarray(a_log, dtype=compute_dtype) + dt_bias_cast = jnp.asarray(dt_bias, dtype=compute_dtype) + beta = jax.nn.sigmoid(b) + g = -jnp.exp(A_log_cast) * jax.nn.softplus(a + dt_bias_cast) + + if num_v_heads > num_k_heads and num_v_heads % num_k_heads == 0: + repeats = num_v_heads // num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query=query, + key=key, + value=value, + g=g, + beta=beta, + chunk_size=chunk_size, + initial_state=recurrent_state, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + next_conv_state = ( + qkv[:, -(conv_kernel_size - 1) :, :] + if seq_len >= conv_kernel_size - 1 + else jnp.zeros((batch, conv_kernel_size - 1, qkv.shape[-1]), dtype=qkv.dtype) + ) + if next_recurrent_state is None: + next_recurrent_state = jnp.zeros((batch, num_v_heads, head_k_dim, head_v_dim), dtype=compute_dtype) + + return core_attn_out.astype(qkv.dtype), (next_conv_state.astype(qkv.dtype), next_recurrent_state.astype(qkv.dtype)) + + +def _compute_forward_conv_and_states( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, jax.Array]: + """Computes convolved QKV and inter-chunk states during forward pass.""" + batch_size, seq_len, dim_size = qkv.shape + num_chunks = seq_len // chunk_size + + # Conv1D + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight.astype(qkv.dtype), + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=dim_size, + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(qkv.dtype) + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(compute_dtype) + + # Chunk states progression + num_kq_heads = num_k_heads + q_size = num_kq_heads * head_k_dim + k_size = num_kq_heads * head_k_dim + repeats = num_v_heads // num_kq_heads + + q = qkv_conv[:, :, :q_size].reshape(batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim) + k = qkv_conv[:, :, q_size : q_size + k_size].reshape(batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim) + v = qkv_conv[:, :, q_size + k_size :].reshape(batch_size, num_chunks, chunk_size, num_v_heads, head_v_dim) + + if use_qk_norm_in_gdn: + from maxtext.layers.normalizations import l2norm + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + + scale = jax.lax.rsqrt(jnp.array(head_k_dim, dtype=jnp.float32)).astype(compute_dtype) + q = q * scale + + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + + if recurrent_state is None: + init_state = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + else: + init_state = recurrent_state.astype(jnp.float32) + + def scan_fn(carry_state, chunk_inputs): + q_i, k_i, v_i, b_i, a_i = chunk_inputs + q_rep = jnp.repeat(q_i, repeats, axis=2) + k_rep = jnp.repeat(k_i, repeats, axis=2) + + beta = jax.nn.sigmoid(b_i) + log_g = -jnp.exp(a_log) * jax.nn.softplus(a_i + dt_bias) + v_beta = v_i * beta[:, :, :, None] + + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.einsum('cd,bhd->bhc', mask_cumsum, log_g.swapaxes(1, 2)).swapaxes(1, 2) + + diff = cumsum_log_g[:, :, None, :] - cumsum_log_g[:, None, :, :] + mask = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype)) + safe_diff = jnp.where(mask[None, :, :, None] == 1.0, diff, -1e4) + G = jnp.exp(safe_diff) * mask[None, :, :, None] + + cross_decay = jnp.exp(cumsum_log_g) + state_decay_end = G[:, chunk_size - 1, :, :] + state_prev_decayed = carry_state * cross_decay[:, -1, :, None, None] + + k_scaled = k_rep * state_decay_end[:, :, :, None] + state_new_intra = jnp.einsum('bchk,bchv->bhkv', k_scaled, v_beta) + + new_state = state_prev_decayed + state_new_intra + return new_state, carry_state + + q_chunks = q.swapaxes(0, 1) + k_chunks = k.swapaxes(0, 1) + v_chunks = v.swapaxes(0, 1) + b_chunks = b_4d.swapaxes(0, 1) + a_chunks = a_4d.swapaxes(0, 1) + + _, chunk_states = jax.lax.scan(scan_fn, init_state, (q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)) + chunk_states = chunk_states.swapaxes(0, 1) + + return qkv_conv, chunk_states + + +def _run_tokamax_fused_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +): + if jax.default_backend() != "tpu": + return pure_jax_fused_conv1d_gdn( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + from tokamax._src.ops.experimental.causal_conv1d_gated_delta_rule import wrapper as tokamax_gdn_wrapper + + batch_size, seq_len, dim_size = qkv.shape + num_seqs = batch_size + + qkv_flat = qkv.reshape(-1, dim_size) + b_flat = b.reshape(-1, b.shape[-1]) + a_flat = a.reshape(-1, a.shape[-1]) + tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2) + + query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32) + state_indices = jnp.arange(num_seqs, dtype=jnp.int32) + seq_lens = jnp.full((num_seqs,), seq_len, dtype=jnp.int32) + distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32) + + if conv_state is None: + tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype) + elif conv_state.shape[0] == num_seqs: + tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0))) + else: + tokamax_conv_state = conv_state + + if recurrent_state is None: + tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=qkv.dtype) + elif recurrent_state.shape[0] == num_seqs: + tokamax_recurrent_state = jnp.pad(recurrent_state, ((1, 0), (0, 0), (0, 0), (0, 0))) + else: + tokamax_recurrent_state = recurrent_state + + (new_conv_state, new_recurrent_state), core_attn_out_flat = tokamax_gdn_wrapper.fused_conv1d_gdn( + qkv=qkv_flat, + b=b_flat, + a=a_flat, + conv_state=tokamax_conv_state, + recurrent_state=tokamax_recurrent_state, + conv_weight=tokamax_conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + query_start_loc=query_start_loc, + state_indices=state_indices, + distribution=distribution, + seq_lens=seq_lens, + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=head_k_dim, + d_v=head_v_dim, + kernel_size=conv_kernel_size, + compute_precision=jnp.dtype(jnp.float32), + ) + + core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim) + return core_attn_out.astype(qkv.dtype), ( + new_conv_state[1:].astype(qkv.dtype), + new_recurrent_state[1:].astype(qkv.dtype), + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16)) +def hybrid_fused_conv1d_gdn( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Hybrid Fused Conv1D + GDN: Tokamax GDN v3 forward + Pallas Custom VJP backward.""" + return _run_tokamax_fused_fwd( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + +def _hybrid_fused_conv1d_gdn_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +): + out, states = _run_tokamax_fused_fwd( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + qkv_conv, chunk_states = _compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + residuals = ( + qkv, + qkv_conv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + chunk_states, + conv_state, + recurrent_state, + ) + return (out, states), residuals + + +def _hybrid_fused_conv1d_gdn_bwd( + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, + residuals: tuple, + cotangents: tuple, +): + ( + pre_conv_qkv, + qkv_conv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + chunk_states, + conv_state, + recurrent_state, + ) = residuals + d_out, d_states = cotangents + d_conv_state, d_recurrent_state = d_states + + if jax.default_backend() == "tpu": + d_qkv, d_b, d_a, d_conv_weight, d_a_log, d_dt_bias = pallas_fused_conv1d_gdn_bwd_computation( + pre_conv_qkv=pre_conv_qkv, + qkv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=d_out, + chunk_states=chunk_states, + conv_weight=conv_weight, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + d_conv_bias = None if conv_bias is None else jnp.zeros_like(conv_bias) + d_conv_state = None if conv_state is None else jnp.zeros_like(conv_state) + d_recurrent_state = None if recurrent_state is None else jnp.zeros_like(recurrent_state) + return (d_qkv, d_b, d_a, d_conv_weight, d_conv_bias, d_a_log, d_dt_bias, d_conv_state, d_recurrent_state) + + # Fallback to JAX Newton-Schulz VJP on non-TPU / CPU + def target_fn(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + return pure_jax_fused_conv1d_gdn( + qkv_, + b_, + a_, + cw_, + cb_, + al_, + dt_, + cs_, + rs_, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + _, vjp_fn = jax.vjp( + target_fn, + pre_conv_qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + ) + return vjp_fn((d_out, (d_conv_state, d_recurrent_state))) + + +hybrid_fused_conv1d_gdn.defvjp(_hybrid_fused_conv1d_gdn_fwd, _hybrid_fused_conv1d_gdn_bwd) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 7cb710bf1e..c16088eac8 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -33,12 +33,16 @@ from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE from maxtext.common.common_types import KV_BATCH, KV_HEAD +from maxtext.common.common_types import HyperConnectionType from maxtext.utils.sharding import logical_to_mesh_axes, get_logical_axis_rules from maxtext.layers import attentions from maxtext.layers import initializers as max_initializers +from maxtext.layers import mhc from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations +from maxtext.layers import nnx_scan +from jax.experimental import xla_metadata from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated from maxtext.layers.quantizations import AqtQuantization as Quant @@ -47,6 +51,7 @@ from maxtext.layers.moe import RoutedMoE from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils from maxtext.inference import kvcache @@ -612,7 +617,7 @@ def __call__( # mixed_qkvz: (B, S, H_k, 2*D_k + 2*D_v*V_per_K) mixed_qkvz = qkvz.reshape(new_shape_qkvz) if self.mesh is not None: - logical_rules = get_logical_axis_rules() + logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules qkvz_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) qkvz_sharding = jax.sharding.NamedSharding(self.mesh, qkvz_pspec) mixed_qkvz = jax.lax.with_sharding_constraint(mixed_qkvz, qkvz_sharding) @@ -658,11 +663,14 @@ def __call__( # vLLM PAGED STATE PATH: use tpu_inference fused conv + ragged delta-rule. # ========================================================================= try: - from tpu_inference.layers.common.gdn_attention import run_jax_gdn_attention # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.layers.common.sharding import ShardingAxisName # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.layers.common.utils import reorder_concatenated_tensor_for_sharding # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel # pytype: disable=import-error + # pylint: disable=import-outside-toplevel + # pytype: disable=import-error + from tpu_inference.layers.common.gdn_attention import GdnAttentionConfig, run_jax_gdn_attention # pylint: disable=import-outside-toplevel + from tpu_inference.layers.common.ragged_gated_delta_rule_wrapper import RaggedGatedDeltaRuleImpl # pylint: disable=import-outside-toplevel + from tpu_inference.layers.common.sharding import ShardingAxisName # pylint: disable=import-outside-toplevel + from tpu_inference.layers.common.utils import reorder_concatenated_tensor_for_sharding # pylint: disable=import-outside-toplevel + from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel + from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel except ImportError as e: raise ImportError( "GDN attention kernel require the vllm-tpu package. Please install it with `pip install vllm-tpu`." @@ -702,6 +710,9 @@ def __call__( conv_state_paged, recurrent_state_paged = kv_cache + # Use REF impl (pure JAX) to avoid Mosaic kernel compilation issues. + gdn_config = GdnAttentionConfig(ragged_gated_delta_rule_impl=RaggedGatedDeltaRuleImpl.REF) + (new_conv_state_paged, new_recurrent_state_paged), gdn_output = run_jax_gdn_attention( mixed_qkv, b_flat, @@ -712,20 +723,22 @@ def __call__( None, # conv_bias: MaxText conv1d uses use_bias=False. jnp.asarray(self.A_log[...], dtype=cfg.dtype), jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), - attention_metadata.mamba_state_indices.astype(jnp.int32), # pyrefly: ignore[missing-attribute] - attention_metadata.query_start_loc, # pyrefly: ignore[missing-attribute] - attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] - attention_metadata.seq_lens, # pyrefly: ignore[missing-attribute] + attention_metadata.mamba_state_indices.astype(jnp.int32), + attention_metadata.query_start_loc, + attention_metadata.request_distribution, + attention_metadata.seq_lens, self.num_k_heads, self.num_v_heads, self.head_k_dim, self.head_v_dim, cfg.gdn_conv_kernel_dim, mesh=self.mesh, + config=gdn_config, ) # Reshape GDN output and apply gated norm + out projection. gdn_output = gdn_output.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) + gdn_output = checkpoint_name(gdn_output, "context") gated_output = self.norm(gdn_output, z) gated_output = gated_output.reshape(batch, seq_len, -1) output = self.out_proj(gated_output) @@ -741,7 +754,7 @@ def __call__( v = value.reshape(batch, seq_len, -1) # ========================================================================= - # STEP B: 1D Convolution + # STEP B & C: 1D Convolution & Gated Delta Rule Recurrence # ========================================================================= qkv = jnp.concatenate([q, k, v], axis=-1) batch, seq_len, _ = qkv.shape @@ -754,7 +767,6 @@ def __call__( recurrent_state, conv_state = active_cache.get_gdn_states() orig_cache_batch = conv_state.shape[0] - # 1. Safely shrink/expand conv_state to match incoming qkv (e.g. 16 -> 1) if conv_state.shape[0] != batch: if conv_state.shape[0] == 1: conv_state = jnp.broadcast_to(conv_state, (batch,) + conv_state.shape[1:]) @@ -764,7 +776,6 @@ def __call__( else: conv_state = conv_state[:batch] - # 2. Safely shrink/expand recurrent_state to match incoming qkv if recurrent_state.shape[0] != batch: if recurrent_state.shape[0] == 1: recurrent_state = jnp.broadcast_to(recurrent_state, (batch,) + recurrent_state.shape[1:]) @@ -774,130 +785,226 @@ def __call__( else: recurrent_state = recurrent_state[:batch] - conv_input = jnp.concatenate([conv_state, qkv], axis=1) - - if decoder_segment_ids is not None: - valid_lens = jnp.sum(decoder_segment_ids != 0, axis=1) - - def extract_state(c_in, v_len): - return jax.lax.dynamic_slice_in_dim(c_in, v_len, conv_kernel_size - 1, axis=0) + if getattr(cfg, "use_gdn_kernel", False) and getattr(cfg, "use_hybrid_gdn", False): + from maxtext.models.hybrid_gdn import hybrid_fused_conv1d_gdn - next_conv_state = jax.vmap(extract_state)(conv_input, valid_lens) - else: - next_conv_state = conv_input[:, -(conv_kernel_size - 1) :, :] - else: - conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) - - # Perform the convolution. - conv_out = self.conv1d(conv_input) - # Slice the output to match the original input sequence length. - conv_out = conv_out[:, -seq_len:, :] - qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) - # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) - q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) - - # Reshape for multi-head processing - # query shape: (B, S, H_k, D_k) - query = q_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) - # key shape: (B, S, H_k, D_k) - key = k_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) - # value shape: (B, S, H_v, D_v) - value = v_conv.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) - - # ========================================================================= - # STEP C: Gated Delta Rule Recurrence - # ========================================================================= - A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) - dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) - # beta shape: (B, S, H_v) - beta = jax.nn.sigmoid(b) - # g shape: (B, S, H_v) - g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) - - if decoder_segment_ids is not None: - mask = decoder_segment_ids != 0 - # Apply mask by broadcasting to respective shapes - key = jnp.where(mask[..., None, None], key, 0.0) - value = jnp.where(mask[..., None, None], value, 0.0) - g = jnp.where(mask[..., None], g, 0.0) - - if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: - repeats = self.num_v_heads // self.num_k_heads - # query shape after repeat: (B, S, H_v, D_k) - query = jnp.repeat(query, repeats, axis=2) - # key shape after repeat: (B, S, H_v, D_k) - key = jnp.repeat(key, repeats, axis=2) - - if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: - core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( - query, - key, - value, - g, - beta, - initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, + conv_state_arg = ( + conv_state + if conv_state is not None + else jnp.zeros((batch, self.config.gdn_conv_kernel_dim - 1, qkv.shape[-1]), dtype=cfg.dtype) ) - elif self.mesh is not None: - logical_rules = get_logical_axis_rules() recurrent_state_arg = ( recurrent_state if recurrent_state is not None else jnp.zeros((batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), dtype=cfg.dtype) ) - qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) - g_beta_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD), mesh=self.mesh, rules=logical_rules) - state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) - - @functools.partial( - jax.shard_map, - mesh=self.mesh, - in_specs=( - qkv_pspec, # query - qkv_pspec, # key - qkv_pspec, # value - g_beta_pspec, # g - g_beta_pspec, # beta - state_pspec, # initial_state - ), - out_specs=( - qkv_pspec, # core_attn_out - state_pspec, # final_state - ), - check_vma=False, + conv_bias_arg = ( + self.conv1d.bias.value + if hasattr(self.conv1d, "bias") and self.conv1d.bias is not None + else jnp.zeros((qkv.shape[-1],), dtype=cfg.dtype) ) - def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): - return jax_chunk_gated_delta_rule( - query=q, - key=k, - value=v, - g=g_val, - beta=beta_val, + + if self.mesh is not None: + logical_rules = get_logical_axis_rules() + batch_pspec3 = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) + batch_pspec4 = logical_to_mesh_axes((KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules) + none_pspec3 = logical_to_mesh_axes((None, None, None), mesh=self.mesh, rules=logical_rules) + none_pspec1 = logical_to_mesh_axes((None,), mesh=self.mesh, rules=logical_rules) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + batch_pspec3, # qkv + batch_pspec3, # b + batch_pspec3, # a + none_pspec3, # conv_weight + none_pspec1, # conv_bias + none_pspec1, # a_log + none_pspec1, # dt_bias + batch_pspec3, # conv_state + batch_pspec4, # recurrent_state + ), + out_specs=( + batch_pspec4, # core_attn_out + (batch_pspec3, batch_pspec4), # (next_conv_state, next_recurrent_state) + ), + check_vma=False, + ) + def shard_mapped_hybrid_gdn(qkv_val, b_val, a_val, cw_val, cb_val, alog_val, dt_val, cs_val, rs_val): + return hybrid_fused_conv1d_gdn( + qkv=qkv_val, + b=b_val, + a=a_val, + conv_weight=cw_val, + conv_bias=cb_val, + a_log=alog_val, + dt_bias=dt_val, + conv_state=cs_val, + recurrent_state=rs_val, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=self.config.dtype, + ) + + core_attn_out, (next_conv_state, next_recurrent_state) = shard_mapped_hybrid_gdn( + qkv, + b, + a, + self.conv1d.kernel.value, + conv_bias_arg, + self.A_log[...], + self.dt_bias[...], + conv_state_arg, + recurrent_state_arg, + ) + else: + core_attn_out, (next_conv_state, next_recurrent_state) = hybrid_fused_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_weight=self.conv1d.kernel.value, + conv_bias=None, + a_log=self.A_log[...], + dt_bias=self.dt_bias[...], + conv_state=conv_state_arg, + recurrent_state=recurrent_state_arg, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=self.config.dtype, + ) + else: + if conv_state is not None: + conv_input = jnp.concatenate([conv_state, qkv], axis=1) + if decoder_segment_ids is not None: + valid_lens = jnp.sum(decoder_segment_ids != 0, axis=1) + + def extract_state(c_in, v_len): + return jax.lax.dynamic_slice_in_dim(c_in, v_len, conv_kernel_size - 1, axis=0) + + next_conv_state = jax.vmap(extract_state)(conv_input, valid_lens) + else: + next_conv_state = conv_input[:, -(conv_kernel_size - 1) :, :] + else: + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + + conv_out = self.conv1d(conv_input) + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) + + query = q_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) + key = k_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) + value = v_conv.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) + + A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) + dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) + beta = jax.nn.sigmoid(b) + g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) + + if decoder_segment_ids is not None: + mask = decoder_segment_ids != 0 + key = jnp.where(mask[..., None, None], key, 0.0) + value = jnp.where(mask[..., None, None], value, 0.0) + g = jnp.where(mask[..., None], g, 0.0) + + if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: + repeats = self.num_v_heads // self.num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: + core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( + query, + key, + value, + g, + beta, + initial_state=recurrent_state, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) + elif getattr(cfg, "use_gdn_kernel", False): + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, chunk_size=cfg.gdn_chunk_size, - initial_state=init_h, + initial_state=recurrent_state, use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=cfg.dtype, ) + elif self.mesh is not None: + logical_rules = self.config.logical_axis_rules + recurrent_state_arg = ( + recurrent_state + if recurrent_state is not None + else jnp.zeros((batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), dtype=cfg.dtype) + ) + qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) + g_beta_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD), mesh=self.mesh, rules=logical_rules) + state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, # query + qkv_pspec, # key + qkv_pspec, # value + g_beta_pspec, # g + g_beta_pspec, # beta + state_pspec, # initial_state + ), + out_specs=( + qkv_pspec, # core_attn_out + state_pspec, # final_state + ), + check_vma=False, + ) + def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): + return jax_chunk_gated_delta_rule( + query=q, + key=k, + value=v, + g=g_val, + beta=beta_val, + chunk_size=cfg.gdn_chunk_size, + initial_state=init_h, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) - core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) - else: - core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( - query, - key, - value, - g, - beta, - chunk_size=cfg.gdn_chunk_size, - initial_state=recurrent_state, - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, - ) + core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) + else: + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, + chunk_size=cfg.gdn_chunk_size, + initial_state=recurrent_state, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) if model_mode != MODEL_MODE_TRAIN and active_cache is not None: assert next_conv_state is not None assert next_recurrent_state is not None - if next_conv_state.shape[0] != orig_cache_batch: # pyrefly: ignore[unbound-name] + if next_conv_state.shape[0] != orig_cache_batch: if next_conv_state.shape[0] == 1: next_conv_state = jnp.broadcast_to(next_conv_state, (orig_cache_batch,) + next_conv_state.shape[1:]) next_recurrent_state = jnp.broadcast_to( @@ -912,7 +1019,9 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): next_recurrent_state = next_recurrent_state[:orig_cache_batch] if model_mode != MODEL_MODE_TRAIN and active_cache is not None: - active_cache.update_gdn_states(next_recurrent_state, next_conv_state) # pyrefly: ignore[bad-argument-type] + active_cache.update_gdn_states(next_recurrent_state, next_conv_state) + + core_attn_out = checkpoint_name(core_attn_out, "context") # ========================================================================= # STEP D: Final Output Stage @@ -948,12 +1057,12 @@ def init_kv_caches(self, batch_size: int): value_heads=self.num_v_heads, key_head_size=self.head_k_dim, value_head_size=self.head_v_dim, - dtype=self.dtype, # pyrefly: ignore[missing-attribute] + dtype=self.dtype, is_gdn=True, conv_kernel_size=conv_kernel_size, conv_dim=conv_dim, - model_mode=self.model_mode, # pyrefly: ignore[missing-attribute] - rngs=self.rngs, # pyrefly: ignore[missing-attribute] + model_mode=self.model_mode, + rngs=self.rngs, ) @@ -1081,12 +1190,13 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - # 2. Instantiate and apply the shared expert. + # 2. Instantiate and apply the shared expert(s). + shared_expert_mlp_dim = maxtext_utils.get_shared_expert_mlp_dim(cfg) self.shared_expert = MlpBlock( config=cfg, mesh=mesh, in_features=cfg.emb_dim, - intermediate_dim=cfg.moe_mlp_dim, + intermediate_dim=cfg.shared_experts * shared_expert_mlp_dim, activations=cfg.mlp_activations, intermediate_dropout_rate=cfg.dropout_rate, dtype=cfg.dtype, @@ -1096,19 +1206,22 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - # 3. Instantiate and apply the gate for the shared expert. - self.shared_expert_gate = DenseGeneral( - in_features_shape=cfg.emb_dim, - out_features_shape=1, - use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias - dtype=cfg.dtype, - kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"), - kernel_axes=("embed", None), - matmul_precision=cfg.matmul_precision, - rngs=rngs, - ) + # 3. Instantiate the (optional) gate for the shared expert. + if cfg.moe_shared_expert_gate: + self.shared_expert_gate = DenseGeneral( + in_features_shape=cfg.emb_dim, + out_features_shape=1, + use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias + dtype=cfg.dtype, + kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"), + kernel_axes=("embed", None), + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + else: + self.shared_expert_gate = None - def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Array | None]: + def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Array | None, Array | None]: """ Applies the sparse MoE block to the input hidden states. @@ -1120,103 +1233,231 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar A tuple containing: - The output array of the MoE block. - The load balancing loss from the routed experts, if applicable during training. + - The aux-loss-free expert-bias updates from the routed experts, if applicable. """ # 1. Apply the routed experts block. - routed_output, load_balance_loss, _ = self.routed_experts(hidden_states) + routed_output, load_balance_loss, moe_bias_updates = self.routed_experts(hidden_states) # 2. Apply the shared expert. shared_expert_output = self.shared_expert(hidden_states, deterministic=deterministic) - # 3. Apply the gate for the shared expert. - shared_gate_output = self.shared_expert_gate(hidden_states) + # 3. Apply the (optional) gate for the shared expert. + if self.shared_expert_gate is not None: + shared_gate_output = self.shared_expert_gate(hidden_states) + shared_expert_output = jax.nn.sigmoid(shared_gate_output) * shared_expert_output # 4. Combine the outputs. - final_output = routed_output + jax.nn.sigmoid(shared_gate_output) * shared_expert_output + final_output = routed_output + shared_expert_output - return final_output, load_balance_loss + return final_output, load_balance_loss, moe_bias_updates class Qwen3NextScannableBlock(nnx.Module): - """A scannable block of Qwen3-Next decoder layers. + """A scannable block of Qwen3-Next decoder layers with hierarchical nested scans. - This module contains a fixed number of heterogeneous decoder layers that form - a repeating pattern, as defined by `config.inhomogeneous_layer_cycle_interval`. It is - intended to be the body of an `nn.scan` transformation to construct the full - decoder stack efficiently. - - Attributes: - config: The model configuration object. - mesh: The device mesh for sharding. - model_mode: The operational mode (e.g., 'train', 'prefill'). - quant: Optional quantization configuration. + Linear attention layers (local) are scanned via `nnx_scan.apply_scanned_layers` + while full attention (global) is scanned via a length-1 `jax.lax.scan`. """ - def __init__(self, config: Config, mesh: Mesh, model_mode: str, quant: None | Quant = None, *, rngs: nnx.Rngs): + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + quant: None | Quant = None, + *, + num_of_layers: int | None = None, + layer_idx_offset: int = 0, + remat_policy_fn: Any | None = None, + apply_internal_remat: bool = False, + rngs: nnx.Rngs, + ): self.config = config self.mesh = mesh self.model_mode = model_mode self.quant = quant self.rngs = rngs + self.remat_policy_fn = remat_policy_fn + self.apply_internal_remat = apply_internal_remat cfg = self.config + if num_of_layers is None: + num_of_layers = cfg.inhomogeneous_layer_cycle_interval + self.num_of_layers = num_of_layers + self.layer_idx_offset = layer_idx_offset + + cycle_interval = cfg.inhomogeneous_layer_cycle_interval + full_attention_offset = getattr(cfg, "full_attention_layer_offset", 0) % cycle_interval + + self.num_local = sum(1 for i in range(num_of_layers) if (layer_idx_offset + i) % cycle_interval != full_attention_offset) + self.num_global = sum(1 for i in range(num_of_layers) if (layer_idx_offset + i) % cycle_interval == full_attention_offset) + + if self.num_local > 0: + self.local_layers = nnx_scan.create_scanned_layers( + lambda layer_rngs: Qwen3NextDecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + layer_idx=0, + is_dense_layer=False, + is_full_attention_layer=False, + rngs=layer_rngs, + ), + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + metadata_axis_name="local_layers", + rngs=self.rngs, + ) + else: + self.local_layers = None - # Instantiate each layer within the block in __init__ - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer_rngs = self.rngs.fork() # Fork RNGs for each layer - layer_name = f"layer_{i}" - layer = Qwen3NextDecoderLayer( + if self.num_global > 0: + self.global_layer = Qwen3NextDecoderLayer( config=self.config, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, - layer_idx=i, - rngs=layer_rngs, + layer_idx=full_attention_offset, + is_dense_layer=False, + is_full_attention_layer=True, + rngs=self.rngs, ) - setattr(self, layer_name, layer) + else: + self.global_layer = None + + def _run_layer(self, layer, y, layer_kwargs, kv_cache=None): + """Invokes one Qwen3NextDecoderLayer, returning (output, updated_kv_cache).""" + out = layer(y, **layer_kwargs, kv_cache=kv_cache) + return out if isinstance(out, tuple) else (out, None) + + @property + def _remat_enabled(self): + """Whether the block rematerializes its own layers.""" + return self.apply_internal_remat and self.config.remat_policy != "none" + + def _scan_local_layers(self, y, layer_kwargs): + """Runs the local (linear attention / GatedDeltaNet) layers via a per-layer rematerialized jax.lax.scan.""" + remat = self._remat_enabled + return nnx_scan.apply_scanned_layers( + self.local_layers, + y, + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + apply_fn=lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0], + remat=remat, + remat_policy=self.remat_policy_fn if remat else None, + prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config) if remat else True, + ) + + def _scan_global_layer(self, y, layer_kwargs): + """Runs the single global-attention layer inside a length-1 jax.lax.scan.""" + cfg = self.config + graphdef_g, intermediate_g, other_g = nnx.split(self.global_layer, nnx.Intermediate, ...) + intermediate_xs = jax.tree.map(lambda x: x[None], intermediate_g) + + def run_global_layer(carry, intermediate_slice): + hidden_states, other = carry + layer = nnx.merge(graphdef_g, intermediate_slice, other) + new_hidden_states = self._run_layer(layer, hidden_states, layer_kwargs)[0] + _, new_intermediate, new_other = nnx.split(layer, nnx.Intermediate, ...) + return (new_hidden_states, new_other), new_intermediate + + global_remat_policy = self.remat_policy_fn + offload_names = maxtext_utils.get_save_and_offload_names(cfg) + if offload_names[0] or offload_names[1]: + save_names, offload_to_device = offload_names + global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device)) + + if self._remat_enabled: + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config) + run_global_layer = jax.checkpoint( + run_global_layer, + policy=global_remat_policy, + prevent_cse=prevent_cse, + ) + + with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}): + (y, final_other), stacked_intermediate = jax.lax.scan( + run_global_layer, + (y, other_g), + intermediate_xs, + length=1, + ) + + intermediate_state = jax.tree.map(lambda x: x[0], stacked_intermediate) + nnx.update(self.global_layer, final_other, intermediate_state) + return y + + def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): + """Runs the block with externally-supplied per-layer kv caches.""" + updated_kvs = [] + if self.local_layers is not None: + graphdef, params, state = nnx.split(self.local_layers, nnx.Param, ...) + scan_axis = self.config.param_scan_axis + if scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params) + per_layer_states = [] + for i in range(self.num_local): + current_params = jax.tree.map(lambda x, i=i: x[i], params) + current_state = jax.tree.map(lambda x, i=i: x[i], state) + layer = nnx.merge(graphdef, current_params, current_state) + current_kv = kv_cache[i] if (kv_cache is not None and i < len(kv_cache)) else None + y, new_kv = self._run_layer(layer, y, layer_kwargs, current_kv) + updated_kvs.append(new_kv) + per_layer_states.append(nnx.state(layer)) + + stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) + if scan_axis != 0: + stacked_params, stacked_other = stacked_state.split(nnx.Param, ...) + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params) + stacked_state = nnx.State.merge(stacked_params, stacked_other) + nnx.update(self.local_layers, stacked_state) + + if self.global_layer is not None: + global_kv = kv_cache[self.num_local] if (kv_cache is not None and self.num_local < len(kv_cache)) else None + y, new_kv = self._run_layer(self.global_layer, y, layer_kwargs, global_kv) + updated_kvs.append(new_kv) + + return y, tuple(updated_kvs) def __call__( self, carry: jnp.ndarray, - decoder_segment_ids: None | jnp.ndarray, - decoder_positions: None | jnp.ndarray, - deterministic: bool, - model_mode: str, + decoder_segment_ids: None | jnp.ndarray = None, + decoder_positions: None | jnp.ndarray = None, + deterministic: bool = False, + model_mode: str = "train", previous_chunk=None, slot: None | int = None, kv_cache=None, attention_metadata=None, ) -> tuple[Array, None]: - """Applies the block of decoder layers to the input carry. + cfg = self.config + inputs = carry + inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + + layer_kwargs = { + "decoder_segment_ids": decoder_segment_ids, + "decoder_positions": decoder_positions, + "deterministic": deterministic, + "model_mode": model_mode, + "slot": slot, + "previous_chunk": previous_chunk, + "attention_metadata": attention_metadata, + } - Args: - carry: The input tensor from the previous scan iteration. - # ... other arguments are broadcasted to each iteration. + if kv_cache is not None: + return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs) - Returns: - A tuple containing the output of the block (the new carry) and an empty - value for the scan's `y` collection. - """ - cfg = self.config - x = carry - - # Loop over the number of sub-layers that make up one repeating pattern. - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer = getattr(self, f"layer_{i}") - # The second return value is kv_cache, which we ignore here because - # it is not passed as a carry in scannable layers. - x, _ = layer( - x, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - ) + y = inputs + if self.local_layers is not None: + y = self._scan_local_layers(y, layer_kwargs) + if self.global_layer is not None: + y = self._scan_global_layer(y, layer_kwargs) - # The output of the block is the carry for the next scan iteration. - return x, None + if cfg.scan_layers: + return y, None + return y class Qwen3NextDecoderLayer(nnx.Module): @@ -1225,9 +1466,9 @@ class Qwen3NextDecoderLayer(nnx.Module): 1. A standard attention + MoE layer. 2. A linear attention + MoE layer. - NOTE: This implementation assumes every layer contains a MoE block, which is true for - models like Qwen3-Next-80B-A3B where `decoder_sparse_step=1`. For models that - interleave dense and sparse MLP layers, conditional logic would be needed here. + The first `config.first_num_dense_layers` layers (by `layer_idx`) use a plain + dense MLP instead of MoE, and always use full attention, mirroring DeepSeek V3's + dense-prefix pattern (see `models/deepseek.py::DeepSeekDenseLayer`). Attributes: config: The model configuration object. @@ -1238,7 +1479,16 @@ class Qwen3NextDecoderLayer(nnx.Module): """ def __init__( - self, config: Config, mesh: Mesh, model_mode: str, layer_idx: int, quant: None | Quant = None, *, rngs: nnx.Rngs + self, + config: Config, + mesh: Mesh, + model_mode: str, + layer_idx: int, + quant: None | Quant = None, + *, + is_dense_layer: bool | None = None, + is_full_attention_layer: bool | None = None, + rngs: nnx.Rngs, ): self.config = config self.mesh = mesh @@ -1247,6 +1497,11 @@ def __init__( self.quant = quant cfg = self.config self.activation_axis_names = ("activation_batch", "activation_norm_length", "activation_embed") + self.is_mhc_enabled = cfg.mhc_expansion_rate > 1 + + if is_dense_layer is None: + is_dense_layer = layer_idx < cfg.first_num_dense_layers + self.is_dense_layer = is_dense_layer # First LayerNorm, applied before the attention block. self.input_layernorm = Qwen3NextRMSNorm( @@ -1257,8 +1512,14 @@ def __init__( rngs=rngs, ) - # Determine the type of attention mechanism for the current layer. - is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 + # Determine the type of attention mechanism for the current layer. Dense layers + # always use full attention (see class docstring). `full_attention_layer_offset` + # picks which position in the cycle is full attention; -1 (Python's negative-modulo + # wraps to cycle-1) reproduces the original "last position in the cycle" schedule. + full_attention_offset = cfg.full_attention_layer_offset % cfg.inhomogeneous_layer_cycle_interval + if is_full_attention_layer is None: + is_full_attention_layer = self.is_dense_layer or self.layer_idx % cfg.inhomogeneous_layer_cycle_interval == full_attention_offset + self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: @@ -1286,8 +1547,78 @@ def __init__( rngs=rngs, ) - # Instantiate our `Qwen3NextSparseMoeBlock`. - self.mlp = Qwen3NextSparseMoeBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + # Dense layers use a plain MLP; all other layers use `Qwen3NextSparseMoeBlock`. + if self.is_dense_layer: + self.mlp = MlpBlock( + in_features=cfg.emb_dim, + intermediate_dim=cfg.mlp_dim, + activations=cfg.mlp_activations, + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + config=cfg, + mesh=self.mesh, + quant=self.quant, + model_mode=model_mode, + rngs=rngs, + ) + else: + self.mlp = Qwen3NextSparseMoeBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + + # Manifold-Constrained Hyper Connections: replaces the plain residual add around the + # attention and MoE branches with a learned multi-stream mixing. See maxtext/layers/mhc.py + # and models/deepseek4.py::DeepSeek4DecoderLayer for the reference implementation. + if self.is_mhc_enabled: + self.mhc_attention = mhc.ManifoldConstrainedHyperConnections(cfg, cfg.emb_dim, self.mesh, rngs) + self.mhc_mlp = mhc.ManifoldConstrainedHyperConnections(cfg, cfg.emb_dim, self.mesh, rngs) + + def pre_attention_norm_op(self, x): + normed = self.input_layernorm(x) + return nn.with_logical_constraint(normed, self.activation_axis_names) + + def post_attention_norm_op(self, x): + normed = self.post_attention_layernorm(x) + return nn.with_logical_constraint(normed, self.activation_axis_names) + + def attention_branch( + self, + inputs_q, + inputs_kv=None, + decoder_segment_ids=None, + inputs_positions=None, + deterministic=None, + model_mode=None, + kv_cache=None, + attention_metadata=None, + **kwargs, + ): + """Adapts Qwen3-Next's two attention variants to mHC's inputs_q/inputs_kv branch_fn convention.""" + del inputs_kv, kwargs + if isinstance(self.attention, Qwen3NextFullAttention): + return self.attention( + inputs_q, + decoder_segment_ids, + inputs_positions, + deterministic, + model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + return self.attention( + inputs_q, + model_mode=model_mode, + kv_cache=kv_cache, + decoder_segment_ids=decoder_segment_ids, + attention_metadata=attention_metadata, + ) + + def mlp_op(self, inputs, deterministic, *args, **kwargs): + """Adapts the dense/MoE MLP's return shape to mHC's MLP_MOE 3-tuple convention.""" + del args, kwargs + if self.is_dense_layer: + return self.mlp(inputs, deterministic=deterministic), None, None + mlp_out, load_balance_loss, moe_bias_updates = self.mlp(inputs, deterministic=deterministic) + return mlp_out, load_balance_loss, moe_bias_updates def __call__( self, @@ -1304,6 +1635,45 @@ def __call__( # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): inputs = inputs[0] + + inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + inputs = checkpoint_name(inputs, "decoder_layer_input") + + if self.is_mhc_enabled: + mhc_expand, mhc_reduce = mhc.get_functions(self.config.mhc_expansion_rate) + inputs = mhc_expand(inputs) + + intermediate_inputs, _ = self.mhc_attention( + self.pre_attention_norm_op, + self.attention_branch, + x=inputs, + mhc_type=HyperConnectionType.ATTENTION, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=deterministic, + model_mode=model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + + layer_output, metadata = self.mhc_mlp( + self.post_attention_norm_op, + self.mlp_op, + x=intermediate_inputs, + mhc_type=HyperConnectionType.MLP_MOE, + deterministic=deterministic, + ) + load_balance_loss = metadata.get("load_balance_loss", None) + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.moe_lb_loss = nnx.Intermediate(load_balance_loss) + moe_bias_updates = metadata.get("moe_bias_updates", None) + if self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: + self.moe_bias_updates = nnx.Intermediate(moe_bias_updates) + + layer_output = mhc_reduce(layer_output) + layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names) + return layer_output, kv_cache + residual = inputs # First LayerNorm, applied before the attention block. @@ -1318,7 +1688,7 @@ def __call__( decoder_positions, deterministic, model_mode, - kv_cache=kv_cache, # pyrefly: ignore[bad-argument-type] + kv_cache=kv_cache, attention_metadata=attention_metadata, ) else: @@ -1341,13 +1711,17 @@ def __call__( hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) - # Instantiate and call our `Qwen3NextSparseMoeBlock`. - mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic) - - # We sow the load balancing loss so it can be collected and added to the total loss - # during training. - if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: - self.moe_lb_loss = nnx.Intermediate(load_balance_loss) + # Apply the dense MLP or `Qwen3NextSparseMoeBlock`. + if self.is_dense_layer: + mlp_output = self.mlp(hidden_states, deterministic=deterministic) + else: + mlp_output, load_balance_loss, moe_bias_updates = self.mlp(hidden_states, deterministic=deterministic) + # We sow the load balancing loss so it can be collected and added to the total loss + # during training. + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.moe_lb_loss = nnx.Intermediate(load_balance_loss) + if self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: + self.moe_bias_updates = nnx.Intermediate(moe_bias_updates) # Final residual connection (after the MoE block) layer_output = residual + mlp_output @@ -1606,8 +1980,8 @@ def update_cache(cache, val): return cache.at[layer_idx].set(val) return cache - stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) # pyrefly: ignore[unbound-name] - return (layer_output, stacked_kv_cache, layer_idx + 1), None # pyrefly: ignore[unbound-name] + stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) + return (layer_output, stacked_kv_cache, layer_idx + 1), None else: return layer_output, kv_cache @@ -1635,7 +2009,7 @@ def __init__( dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init(1.0, "fan_in", "normal"), - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision patch merger. @@ -1749,7 +2123,7 @@ def __init__( dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init(1.0, "fan_in", "normal"), - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision MLP. @@ -1826,7 +2200,7 @@ def __init__( # Default to float32 for numerical stability in 3D convolutions on image/video inputs dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision patch embedding. @@ -1878,8 +2252,8 @@ def __call__(self, hidden_states: Array, video_mask: Array | None = None) -> tup attention_mask = None if video_mask is not None: - mask_patch_elements = self.temporal_patch_size * self.patch_size * self.patch_size - attention_mask = video_mask.reshape(video_mask.shape[0], -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) + patch_mask = video_mask[:, 0, :: self.temporal_patch_size, :: self.patch_size, :: self.patch_size] + attention_mask = patch_mask.reshape(video_mask.shape[0], -1).astype(jnp.int32) return hidden_states, attention_mask @@ -1892,7 +2266,7 @@ class Qwen3OmniMoeVisionAttention(nnx.Module): attn: Underlying attention module """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision attention layer. Args: @@ -1916,7 +2290,7 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyre float32_logits=self.config.float32_logits, dtype=self.config.dtype_mm, weight_dtype=self.config.weight_dtype, - mesh=mesh, # pyrefly: ignore[bad-argument-type] + mesh=mesh, dropout_rate=0.0, attention_type=AttentionType.FULL, is_nope_layer=False, @@ -1934,8 +2308,6 @@ def __call__( num_frames: int, height: int, width: int, - attention_mask: Array | None = None, - valid_grid: tuple[int, int, int] | None = None, deterministic: bool = True, ) -> Array: """ @@ -1944,8 +2316,6 @@ def __call__( num_frames: Number of temporal frames (static) height: Height in patches (static) width: Width in patches (static) - attention_mask: Optional mask identifying valid tokens in the padded sequence. - valid_grid: Optional unpadded `(frames, height, width)` grid used for vision RoPE. deterministic: Whether to use deterministic mode (disable dropout) Returns: @@ -1956,14 +2326,11 @@ def __call__( "num_frames": num_frames, "height": height, "width": width, - "token_mask": attention_mask, - "valid_grid": valid_grid, } output, _ = self.attn( inputs_q=hidden_states, inputs_kv=hidden_states, deterministic=deterministic, - decoder_segment_ids=attention_mask, rope_kwargs=rope_kwargs, ) @@ -1982,7 +2349,7 @@ class Qwen3OmniMoeVisionBlock(nnx.Module): mlp_out: Second MLP layer """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision transformer block. Args: @@ -2016,8 +2383,6 @@ def __call__( num_frames: int, height: int, width: int, - attention_mask: Array | None = None, - valid_grid: tuple[int, int, int] | None = None, ) -> Array: """ Args: @@ -2029,14 +2394,7 @@ def __call__( Returns: Output tensor of shape (batch, T*H*W, hidden_size) """ - x = x + self.attn( - self.ln1(x), - num_frames=num_frames, - height=height, - width=width, - attention_mask=attention_mask, - valid_grid=valid_grid, - ) + x = x + self.attn(self.ln1(x), num_frames=num_frames, height=height, width=width) y = self.ln2(x) y = self.mlp(y) y = jax.nn.gelu(y) @@ -2057,7 +2415,7 @@ class Qwen3OmniMoeVisionEncoder(nnx.Module): deep_idx: Indices of layers to extract deep features from """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision encoder. Args: @@ -2097,8 +2455,6 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyre def __call__( self, hidden_states: Array, - video_mask: Array | None = None, - video_grid_thw: Array | tuple[int, int, int] | None = None, deterministic: bool = True, ): """ @@ -2115,12 +2471,6 @@ def __call__( num_frames = num_frames // self.config.temporal_patch_size_for_vit height = height // self.config.patch_size_for_vit width = width // self.config.patch_size_for_vit - attention_mask = None - if video_mask is not None: - mask_patch_elements = ( - self.config.temporal_patch_size_for_vit * self.config.patch_size_for_vit * self.config.patch_size_for_vit - ) - attention_mask = video_mask.reshape(batch_size, -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) hidden_states = hidden_states.reshape( -1, self.config.num_channels_for_vit, @@ -2131,30 +2481,16 @@ def __call__( x, _ = self.patch_embed(hidden_states) x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit) - if attention_mask is not None and video_grid_thw is None: - raise ValueError("video_grid_thw is required when video_mask is provided.") - pos = self.pos_embed_interpolate( - num_frames, - height, - width, - video_grid_thw=video_grid_thw, # pyrefly: ignore[bad-argument-type] - attention_mask=attention_mask, - ) + pos = self.pos_embed_interpolate(num_frames, height, width) + + pos = pos[jnp.newaxis, :, :] x = x + pos - valid_grid = video_grid_thw h_traj = [] for i in range(self.depth): block_name = f"blocks_{i}" blk = getattr(self, block_name) - x = blk( - x, - num_frames=num_frames, - height=height, - width=width, - attention_mask=attention_mask, - valid_grid=valid_grid, - ) + x = blk(x, num_frames=num_frames, height=height, width=width) h_traj.append(x) deep_feats = [] @@ -2176,7 +2512,7 @@ class Qwen3OmniMoeVisionProjector(nnx.Module): merger: Patch merger for spatial reduction """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision projector. Args: @@ -2224,7 +2560,7 @@ def qwen3omni_visionprojector_as_linen(config: Config, mesh: Mesh) -> nn.Module: class Qwen3OmniAudioEncoderLayer(nnx.Module): """Transformer encoder layer for audio model.""" - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): self.config = config self.mesh = mesh self.rngs = rngs @@ -2329,7 +2665,7 @@ class Qwen3OmniAudioEncoder(nnx.Module): mesh: Mesh, JAX device mesh (used for sharding) """ - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): self.config = config self.mesh = mesh self.rngs = rngs @@ -2480,7 +2816,7 @@ def __call__( class Qwen3OmniAudioProjector(nnx.Module): """Projection layer that converts audio encoder output to model embedding space.""" - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): self.config = config self.proj1 = DenseGeneral( in_features_shape=config.d_model_for_audio, diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 331df17f41..0dc9d7ce3f 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -144,6 +144,7 @@ def __init__( # Determine the type of attention mechanism for the current layer. is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 + self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 67e1f589ca..de288b24a2 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -204,7 +204,7 @@ def get_optimizer(config, learning_rate_schedule, model=None): ns_steps = 10 else: ns_coeffs = (3.4445, -4.7750, 2.0315) - ns_steps = 5 + ns_steps = getattr(config, "muon_ns_steps", 5) muon_kwargs = { # Shared parameters: "nesterov" uses default diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 7729fffd0a..1379adeabe 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -309,20 +309,25 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, nested_key, None) else: # NNX intermediates are model-rooted (no "intermediates" prefix), so match by - # suffix instead. Unlike collect_intermediates_by_suffix we must not ravel: - # the update is a 2-D matrix that's transposed at the apply site below. - moe_bias_updates = next( - ( - val - for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs) - if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ("moe_bias_updates",) - ), - None, - ) - if moe_bias_updates is not None: - # The Linen path returns the sow tuple and indexes [0] downstream; tree_leaves - # already descended that tuple, so wrap it back so the apply site is uniform. - moe_bias_updates = (moe_bias_updates,) + # suffix instead. A decoder block may sow more than one moe_bias_updates leaf + # (e.g. Qwen3-Next sows one per layer position inside each scanned block, since + # its MoE gates aren't a single homogeneous scanned collection like DeepSeek's), + # so collect every match keyed by its path, joined into a single string -- the + # path is used downstream (maxtext_utils_nnx.apply_moe_bias_updates) to locate + # the matching gate.bias parameter generically. This must be a dict (path string + # -> array), not a list of (path, array) pairs: this aux value flows through + # jax.lax.scan under gradient accumulation (gradient_accumulation.py), which + # requires every leaf to be a valid JAX array -- a plain string leaf (as part of + # a path tuple) would break that, whereas dict keys are pytree structure, not + # leaves, so they pass through untouched. Unlike collect_intermediates_by_suffix + # we must not ravel: each update is a 2-D matrix transposed at the apply site. + moe_bias_updates = { + "/".join(tuple(k.key for k in path if hasattr(k, "key"))): val + for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs) + if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ("moe_bias_updates",) + } + if not moe_bias_updates: + moe_bias_updates = None # Add the model's primary output to the intermediates dict so it can be used # by the acceptance rate calculation in eval_step. @@ -562,7 +567,7 @@ def move(path, value): state.apply_gradients(grads) new_state = state - # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family + # Apply updates for Auxiliary-Loss-Free load balancing # pylint: disable=too-many-nested-blocks if config.routed_bias and config.routed_bias_update_rate > 0.0: if config.model_name.startswith("deepseek4"): @@ -593,8 +598,7 @@ def move(path, value): if getattr(config, "log_moe_bias_norms", False): bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) elif moe_bias_updates is not None: - target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias - target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose() + maxtext_utils_nnx.apply_moe_bias_updates(new_state.model, moe_bias_updates) lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index f162fc8275..916283618d 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -748,7 +748,11 @@ def get_dense_moe_layers(config): num_moe_layers = config.num_decoder_layers // config.interleave_moe_layer_step num_dense_layers = config.num_decoder_layers - num_moe_layers return num_dense_layers, num_moe_layers - elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): + elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): + num_dense_layers = config.first_num_dense_layers + num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers + return num_dense_layers, num_moe_layers + elif config.decoder_block == DecoderBlockType.DEEPSEEK4: return 0, config.num_decoder_layers elif config.decoder_block == DecoderBlockType.DEFAULT: raise ValueError("Unsupported decoder block for dense/MoE layer calculation") diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 1b67c2592c..58c0ad58f9 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -164,7 +164,7 @@ def create_nnx_sharded_model( named_sharding = nnx_extract_named_sharding(abstract_state) if mesh is None: - mesh = abstract_model.mesh # pyrefly: ignore[missing-attribute] + mesh = abstract_model.mesh # JIT a function that creates the model state with proper sharding from the start. # By providing out_shardings, we instruct JAX to produce sharded output directly, @@ -218,6 +218,33 @@ def nnx_update_sharding_meta(variable, transform_fn): return variable +def nnx_sync_moveaxis(tree, from_axis, to_axis): + """Moves an axis in both values and sharding metadata of nnx.Variables.""" + if from_axis == to_axis: + return tree + + def _op(x): + is_var = isinstance(x, nnx.Variable) + val = x.get_value() if is_var else x + if not hasattr(val, "shape"): + return x + + new_val = jnp.moveaxis(val, from_axis, to_axis) + if not is_var: + return new_val + + def move_fn(l): + while len(l) < val.ndim: + l.append(None) + if len(l) > max(from_axis, to_axis): + l.insert(to_axis, l.pop(from_axis)) + return l + + return nnx_update_sharding_meta(x.replace(value=new_val), move_fn) + + return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable) or hasattr(x, "shape")) + + def nnx_remove_scan_axis(tree, name="layers"): """Removes the given scan axis from the PartitionSpec.""" @@ -225,25 +252,11 @@ def _op(x): if not isinstance(x, nnx.Variable): return x - # Scanned stacks record their own axis name, such as "dense_layers" or "moe_layers", - # so prefer it over the caller's default. Otherwise the name never matches and the - # check below strips a real logical axis instead of the scan axis. - axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name) - def remove_fn(l): - removed = axis_name in l - if removed: - l.remove(axis_name) - if len(l) > x.get_value().ndim: - if removed: - raise ValueError( - f"Sharding names {l} still exceed value rank {x.get_value().ndim} after removing scan axis " - f"{axis_name!r}; the partition metadata is inconsistent." - ) - raise ValueError( - f"Scan axis {axis_name!r} not found in sharding names {l} for a rank-{x.get_value().ndim} value; " - "the partition metadata is inconsistent." - ) + if name in l: + l.remove(name) + while len(l) > x.get_value().ndim: + l.pop(0) return l return nnx_update_sharding_meta(x, remove_fn) @@ -251,31 +264,18 @@ def remove_fn(l): return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable)) -def nnx_add_and_sync_scan_axis(tree, name="layers", pos=0): - """Restores the scan axis on each variable's value and sharding metadata. - - jax.lax.scan stacks its outputs with the scan axis at position 0. For each - variable this moves that axis to the variable's own param_scan_axis (falling - back to pos when the metadata is absent) and inserts the matching axis name at - the same position, so the value and its sharding metadata stay aligned. - """ +def nnx_add_scan_axis(tree, name="layers", pos=0): + """Adds the given scan axis to the PartitionSpec at the specified position.""" def _op(x): if not isinstance(x, nnx.Variable): return x - axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name) - target = x.get_metadata().get("param_scan_axis", pos) - - val = x.get_value() - if target != 0 and hasattr(val, "ndim") and val.ndim > target: - x = x.replace(value=jnp.moveaxis(val, 0, target)) - def add_fn(l): - if axis_name not in l: + if name not in l: while len(l) < x.get_value().ndim - 1: l.append(None) - l.insert(target, axis_name) + l.insert(pos, name) else: while len(l) < x.get_value().ndim: l.append(None) @@ -284,3 +284,73 @@ def add_fn(l): return nnx_update_sharding_meta(x, add_fn) return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable)) + + +def nnx_add_and_sync_scan_axis(tree, name="layers", scan_axis=0): + """Adds the given scan axis to PartitionSpec and moves axis if scan_axis != 0.""" + tree = nnx_add_scan_axis(tree, name, 0) + if scan_axis != 0: + new_params, new_rest = tree.split(nnx.Param, ...) + new_params = nnx_sync_moveaxis(new_params, 0, scan_axis) + tree = nnx.merge_state(new_params, new_rest) + return tree + + +def apply_moe_bias_updates(model: nnx.Module, moe_bias_updates: dict[str, jax.Array]) -> None: + """Applies aux-loss-free load-balancing bias updates to each MoE gate's bias, in place. + + Different decoder blocks sow `moe_bias_updates` intermediates at different depths: DeepSeek + sows once per homogeneous scanned `moe_layers` collection, while Qwen3-Next sows once per + layer position inside each scanned block (its MoE gates aren't a single uniform scanned + collection). Rather than hardcoding an absolute attribute path for one decoder block, this + locates each target `gate.bias` parameter generically: it's the unique `Param` leaf whose + path ends in `("gate", "bias")` and lives under the same parent module that produced the + corresponding sow, found by matching path prefixes against the live parameter tree. + + Args: + model: The NNX model whose gate-bias parameters will be updated in place. + moe_bias_updates: A dict mapping each sow path (its components joined with "/") to its + update array, as collected in `train.py`'s `train_step` from the model's sown + `moe_bias_updates` intermediates. A dict (rather than a list of (path, array) pairs) + because this value flows through `jax.lax.scan` under gradient accumulation, which + requires every leaf to be an array -- dict keys are pytree structure, not leaves, so + the string paths survive that unscathed. + """ + param_paths = [ + tuple(k.key for k in path if hasattr(k, "key")) + for path, _ in jax.tree_util.tree_leaves_with_path(nnx.state(model, nnx.Param).to_pure_dict()) + ] + gate_bias_paths = [p for p in param_paths if p[-2:] == ("gate", "bias")] + + for sow_path_str, update in moe_bias_updates.items(): + sow_path = tuple(sow_path_str.split("/")) + parent_path = sow_path[:-1] + matches = [p for p in gate_bias_paths if p[: len(parent_path)] == parent_path] + if len(matches) != 1: + raise ValueError( + f"Expected exactly one gate.bias parameter under {parent_path} for the " + f"moe_bias_updates sown at {sow_path}, found {len(matches)}: {matches}" + ) + target_bias = model + for key in matches[0]: + target_bias = getattr(target_bias, key) + update_arr = jnp.array(update) + if target_bias.value.shape == update_arr.shape: + target_bias.value = target_bias.value + update_arr + elif target_bias.value.ndim == update_arr.ndim and sorted(target_bias.value.shape) == sorted(update_arr.shape): + target_shape = target_bias.value.shape + update_shape = update_arr.shape + used = [False] * len(update_shape) + perm = [] + for s in target_shape: + for i, us in enumerate(update_shape): + if us == s and not used[i]: + perm.append(i) + used[i] = True + break + if len(perm) == len(target_shape): + target_bias.value = target_bias.value + jnp.transpose(update_arr, perm) + else: + target_bias.value = target_bias.value + update_arr + else: + target_bias.value = target_bias.value + update_arr diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index ff77c57807..202495d142 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -91,6 +91,14 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: "hc_base", "sinks", "tid2eid", + "A_log", + "dt_bias", + "conv1d", + "gate", + "shared_expert_gate", + "post_alpha", + "pre_alpha", + "res_alpha", ) ) or segment == "bias" @@ -101,7 +109,7 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: # 2 Special weights # 2.1 Special weights: MoE, [0, L, -2, -1] # L (optional) stands for layer when scan_layers=True - if "MoeBlock_0" in path: + if _is_path_contain_any(("MoeBlock_0", "routed_experts"), path): # exclude gate if _is_path_contain_any(("wi_0", "wi_1", "wo"), path): return mdn((-2,), (-1,)) @@ -119,6 +127,13 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b", "wkv"), path): return mdn((0,), (-2, -1)) + # 2.3 Special weights: Gated Delta Net (GDN) + elif _is_path_contain_any(("in_proj_qkvz", "in_proj_ba", "out_proj"), path): + if "out_proj" in path: + return mdn((0, -2), (-1,)) + elif _is_path_contain_any(("in_proj_qkvz", "in_proj_ba"), path): + return mdn((0,), (-2, -1)) + # 3 Standard weights, [0, L, -1] return mdn((0,), (-1,)) diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 3ebaa21610..4e3a29f880 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -381,7 +381,10 @@ def logical_to_mesh_sharding(tree, mesh, rules=None): def create_sharding(mesh, logical_names, rules=None): """Create NamedSharding with given logical names.""" - return NamedSharding(mesh, logical_to_mesh_axes(logical_names, mesh, rules=rules)) + spec = logical_to_mesh_axes(logical_names, mesh, rules=rules) + if spec is None: + spec = P() + return NamedSharding(mesh, spec) def truncate_out_sharding(out_sharding, out_ndim: int): @@ -758,15 +761,24 @@ def _extract_param_only(state): # state_mesh_shardings.optimizer contains the sharding for the nnx.Optimizer opt_state = state_mesh_shardings.optimizer.opt_state - def find_adam_mu(obj): + def collect_mu_trees(obj, out): + # Leaf variable (e.g. a bare count/step OptArray): nothing to collect below + # it, and its __contains__ delegates to the wrapped value (a NamedSharding), + # so the dict-style 'mu' membership test below would raise on it. + if isinstance(obj, nnx.Variable): + return + # 1. Direct hit on ScaleByAdamState (Linen path or unflattened NNX) if isinstance(obj, optax.ScaleByAdamState): - return obj.mu + out.append(obj.mu) + return - # 2. Check for flattened ScaleByAdamState (nnx.State/dict) - # These nodes contain 'mu', 'nu', and 'count' as keys. - if hasattr(obj, "__getitem__") and "mu" in obj and "nu" in obj: - return obj["mu"] + # 2. Flattened optax moment state (nnx.State/dict) holding a param-mirroring + # 'mu' tree. Matching on 'mu' alone (not 'mu' and 'nu') also catches Muon's + # scale_by_muon state (mu + ns_coeffs, no nu) alongside Adam's (mu + nu). + if hasattr(obj, "__getitem__") and "mu" in obj: + out.append(obj["mu"]) + return # 3. Recursive search through containers (nnx.State, dict, list, tuple) values = None @@ -777,24 +789,35 @@ def find_adam_mu(obj): if values: for v in values: - res = find_adam_mu(v) - if res is not None: - return res - return None + collect_mu_trees(v, out) - sharded_fp32_params = find_adam_mu(opt_state) + mu_trees = [] + collect_mu_trees(opt_state, mu_trees) + sharded_fp32_params = mu_trees or None if sharded_fp32_params is None: actual_type = type(state_mesh_shardings.optimizer.get("opt_state", "None")) raise NotImplementedError(f"Could not find Adam optimizer state in: {actual_type}") # Update model parameter sharding to match the mu (first moment) sharding. # This ensures parameter sharding is consistent with the Zero-1 distributed layout. - # Build a path β†’ new_PS lookup from sharded_fp32_params (mu), then update model_shardings - # at those paths while preserving rngs and any other non-Param variables. - mu_leaves_with_paths = list( - jax.tree_util.tree_leaves_with_path(sharded_fp32_params, is_leaf=lambda x: isinstance(x, nnx.Variable)) - ) - mu_lookup = {path: mu_var.get_value() for path, mu_var in mu_leaves_with_paths} + # Build a path β†’ new_PS lookup across every collected mu tree, then update + # model_shardings at those paths while preserving rngs and any other non-Param + # variables. A partitioned optimizer (e.g. Muon: a 'muon' branch owning the 2D + # weights plus an 'adam' fallback owning the rest) holds one param-mirroring mu + # per branch, with optax MaskedNode at every position the branch does NOT own β€” + # so each param's sharding must come from whichever branch actually tracks it, + # and MaskedNode placeholders must never overwrite a real param sharding (they + # otherwise clobber the shardings tree and crash with_sharding_constraint with + # a MaskedNode-vs-leaf pytree structure error). + mu_lookup = {} + for mu_tree in mu_trees: + for path, mu_var in jax.tree_util.tree_leaves_with_path(mu_tree, is_leaf=lambda x: isinstance(x, nnx.Variable)): + if not isinstance(mu_var, nnx.Variable): + continue + mu_value = mu_var.get_value() + if not isinstance(mu_value, NamedSharding): + continue # skip MaskedNode / other non-sharding placeholders + mu_lookup.setdefault(path, mu_value) def _update_model_var(path, var): if path in mu_lookup: diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index 58bfadf29a..dd1a6c0a1c 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -19,6 +19,7 @@ import io import contextlib import unittest +import pytest from unittest import mock import jax @@ -69,9 +70,9 @@ def test_moe_wi_1_uses_last_two_axes(self): def test_moe_wo_uses_last_two_axes(self): self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "wo")), mdn((-2,), (-1,))) - def test_moe_gate_falls_through_to_standard(self): - # 'gate' is inside MoeBlock_0 but not one of (wi_0, wi_1, wo) β†’ standard. - self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate", "kernel")), mdn((0,), (-1,))) + def test_moe_gate_is_excluded(self): + # 'gate' is excluded from Muon (optimized with standard AdamW). + self.assertIsNone(muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate", "kernel"))) # --- 2.2 Self-attention --- def test_self_attention_out_projection(self): @@ -119,6 +120,27 @@ def test_deepseek_v4_self_attention_grouped_projection(self): # and output on out_features_per_group (-1) self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,))) + # --- 5. Qwen3-Next Specific --- + @pytest.mark.tpu_only + def test_qwen3_next_moe_routed_experts(self): + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_0")), mdn((-2,), (-1,))) + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_1")), mdn((-2,), (-1,))) + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wo")), mdn((-2,), (-1,))) + + @pytest.mark.tpu_only + def test_qwen3_next_gdn_projections(self): + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "in_proj_qkvz")), mdn((0,), (-2, -1))) + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "in_proj_ba")), mdn((0,), (-2, -1))) + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "out_proj")), mdn((0, -2), (-1,))) + + @pytest.mark.tpu_only + def test_qwen3_next_exclusions(self): + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "A_log"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "dt_bias"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "conv1d"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "gate"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "shared_expert_gate"))) + class TestGetTransformTree(unittest.TestCase): """Tests for get_transform_tree: recursive dict walk that applies transform_logic.""" diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index abdfb30495..b9c1242973 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -51,7 +51,7 @@ from maxtext.layers.embeddings import Embed from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process from maxtext.layers.normalizations import RMSNorm -from maxtext.models import gemma4, gemma4_small +from maxtext.models import gemma4, gemma4_small, qwen3 from maxtext.models.gpt3 import Gpt3LayerNorm from maxtext.models.llama2 import LlamaDecoderLayer from maxtext.utils import maxtext_utils, maxtext_utils_nnx @@ -716,6 +716,93 @@ def test_scan_layers(self): self.assertEqual(logits.shape, (batch, seq_len, cfg.vocab_size)) +class _StatefulQwen3NextDecoderLayer(nnx.Module): + """Small stand-in that exposes cache ordering and mutable-state updates for Qwen3-Next.""" + + def __init__(self, *, layer_idx, **unused_kwargs): + is_global = (layer_idx + 1) % 4 == 0 + self.increment = 10 if is_global else 1 + self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32)) + self.received_attention_metadata = nnx.Intermediate(jnp.array(False)) + + def __call__( + self, + inputs, + *unused_args, + kv_cache=None, + attention_metadata=None, + **unused_kwargs, + ): + self.call_count.value += 1 + self.received_attention_metadata.value = attention_metadata is not None + output = inputs + self.increment + if kv_cache is None: + return output + return output, kv_cache + self.increment + + +class TestQwen3NextScannableBlock(unittest.TestCase): + """Tests Qwen3-Next's nested local/global decoder block behavior.""" + + def setUp(self): + super().setUp() + self.config = SimpleNamespace( + dtype=jnp.float32, + param_scan_axis=1, + remat_policy="none", + scan_layers=True, + inhomogeneous_layer_cycle_interval=4, + full_attention_layer_offset=3, + ) + + def _make_block(self): + return qwen3.Qwen3NextScannableBlock( + config=self.config, + mesh=None, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(0), + ) + + def test_updates_state_through_global_single_iteration_scan(self): + with mock.patch.object(qwen3, "Qwen3NextDecoderLayer", _StatefulQwen3NextDecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) + self.assertIsNone(updated_kvs) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(3, dtype=jnp.int32)) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + + def test_restores_local_state_and_preserves_kv_order(self): + attention_metadata = object() + + with mock.patch.object(qwen3, "Qwen3NextDecoderLayer", _StatefulQwen3NextDecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + kv_cache=tuple(jnp.array(i) for i in range(4)), + attention_metadata=attention_metadata, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) + np.testing.assert_array_equal(jnp.stack(updated_kvs), jnp.array([1, 2, 3, 13])) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(3, dtype=jnp.int32)) + np.testing.assert_array_equal( + block.local_layers.received_attention_metadata.value, + jnp.ones(3, dtype=jnp.bool_), + ) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True) class _StatefulGemma4DecoderLayer(nnx.Module): """Small stand-in that exposes cache ordering and mutable-state updates.""" @@ -1076,6 +1163,53 @@ def make_block(): self.assertEqual(len(updated_kvs), num_layers) np.testing.assert_allclose(y_external, y_scanned, rtol=1e-5, atol=1e-5) + def test_qwen3_next_scanned_layers(self): + """Test NNXDecoder with qwen3_next block, dense prefix, and scan_layers=True.""" + cfg = _make_config( + decoder_block="qwen3_next", + scan_layers=True, + first_num_dense_layers=1, + inhomogeneous_layer_cycle_interval=3, + full_attention_layer_offset=0, + num_decoder_layers=4, + base_emb_dim=128, + base_num_query_heads=4, + base_num_kv_heads=4, + base_mlp_dim=256, + base_moe_mlp_dim=128, + shared_experts=1, + num_experts=4, + num_experts_per_tok=2, + gdn_num_key_heads=4, + gdn_num_value_heads=4, + gdn_key_head_dim=32, + gdn_value_head_dim=32, + vocab_size=256, + max_target_length=16, + ) + decoder = NNXDecoder( + config=cfg, + mesh=self.mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=self.rngs, + ) + shared_embedding = self._make_shared_embedding(cfg) + ids, segment_ids, positions = self._make_token_inputs(cfg) + + logits, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + self.assertEqual( + logits.shape, + (cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.vocab_size), + ) + + @pytest.mark.tpu_only class TestGemma4SmallNNXDecoder(unittest.TestCase): @@ -1315,7 +1449,5 @@ def __call__(self, x, **kwargs): ) finally: maxtext_utils_nnx.nnx_add_and_sync_scan_axis = original_add_scan_axis - - if __name__ == "__main__": unittest.main() diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index cea6485817..f8fd22c5d1 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -109,7 +109,14 @@ def test_qwen3_next_mapping_scanned(self): maxtext_config = mock.Mock() maxtext_config.inhomogeneous_layer_cycle_interval = 2 mapping = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) - self.assertIn("params-decoder-layers-layer_0-input_layernorm-scale", mapping) + self.assertIn("params-decoder-scanned_blocks-local_layers-input_layernorm-scale", mapping) + self.assertIn("params-decoder-scanned_blocks-global_layer-input_layernorm-scale", mapping) + num_blocks = config["num_hidden_layers"] // maxtext_config.inhomogeneous_layer_cycle_interval + local_val = mapping["params-decoder-scanned_blocks-local_layers-input_layernorm-scale"] + global_val = mapping["params-decoder-scanned_blocks-global_layer-input_layernorm-scale"] + self.assertEqual(len(local_val), num_blocks) + self.assertEqual(len(local_val[0]), 1) + self.assertEqual(len(global_val), num_blocks) def test_deepseek_mapping(self): config = { From ec457409d904a781aa01f0b53ea2ac779ef7e75e Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Wed, 19 Aug 2026 14:31:55 +0000 Subject: [PATCH 2/3] my recipe --- run_custom_qwen3_next_on_xpk.sh | 191 ++++++++++++++++++++++++++++++++ run_qwen3_80b_xpk.sh | 58 +++++----- 2 files changed, 220 insertions(+), 29 deletions(-) create mode 100644 run_custom_qwen3_next_on_xpk.sh diff --git a/run_custom_qwen3_next_on_xpk.sh b/run_custom_qwen3_next_on_xpk.sh new file mode 100644 index 0000000000..7b1cc0c5a6 --- /dev/null +++ b/run_custom_qwen3_next_on_xpk.sh @@ -0,0 +1,191 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment +source /usr/local/google/home/muskansh/maxtext_env/bin/activate + +# --- Environment Variables --- +export PROJECT_ID="tpu-prod-env-one-vm" +export CLUSTER_NAME="v6e-256-c2b3-b478935789" +export ZONE="us-central2-b" + +# --- Configuration --- +TIMESTAMP=$(date +%m%d%H%M%S) +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:muskansh_${TIMESTAMP}" +export WORKLOAD_NAME="muskansh-qn80b-${TIMESTAMP}" +export DEVICE_TYPE="v6e-256" +export NUM_SLICES=1 +export PRIORITY="very-high" +export MAX_RESTARTS=0 +export NUM_STEPS=15 +export MODEL_NAME="qwen3-next-80b-a3b" +export BASE_OUTPUT_DIR="gs://darisoy-hlo-dumps/qwen3-next-80b-profiles/run-${TIMESTAMP}" + +echo "========================================================================" +echo "Using pre-built Docker runner image: ${WORKLOAD_IMAGE}" +echo "========================================================================" + +# --- XLA Flags --- +XLA_FLAGS_ARRAY=( + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_enable_concurrent_sparse_core_offloading=true" + "--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" + "--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" + "--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" + "--xla_tpu_sparse_core_all_gather_latency_multiplier=1" + "--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_use_tc_device_shape_on_sc=true" + "--xla_sc_enable_instruction_fusion=false" + "--xla_sc_disable_megacore_partitioning=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" + "--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" + "--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + "--xla_tpu_enable_all_experimental_scheduler_features=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=4" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=ENABLED" + "--xla_tpu_enable_ag_backward_pipelining=true" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" + "--xla_tpu_scheduler_percent_shared_memory_limit=150" +) +export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}" + +# --- MaxText Workload Overrides --- +MAXTEXT_ARGS_ARRAY=( + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=True" + "ici_expert_parallelism=4" + "use_ring_of_experts=True" + "custom_mesh=hybrid_ring_64x4" + "use_ragged_sort=True" + "use_random_routing=True" + "num_moe_token_chunks=2" + "per_device_batch_size=8" + "opt_type=adamw" + "max_target_length=2048" + "ragged_buffer_factor=1.5" + "remat_policy=full" + "reuse_example_batch=1" + "decoder_layer_input=device" + "ici_fsdp_parallelism=-1" + "steps=15" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=1024" + "sa_block_kv=1024" + "sa_block_kv_compute=512" + "sa_block_q_dkv=1024" + "sa_block_kv_dkv=1024" + "sa_block_kv_dkv_compute=1024" + "sa_fuse_reciprocal=false" + "use_splash_scheduler=true" + "sa_use_base2_exp=true" + "dq_reduction_steps=3" + "hardware=tpu" + "skip_jax_distributed_system=False" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=512" + "wi_tile_dlhs_batch_seq=512" + "wi_tile_drhs_batch_seq=512" + "wo_tile_fwd_batch_seq=512" + "wo_tile_dlhs_batch_seq=512" + "wo_tile_drhs_batch_seq=512" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=huggingface" + "tokenizer_path=src/maxtext/assets/tokenizers/qwen3-tokenizer" + "override_model_config=true" + "mhc_expansion_rate=4" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "profiler=xplane" + "profiler_steps=4" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" +) +MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}" + +# The command to run inside the container +RUN_COMMAND="set -e && \ +export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && \ +export JAX_PLATFORMS='tpu,cpu' && \ +export ENABLE_PJRT_COMPATIBILITY='true' && \ +export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && \ +export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && \ +python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS}" + +# --- XPK Workload Creation --- +echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" + +/usr/local/google/home/muskansh/maxtext_env/bin/python3 -m xpk.main workload create \ + --cluster="${CLUSTER_NAME}" \ + --project="${PROJECT_ID}" \ + --zone="${ZONE}" \ + --priority="${PRIORITY}" \ + --max-restarts="${MAX_RESTARTS}" \ + --device-type="${DEVICE_TYPE}" \ + --num-slices="${NUM_SLICES}" \ + --docker-image="${WORKLOAD_IMAGE}" \ + --enable-debug-logs \ + --workload="${WORKLOAD_NAME}" \ + --command="${RUN_COMMAND}" + +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22us-central2%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/us-central2/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" + +echo "========================================================================" +echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" +echo "${LOGS_URL}" +echo "" +echo "☸️ GKE Workload Details:" +echo "${GKE_URL}" +echo "" +echo "========================================================================" diff --git a/run_qwen3_80b_xpk.sh b/run_qwen3_80b_xpk.sh index 8a835caf35..9a752c7a9d 100755 --- a/run_qwen3_80b_xpk.sh +++ b/run_qwen3_80b_xpk.sh @@ -2,32 +2,32 @@ set -e # Activate Python virtual environment -source /usr/local/google/home/chengnuojin/.venv/bin/activate +source /usr/local/google/home/muskansh/maxtext_env/bin/activate # --- Environment Variables --- export PROJECT_ID="tpu-prod-env-one-vm" -export CLUSTER_NAME="bodaborg-v6e-256-lcscld-c" -export ZONE="southamerica-west1-a" +export CLUSTER_NAME="v6e-256-c2b3-b478935789" +export ZONE="us-central2-b" # --- Configuration & Automated Image Build --- TIMESTAMP=$(date +%m%d%H%M%S) -export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:chengnuojin_${TIMESTAMP}" -export WORKLOAD_NAME="chengnuojin-qn80b-${TIMESTAMP}" +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:muskansh_${TIMESTAMP}" +export WORKLOAD_NAME="muskansh-qn80b-${TIMESTAMP}" export DEVICE_TYPE="v6e-256" export NUM_SLICES=1 export PRIORITY="very-high" export NUM_STEPS=15 export MAX_RESTARTS=0 export MODEL_NAME="qwen3-next-80b-a3b" -export BASE_OUTPUT_DIR="gs://chengnuojin-maxtext-logs/qwen3-next-80b-profiles/run-${TIMESTAMP}" +export BASE_OUTPUT_DIR="gs://darisoy-hlo-dumps/qwen3-next-80b-profiles/run-${TIMESTAMP}" echo "========================================================================" -echo "Building and uploading Docker runner image from /usr/local/google/home/chengnuojin/maxtext" +echo "Building and uploading Docker runner image from /usr/local/google/home/muskansh/maxtext" echo "Target Image: ${WORKLOAD_IMAGE}" echo "========================================================================" ( - cd /usr/local/google/home/chengnuojin/maxtext && \ + cd /usr/local/google/home/muskansh/maxtext && \ if ! docker image inspect maxtext_base_image &> /dev/null; then echo "Local image 'maxtext_base_image' not found. Pulling gcr.io/tpu-prod-env-one-vm/param3_21jul:latest and tagging as maxtext_base_image..." docker pull gcr.io/tpu-prod-env-one-vm/param3_21jul:latest @@ -134,24 +134,24 @@ MAXTEXT_ARGS_ARRAY=( "sa_use_fused_bwd_kernel=True" "sparse_matmul=True" "megablox=True" - "wi_tile_fwd_batch_seq=128" - "wi_tile_dlhs_batch_seq=128" - "wi_tile_drhs_batch_seq=128" - "wo_tile_fwd_batch_seq=128" - "wo_tile_dlhs_batch_seq=128" - "wo_tile_drhs_batch_seq=128" - "wi_tile_fwd_embed_dim=3072" - "wi_tile_fwd_mlp_dim=1536" - "wi_tile_dlhs_embed_dim=3072" - "wi_tile_dlhs_mlp_dim=1536" - "wi_tile_drhs_embed_dim=3072" - "wi_tile_drhs_mlp_dim=1536" - "wo_tile_fwd_embed_dim=3072" - "wo_tile_fwd_mlp_dim=1536" - "wo_tile_dlhs_embed_dim=3072" - "wo_tile_dlhs_mlp_dim=1536" - "wo_tile_drhs_embed_dim=3072" - "wo_tile_drhs_mlp_dim=1536" + "wi_tile_fwd_batch_seq=256" + "wi_tile_dlhs_batch_seq=256" + "wi_tile_drhs_batch_seq=256" + "wo_tile_fwd_batch_seq=256" + "wo_tile_dlhs_batch_seq=256" + "wo_tile_drhs_batch_seq=256" + "wi_tile_fwd_embed_dim=512" + "wi_tile_fwd_mlp_dim=512" + "wi_tile_dlhs_embed_dim=512" + "wi_tile_dlhs_mlp_dim=512" + "wi_tile_drhs_embed_dim=512" + "wi_tile_drhs_mlp_dim=512" + "wo_tile_fwd_embed_dim=512" + "wo_tile_fwd_mlp_dim=512" + "wo_tile_dlhs_embed_dim=512" + "wo_tile_dlhs_mlp_dim=512" + "wo_tile_drhs_embed_dim=512" + "wo_tile_drhs_mlp_dim=512" "use_tokamax_gmm=True" "use_gmm_v2=True" "optimizer_memory_host_offload=False" @@ -183,7 +183,7 @@ python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${M # --- XPK Workload Creation --- echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" -PYTHONPATH=/usr/local/google/home/chengnuojin/xpk/src python3 -m xpk.main workload create \ +/usr/local/google/home/muskansh/maxtext_env/bin/python3 -m xpk.main workload create \ --cluster="${CLUSTER_NAME}" \ --project="${PROJECT_ID}" \ --zone="${ZONE}" \ @@ -195,8 +195,8 @@ PYTHONPATH=/usr/local/google/home/chengnuojin/xpk/src python3 -m xpk.main worklo --workload="${WORKLOAD_NAME}" \ --command="${RUN_COMMAND}" -LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" -GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22us-central2%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/us-central2/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" TB_URL="https://tensorboard.corp.google.com/?logdir=${BASE_OUTPUT_DIR}/param-3/tensorboard" echo "========================================================================" From 2acb7ada502eddc93133b72afb85857a24094340 Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Wed, 19 Aug 2026 16:13:36 +0000 Subject: [PATCH 3/3] use l2norm func --- run_custom_qwen3_next_on_xpk.sh | 44 ++++++++++++++++---------------- run_qwen3_80b_xpk.sh | 38 +++++++++++++-------------- src/maxtext/models/hybrid_gdn.py | 6 +++-- 3 files changed, 45 insertions(+), 43 deletions(-) mode change 100644 => 100755 run_custom_qwen3_next_on_xpk.sh diff --git a/run_custom_qwen3_next_on_xpk.sh b/run_custom_qwen3_next_on_xpk.sh old mode 100644 new mode 100755 index 7b1cc0c5a6..5a916bdabe --- a/run_custom_qwen3_next_on_xpk.sh +++ b/run_custom_qwen3_next_on_xpk.sh @@ -7,11 +7,11 @@ source /usr/local/google/home/muskansh/maxtext_env/bin/activate # --- Environment Variables --- export PROJECT_ID="tpu-prod-env-one-vm" export CLUSTER_NAME="v6e-256-c2b3-b478935789" -export ZONE="us-central2-b" +export ZONE="us-central2" # --- Configuration --- TIMESTAMP=$(date +%m%d%H%M%S) -export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:muskansh_${TIMESTAMP}" +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:darisoy_0819071448" export WORKLOAD_NAME="muskansh-qn80b-${TIMESTAMP}" export DEVICE_TYPE="v6e-256" export NUM_SLICES=1 @@ -115,24 +115,24 @@ MAXTEXT_ARGS_ARRAY=( "sa_use_fused_bwd_kernel=True" "sparse_matmul=True" "megablox=True" - "wi_tile_fwd_batch_seq=512" - "wi_tile_dlhs_batch_seq=512" - "wi_tile_drhs_batch_seq=512" - "wo_tile_fwd_batch_seq=512" - "wo_tile_dlhs_batch_seq=512" - "wo_tile_drhs_batch_seq=512" - "wi_tile_fwd_embed_dim=3072" - "wi_tile_fwd_mlp_dim=1536" - "wi_tile_dlhs_embed_dim=3072" - "wi_tile_dlhs_mlp_dim=1536" - "wi_tile_drhs_embed_dim=3072" - "wi_tile_drhs_mlp_dim=1536" - "wo_tile_fwd_embed_dim=3072" - "wo_tile_fwd_mlp_dim=1536" - "wo_tile_dlhs_embed_dim=3072" - "wo_tile_dlhs_mlp_dim=1536" - "wo_tile_drhs_embed_dim=3072" - "wo_tile_drhs_mlp_dim=1536" + "wi_tile_fwd_batch_seq=256" + "wi_tile_dlhs_batch_seq=256" + "wi_tile_drhs_batch_seq=256" + "wo_tile_fwd_batch_seq=256" + "wo_tile_dlhs_batch_seq=256" + "wo_tile_drhs_batch_seq=256" + "wi_tile_fwd_embed_dim=512" + "wi_tile_fwd_mlp_dim=512" + "wi_tile_dlhs_embed_dim=512" + "wi_tile_dlhs_mlp_dim=512" + "wi_tile_drhs_embed_dim=512" + "wi_tile_drhs_mlp_dim=512" + "wo_tile_fwd_embed_dim=512" + "wo_tile_fwd_mlp_dim=512" + "wo_tile_dlhs_embed_dim=512" + "wo_tile_dlhs_mlp_dim=512" + "wo_tile_drhs_embed_dim=512" + "wo_tile_drhs_mlp_dim=512" "use_tokamax_gmm=True" "use_gmm_v2=True" "optimizer_memory_host_offload=False" @@ -178,8 +178,8 @@ echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" --workload="${WORKLOAD_NAME}" \ --command="${RUN_COMMAND}" -LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22us-central2%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" -GKE_URL="https://console.cloud.google.com/kubernetes/service/us-central2/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" echo "========================================================================" echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" diff --git a/run_qwen3_80b_xpk.sh b/run_qwen3_80b_xpk.sh index 9a752c7a9d..a5fa5ce961 100755 --- a/run_qwen3_80b_xpk.sh +++ b/run_qwen3_80b_xpk.sh @@ -19,7 +19,7 @@ export PRIORITY="very-high" export NUM_STEPS=15 export MAX_RESTARTS=0 export MODEL_NAME="qwen3-next-80b-a3b" -export BASE_OUTPUT_DIR="gs://darisoy-hlo-dumps/qwen3-next-80b-profiles/run-${TIMESTAMP}" +export BASE_OUTPUT_DIR="/tmp/maxtext-profiles/run-${TIMESTAMP}" echo "========================================================================" echo "Building and uploading Docker runner image from /usr/local/google/home/muskansh/maxtext" @@ -134,24 +134,24 @@ MAXTEXT_ARGS_ARRAY=( "sa_use_fused_bwd_kernel=True" "sparse_matmul=True" "megablox=True" - "wi_tile_fwd_batch_seq=256" - "wi_tile_dlhs_batch_seq=256" - "wi_tile_drhs_batch_seq=256" - "wo_tile_fwd_batch_seq=256" - "wo_tile_dlhs_batch_seq=256" - "wo_tile_drhs_batch_seq=256" - "wi_tile_fwd_embed_dim=512" - "wi_tile_fwd_mlp_dim=512" - "wi_tile_dlhs_embed_dim=512" - "wi_tile_dlhs_mlp_dim=512" - "wi_tile_drhs_embed_dim=512" - "wi_tile_drhs_mlp_dim=512" - "wo_tile_fwd_embed_dim=512" - "wo_tile_fwd_mlp_dim=512" - "wo_tile_dlhs_embed_dim=512" - "wo_tile_dlhs_mlp_dim=512" - "wo_tile_drhs_embed_dim=512" - "wo_tile_drhs_mlp_dim=512" + "wi_tile_fwd_batch_seq=128" + "wi_tile_dlhs_batch_seq=128" + "wi_tile_drhs_batch_seq=128" + "wo_tile_fwd_batch_seq=128" + "wo_tile_dlhs_batch_seq=128" + "wo_tile_drhs_batch_seq=128" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" "use_tokamax_gmm=True" "use_gmm_v2=True" "optimizer_memory_host_offload=False" diff --git a/src/maxtext/models/hybrid_gdn.py b/src/maxtext/models/hybrid_gdn.py index 8f2ca59d39..74a12a509f 100644 --- a/src/maxtext/models/hybrid_gdn.py +++ b/src/maxtext/models/hybrid_gdn.py @@ -95,8 +95,10 @@ def chunk_forward(q, k, v, b_val, a_val, a_log_val, dt_bias_val, state_prev): k = k.astype(jnp.float32) v = v.astype(jnp.float32) if use_qk_norm_in_gdn: - q = q / (jnp.linalg.norm(q, axis=-1, keepdims=True) + 1e-6) - k = k / (jnp.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + from maxtext.layers.normalizations import l2norm + + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) scale = 1.0 / jnp.sqrt(kq_head_dim) q = q * scale b_val = b_val.astype(jnp.float32)