บทความนี้อธิบายความหมายและการทำงานเบื้องต้นของ NCCL (NVIDIA Collective Communications Library) เพื่อเพิ่มประสิทธิภาพการสื่อสารระหว่าง GPU พร้อมทั้งนำเสนอผลการทดลองของการตั้งค่า parameter เพื่อให้เข้าใจและเป็นแนวทางในการนำไปใช้งาน โดยหัวข้อต่อไปนี้ให้ข้อมูลสรุปเนื้อหาของบทความ เพื่อให้ผู้อ่านสามารถระบุส่วนที่ต้องการอ่านได้อย่างรวดเร็ว

ภาพรวมของ NCCL

NCCL (NVIDIA Collective Communications Library) เป็น library การสื่อสารแบบกลุ่มที่พัฒนาโดย NVIDIA สำหรับการทำ Parallel processing บน GPU โดยเฉพาะอย่างยิ่งในการเทรนโมเดล deep learning ที่ใช้ GPU หลายใบ NCCL ถูกออกแบบมาเพื่อเพิ่มประสิทธิภาพการสื่อสารระหว่าง GPU ในระบบที่มี GPU หลายใบ ทั้งภายใน single node และ multiple nodes

การตั้งค่า NCCL parameters อย่างเหมาะสมมีความสำคัญในการเพิ่มประสิทธิภาพการทำงานของระบบ โดยเฉพาะใน cluster พารามิเตอร์ต่างๆ สามารถส่งผลกระทบต่อประสิทธิภาพการสื่อสารและการใช้ทรัพยากรของระบบ

การทดสอบหา NCCL parameters ที่เหมาะสม

ในการทดสอบ เราได้ทำการทดสอบ parameter 2 ตัวได้แก่ NCCL_SOCKET_NTHREADS และ NCCL_NSOCKS_PERTHREAD

NCCL_SOCKET_NTHREADS

NCCL_SOCKET_NTHREADS คือจำนวนการใช้งาน thread ของ CPU ที่จะใช้สำหรับเปิดการเชื่อมต่อระหว่าง node ผ่าน network socket

NCCL_NSOCKS_PERTHREAD

NCCL_NSOCKS_PERTHREAD คือจำนวน network socket ที่ถูกใช้สำหรับแต่ละ thread ของ CPU

รายละเอียดในการตั้งค่า

หากกำหนด NCCL_SOCKET_NTHREADS เป็น 8 ; CPU จะใช้งาน 8 Threads และตั้งค่า NCCL_NSOCKS_PERTHREAD เป็น 2 ; แต่ละ Thread จะใช้ Network Socket 2 ตัวในการสื่อสารระหว่าง GPU/node

ดังนั้นจะใช้ทั้งสิ้น 16 Network Socket ในการสื่อสารระหว่าง GPU/node

สำหรับเครื่อง LANTA ที่ปิด Hyper-threading ไว้ การระบุจำนวน thread จะเทียบเท่ากับการใช้งาน CPU core โดยสามารถอ่านรายละเอียดของเครื่อง LANTA เพิ่มเติมได้ที่ https://thaisc.io/thaisc-resorces/lanta

นอกจากนี้ parameter 2 ตัวนี้ จะใช้ทรัพยากรส่วนหนึ่งของ CPU ในการสื่อสารกันระหว่าง GPU/node ซึ่งในการประมวลผลของ main task จะต้องใช้ทรัพยากรของ CPU เช่นกันจึงต้อง balance เพื่อหาค่า parameter ที่เหมาะสม

โดยปกติแล้วผลคูณของ NCCL_SOCKET_NTHREADS และ NCCL_NSOCKS_PERTHREAD จะต้องไม่เกิน 64 เพื่อป้องกันผลกระทบต่อการประมวลผลของ main task

ผลลัพธ์จากการทดสอบ

เราได้ทดลอง Train Model Llama2-13b โดยใช้ Dataset alpaca (52k) เป็นชุดข้อมูล

ใช้ 32 node เทรนจำนวน 1 epoch และ Deepspeed - ZeRO Stage 3 ได้ผลลัพธ์ดังนี้

NCCL_SOCKET_NTHREADS	NCCL_NSOCKS_PERTHREAD	Testing Rounds	Average time (sec)	Standard Deviation
8	2	4	307.91	3.21
16	2	4	309.88	7.19
8	4	4	311.61	2.29
16	4	4	313.36	0.98
4	4	4	316.64	3.39
8	8	4	321.70	3.98
4	8	1	321.78	N/A
2	2	4	323.78	2.90
2	8	1	326.43	N/A
2	16	1	348.85	N/A
1	1	4	373.15	6.39

หมายเหตุ: ในการทดสอบโดยใช้ model หรือ dataset อื่นอาจได้ผลลัพธ์ที่แตกต่างไป

ตัวอย่างการตั้งค่า NCCL Parameters

ในตัวอย่างนี้ เป็นการตั้งค่า NCCL parameters โดยใช้ NCCL_SOCKET_NTHREADS=8 และ NCCL_NSOCKS_PERTHREAD=2 โดยสามารถดูรายละเอียดได้จากไฟล์ต่อไปนี้

ไฟล์ submit-batch.sh

#!/bin/bash
#SBATCH -p gpu                           # Specify partition [Compute/Memory/GPU]
#SBATCH -c 64                            # Specify number of processors per task
#SBATCH --ntasks-per-node=1		         # Specify number of tasks per node
#SBATCH --gpus-per-node=4		         # Specify total number of GPUs
#SBATCH -t 1:00:00                       # Specify maximum time limit (hour: minute: second)
#SBATCH -A xxyyyyyy                      # Specify project name
#SBATCH -J nccl-test                     # Specify job name
#SBATCH -o ./logs/finetune-%j.out        # Specify output file

module restore
module load Mamba
module load PrgEnv-gnu
module load cpe-cuda/23.03
module load cudatoolkit/23.3_11.8

: "${NTHREADS:=8}"
: "${PTHREADS:=2}"
: "${BATCH_SIZE:=4}"
: "${DEEPSPEED_STAGE:=3}"
: "${MODEL_SIZE:=7b}"
: "${TASK:=finetune}"
: "${RUN_WITH:=conda}"
: "${ENV_PATH:=}"
: "${SCALING_TYPE:=}"
: "${WO_LORA:=NO}"
: "${PROJ_PATH:=}"
: "${SHARED_PATH:=}"
: "${CACHE_PATH:=}"
: "${ENV_PATH:=}"

conda deactivate
conda activate $ENV_PATH

export WANDB_MODE="offline"
export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12802
export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)

LOG_DIR="./logs/${NTHREADS}nth-${PTHREADS}pth-${SLURM_JOB_ID}"

mkdir -p $LOG_DIR/node_log

export LOG_DIR=$LOG_DIR

export NCCL_TIMEOUT=3600000
export NCCL_DEBUG=DEBUG
export NCCL_SOCKET_IFNAME=hsn
export NCCL_SOCKET_NTHREADS=$NTHREADS
export NCCL_NSOCKS_PERTHREAD=$PTHREADS
export NCCL_DEBUG_FILE=${LOG_DIR}/nccl-${SLURM_JOB_ID}.log
export NCCL_TOPO_DUMP_FILE=${LOG_DIR}/nccl-topo-${SLURM_JOB_ID}.log
export BATCH_SIZE=$BATCH_SIZE
export DEEPSPEED_STAGE=$DEEPSPEED_STAGE
export MODEL_SIZE=$MODEL_SIZE

export TORCH_NCCL_BLOCKING_WAIT=0
export TORCH_EXTENSIONS_DIR=$CACHE_PATH
export HF_HUB_CACHE="$CACHE_PATH/huggingface"
export HF_HOME="$CACHE_PATH/huggingface"
export HF_DATASETS_CACHE="$CACHE_PATH/huggingface"
export TORCH_HOME=$CACHE_PATH
export XDG_CACHE_HOME=$CACHE_PATH
export HF_DATASETS_OFFLINE=1 
export HF_HUB_OFFLINE=1

echo -------ENVIRONMENT-------
echo Python Path: $(which python)
echo Batch Size: $BATCH_SIZE
echo Deepspeed Stage: $DEEPSPEED_STAGE
echo Model Size: $MODEL_SIZE
echo Train with LoRA: $WO_LORA
echo -------------------------
echo NTHREADS: $NTHREADS
echo PTHREADS: $PTHREADS
echo NODES: $COUNT_NODE
echo HOSTNAMES: $HOSTNAMES
echo MASTER_ADDR: $MASTER_ADDR
echo MASTER_PORT: $MASTER_PORT
echo -------------------------

srun --output=${LOG_DIR}/node_log/node-%t.out sh submit-node.sh

ไฟล์ submit-node.sh

#!/bin/bash

echo "------INFO--------"
echo node_number: $SLURM_PROCID
echo hostname: $(hostname)
echo "------------------"
echo ""

echo "----LAUNCHING TRAINING----"

accelerate launch \
    --num_processes $((4 * $COUNT_NODE)) \
    --num_machines $COUNT_NODE \
    --multi_gpu \
    --mixed_precision bf16 \
    --machine_rank $SLURM_PROCID \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --dynamo_backend inductor \
    $PROJ_PATH/scripts/train_wo_lora.py \
        --pretrained_model_name_or_path "meta-llama/Llama-2-$MODEL_SIZE-chat-hf" \
        --train_file $SHARED_PATH/datasets/alpaca_json/alpaca_all.json \
        --validation_file $SHARED_PATH/datasets/alpaca_json/alpaca_validation.json \
        --seed 42 \
        --max_seq_length 1300 \
        --output_dir $PROJ_PATH/checkpoint \
        --num_train_epochs 1 \
        --per_device_train_batch_size $BATCH_SIZE \
        --per_device_eval_batch_size $BATCH_SIZE \
        --save_steps 700 \
        --save_total_limit 5 \
        --learning_rate 8e-5 \
        --weight_decay 0.01 \
        --warmup_ratio 0.05 \
        --lr_scheduler_type cosine \
        --gradient_accumulation_steps 1 \
        --deepspeed "$PROJ_PATH/deepspeed_config/deepspeed_$DEEPSPEED_STAGE.json" \
        --gradient_checkpointing True \
        --tf32 True \
        --bf16 True \
        --max_grad_norm 1.0 \
        --logging_steps 10 \
        --dataloader_num_workers 16 \
        --ddp_find_unused_parameters False \
        --log_dir $LOG_DIR \
        --node_number $SLURM_PROCID

ข้อควรระวังในการปรับค่า Parameters

ในการปรับตัวแปร NCCL ควรระวังการตั้งค่าที่สูงเกินไป อาจส่งผลให้ระบบนั้นใช้ทรัพยากร CPU และ RAM มากเกินความจำเป็น และเกิดการแย่งทรัพยากรกับกระบวนการอื่นๆ ใน Task หลัก เช่น

การ load data
การปรับ weight
เกิด overhead ในการจัดการ threads และ sockets ที่มากเกินไป

แหล่งอ้างอิง

คู่มือการใช้งาน NCCL - NVIDIA Collective Communication Library (NCCL) Documentation
คู่มือการใช้งาน Deepspeed โดย Huggingface - https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed
Repository ที่ใช้ในการทดสอบ - GitHub - Kentakoong/optimize-finetuning-llm

LANTA user guide

การตั้งค่า NCCL Parameters สำหรับการเชื่อมต่อแบบ Multinode