Bench Shuffle Pinned Memory Debugging

[quasiben]

When running current bench_shuffle benchmarks on an NVL72 I'm seeing an error then pinned memory is enabled (this is the default)

APP_ARGS="-C ucxx -w 3 -r 10 -g -s -x -l $((40 * 1024)) -p ${NUM_PARTITIONS} -o 4 -c ${NUM_COLUMNS} -n ${NUM_ROWS} -x -m async"

  -c ucxx (communicator)
  -r 10 (number of runs)
  -w 3 (number of warmup runs)
  -c 4 (number of columns)
  -n 67108864 (number of rows per rank)
  -p 50 (number of input partitions per rank)
  -o 4 (number of output partitions per rank)
  -m async (RMM memory resource)
  -l 40960 (device memory limit in MiB)
  -s (enable output discard to simulate streaming)
  -x (enable memory profiling)
  -g (use pre-partitioned input tables)
Local size: 50 GiB

terminate called after throwing an instance of 'cuda::__4::cuda_error'
  what():  /opt/conda/envs/rapidsmpf/include/rapids/cuda/__driver/driver_api.h:152 out of memory(2): Failed to allocate memory from a memory pool
[presto-gb200-gcn-07:1222606] *** Process received signal ***
[presto-gb200-gcn-07:1222606] Signal: Aborted (6)

When I disable pinned memory with -L the job completes without issue. One interesting note is that this does not seem to reproduce on a cluster made up of DGXH100s

[nirandaperera]

@quasiben I will look into this. I have seen exceptions when I was testing pinned memory blowup as well.
On a side note, in NVL72. this issue likely affects this as well (processes are not properly pinned to a numa node). For host memory, this is not a problem, but it would be suboptimal.

[pentschev]

To add more details, here are some observations I collected today running on 2xNVL4 nodes:

Spilling starts at 80GiB/rank (640GiB global) with 40GiB device spill limit (seemingly indicating we spill roughly when data is roughly 2x the device spill limit).

Large sizes with small device memory limit (e.g., 120GiB/rank and 40GiB device spill limit, 150 GiB/rank and 40GiB or 80GiB device spill limit) we hit OOMs:

terminate called after throwing an instance of 'cuda::__4::cuda_error'
  what():  /opt/conda/envs/rapidsmpf/include/rapids/cuda/__driver/driver_api.h:152 out of memory(2): Failed to allocate memory from a memory pool

Sizes >=120GiB/rank (960GiB global) we observe rows/column size limit errors when the benchmark does not OOM first (i.e., has a high device spill limit):

terminate called after throwing an instance of 'std::overflow_error'
  what():  CUDF failure at: /tmp/conda-bld-output/bld/rattler-build_libcudf/work/cpp/src/copying/concatenate.cu:474: Total number of concatenated rows exceeds the column size limit

[nirandaperera]

@quasiben @pentschev I think I found the issue. From what I can see, slurm in the NVL72 cluster does not pin processes to proper numa nodes. If we set --tasks-per-node 4, all processes are not given a numa binding. However, each process correctly receives a unique GPU. See the following,

srun \
  --nodes=1 \
  --ntasks-per-node=4 \
  --gpus-per-node=4 \
  --gpus-per-task=1 \
  --cpus-per-task=36 \
  --cpu-bind=verbose,cores \
  --mem-bind=verbose,local \
  --hint=nomultithread \
  --exclusive \
  --label \
  bash -c '
    echo "=== rank=$SLURM_PROCID localid=$SLURM_LOCALID host=$(hostname) ==="
    echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
    echo "CPU affinity        : $(taskset -pc $$ | cut -d: -f2 | xargs)"
    echo "NUMA nodes (Mems)   : $(grep Mems_allowed_list /proc/self/status | cut -f2)"
    echo "NUMA nodes (Cpus->) : $(grep Cpus_allowed_list /proc/self/status | cut -f2)"
    nvidia-smi --query-gpu=index,uuid,pci.bus_id --format=csv,noheader
    for bus in $(nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader); do
      short=$(echo $bus | tr "[:upper:]" "[:lower:]" | sed "s/^0000//")
      echo "GPU $bus -> NUMA $(cat /sys/bus/pci/devices/$short/numa_node 2>/dev/null)"
    done
    echo
  '

output 

0: 
0: 0, GPU-6b83f80d-2369-92fa-5df8-0d1f78c71561, 00000008:01:00.0
0: === rank=0 localid=0 host=presto-gb200-gcn-05 ===
0: CPU affinity        : 0-143
0: CUDA_VISIBLE_DEVICES=0
0: GPU 00000008:01:00.0 -> NUMA 0
0: NUMA nodes (Cpus->) : 0-143
0: NUMA nodes (Mems)   : 0-2,10,18,26
1: 
1: 0, GPU-283a0810-d2b9-79e2-66d4-06218d92c3ca, 00000009:01:00.0
1: === rank=1 localid=1 host=presto-gb200-gcn-05 ===
1: CPU affinity        : 0-143
1: CUDA_VISIBLE_DEVICES=0
1: GPU 00000009:01:00.0 -> NUMA 0
1: NUMA nodes (Cpus->) : 0-143
1: NUMA nodes (Mems)   : 0-2,10,18,26
2: 
2: 0, GPU-cd08b629-d517-d919-b365-b22e21e1cf95, 00000018:01:00.0
2: === rank=2 localid=2 host=presto-gb200-gcn-05 ===
2: CPU affinity        : 0-143
2: CUDA_VISIBLE_DEVICES=0
2: GPU 00000018:01:00.0 -> NUMA 1
2: NUMA nodes (Cpus->) : 0-143
2: NUMA nodes (Mems)   : 0-2,10,18,26
3: 
3: 0, GPU-5a7cf4eb-ccf9-e045-4863-4c30236040d2, 00000019:01:00.0
3: === rank=3 localid=3 host=presto-gb200-gcn-05 ===
3: CPU affinity        : 0-143
3: CUDA_VISIBLE_DEVICES=0
3: GPU 00000019:01:00.0 -> NUMA 1
3: NUMA nodes (Cpus->) : 0-143
3: NUMA nodes (Mems)   : 0-2,10,18,26

So, no CPU affinity is set. But if we manually set the CPU affinity using numactl, we can run the benchmark successfully. Attached sbatch runs properly in NVL4 node.

#!/bin/bash
#SBATCH --job-name=nperera-bench
#SBATCH --output=logs/%x_%j.out
#SBATCH --error=logs/%x_%j.err
#SBATCH --time=04:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --gres=gpu:4
#SBATCH --mem=0
#SBATCH --mem-bind=local
#SBATCH --exclusive


CONTAINER_PATH=/scratch/prestouser/images/mg-polars-tpch-nvl72-2026-04-30-arm64-cuda-13-1-ucxthrottling-nvcr_nperera.sqsh
BENCH_PATH=/scratch/prestouser/nperera/rapidsmpf/cpp/build/benchmarks/bench_shuffle
ARGS="-C ucxx -r 10 -w 3 -c 4 -n 67108864 -p 50 -o 4 -m async -l 40960 -s -x -g"
# ARGS="-C ucxx -r 10 -w 3 -c 4 -n 1000 -p 50 -o 4 -m async -l 40960 -s -x -g"

export UCX_TLS=^ib,ud:aux
export UCX_NET_DEVICES=bond0
export UCX_MAX_RNDV_RAILS=1
export UCX_RNDV_PIPELINE_ERROR_HANDLING=y

# Maximum amount of memory UCX may use as staging buffers, must be a multiple of 512M
export UCX_RNDV_MTYPE_WORKER_MAX_MEM=1G
export UCX_RNDV_MTYPE_WORKER_FC_ENABLE=y # Enable staging buffer throttling
export UCX_RNDV_FRAG_MEM_TYPES=cuda # Needs to be set due to a bug in UCX master

export POLARS_MAX_THREADS=1
export RAPIDSMPF_num_streaming_threads=8
export OMP_NUM_THREADS=1
export LIBCUDF_NUM_HOST_WORKERS=4
export KVIKIO_NTHREADS=8

echo "========================================"
echo "Job ID: $SLURM_JOB_ID"
echo "Task count: $SLURM_NTASKS"
echo "Nodes: $SLURM_JOB_NUM_NODES"
echo "Node list: $SLURM_JOB_NODELIST"
echo "Container path: $CONTAINER_PATH"
echo "Bench path: $BENCH_PATH"
echo "Args: $ARGS"
echo "========================================"


srun \
  --export=ALL \
  --mpi=pmix_v4 \
  --container-mounts=/scratch:/scratch \
  --container-image=${CONTAINER_PATH} \
  --container-remap-root \
  --container-mount-home \
  --container-writable \
  bash -c "
    # ----- per-rank CPU/NUMA pinning (slurm task/cgroup does not bind cores) -----
    # Layout (presto GB200, 144 cores, 4 GPUs):
    #   localid 0 -> GPU bus 0008 -> NUMA 0 -> cores   0- 35
    #   localid 1 -> GPU bus 0009 -> NUMA 0 -> cores  36- 71
    #   localid 2 -> GPU bus 0018 -> NUMA 1 -> cores  72-107
    #   localid 3 -> GPU bus 0019 -> NUMA 1 -> cores 108-143
    CPUS_PER_TASK=36
    CPU_START=\$(( SLURM_LOCALID * CPUS_PER_TASK ))
    CPU_END=\$(( CPU_START + CPUS_PER_TASK - 1 ))
    case \$SLURM_LOCALID in
      0|1) NUMA=0 ;;
      2|3) NUMA=1 ;;
    esac

    export CUDA_VISIBLE_DEVICES=\$SLURM_LOCALID
    source /opt/conda/etc/profile.d/conda.sh
    conda activate rapidsmpf

    # ----- sanity check: prove the pinning actually took effect -----
    {
      echo \"--- rank=\$SLURM_PROCID localid=\$SLURM_LOCALID host=\$(hostname) ---\"
      echo \"  requested        : cpus=\$CPU_START-\$CPU_END  numa=\$NUMA  CVD=\$CUDA_VISIBLE_DEVICES\"
      numactl --physcpubind=\$CPU_START-\$CPU_END --membind=\$NUMA \
        bash -c '
          echo \"  kernel CPUs      : \$(grep Cpus_allowed_list /proc/self/status | cut -f2)\"
          echo \"  kernel Mems      : \$(grep Mems_allowed_list /proc/self/status | cut -f2)\"
          echo \"  taskset -pc      : \$(taskset -pc \$\$ | cut -d: -f2 | xargs)\"
          nvidia-smi --query-gpu=index,uuid,pci.bus_id --format=csv,noheader \
            | sed \"s/^/  nvidia-smi: /\"
        '
      echo
    } 1>&2

    # ----- actual benchmark, with the same numactl wrapper -----
    exec numactl --physcpubind=\$CPU_START-\$CPU_END --membind=\$NUMA $BENCH_PATH $ARGS
  "

[nirandaperera]

Either, we need to add the new rrun bindings to the bench shuffle or I think we could use rrun executable to run this inside slurm allocation. @pentschev WDYT?

[nirandaperera]

This is how PDX cluster allocates resources,

srun \
  --account=datascience_rapids_testing   --partition=interactive \
  --nodes=1 \
  --ntasks-per-node=4 \
  --gpus-per-node=4 \
  --gpus-per-task=1 \
  --cpus-per-task=32 \
  --cpu-bind=verbose,cores \
  --mem-bind=verbose,local \
  --hint=nomultithread \
  --exclusive \
  --label \
  bash -c '
    echo "=== rank=$SLURM_PROCID localid=$SLURM_LOCALID host=$(hostname) ==="
    echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
    echo "CPU affinity        : $(taskset -pc $$ | cut -d: -f2 | xargs)"
    echo "NUMA nodes (Mems)   : $(grep Mems_allowed_list /proc/self/status | cut -f2)"
    echo "NUMA nodes (Cpus->) : $(grep Cpus_allowed_list /proc/self/status | cut -f2)"
    nvidia-smi --query-gpu=index,uuid,pci.bus_id --format=csv,noheader
    for bus in $(nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader); do
      short=$(echo $bus | tr "[:upper:]" "[:lower:]" | sed "s/^0000//")
      echo "GPU $bus -> NUMA $(cat /sys/bus/pci/devices/$short/numa_node 2>/dev/null)"
    done
    echo
  '

0: 
0: 0, GPU-2bc145d8-ee64-47df-a132-b60694783f83, 00000000:19:00.0
0: === rank=0 localid=0 host=gpu-h100-0283 ===
0: CPU affinity        : 0-15,64-79
0: cpu-bind=MASK - gpu-h100-0283, task  0  0 [854749]: mask 0xffff000000000000ffff set
0: cpu-bind=MASK - gpu-h100-0283, task  1  1 [854750]: mask 0xffff000000000000ffff0000 set
0: cpu-bind=MASK - gpu-h100-0283, task  2  2 [854751]: mask 0xffff000000000000ffff00000000 set
0: cpu-bind=MASK - gpu-h100-0283, task  3  3 [854752]: mask 0xffff000000000000ffff000000000000 set
0: CUDA_VISIBLE_DEVICES=0
0: GPU 00000000:19:00.0 -> NUMA 0
0: GpuFreq=control_disabled
0: mem-bind=LOC - gpu-h100-0283, task  0  0 [854749]: mask 0x1 set
0: NUMA nodes (Cpus->) : 0-15,64-79
0: NUMA nodes (Mems)   : 0-1
1: 
1: 0, GPU-08171093-acd8-6c36-8d4d-2f76cc69546c, 00000000:2D:00.0
1: === rank=1 localid=1 host=gpu-h100-0283 ===
1: CPU affinity        : 16-31,80-95
1: CUDA_VISIBLE_DEVICES=0
1: GPU 00000000:2D:00.0 -> NUMA 0
1: mem-bind=LOC - gpu-h100-0283, task  1  1 [854750]: mask 0x1 set
1: NUMA nodes (Cpus->) : 16-31,80-95
1: NUMA nodes (Mems)   : 0-1
2: 
2: 0, GPU-3cd2f6e7-b7cd-05c5-d3af-2d47d6cb8734, 00000000:3F:00.0
2: === rank=2 localid=2 host=gpu-h100-0283 ===
2: CPU affinity        : 32-47,96-111
2: CUDA_VISIBLE_DEVICES=0
2: GPU 00000000:3F:00.0 -> NUMA 0
2: mem-bind=LOC - gpu-h100-0283, task  2  2 [854751]: mask 0x2 set
2: NUMA nodes (Cpus->) : 32-47,96-111
2: NUMA nodes (Mems)   : 0-1
3: 
3: 0, GPU-bf3ddabc-7924-432a-fc80-0812bcbcfc0a, 00000000:66:00.0
3: === rank=3 localid=3 host=gpu-h100-0283 ===
3: CPU affinity        : 48-63,112-127
3: CUDA_VISIBLE_DEVICES=0
3: GPU 00000000:66:00.0 -> NUMA 0
3: mem-bind=LOC - gpu-h100-0283, task  3  3 [854752]: mask 0x2 set
3: NUMA nodes (Cpus->) : 48-63,112-127
3: NUMA nodes (Mems)   : 0-1

[nirandaperera]

Hmm. This is a bit strange. When I add this,

diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp
index 1744f488..336fa26e 100644
--- a/cpp/benchmarks/bench_shuffle.cpp
+++ b/cpp/benchmarks/bench_shuffle.cpp
@@ -25,6 +25,8 @@
 #include <rapidsmpf/statistics.hpp>
 #include <rapidsmpf/utils/string.hpp>
 
+#include <rrun/rrun.hpp>
+
 #ifdef RAPIDSMPF_HAVE_CUPTI
 #include <rapidsmpf/cupti.hpp>
 #endif
@@ -534,6 +536,13 @@ int main(int argc, char** argv) {
         );
     }
 
+    rapidsmpf::rrun::bind(std::nullopt, rapidsmpf::rrun::bind_options{
+        .cpu = true,
+        .memory = true,
+        .network = true,
+        .verify = true
+    });
+
     ArgumentParser args{argc, argv, !use_bootstrap};
 
     // Initialize configuration options from environment variables.

I get the following error

terminate called after throwing an instance of 'std::runtime_error'
  what():  rapidsmpf::rrun::bind(): GPU 0 not found in topology

Why would GPU 0 not found in topology? 🤔

[pentschev]

@nirandaperera #844 is useful to work around all this. Just launch one task with all 4 GPUs per node (--ntasks-per-node=1 --gpus=4) and use rrun -n 4 bench_shuffle ..., it will launch 4 processes like in baremetal and bind resources properly.

EDIT: What you're experiencing is not the same issue from the original post, I have confirmed that independently.