slurm_example.sub

#!/bin/bash
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH --overcommit
#SBATCH --parsable

# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
# Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.


#
# This Slurm script demonstrates how to use the LDDL preprocessor, load balancer
# and data loader and scale them to multi-nodes on Slurm clusters for (mock)
# BERT Phase 2 pretraining with static masking and sequence binning enabled.
#

set -eux

#
# The following configurations might need to be customized based on the setup
# of the Slurm cluster you are using.
#

# The URL of the container image built via `bash docker/build.sh`.
# For example, if you build the container image by 
# `bash docker/build.sh ngc_pyt 21.11-py3 lddl:latest push`, 
# then the URL would be "lddl:latest":
readonly docker_image=${DOCKER_IMAGE:-"lddl:latest"}

# Create a directory to store data.
mkdir -p data/

# Assume the Wikipedia dump is already downloaded and moved to the following
# location in the NFS of your Slurm cluster.
#
# Please refer to examples/local_example.sh on how to use the LDDL downloader
# to download the Wikipedia dump.
readonly wikipedia_path=data/wikipedia

# Download the vocab file from NVIDIA Deep Learning Examples (but you can
# certainly get it from other sources as well).
readonly vocab_source_url=https://raw.githubusercontent.com/NVIDIA/DeepLearningExamples/master/PyTorch/LanguageModeling/BERT/vocab/vocab
mkdir -p data/vocab/
readonly vocab_path=data/vocab/bert-en-uncased.txt
wget ${vocab_source_url} -O ${vocab_path}

# Run the LDDL preprocessor for BERT Phase 2 pretraining with static masking and
# sequence binning enabled (where the bin size is 64).
readonly mounts=$(realpath data/):/workspace/lddl/data
readonly workdir=/workspace/lddl
readonly num_shards=4096
readonly bin_size=64
readonly tasks_per_node=128
readonly pretrain_input_path=data/bert/pretrain/phase2/bin_size_${bin_size}/
srun \
  -l \
  --mpi=pmix \
  --container-image="${docker_image}" \
  --container-mounts="${mounts}" \
  --container-workdir=${workdir} \
  --ntasks-per-node=${tasks_per_node} \
  --export=ALL,LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so \
    preprocess_bert_pretrain \
      --schedule mpi \
      --vocab-file ${vocab_path} \
      --wikipedia ${wikipedia_path}/source/ \
      --sink ${pretrain_input_path} \
      --target-seq-length 512 \
      --num-blocks ${num_shards} \
      --bin-size ${bin_size} \
      --masking

# Run the LDDL load balancer to balance the parquet shards generated by the LDDL
# preprocessor.
srun \
  -l \
  --mpi=pmix \
  --container-image="${docker_image}" \
  --container-mounts="${mounts}" \
  --container-workdir=${workdir} \
  --ntasks-per-node=${tasks_per_node} \
    balance_dask_output \
      --indir ${pretrain_input_path} \
      --num-shards ${num_shards}

# Run a mock PyTorch training script that loads the input from the balanced
# parquet shards using the LDDL data loader.
# Once these training processes is up and running (as you can see from the
# stdout printing), it simply emulates training and you can kill it at any time.
readonly gpus_per_node=8
srun \
  -l \
  --container-image="${docker_image}" \
  --container-mounts="${mounts}" \
  --container-workdir=${workdir} \
  --ntasks-per-node=${gpus_per_node} \
    python benchmarks/torch_train.py \
      --path ${pretrain_input_path} \
      --vocab-file ${vocab_path}