Replace redundant preprocess_data.py with preprocess_data_partitions.py

Author: jon-barker

.gitlab-ci.yml

@@ -40,7 +40,7 @@ unit_tests:
     - export BUILD_DIR=`pwd`
     - export RUN_NAME=resume_${RUN_MODEL}_tp${TP_SIZE}_pp${PP_SIZE}_${NUM_NODES}nodes
     - echo "In case of error check ${SELENE_ADLR_CI_PATH}/${CI_PIPELINE_ID}/${RUN_NAME}/results directory for result logs."
-    - export TP_SIZE PP_SIZE NUM_NODES MAX_STEPS 
+    - export TP_SIZE PP_SIZE NUM_NODES MAX_STEPS
     - export DATA_DIR=$DATA_DIR
     - echo "Run name is $RUN_NAME"
     - mkdir -p $SELENE_ADLR_CI_PATH/$CI_PIPELINE_ID/$RUN_NAME/checkpoints
@@ -65,7 +65,7 @@ unit_tests:
     # Gitlab logs collapsible section markers
     - echo -e "\e[0Ksection_end:`date +%s`:slurm_setup\r\e[0K"
     # Follow output of the job
-    - echo "Finished job"  
+    - echo "Finished job"
     - export SLURM_STATE=$(sacct -j "${SLURM_JOBID}" --format State --parsable2 --noheader |& head -n 1)
     - echo "Slurm job state $SLURM_STATE"
     - if [[ "$SLURM_STATE" != "COMPLETED" ]]; then echo "Slurm job did not complete. See ${SELENE_ADLR_CI_PATH}/${CI_PIPELINE_ID}/${RUN_NAME}/results directory for result logs. Skipping pytest."; exit 1; fi
@@ -79,7 +79,7 @@ unit_tests:
     - if: '$CI_COMMIT_REF_NAME == $CI_DEFAULT_BRANCH && $TEST_LEVEL =~ $TESTS_TO_RUN_AFTER_MERGING'
       when: always
     - if: $CI_MERGE_REQUEST_APPROVED && $TEST_LEVEL =~ $TESTS_TO_RUN_AFTER_MERGE_REQ_APPROVED
-      when: always      
+      when: always
   allow_failure: false
 
 .selene_test_launcher: &selene-test-launcher
@@ -146,7 +146,7 @@ unit_tests:
     - if: '$CI_COMMIT_REF_NAME == $CI_DEFAULT_BRANCH && $TEST_LEVEL =~ $TESTS_TO_RUN_AFTER_MERGING'
       when: always
     - if: $CI_MERGE_REQUEST_APPROVED && $TEST_LEVEL =~ $TESTS_TO_RUN_AFTER_MERGE_REQ_APPROVED
-      when: always      
+      when: always
   allow_failure: false
 
 train.te_gpt3.345m_tp2_pp2_1node_50steps:
@@ -199,7 +199,7 @@ train.gpt3.345m_tp1_pp2_1node_50steps:
     NUM_NODES: 1
     MAX_STEPS: 50
     TIME_LIMIT: "20:00"
-    TEST_LEVEL: L0 
+    TEST_LEVEL: L0
 
 train.gpt3.345m_tp1_pp4_1node_50steps:
   <<: *selene-test-launcher
@@ -224,7 +224,7 @@ resume.checkpoint.gpt3.345m_tp1_pp2_1node:
     PP_SIZE: 2
     NUM_NODES: 1
     TIME_LIMIT: "30:00"
-    TEST_LEVEL: L0 
+    TEST_LEVEL: L0
 
 train.bert.345m_tp4_pp1_1node_50steps:
   <<: *selene-test-launcher
@@ -260,7 +260,7 @@ train.bert.345m_tp1_pp2_1node_50steps:
     NUM_NODES: 1
     MAX_STEPS: 50
     TIME_LIMIT: "20:00"
-    TEST_LEVEL: L0 
+    TEST_LEVEL: L0
 
 train.bert.345m_tp1_pp4_1node_50steps:
   <<: *selene-test-launcher
@@ -284,7 +284,7 @@ resume.checkpoint.bert.345m_tp1_pp2_1node:
     PP_SIZE: 2
     NUM_NODES: 1
     TIME_LIMIT: "30:00"
-    TEST_LEVEL: L0       
+    TEST_LEVEL: L0
 
 cleanup.selene:
   tags:

README.md

@@ -102,13 +102,12 @@ The training data requires preprocessing. First, place your training data in a l
 
 The name of the `text` field of the json can be changed by using the `--json-key` flag in [`preprocess_data.py`](./tools/preprocess_data.py) The other metadata are optional and are not used in training.
 
-The loose json is then processed into a binary format for training. To convert the json into mmap, cached index file, or the lazy loader format use `preprocess_data.py`. Set the `--dataset-impl` flag to `mmap`, `cached`, or `lazy`, respectively (default is `mmap`). An example script to prepare data for BERT training is:
+The loose json is then processed into a binary format for training. To convert the json into mmap format use `preprocess_data.py`. An example script to prepare data for BERT training is:
 <pre>
 python tools/preprocess_data.py \
        --input my-corpus.json \
        --output-prefix my-bert \
-       --vocab bert-vocab.txt \
-       --dataset-impl mmap \
+       --vocab-file bert-vocab.txt \
        --tokenizer-type BertWordPieceLowerCase \
        --split-sentences
 </pre>
@@ -125,7 +124,7 @@ Some minor modifications are required for GPT data preprocessing, namely, the ad
 python tools/preprocess_data.py \
        --input my-corpus.json \
        --output-prefix my-gpt2 \
-       --vocab gpt2-vocab.json \
+       --vocab-file gpt2-vocab.json \
        --dataset-impl mmap \
        --tokenizer-type GPT2BPETokenizer \
        --merge-file gpt2-merges.txt \

cluster_scripts/debug_gpt3.sh

@@ -0,0 +1,69 @@
+#! /bin/bash
+
+
+NAME=gpt3_126m_2_2_debug
+BASE_DIR=/lustre/fsw/adlr/adlr-nlp/jbarker/next-llm/source
+SCRIPTS=${BASE_DIR}/scripts
+MEGATRON=${BASE_DIR}/megatron-lm
+OUTPUT_DIR=${BASE_DIR}/output/debug
+LOGDIR=${OUTPUT_DIR}/logs/${NAME}
+CHECKPOINT_DIR=/lustre/fsw/adlr/adlr-nlp/jbarker/checkpoints/${NAME}
+TENSORBOARD_DIR=${OUTPUT_DIR}/tensorboard/${NAME}
+
+WORLD_SIZE=8
+
+# Get the data blend
+. /lustre/fsw/adlr/adlr-nlp-large/data/gpt3/gpt3_blend.sh
+
+TRAIN_COMMAND=(
+    ${MEGATRON}/pretrain_gpt.py
+    --exit-duration-in-mins 230
+    --tensor-model-parallel-size 1
+    --pipeline-model-parallel-size 8
+    --num-layers 24
+    --hidden-size 768
+    --num-attention-heads 12
+    --seq-length 2048
+    --max-position-embeddings 2048
+    --micro-batch-size 1
+    --global-batch-size 8
+    --train-samples 192000000
+    --lr-decay-samples 166400000
+    --lr-warmup-samples 162761
+    --lr 6.0e-4
+    --min-lr 6.0e-5
+    --lr-decay-style cosine
+    --log-interval 10
+    --exit-interval 1000
+    --log-num-zeros-in-grad
+    --eval-iters 200
+    --eval-interval 2000
+    --data-path ${DATA_BLEND}
+    --vocab-file /lustre/fsw/adlr/adlr-nlp-large/data/bpe/gpt2-vocab.json
+    --merge-file /lustre/fsw/adlr/adlr-nlp-large/data/bpe/gpt2-merges.txt
+    --split 98,2,0
+    --clip-grad 1.0
+    --weight-decay 0.1
+    --adam-beta1 0.9
+    --adam-beta2 0.95
+    --init-method-std 0.023
+    --log-params-norm
+    --log-num-zeros-in-grad
+    --timing-log-level 0
+    --bf16
+    --DDP-impl local
+    --save-interval 1000
+    --save ${CHECKPOINT_DIR}
+)
+
+#    --num-layers-per-virtual-pipeline-stage 1
+
+#    --use-flash-attn
+
+#    --load ${CHECKPOINT_DIR}
+
+CUDA_DEVICE_MAX_CONNECTIONS=1 \
+torchrun --nproc_per_node ${WORLD_SIZE} ${TRAIN_COMMAND[*]}
+
+#    --global-batch-size 256
+#    --rampup-batch-size 32 32 1953125

cluster_scripts/debug_nextllm.sh

@@ -0,0 +1,78 @@
+#! /bin/bash
+
+export CUBLAS_WORKSPACE_CONFIG=:16:8
+
+NAME=nextllm_determinism_debug
+BASE_DIR=/lustre/fsw/adlr/adlr-nlp/jbarker/next-llm
+SCRIPTS=${BASE_DIR}/scripts
+MEGATRON=${BASE_DIR}/source/megatron-lm
+OUTPUT_DIR=${BASE_DIR}/output/debug
+LOGDIR=${OUTPUT_DIR}/logs/${NAME}
+CHECKPOINT_DIR=/lustre/fsw/adlr/adlr-nlp/jbarker/checkpoints/${NAME}
+TENSORBOARD_DIR=${OUTPUT_DIR}/tensorboard/${NAME}
+
+WORLD_SIZE=8
+
+# Get the data blend
+. /lustre/fsw/adlr/adlr-nlp/data/pile-cc1-cc2-shuf/gpt3_blend.sh
+
+BPE_DIR="/lustre/fsw/adlr/adlr-nlp/data/pile-cc1-cc2-shuf/bpe"
+
+TRAIN_COMMAND=(
+    ${MEGATRON}/pretrain_gpt.py
+    --exit-duration-in-mins 230 \
+    --tensor-model-parallel-size 8 \
+    --pipeline-model-parallel-size 8 \
+    #--num-layers-per-virtual-pipeline-stage 1 \
+    --recompute-activations \
+    --sequence-parallel \
+    --num-layers 24 \
+    --hidden-size 768 \
+    --num-attention-heads 24 \
+    --seq-length 2048 \
+    --max-position-embeddings 2048 \
+    --micro-batch-size 1 \
+    --global-batch-size 8 \
+    --train-samples 192000000 \
+    --lr-decay-samples 166400000 \
+    --lr-warmup-samples 244141 \
+    --lr 1.0e-4 \
+    --min-lr 1.0e-5 \
+    --lr-decay-style cosine \
+    --log-interval 1 \
+    --eval-iters 50 \
+    --eval-interval 2000 \
+    --data-path ${DATA_BLEND} \
+    --vocab-file ${BPE_DIR}/gpt2-vocab.json \
+    --merge-file ${BPE_DIR}/gpt2-merges.txt \
+    --save-interval 20000 \
+    --save ${CHECKPOINT_DIR} \
+    --load ${CHECKPOINT_DIR} \
+    --exit-interval 1 \
+    --split 98,2,0 \
+    --clip-grad 1.0 \
+    --weight-decay 0.1 \
+    --adam-beta1 0.9 \
+    --adam-beta2 0.95 \
+    --init-method-std 0.01 \
+    --log-params-norm \
+    --log-num-zeros-in-grad \
+    --bf16 \
+    --DDP-impl local \
+    --tensorboard-dir ${TENSORBOARD_DIR} \
+    --timing-log-level 1 \
+    --timing-log-option minmax \
+)
+
+#    --num-layers-per-virtual-pipeline-stage 1
+
+#    --use-flash-attn
+
+#    --load ${CHECKPOINT_DIR}
+
+CUDA_DEVICE_MAX_CONNECTIONS=1 \
+CUBLAS_WORKSPACE_CONFIG=:16:8 \
+torchrun --nproc_per_node ${WORLD_SIZE} ${TRAIN_COMMAND[*]}
+
+#    --global-batch-size 256
+#    --rampup-batch-size 32 32 1953125

cluster_scripts/run_foundation_model_medium_dp1_adaptive_routing-22.12-noflash.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+
+#SBATCH -p luna -A adlr -t 04:00:00 --dependency=singleton --nodes=1 --exclusive --mem=0 --overcommit --ntasks-per-node=8 --job-name=adlr-nlp:foundation-model-medium_dp1_adaptve_routing-22.12-noflash-repeat
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_SL=1
+
+BRANCH=${1}
+COMMIT=${2}
+CONTAINER=${3}
+NUMBER=${4}
+
+NAME="foundation-model-medium_dp1_adaptive_routing-22.12-noflash-${NUMBER}"
+
+SOURCE="/lustre/fsw/adlr/adlr-nlp/jbarker/next-llm/source/megatron-lm"
+OUTPUT="/lustre/fsw/adlr/adlr-nlp/jbarker/next-llm/output/pretraining.${BRANCH}.${COMMIT}.${CONTAINER}/${NAME}/"
+
+SCRIPTS_DIR="/lustre/fsw/adlr/adlr-nlp/jbarker/next-llm/source/"
+
+CHECKPOINTS_DIR="${OUTPUT}/checkpoints"
+TENSORBOARD_DIR="${OUTPUT}/tensorboard"
+LOGS_DIR="${OUTPUT}/logs"
+
+mkdir -p ${CHECKPOINTS_DIR}
+mkdir -p ${TENSORBOARD_DIR}
+mkdir -p ${LOGS_DIR}
+
+# CHECKPOINT_DIR="/lustre/fsw/adlr/adlr-nlp/mshoeybi/checkpoints/foundation_model/speed/${NAME}"
+
+# Get the data blend
+. /lustre/fsw/adlr/adlr-nlp/data/pile-cc1-cc2-shuf/gpt3_blend.sh
+
+BPE_DIR="/lustre/fsw/adlr/adlr-nlp/data/pile-cc1-cc2-shuf/bpe"
+
+#    --num-layers-per-virtual-pipeline-stage 3 \
+
+options=" \
+    --exit-duration-in-mins 230 \
+    --exit-interval 100000 \
+    --tensor-model-parallel-size 8 \
+    --pipeline-model-parallel-size 1 \
+    --recompute-activations \
+    --sequence-parallel \
+    --num-layers 12 \
+    --hidden-size 8192 \
+    --num-attention-heads 64 \
+    --seq-length 2048 \
+    --max-position-embeddings 2048 \
+    --micro-batch-size 1 \
+    --global-batch-size 16 \
+    --train-samples 192000000 \
+    --lr-decay-samples 166400000 \
+    --lr-warmup-samples 244141 \
+    --lr 1.0e-4 \
+    --min-lr 1.0e-5 \
+    --lr-decay-style cosine \
+    --log-interval 1 \
+    --eval-iters 50 \
+    --eval-interval 2000 \
+    --data-path ${DATA_BLEND} \
+    --vocab-file ${BPE_DIR}/gpt2-vocab.json \
+    --merge-file ${BPE_DIR}/gpt2-merges.txt \
+    --save-interval 2000 \
+    --save ${CHECKPOINTS_DIR} \
+    --load ${CHECKPOINTS_DIR} \
+    --split 98,2,0 \
+    --clip-grad 1.0 \
+    --weight-decay 0.1 \
+    --adam-beta1 0.9 \
+    --adam-beta2 0.95 \
+    --init-method-std 0.01 \
+    --log-params-norm \
+    --log-num-zeros-in-grad \
+    --bf16 \
+    --DDP-impl local \
+    --tensorboard-dir ${TENSORBOARD_DIR} \
+    --timing-log-level 1 \
+    --timing-log-option minmax \
+"
+
+run_cmd="${SCRIPTS_DIR}/bind.sh --cpu=${SCRIPTS_DIR}/dgxa100_ccx.sh --mem=${SCRIPTS_DIR}/dgxa100_ccx.sh python -u ${SOURCE}/pretrain_gpt.py ${options}"
+
+DATETIME=`date +'date_%y-%m-%d_time_%H-%M-%S'`
+
+#     --container-image "gitlab-master.nvidia.com/adlr/megatron-lm/pytorch_flash_att:22.12-py3" \
+
+srun -l \
+     --container-image nvcr.io#nvidia/pytorch:22.09-py3 \
+     --container-mounts "/lustre/fsw/adlr:/lustre/fsw/adlr" \
+     --output=${LOGS_DIR}/%x_%j_$DATETIME.log sh -c "${run_cmd}"
+
+set +x
+