Initial Commit

scripts/kaiyuan2b-training/training_smoke_gpt2.sh

@@ -0,0 +1,66 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+MODE=${1:-smoke}
+TRAIN_NAME=${2:-smoke_gpt}
+
+MEGATRON_PATH=/apps/yi/model_training/Megatron-LM
+ARTIFACT_ROOT=/apps/yi/model_training/artifacts
+TB_DIR="${ARTIFACT_ROOT}/tb_logs/${TRAIN_NAME}"
+CKPT_DIR="${ARTIFACT_ROOT}/checkpoints/${TRAIN_NAME}"
+
+source params/optim_common.sh
+source params/gpt_smoke/model.sh
+source params/gpt_smoke/data.sh
+source params/gpt_smoke/hparams.sh
+
+RUN_ARGS=$HPARAM_ARGS
+
+LOGGING_ARGS="
+  --tensorboard-dir ${TB_DIR}
+  --tensorboard-log-interval 1
+  --log-interval 1
+  --log-timers-to-tensorboard
+  --log-validation-ppl-to-tensorboard
+  --log-memory-to-tensorboard
+  --log-world-size-to-tensorboard
+  --log-num-zeros-in-grad
+  --log-device-memory-used
+  --log-throughput
+  --log-params-norm
+"
+
+PRECISION_ARGS="--bf16"
+PARALLEL_ARGS="
+    --tensor-model-parallel-size 1
+    --pipeline-model-parallel-size 1
+"
+# PARALLEL_ARGS="
+#     --tensor-model-parallel-size 2
+#     --sequence-parallel
+# "
+
+mkdir -p "$CKPT_DIR" "$TB_DIR"
+
+DISTRIBUTED_ARGS="
+  --nproc_per_node 8
+  --nnodes 1
+  --node_rank 0
+  --master_addr localhost
+  --master_port 6000
+"
+
+torchrun $DISTRIBUTED_ARGS \
+  $MEGATRON_PATH/pretrain_gpt.py \
+  $MODEL_ARGS \
+  $OPTIM_ARGS \
+  $PRECISION_ARGS \
+  $PARALLEL_ARGS \
+  $DATA_ARGS \
+  $RUN_ARGS \
+  $LOGGING_ARGS\
+  --save "$CKPT_DIR" \
+  --load "$CKPT_DIR" \