Skip to content

Commit a9df444

Browse files
aireenmeiaireenmei
committed
Use synthetic data for scripts under src/maxtext/configs/tpu
1 parent c91ca12 commit a9df444

24 files changed

Lines changed: 38 additions & 62 deletions

‎src/maxtext/configs/tpu/v4/22b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,10 @@ echo "Running 22b.sh"
1818
#
1919
# Command Flags:
2020
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
21-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
2221
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
2322
#
2423
# Example to invoke this script for training:
25-
# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
24+
# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
2625
#
2726
# Example to AOT compile:
2827
# 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
5857
python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml\
5958
ici_fsdp_parallelism=64 steps=10 per_device_batch_size=13 profiler=xplane remat_policy=full attention=flash num_vocab_tiling=8\
6059
base_emb_dim=6144 base_num_kv_heads=24 base_num_query_heads=24 base_mlp_dim=24576 base_num_decoder_layers=48\
61-
base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH
60+
base_output_directory=$OUTPUT_PATH dataset_type=synthetic

‎src/maxtext/configs/tpu/v4/52b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,10 @@ echo "Running 52b.sh"
1818
#
1919
# Command Flags:
2020
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
21-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
2221
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
2322
#
2423
# Example to invoke this script:
25-
# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
24+
# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
2625
#
2726
# Example to AOT compile:
2827
# 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
5958
profiler=xplane enable_checkpointing=false steps=10\
6059
ici_fsdp_parallelism=192 ici_tensor_parallelism=1 per_device_batch_size=7 remat_policy=full attention=flash num_vocab_tiling=8\
6160
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\
62-
base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH
61+
base_output_directory=$OUTPUT_PATH dataset_type=synthetic

‎src/maxtext/configs/tpu/v5e/128b.sh‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@ echo "Running 128b.sh"
44
#
55
# Command Flags:
66
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
7-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
87
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
98
#
109
# Example to invoke this script:
11-
# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
10+
# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1211
#
1312
# Example to AOT compile:
1413
# bash src/maxtext/configs/tpu/v5e/128b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2

‎src/maxtext/configs/tpu/v5e/16b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@ echo "Running 16b.sh"
44
#
55
# Command Flags:
66
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
7-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
87
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
98
#
109
# Example to invoke this script:
11-
# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
10+
# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1211
#
1312
# Example to AOT compile:
1413
# 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
4544
steps=15 per_device_batch_size=6 enable_checkpointing=false\
4645
remat_policy=full global_parameter_scale=16\
4746
max_target_length=2048 base_output_directory=$OUTPUT_PATH\
48-
dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\
47+
use_iota_embed=true reuse_example_batch=1\
4948
dataset_type=synthetic attention='flash' gcs_metrics=true

‎src/maxtext/configs/tpu/v5e/32b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@ echo "Running 32b.sh"
44
#
55
# Command Flags:
66
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
7-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
87
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
98
#
109
# Example to invoke this script:
11-
# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
10+
# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1211
#
1312
# Example to AOT compile:
1413
# 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
4544
steps=15 per_device_batch_size=4 enable_checkpointing=false\
4645
remat_policy=full global_parameter_scale=32\
4746
max_target_length=2048 base_output_directory=$OUTPUT_PATH\
48-
dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\
47+
use_iota_embed=true reuse_example_batch=1\
4948
dataset_type=synthetic attention='flash' gcs_metrics=true

‎src/maxtext/configs/tpu/v5e/64b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@ echo "Running 64b.sh"
44
#
55
# Command Flags:
66
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
7-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
87
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
98
#
109
# Example to invoke this script:
11-
# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
10+
# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1211
#
1312
# Example to AOT compile:
1413
# 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
4544
steps=15 per_device_batch_size=2 enable_checkpointing=false\
4645
remat_policy=full global_parameter_scale=64\
4746
max_target_length=2048 base_output_directory=$OUTPUT_PATH\
48-
dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\
47+
use_iota_embed=true reuse_example_batch=1\
4948
dataset_type=synthetic attention='flash' gcs_metrics=true

‎src/maxtext/configs/tpu/v5e/gpt3_175b.sh‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@
33
#
44
# Command Flags:
55
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
6-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
76
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
87
#
98
# Example to invoke this script:
10-
# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
9+
# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1110
#
1211
# Example to AOT compile:
1312
# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=8

‎src/maxtext/configs/tpu/v5e/llama2_13b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@
33
#
44
# Command Flags:
55
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
6-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
76
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
87
#
98
# Example to invoke this script:
10-
# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
9+
# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1110
#
1211
# Example to AOT compile:
1312
# 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
4241
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"
4342

4443
python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-13b\
45-
base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\
44+
base_output_directory=$OUTPUT_PATH dataset_type=synthetic\
4645
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\
4746
steps=15 enable_checkpointing=false use_iota_embed=true

‎src/maxtext/configs/tpu/v5e/llama2_70b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@
33
#
44
# Command Flags:
55
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
6-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
76
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
87
#
98
# Example to invoke this script:
10-
# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
9+
# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1110
#
1211
# Example to AOT compile:
1312
# 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
4241
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"
4342

4443
python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-70b\
45-
base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\
44+
base_output_directory=$OUTPUT_PATH dataset_type=synthetic\
4645
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\
4746
steps=15 enable_checkpointing=false use_iota_embed=true

‎src/maxtext/configs/tpu/v5e/llama2_7b.sh‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@
33
#
44
# Command Flags:
55
# OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml)
6-
# DATASET_PATH (Required, unless dataset_path is already set in base.yml)
76
# RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE)
87
#
98
# Example to invoke this script:
10-
# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>" DATASET_PATH="gs://<your_dataset_path>"
9+
# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="<your_run_name>" OUTPUT_PATH="gs://<your_output_path>"
1110
#
1211
# Example to AOT compile:
1312
# 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
4241
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"
4342

4443
python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-7b\
45-
base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\
44+
base_output_directory=$OUTPUT_PATH dataset_type=synthetic\
4645
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\
4746
steps=15 enable_checkpointing=false use_iota_embed=true

0 commit comments

Comments
 (0)