offline_feature_extract.sh

#!/bin/bash

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"

# Common arguments and base model specific arguments
source "${SCRIPT_DIR}/conf/arguments.sh"


# Set up cache dir for HF to avoid out of space error
export HF_DATASETS_CACHE="/tmp/hf_datasets_cache"

# Extra arguments of this script
MLM_DEFAULT_ARGS=" \
    --distributed-timeout-minutes 30 \
    --auto-detect-ckpt-format \
    --export-te-mcore-model \
    --finetune \
"


if [ -z ${MLM_DATA_ARGS} ]; then
    MLM_DATA_ARGS=" \
        --num-samples 128000 \
        --finetune-hf-dataset nvidia/Daring-Anteater \
    "
fi


${LAUNCH_SCRIPT} ${SCRIPT_DIR}/offline_feature_extract.py \
    ${MODEL_ARGS} \
    --tensor-model-parallel-size ${TP} \
    --expert-tensor-parallel-size ${ETP} \
    --expert-model-parallel-size ${EP} \
    --pipeline-model-parallel-size ${PP} \
    --tokenizer-model ${TOKENIZER_MODEL} \
    --load ${MLM_MODEL_CKPT} \
    ${MLM_DATA_ARGS} \
    ${MLM_DEFAULT_ARGS} ${MLM_EXTRA_ARGS}