NVIDIA · nv-kkudrynski · Aug 17, 2020 · Aug 17, 2020
diff --git a/PyTorch/Recommendation/DLRM/Dockerfile b/PyTorch/Recommendation/DLRM/Dockerfile
@@ -12,15 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:20.06-py3
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:20.07-py3
 FROM ${FROM_IMAGE_NAME}
 
-RUN apt update && \
-    apt install -y openjdk-8-jdk && \
-    curl http://archive.apache.org/dist/spark/spark-2.4.5/spark-2.4.5-bin-hadoop2.7.tgz -o /opt/spark-2.4.5-bin-hadoop2.7.tgz && \
-    tar zxf /opt/spark-2.4.5-bin-hadoop2.7.tgz -C /opt/ && \
-    rm /opt/spark-2.4.5-bin-hadoop2.7.tgz
-
 ADD requirements.txt .
 RUN pip install -r requirements.txt
 

diff --git a/PyTorch/Recommendation/DLRM/Dockerfile_spark b/PyTorch/Recommendation/DLRM/Dockerfile_spark
@@ -0,0 +1,39 @@
+# Copyright (c) 2020 NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:20.03-py3
+FROM ${FROM_IMAGE_NAME}
+
+RUN apt update && \
+    apt install -y openjdk-8-jdk && \
+    curl https://downloads.apache.org/spark/spark-3.0.0/spark-3.0.0-bin-hadoop3.2.tgz -o /opt/spark-3.0.0-bin-hadoop3.2.tgz && \
+    tar zxf /opt/spark-3.0.0-bin-hadoop3.2.tgz -C /opt/ && \
+    rm /opt/spark-3.0.0-bin-hadoop3.2.tgz && \
+    curl https://repo1.maven.org/maven2/ai/rapids/cudf/0.14/cudf-0.14-cuda10-2.jar -o /opt/cudf-0.14-cuda10-2.jar && \
+    curl https://repo1.maven.org/maven2/com/nvidia/rapids-4-spark_2.12/0.1.0/rapids-4-spark_2.12-0.1.0.jar -o /opt/rapids-4-spark_2.12-0.1.0.jar
+
+ADD requirements.txt .
+RUN pip install -r requirements.txt
+
+WORKDIR /workspace/dlrm
+
+COPY . .
+
+RUN mv /opt/cudf-0.14-cuda10-2.jar  /opt/spark-3.0.0-bin-hadoop3.2/jars && \
+    mv /opt/rapids-4-spark_2.12-0.1.0.jar /opt/spark-3.0.0-bin-hadoop3.2/jars/ && \
+    mv /workspace/dlrm/preproc/gpu/get_gpu_resources.sh /opt/spark-3.0.0-bin-hadoop3.2/conf/ && \
+    mv /workspace/dlrm/preproc/gpu/spark-defaults.conf /opt/spark-3.0.0-bin-hadoop3.2/conf/ && \
+    rm -fr /workspace/dlrm/preproc/gpu
+
+RUN chmod +x /opt/spark-3.0.0-bin-hadoop3.2/conf/get_gpu_resources.sh
diff --git a/PyTorch/Recommendation/DLRM/README.md b/PyTorch/Recommendation/DLRM/README.md
diff --git a/PyTorch/Recommendation/DLRM/bind.sh b/PyTorch/Recommendation/DLRM/bind.sh
@@ -0,0 +1,203 @@
+# Copyright (c) 2020 NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+#! /bin/bash
+set -euo pipefail
+
+print_usage() {
+    cat << EOF
+${0} [options] [--] COMMAND [ARG...]
+
+Control binding policy for each task. Assumes one rank will be launched for each GPU.
+
+Options:
+    --cpu=MODE
+        * exclusive -- bind each rank to an exclusive set of cores near its GPU
+        * exclusive,nosmt -- bind each rank to an exclusive set of cores near its GPU, without hyperthreading
+        * node -- bind each rank to all cores in the NUMA node nearest its GPU [default]
+	* *.sh -- bind each rank using the bash associative array bind_cpu_cores or bind_cpu_nodes from a file
+        * off -- don't bind
+    --mem=MODE
+        * node -- bind each rank to the nearest NUMA node [default]
+	* *.sh -- bind each rank using the bash associative array bind_mem from a file
+        * off -- don't bind
+    --ib=MODE
+        * single -- bind each rank to a single IB device near its GPU
+        * off -- don't bind [default]
+EOF
+}
+
+################################################################################
+# Argument parsing
+################################################################################
+
+cpu_mode='node'
+mem_mode='node'
+ib_mode='off'
+while [ $# -gt 0 ]; do
+    case "$1" in
+        -h|--help) print_usage ; exit 0 ;;
+        --cpu=*) cpu_mode="${1/*=/}"; shift ;;
+        --cpu)   cpu_mode="$2"; shift 2 ;;
+        --mem=*) mem_mode="${1/*=/}"; shift ;;
+        --mem)   mem_mode="$2"; shift 2 ;;
+        --ib=*) ib_mode="${1/*=/}"; shift ;;
+        --ib)   ib_mode="$2"; shift 2 ;;
+        --) shift; break ;;
+        *) break ;;
+    esac
+done
+if [ $# -lt 1 ]; then
+    echo 'ERROR: no command given' 2>&1
+    print_usage
+    exit 1
+fi
+
+################################################################################
+# Get system params
+################################################################################
+
+# LOCAL_RANK is set with an enroot hook for Pytorch containers
+# SLURM_LOCALID is set by Slurm
+# OMPI_COMM_WORLD_LOCAL_RANK is set by mpirun
+readonly local_rank="${LOCAL_RANK:=${SLURM_LOCALID:=${OMPI_COMM_WORLD_LOCAL_RANK:-}}}"
+if [ -z "${local_rank}" ]; then
+    echo 'ERROR: cannot read LOCAL_RANK from env' >&2
+    exit 1
+fi
+
+num_gpus=$(nvidia-smi -i 0 --query-gpu=count --format=csv,noheader,nounits)
+if [ "${local_rank}" -ge "${num_gpus}" ]; then
+    echo "ERROR: local rank is ${local_rank}, but there are only ${num_gpus} gpus available" >&2
+    exit 1
+fi
+
+get_lscpu_value() {
+    awk -F: "(\$1 == \"${1}\"){gsub(/ /, \"\", \$2); print \$2; found=1} END{exit found!=1}"
+}
+lscpu_out=$(lscpu)
+num_sockets=$(get_lscpu_value 'Socket(s)' <<< "${lscpu_out}")
+num_nodes=$(get_lscpu_value 'NUMA node(s)' <<< "${lscpu_out}")
+cores_per_socket=$(get_lscpu_value 'Core(s) per socket' <<< "${lscpu_out}")
+
+echo "num_sockets = ${num_sockets} num_nodes=${num_nodes} cores_per_socket=${cores_per_socket}"
+
+readonly cores_per_node=$(( (num_sockets * cores_per_socket) / num_nodes ))
+if [ ${num_gpus} -gt 1 ]; then
+    readonly gpus_per_node=$(( num_gpus / num_nodes ))
+else
+    readonly gpus_per_node=1
+fi
+readonly cores_per_gpu=$(( cores_per_node / gpus_per_node ))
+readonly local_node=$(( local_rank / gpus_per_node ))
+
+
+declare -a ibdevs=()
+readonly num_ibdevs="${#ibdevs[@]}"
+
+################################################################################
+# Setup for exec
+################################################################################
+
+declare -a numactl_args=()
+
+case "${cpu_mode}" in
+    exclusive)
+        numactl_args+=( "$(printf -- "--physcpubind=%u-%u,%u-%u" \
+            $(( local_rank * cores_per_gpu )) \
+            $(( (local_rank + 1) * cores_per_gpu - 1 )) \
+            $(( local_rank * cores_per_gpu + (cores_per_gpu * gpus_per_node * num_nodes) )) \
+            $(( (local_rank + 1) * cores_per_gpu + (cores_per_gpu * gpus_per_node * num_nodes) - 1 )) \
+        )" )
+        ;;
+    exclusive,nosmt)
+        numactl_args+=( "$(printf -- "--physcpubind=%u-%u" \
+            $(( local_rank * cores_per_gpu )) \
+            $(( (local_rank + 1) * cores_per_gpu - 1 )) \
+        )" )
+        ;;
+    node)
+        numactl_args+=( "--cpunodebind=${local_node}" )
+        ;;
+    *.sh)
+    source "${cpu_mode}"
+    if [ -n "${bind_cpu_cores:-}" ]; then
+        numactl_args+=( "--physcpubind=${bind_cpu_cores[${local_rank}]}" )
+    elif [ -n "${bind_cpu_nodes:-}" ]; then
+        numactl_args+=( "--cpunodebind=${bind_cpu_nodes[${local_rank}]}" )
+    else
+        echo "ERROR: invalid CPU affinity file ${cpu_mode}." >&2
+        exit 1
+    fi
+    ;;
+    off|'')
+        ;;
+    *)
+        echo "ERROR: invalid cpu mode '${cpu_mode}'" 2>&1
+        print_usage
+        exit 1
+        ;;
+esac
+
+case "${mem_mode}" in
+    node)
+        numactl_args+=( "--membind=${local_node}" )
+        ;;
+    *.sh)
+    source "${mem_mode}"
+    if [ -z "${bind_mem:-}" ]; then
+        echo "ERROR: invalid memory affinity file ${mem_mode}." >&2
+        exit 1
+    fi
+    numactl_args+=( "--membind=${bind_mem[${local_rank}]}" )
+    ;;
+    off|'')
+        ;;
+    *)
+        echo "ERROR: invalid mem mode '${mem_mode}'" 2>&1
+        print_usage
+        exit 1
+        ;;
+esac
+
+case "${ib_mode}" in
+    single)
+        if [ "${num_ibdevs}" -eq 0 ]; then
+            echo "WARNING: used '$0 --ib=single', but there are 0 IB devices available; skipping IB binding." 2>&1
+        else
+            readonly ibdev="${ibdevs[$(( local_rank * num_ibdevs / num_gpus ))]}"
+            export OMPI_MCA_btl_openib_if_include="${OMPI_MCA_btl_openib_if_include-$ibdev}"
+            export UCX_NET_DEVICES="${UCX_NET_DEVICES-$ibdev:1}"
+        fi
+        ;;
+    off|'')
+        ;;
+    *)
+        echo "ERROR: invalid ib mode '${ib_mode}'" 2>&1
+        print_usage
+        exit 1
+        ;;
+esac
+
+################################################################################
+# Exec
+################################################################################
+
+if [ "${#numactl_args[@]}" -gt 0 ] ; then
+    set -x
+    exec numactl "${numactl_args[@]}" -- "${@}"
+else
+    exec "${@}"
+fi
diff --git a/PyTorch/Recommendation/DLRM/dgxa100_ccx.sh b/PyTorch/Recommendation/DLRM/dgxa100_ccx.sh
@@ -0,0 +1,22 @@
+# Copyright (c) 2020 NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#! /bin/bash
+
+
+bind_cpu_cores=([0]="48-51,176-179" [1]="60-63,188-191" [2]="16-19,144-147" [3]="28-31,156-159"
+                [4]="112-115,240-243" [5]="124-127,252-255" [6]="80-83,208-211" [7]="92-95,220-223")
+
+bind_mem=([0]="3" [1]="3" [2]="1" [3]="1"
+          [4]="7" [5]="7" [6]="5" [7]="5")
diff --git a/PyTorch/Recommendation/DLRM/dlrm/data/datasets.py b/PyTorch/Recommendation/DLRM/dlrm/data/datasets.py
@@ -236,7 +236,10 @@ def _get_categorical_features(self, idx: int) -> Optional[torch.Tensor]:
         return torch.cat(categorical_features, dim=1)
 
     def __del__(self):
-        data_files = [self._label_file, self._numerical_features_file] + self._categorical_features_files
+        data_files = [self._label_file, self._numerical_features_file]
+        if self._categorical_features_files is not None:
+            data_files += self._categorical_features_files
+
         for data_file in data_files:
             if data_file is not None:
                 os.close(data_file)
diff --git a/PyTorch/Recommendation/DLRM/dlrm/data/factories.py b/PyTorch/Recommendation/DLRM/dlrm/data/factories.py
@@ -163,19 +163,24 @@ def create_datasets(self) -> Tuple[Dataset, Dataset]:
         test_dataset_path = os.path.join(self._flags.dataset, "test")
         categorical_sizes = get_categorical_feature_sizes(self._flags)
 
+        # prefetching is currently unsupported if using the batch-wise shuffle
+        prefetch_depth = 0 if self._flags.shuffle_batch_order else 10
+
         dataset_train = SplitCriteoDataset(
             data_path=train_dataset_path,
             batch_size=self._flags.batch_size,
             numerical_features=self._numerical_features,
             categorical_features=self._categorical_features,
-            categorical_feature_sizes=categorical_sizes
+            categorical_feature_sizes=categorical_sizes,
+            prefetch_depth=prefetch_depth
         )
         dataset_test = SplitCriteoDataset(
             data_path=test_dataset_path,
             batch_size=self._flags.test_batch_size,
             numerical_features=self._numerical_features,
             categorical_features=self._categorical_features,
-            categorical_feature_sizes=categorical_sizes
+            categorical_feature_sizes=categorical_sizes,
+            prefetch_depth=prefetch_depth
         )
         return dataset_train, dataset_test
 

diff --git a/PyTorch/Recommendation/DLRM/dlrm/data/samplers.py b/PyTorch/Recommendation/DLRM/dlrm/data/samplers.py
@@ -30,7 +30,7 @@ def __iter__(self):
         to other rank by writing to disk
         """
         if get_local_rank() == 0:
-            np.save(self._SAMPLE_FILE, np.array(super().__iter__()))
+            np.save(self._SAMPLE_FILE, np.array(list(super().__iter__())))
         torch.distributed.barrier()
 
         sample = np.load(self._SAMPLE_FILE)

diff --git a/PyTorch/Recommendation/DLRM/dlrm/nn/parts.py b/PyTorch/Recommendation/DLRM/dlrm/nn/parts.py
@@ -85,7 +85,7 @@ def forward(self, numerical_input, categorical_inputs) -> Tuple[torch.Tensor, Op
         Returns:
             Tensor: Concatenated bottom mlp and embedding output in shape [batch, 1 + #embedding, embedding_dim]
         """
-        batch_size = categorical_inputs.size()[0]
+        batch_size = len(numerical_input) if numerical_input is not None else len(categorical_inputs)
         bottom_output = []
         bottom_mlp_output = None
 
@@ -97,7 +97,8 @@ def forward(self, numerical_input, categorical_inputs) -> Tuple[torch.Tensor, Op
             # reshape bottom mlp to concatenate with embeddings
             bottom_output.append(bottom_mlp_output.view(batch_size, 1, -1))
 
-        bottom_output += self.embeddings(categorical_inputs)
+        if self.num_categorical_features > 0:
+            bottom_output += self.embeddings(categorical_inputs)
 
         if self._fp16:
             bottom_output = [x.half() if x.dtype != torch.half else x for x in bottom_output]

diff --git a/PyTorch/Recommendation/DLRM/dlrm/scripts/dist_main.py b/PyTorch/Recommendation/DLRM/dlrm/scripts/dist_main.py
@@ -32,7 +32,7 @@
 from dlrm.scripts.main import FLAGS, get_categorical_feature_sizes
 from dlrm.utils import distributed as dist
 from dlrm.utils.checkpointing.distributed import make_distributed_checkpoint_writer, make_distributed_checkpoint_loader
-from dlrm.utils.distributed import get_gpu_batch_sizes, get_criteo_device_mapping, is_main_process, is_distributed
+from dlrm.utils.distributed import get_gpu_batch_sizes, get_device_mapping, is_main_process, is_distributed
 
 # Training schedule flags
 FLAGS.set_default("batch_size", 65536)
@@ -75,7 +75,7 @@ def main(argv):
 
     categorical_feature_sizes = get_categorical_feature_sizes(FLAGS)
     world_categorical_feature_sizes = np.asarray(categorical_feature_sizes)
-    device_mapping = get_criteo_device_mapping(world_size)
+    device_mapping = get_device_mapping(categorical_feature_sizes, num_gpus=world_size)
 
     batch_sizes_per_gpu = get_gpu_batch_sizes(FLAGS.batch_size, num_gpus=world_size)
     batch_indices = tuple(np.cumsum([0] + list(batch_sizes_per_gpu)))

diff --git a/PyTorch/Recommendation/DLRM/dlrm/scripts/main.py b/PyTorch/Recommendation/DLRM/dlrm/scripts/main.py
@@ -139,6 +139,9 @@ def is_data_prefetching_enabled() -> bool:
 def create_model():
     print("Creating model")
 
+    FLAGS.top_mlp_sizes = [int(s) for s in FLAGS.top_mlp_sizes]
+    FLAGS.bottom_mlp_sizes = [int(s) for s in FLAGS.bottom_mlp_sizes]
+
     model_config = {
         'top_mlp_sizes': FLAGS.top_mlp_sizes,
         'bottom_mlp_sizes': FLAGS.bottom_mlp_sizes,