Skip to content

Fix NaN loss by using l2Norm in GDN v3 kernel #4946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
muskansh-google wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Fix NaN loss by using l2Norm in GDN v3 kernel #4946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
31951ab
feat(qwen3-next): support Qwen3-Next 80B with GDNv3, fine-grained rem…
NuojCheng Aug 13, 2026
ec45740
my recipe
muskansh-google Aug 19, 2026
2acb7ad
use l2norm func
muskansh-google Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .dockerignore
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1,4 @@
.git
maxtext_venv
.venv
venv13
191 changes: 191 additions & 0 deletions run_custom_qwen3_next_on_xpk.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"

# --- Configuration ---
TIMESTAMP=$(date +%m%d%H%M%S)
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
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=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"
"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%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):"
echo "${LOGS_URL}"
echo ""
echo "☸️ GKE Workload Details:"
echo "${GKE_URL}"
echo ""
echo "========================================================================"
158 changes: 158 additions & 0 deletions run_qwen3_80b_aot.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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[@]}"
)
Loading
Loading