Maximize per_device_batch_size in 1024b config

google · Nov 14, 2023 · 54bdd80 · 54bdd80
1 parent b3e1439
commit 54bdd80
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/configs/experimental/1024b.sh b/MaxText/configs/experimental/1024b.sh
@@ -31,7 +31,7 @@ fi
 # Train
 export LIBTPU_INIT_ARGS="--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 MaxText/train.py MaxText/configs/base.yml run_name=$RUN_NAME\
- steps=20 per_device_batch_size=1 enable_checkpointing=false\
+ steps=20 per_device_batch_size=2 enable_checkpointing=false\
  enable_profiler=false remat_policy=full global_parameter_scale=1024\
  ici_fsdp_parallelism=-1 ici_tensor_parallelism=16\
  max_target_length=2048 base_output_directory=$OUTPUT_PATH\