pretrain_base_benchmark.sh 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. #!/usr/bin/env bash
  2. # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. set -a
  16. : ${NUM_WARMUP_EPOCHS:=2} # Number of warmup epochs
  17. : ${NUM_EPOCHS:=5} # Number of epochs for collecting perf measurements
  18. : ${TRAIN_SUBSET:="train-full-960"}
  19. : ${FP16:=false}
  20. : ${BF16:=false}
  21. : ${NUM_GPUS:=8}
  22. : ${MAX_TOKENS:=1400000}
  23. : ${NUM_CONCAT_BATCHES:=8}
  24. : ${UPDATE_FREQ:=1}
  25. if [ "$FP16" = true ]; then PREC=fp16; elif [ "$BF16" = true ]; then PREC=bf16; else PREC=fp32; fi
  26. : ${OUTPUT_DIR:="results/pretrain_base_benchmark_${NUM_GPUS}x${MAX_TOKENS}x${NUM_CONCAT_BATCHES}x${UPDATE_FREQ}_${PREC}"}
  27. NO_SAVE=true
  28. EPOCHS_THIS_JOB=$(($NUM_EPOCHS + $NUM_WARMUP_EPOCHS))
  29. ARGS+=" --benchmark_epochs_num $NUM_EPOCHS"
  30. bash scripts/pretrain_base.sh