Skip to content

[No merge] Draft test #4747

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
NuojCheng wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

[No merge] Draft test #4747

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0dae68e
Update qwen3-next-80b-a3b.yml config parameters
darisoy Jul 29, 2026
d3d7a3a
Integrate mHC into Qwen3NextDecoderLayer (mhc_expansion_rate=4)
darisoy Jul 30, 2026
d50c557
Implement Newton-Schulz log-depth custom VJP triangular matrix solver…
darisoy Jul 29, 2026
63fc9e1
add scan support for qwen3-next model
NuojCheng Jul 30, 2026
8e94b20
Implement Hybrid GDN v3 Tokamax forward + Newton-Schulz Custom VJP ba…
darisoy Jul 31, 2026
d7fed6e
Checkpoint [B, S, E] hidden states and recompute [B, S, 4, E] 4-manif…
darisoy Aug 3, 2026
b001b94
update scan logic and muon
NuojCheng Aug 5, 2026
5d1a147
add pure fsdp scripts
NuojCheng Aug 5, 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
156 changes: 156 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,156 @@
#!/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"
)
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=8"
"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"
"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[@]}"
)
176 changes: 176 additions & 0 deletions run_qwen3_80b_aot_fsdp.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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[@]}"
)
Loading
Loading