From a9df4440119b9bca73d11501ee2b8ae7884a591e Mon Sep 17 00:00:00 2001 From: aireenmei Date: Wed, 1 Jul 2026 23:11:27 +0000 Subject: [PATCH] Use synthetic data for scripts under src/maxtext/configs/tpu --- src/maxtext/configs/tpu/v4/22b.sh | 5 ++--- src/maxtext/configs/tpu/v4/52b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/128b.sh | 3 +-- src/maxtext/configs/tpu/v5e/16b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/32b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/64b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/gpt3_175b.sh | 3 +-- src/maxtext/configs/tpu/v5e/llama2_13b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/llama2_70b.sh | 5 ++--- src/maxtext/configs/tpu/v5e/llama2_7b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/1024b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/128b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/256b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/32b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/512b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/64b.sh | 5 ++--- src/maxtext/configs/tpu/v5p/llama2_70b.sh | 3 +-- src/maxtext/configs/tpu/v5p/llama2_7b.sh | 3 +-- src/maxtext/configs/tpu/v6e/gemma2_27b.sh | 3 +-- src/maxtext/configs/tpu/v6e/gemma2_9b.sh | 3 +-- src/maxtext/configs/tpu/v6e/gemma3_27b.sh | 3 +-- src/maxtext/configs/tpu/v6e/gpt3_175b.sh | 3 +-- src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh | 3 +-- src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh | 3 +-- 24 files changed, 38 insertions(+), 62 deletions(-) diff --git a/src/maxtext/configs/tpu/v4/22b.sh b/src/maxtext/configs/tpu/v4/22b.sh index 35798dbdc1..83fad6c38e 100644 --- a/src/maxtext/configs/tpu/v4/22b.sh +++ b/src/maxtext/configs/tpu/v4/22b.sh @@ -18,11 +18,10 @@ echo "Running 22b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script for training: -# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v4/22b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v4-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -58,4 +57,4 @@ export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true TPU_MEGACORE=MEGACOR python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml\ ici_fsdp_parallelism=64 steps=10 per_device_batch_size=13 profiler=xplane remat_policy=full attention=flash num_vocab_tiling=8\ base_emb_dim=6144 base_num_kv_heads=24 base_num_query_heads=24 base_mlp_dim=24576 base_num_decoder_layers=48\ - base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH + base_output_directory=$OUTPUT_PATH dataset_type=synthetic diff --git a/src/maxtext/configs/tpu/v4/52b.sh b/src/maxtext/configs/tpu/v4/52b.sh index 702a4812ef..f774679c50 100644 --- a/src/maxtext/configs/tpu/v4/52b.sh +++ b/src/maxtext/configs/tpu/v4/52b.sh @@ -18,11 +18,10 @@ echo "Running 52b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v4/52b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v4-384 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -59,4 +58,4 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT profiler=xplane enable_checkpointing=false steps=10\ ici_fsdp_parallelism=192 ici_tensor_parallelism=1 per_device_batch_size=7 remat_policy=full attention=flash num_vocab_tiling=8\ base_num_decoder_layers=32 base_emb_dim=12288 base_mlp_dim=49152 base_num_query_heads=32 base_num_kv_heads=32 learning_rate=1e-8\ - base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH + base_output_directory=$OUTPUT_PATH dataset_type=synthetic diff --git a/src/maxtext/configs/tpu/v5e/128b.sh b/src/maxtext/configs/tpu/v5e/128b.sh index 069abf376e..2fdc307b25 100644 --- a/src/maxtext/configs/tpu/v5e/128b.sh +++ b/src/maxtext/configs/tpu/v5e/128b.sh @@ -4,11 +4,10 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/128b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 diff --git a/src/maxtext/configs/tpu/v5e/16b.sh b/src/maxtext/configs/tpu/v5e/16b.sh index 4c45661146..da66a72adf 100644 --- a/src/maxtext/configs/tpu/v5e/16b.sh +++ b/src/maxtext/configs/tpu/v5e/16b.sh @@ -4,11 +4,10 @@ echo "Running 16b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/16b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=6 enable_checkpointing=false\ remat_policy=full global_parameter_scale=16\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/32b.sh b/src/maxtext/configs/tpu/v5e/32b.sh index 2e72b5da2e..ff870948c3 100644 --- a/src/maxtext/configs/tpu/v5e/32b.sh +++ b/src/maxtext/configs/tpu/v5e/32b.sh @@ -4,11 +4,10 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/32b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=4 enable_checkpointing=false\ remat_policy=full global_parameter_scale=32\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/64b.sh b/src/maxtext/configs/tpu/v5e/64b.sh index f28eeb6394..faec50b553 100644 --- a/src/maxtext/configs/tpu/v5e/64b.sh +++ b/src/maxtext/configs/tpu/v5e/64b.sh @@ -4,11 +4,10 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/64b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=2 enable_checkpointing=false\ remat_policy=full global_parameter_scale=64\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh index d71f9b1a78..d9c95ce2c6 100644 --- a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=8 diff --git a/src/maxtext/configs/tpu/v5e/llama2_13b.sh b/src/maxtext/configs/tpu/v5e/llama2_13b.sh index 8dca05f36f..e79997da64 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_13b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_13b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_13b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-13b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=8 remat_policy=qkv_proj_offloaded\ steps=15 enable_checkpointing=false use_iota_embed=true diff --git a/src/maxtext/configs/tpu/v5e/llama2_70b.sh b/src/maxtext/configs/tpu/v5e/llama2_70b.sh index 528970ccc2..aee3babe44 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_70b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_70b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-70b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=2 remat_policy=qkv_proj_offloaded\ steps=15 enable_checkpointing=false use_iota_embed=true diff --git a/src/maxtext/configs/tpu/v5e/llama2_7b.sh b/src/maxtext/configs/tpu/v5e/llama2_7b.sh index 6b68297a1e..e4aa662b2c 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_7b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_7b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-7b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=4 remat_policy=save_qkv_proj\ steps=15 enable_checkpointing=false use_iota_embed=true \ No newline at end of file diff --git a/src/maxtext/configs/tpu/v5p/1024b.sh b/src/maxtext/configs/tpu/v5p/1024b.sh index 14bdcc6692..d42bc8a70e 100644 --- a/src/maxtext/configs/tpu/v5p/1024b.sh +++ b/src/maxtext/configs/tpu/v5p/1024b.sh @@ -4,11 +4,10 @@ echo "Running 1024b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/1024b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/1024b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/1024b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-2048 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=full global_parameter_scale=1024\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=16\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/128b.sh b/src/maxtext/configs/tpu/v5p/128b.sh index ce068f7fde..e53c6a9ede 100644 --- a/src/maxtext/configs/tpu/v5p/128b.sh +++ b/src/maxtext/configs/tpu/v5p/128b.sh @@ -4,11 +4,10 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/128b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=128\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/256b.sh b/src/maxtext/configs/tpu/v5p/256b.sh index 04e16ca0cc..984f281e32 100644 --- a/src/maxtext/configs/tpu/v5p/256b.sh +++ b/src/maxtext/configs/tpu/v5p/256b.sh @@ -4,12 +4,11 @@ echo "Running 256b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/256b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/256b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/256b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-1024 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -47,5 +46,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=256\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/32b.sh b/src/maxtext/configs/tpu/v5p/32b.sh index 813f05d3a8..de129f0ec5 100644 --- a/src/maxtext/configs/tpu/v5p/32b.sh +++ b/src/maxtext/configs/tpu/v5p/32b.sh @@ -4,11 +4,10 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/32b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=32\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=4\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/512b.sh b/src/maxtext/configs/tpu/v5p/512b.sh index 293c0b65eb..42e6c7cec7 100644 --- a/src/maxtext/configs/tpu/v5p/512b.sh +++ b/src/maxtext/configs/tpu/v5p/512b.sh @@ -4,12 +4,11 @@ echo "Running 512b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/512b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/512b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/512b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-1024 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -47,5 +46,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=full global_parameter_scale=512\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/64b.sh b/src/maxtext/configs/tpu/v5p/64b.sh index ec44e130f3..90c01c2c74 100644 --- a/src/maxtext/configs/tpu/v5p/64b.sh +++ b/src/maxtext/configs/tpu/v5p/64b.sh @@ -4,11 +4,10 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/64b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=64\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=4\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/llama2_70b.sh b/src/maxtext/configs/tpu/v5p/llama2_70b.sh index b0e1836ea1..2228794617 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_70b.sh @@ -5,7 +5,6 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: @@ -45,7 +44,7 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-70b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 remat_policy=save_dot_except_mlpwi per_device_batch_size=4\ steps=30 enable_checkpointing=false use_iota_embed=true max_target_length=4096\ profiler=xplane skip_first_n_steps_for_profiler=10 profiler_steps=5 gcs_metrics=true\ diff --git a/src/maxtext/configs/tpu/v5p/llama2_7b.sh b/src/maxtext/configs/tpu/v5p/llama2_7b.sh index 08683ace65..774b7ffe61 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_7b.sh @@ -5,7 +5,6 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml or using synthetic dataset_type) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: @@ -45,7 +44,7 @@ fi # Train export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-7b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 remat_policy=minimal per_device_batch_size=4\ steps=30 enable_checkpointing=false use_iota_embed=true max_target_length=4096\ profiler=xplane skip_first_n_steps_for_profiler=10 profiler_steps=5 gcs_metrics=true\ diff --git a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh index 4bbc58643e..970722c6f8 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma2_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma2_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh index 2a7deb01aa..2728f46aed 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma2_9b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma2_9b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh index 7441dc358a..80824e9f46 100644 --- a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma3_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma3_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh index 79ca9d30fc..b7f3b40454 100644 --- a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh index 4e08e3d76f..8279acede6 100644 --- a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh +++ b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh index 2eac06d68c..c5bd75d94f 100644 --- a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh +++ b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh RUN_NAME="" OUTPUT_PATH="gs://" #