71 lines
2.1 KiB
Bash
Executable File
71 lines
2.1 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
|
|
IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
|
|
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
|
|
DATA_PREFIX="${DATA_PREFIX:-}"
|
|
DATA_MANIFEST="${DATA_MANIFEST-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_65536_direct_v1/manifest.json}"
|
|
TRAIN_ITERS="${TRAIN_ITERS:-10}"
|
|
SEQ_LENGTH="${SEQ_LENGTH:-8192}"
|
|
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
|
|
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-1024}"
|
|
TP="${TP:-1}"
|
|
PP="${PP:-1}"
|
|
EP="${EP:-1}"
|
|
CP="${CP:-1}"
|
|
DRY_RUN="${DRY_RUN:-0}"
|
|
|
|
if [[ -n "$DATA_MANIFEST" ]]; then
|
|
if [[ ! -f "$DATA_MANIFEST" ]]; then
|
|
echo "missing Megatron indexed dataset manifest: $DATA_MANIFEST" >&2
|
|
echo "run scripts/preprocess_megatron_bridge_pretrain_direct.sh first, or set DATA_MANIFEST" >&2
|
|
exit 1
|
|
fi
|
|
elif [[ "$DRY_RUN" != "1" && ( -z "$DATA_PREFIX" || ! -f "${DATA_PREFIX}.idx" ) ]]; then
|
|
echo "missing Megatron indexed data prefix: ${DATA_PREFIX:-<empty>}" >&2
|
|
echo "set DATA_MANIFEST or DATA_PREFIX" >&2
|
|
exit 1
|
|
fi
|
|
|
|
DRY_RUN_ARG=""
|
|
if [[ "$DRY_RUN" == "1" ]]; then
|
|
DRY_RUN_ARG="--dry-run"
|
|
fi
|
|
|
|
DATA_ARGS=()
|
|
if [[ -n "$DATA_MANIFEST" ]]; then
|
|
DATA_ARGS=(--data-manifest "$DATA_MANIFEST")
|
|
elif [[ -n "$DATA_PREFIX" ]]; then
|
|
DATA_ARGS=(--data-prefix "$DATA_PREFIX")
|
|
fi
|
|
|
|
DOCKER_MOUNTS=()
|
|
if [[ -d /mnt/beegfs ]]; then
|
|
DOCKER_MOUNTS+=(-v /mnt/beegfs:/mnt/beegfs)
|
|
fi
|
|
if [[ -d /ssd ]]; then
|
|
DOCKER_MOUNTS+=(-v /ssd:/ssd)
|
|
fi
|
|
|
|
docker run --rm --gpus all --ipc=host --network=host \
|
|
--ulimit memlock=-1 --ulimit stack=67108864 \
|
|
"${DOCKER_MOUNTS[@]}" \
|
|
-w "$REPO_ROOT" \
|
|
"$IMAGE" \
|
|
bash -lc "
|
|
set -euo pipefail
|
|
torchrun --nproc_per_node='$NPROC_PER_NODE' \
|
|
training/megatron_bridge/laoyao_2b_moe_pretrain.py \
|
|
${DATA_ARGS[*]} \
|
|
--seq-length '$SEQ_LENGTH' \
|
|
--train-iters '$TRAIN_ITERS' \
|
|
--micro-batch-size '$MICRO_BATCH_SIZE' \
|
|
--global-batch-size '$GLOBAL_BATCH_SIZE' \
|
|
--tensor-parallel '$TP' \
|
|
--pipeline-parallel '$PP' \
|
|
--expert-parallel '$EP' \
|
|
--context-parallel '$CP' \
|
|
$DRY_RUN_ARG
|
|
"
|