diff --git a/training/slurm/quickstart/launch_training.sbatch b/training/slurm/quickstart/launch_training.sbatch new file mode 100644 index 0000000..111d5c1 --- /dev/null +++ b/training/slurm/quickstart/launch_training.sbatch @@ -0,0 +1,162 @@ +#!/bin/bash + +#SBATCH --nodes 1 +#SBATCH --ntasks-per-node 8 +#SBATCH --gpus-per-node 8 +#SBATCH --constraint gpu +#SBATCH --job-name test +#SBATCH --output test.%j +#SBATCH --export all +#SBATCH --exclusive + +export NCCL_SOCKET_IFNAME=eth0 +export SHARP_COLL_ENABLE_PCI_RELAXED_ORDERING=1 +export NCCL_COLLNET_ENABLE=0 +export NCCL_IB_HCA=ibp +export UCX_NET_DEVICES=ibp0:1,ibp1:1,ibp2:1,ibp3:1,ibp4:1,ibp5:1,ibp6:1,ibp7:1 + +export MASTER_PORT="$(expr 10000 + "$(echo -n "${SLURM_JOB_ID:?}" | tail -c 4)")" +export MASTER_ADDR="$(scontrol show hostnames "${SLURM_JOB_NODELIST:?}" | head -n 1)" + +# This will be used for both the output directory as well as the wandb project +export PROJECT_NAME="test" + +# If you change this, make sure the --container-mount can find it +export OUTDIR="/mnt/data/coreweave/${PROJECT_NAME}" + +# Megatron's dataset indices will go here once built +export CACHEDIR="${OUTDIR}/cache" + +# save images here +export CONTAINER_DIR="/mnt/data/coreweave/images" + +# uncomment if running on H100/H200 +#CPU_BIND='map_ldom:0,0,0,0,1,1,1,1' + + +# Use squished container if available, otw download and save it. +export REMOTE_IMAGE_URI="ghcr.io#coreweave/ml-containers/megatron:dmarx-megatron-update-0c68584" +export IMAGE_NAME="training-test.sqsh" +export LOCAL_IMAGE_PATH="${CONTAINER_DIR}/${IMAGE_NAME}" + +CONTAINER_SAVE="" +if [ -f "$LOCAL_IMAGE_PATH" ]; then + CONTAINER_IMAGE="${LOCAL_IMAGE_PATH}" +else + CONTAINER_IMAGE="${REMOTE_IMAGE_URI}" + CONTAINER_SAVE="--container-save=${LOCAL_IMAGE_PATH}" +fi + + +srun --container-image "${CONTAINER_IMAGE}" ${CONTAINER_SAVE:+"$CONTAINER_SAVE"} \ + --container-mounts /mnt/data:/mnt/data \ + --export=ALL \ + --mpi=pmix \ + --kill-on-bad-exit=1 \ + ${CPU_BIND:+"--cpu-bind=$CPU_BIND"} \ + bash <<'EOF' +#### in-line srun script #### +#!/bin/bash + +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +export WORLD_SIZE="${SLURM_NTASKS:?}" +export RANK="${SLURM_PROCID:?}" +export LOCAL_RANK="${SLURM_LOCALID:?}" +export CUDA_DEVICE_ORDER='PCI_BUS_ID' + + +CKPTDIR_LOAD="${CKPTDIR_LOAD:-${OUTDIR}/checkpoints}" +CKPTDIR_SAVE="${CKPTDIR_SAVE:-${OUTDIR}/checkpoints}" +mkdir -p "${CKPTDIR_SAVE}" + + +touch "${CKPTDIR_SAVE}/progress.txt" + + +cd /usr/src/app/megatron-lm + +WARNING_FILTERS=( +'-Wignore::DeprecationWarning' +'-Wignore::FutureWarning' +'-Wignore::UserWarning:megatron.core.tensor_parallel.layers' # "async_grad_allreduce is deprecated" +'-Wignore::UserWarning:megatron.core.optimizer.distrib_optimizer' # "pre_hook" method deprecations +) + +WANDB_ARGS=() +# Enable W&B flags only if credentials are available via ~/.netrc or env var +if [ -f "${HOME}/.netrc" ] || [ -n "${WANDB_API_KEY:-}" ] || [ -n "${WANDB_APIP_KEY:-}" ]; then + WANDB_ARGS+=("--wandb-exp-name" "${SLURM_JOB_ID:?}/megatron-test") + WANDB_ARGS+=("--wandb-project" "${PROJECT_NAME}") + WANDB_ARGS+=("--wandb-save-dir" "${OUTDIR}/logs") +fi + +python3 "${WARNING_FILTERS[@]:?}" \ + "/usr/src/app/megatron-lm/pretrain_gpt.py" \ + ${WANDB_ARGS[@]+"${WANDB_ARGS[@]}"} \ + --tensorboard-dir "${OUTDIR}/tensorboard" \ + --data-cache-path ${CACHEDIR} \ + --load "${CKPTDIR_LOAD}" \ + --save "${CKPTDIR_SAVE}" \ + --save-interval 3500 \ + --tensor-model-parallel-size 4 \ + --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --sequence-parallel \ + --overlap-grad-reduce \ + --overlap-param-gather \ + --train-iters 2507400 \ + --lr 4e-06 \ + --min-lr 4e-07 \ + --lr-warmup-iters 2000 \ + --lr-decay-iters 2507400 \ + --lr-decay-style cosine \ + --clip-grad 1.0 \ + --bf16 \ + --use-flash-attn \ + --rotary-seq-len-interpolation-factor 32 \ + --no-fp8-wgrad \ + --use-distributed-optimizer \ + --distributed-backend nccl \ + --split 949,50,1 \ + --seed 42 \ + --use-checkpoint-args \ + --no-masked-softmax-fusion \ + --attention-softmax-in-fp32 \ + --transformer-impl transformer_engine \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --rotary-base 500000 \ + --rotary-percent 1.0 \ + --use-rope-scaling \ + --micro-batch-size 1 \ + --log-interval 1 \ + --tensorboard-log-interval 1 \ + --eval-interval 100 \ + --eval-iters 10 \ + --logging-level 20 \ + --log-params-norm \ + --log-num-zeros-in-grad \ + --log-throughput \ + --log-progress \ + --timing-log-level 0 \ + --timing-log-option all \ + --log-timers-to-tensorboard \ + --log-validation-ppl-to-tensorboard \ + --log-memory-to-tensorboard \ + --log-world-size-to-tensorboard \ + --ffn-hidden-size 11008 \ + --num-attention-heads 32 \ + --num-layers 32 \ + --hidden-size 4096 \ + --seq-length 8192 \ + --max-position-embeddings 8192 \ + --untie-embeddings-and-output-weights \ + --normalization RMSNorm \ + --swiglu \ + --position-embedding-type rope \ + --disable-bias-linear \ + --group-query-attention \ + --num-query-groups 8 \ + --mock-data +EOF