บทความนี้เขียนขึ้นโดยมีสมมุติฐานว่าท่านมีประสบการณ์การใช้งาน HPC cluster มาก่อน เช่น TARA

และควรมีชุดประสบการณ์ดังต่อไปนี้

เข้าใจความแตกต่าง และวัตถุประสงค์การใช้งานเบื้องต้นของ frontend-node และ compute-node อ่านเพิ่มเติม

ส่วนสำคัญคือ ท่านทราบว่าไม่สามารถ download ใด ๆ ได้ใน compute node ต้อง pre-download files ที่ Frontend-node เท่านั้น

สามารถรัน batch job บน HPC cluster โดยใช้คำสั่ง sbatch ตามด้วย Slurm script ได้ อ่านเพิ่มเติม

ทำการติดตั้งและ activate conda environment ใน home หรือ project directory ของท่านบน HPC cluster ได้ อ่านเพิ่มเติม

มีความเข้าใจในโปรแกรมที่ท่านกำลังใช้งานอยู่เป็นอย่างดี

บทความนี้แบ่งเนื้อหาเป็นสามส่วนหลัก คือ

ส่วนแนะนำภาพกว้างต่าง ๆ (ดังที่ท่านได้อ่านแล้วบางส่วน)
ส่วน Setup หลักเพื่อใช้ DDP PyTorch บน HPC
ส่วน ตัวอย่างการใช้งาน DDP เพื่อประยุกต์กับงานของท่านบน TARA หรือ LANTA

ซึ่งท่านสามารถดูภาพรวมเนื้อหาได้จาก Table of Contents ด้านล่างนี้ เพื่อประโยชน์ในการไปดูส่วนที่สนใจได้ทันที ตัวอย่างทั้งหมดได้ผ่านการทดลองบน TARA แล้วทั้งสิ้น

1 เกริ่นนำ
2 Distributed Data-Parallel (DDP) ใน PyTorch
3 Setup หลักเพื่อใช้ DDP PyTorch บน HPC
- 3.1 Setup ส่วนของ PyTorch
  - 3.1.1 dist.init_process_group()
    - 3.1.1.1 ตัวอย่างของ single-gpu training
    - 3.1.1.2 ตัวอย่างของ multi-GPUs training with DDP
  - 3.1.2 data.distributed.DistributedSampler & data.DataLoader
- 3.2 Setup ส่วนของ HPC (Slurm configuration)
4 Examples
- 4.1 Simple DDP
  - 4.1.1 SimpleDDP.py
  - 4.1.2 rank, world_size, gpus_per_node, local_rank
  - 4.1.3 Output SimpleDDP
  - 4.1.4 SimpleScriptDDP.sh
- 4.2 MNIST DDP
  - 4.2.1 DDP.py
  - 4.2.2 --cpus-per-task หรือ SLURM_CPUS_PER_TASK
  - 4.2.3 script-N-1-worldsize-8.sh
  - 4.2.4 Output DDP กับ 1 Node DGX-1 (8 GPUs)
  - 4.2.5 Output DDP กับ 2 Nodes DGX-1 (16 GPUs)
5 การติดตั้งอื่น ๆ ที่จำเป็น
- 5.1 อย่าลืม Pre-download ข้อมูลที่ frontend-node ก่อนรัน
6 Further Reading
7 Acknowledgment

เกริ่นนำ

โดยทั่วไปแล้วมีเพียงสองเหตุผลที่ทำให้เราต้องการใช้ multiple GPUs ในการเทรน neural networks:

execution time บน single GPU นั้นนานเกินกว่าทีมวิจัยจะรับได้
โมเดลที่เลือกใช้นั้นมีขนาดใหญ่เกินไปที่จะฟิตใน single GPU

ทั้งนี้ การเลือก GPUs จำนวนมากขึ้นย่อมทำให้การรอคิวของท่านใน HPC cluster นานมากขึ้นตามไปด้วย โดยเฉพาะเมื่อคลัสเตอร์มีความหนาแน่นของ Job ที่เข้าใช้งานพร้อมกันจำนวนมาก

ตรวจสอบความหนาแน่นของคิวได้จากคำสั่ง squeue หรือตรวจสอบเครื่องใน partition ต่าง ๆ ว่าว่างหรือไม่ด้วยคำสั่ง sinfo

Distributed Data-Parallel (DDP) ใน PyTorch

SPMD หรือที่ย่อมาจาก Single-Program, Multiple Data คือไอเดียที่ว่า “โมเดล” จะถูกทำสำเนาเก็บไว้ที่ GPUs ทุกตัวที่จะใช้ และข้อมูลที่จะใช้กับโมเดลดังกล่าวก็จะถูกแบ่งออกเป็นจำนวนเท่า ๆ กันสำหรับแต่ละ GPUs เมื่อ gradients ถูกคำนวณจากทุก GPUs แล้วจะนำมารวมกันเพื่อหาค่าเฉลี่ย และค่า weights ก็จะถูกอัพเดททั้งชุดผ่าน gradient all-reduced โดยที่กระบวนการนี้จะถูกทำซ้ำด้วย mini-batches ชุดใหม่ที่ถูกส่งเข้า GPUs แต่ละตัวอีกครั้งหนึ่ง

Setup หลักเพื่อใช้ DDP PyTorch บน HPC

Setup ส่วนของ PyTorch

dist.init_process_group()

def setup(rank, world_size):
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

Code block นี้คือจุดสังเกตของการใช้งาน DDP โดยเป็นการสร้าง process group โดยมีฟังก์ชั่น dist.init_process_group() เป็นจุด check point ที่จะ “บล็อก” การทำงานเอาไว้เพื่อรอให้ processes ทั้งหมดที่กระจายกันไปได้ทำงานให้เสร็จกันก่อนที่นี่ (ดูการใช้งานภาพรวมได้ใน Code block ของ SimpleDDP.py ด้านล่าง)

ตัวอย่างของ single-gpu training

model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

ตัวอย่างของ multi-GPUs training with DDP

มีโมเดลที่ส่งให้ devices ตามจำนวน local_rank โดยที่โมเดลถูกส่งเข้า DDP อีกทอดหนึ่ง

model = Net().to(local_rank)
ddp_model = DDP(model, device_ids=[local_rank])
optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr)

ซึ่ง local_rank ก็คือ gpu index ในเครื่องที่เราจะเห็นเป็นเลข 0-x เวลาเรา nvidia-smi หรือเรียก cuda ดู เช่น ใน DGX เราจะเห็น 0, 1, …, 7 เป็นต้น

นอกจากนี้เรายังต้องทำให้แน่ใจว่า แต่ละ batch จะถูกส่งไปที่ gpu ทุกตัวใน node

data.distributed.DistributedSampler & data.DataLoader

Setup ส่วนของ HPC (Slurm configuration)

สำหรับ Slurm script ที่ใช้รันจริงดูได้จาก code block ที่แสดงให้ดูเป็นตัวอย่างที่สามารถทดลองทำซ้ำได้ เช่น SimpleScriptDDP.sh, Script-N-1-worldsize-8.sh เป็นต้น ใน code block ด้านล่างนี้คือ setup หลักที่อยากแสดงให้เห็นก่อน นั่นคือ การเลือก partition ที่จะใช้งาน (-p) การเลือกจำนวนโหนดที่ต้องการใช้งาน (-N) การเลือกจำนวน tasks หรือ process ที่ให้รันในหนึ่งโหนด (--ntasks-per-node) ซึ่งโดยทั่วไปเพื่อให้เกิดประสิทธิภาพการใช้งานสูงสุดเราจะเลือกใช้เต็มจำนวน gpus ที่มี และการเลือกจำนวน cpu ต่อ tasks หรือ process (--cpus-per-task) ซึ่งโดยปกติเราก็จะเลือกใช้เต็มจำนวน cpus ที่มี

โดยการคำนวณค่า --cpus-per-task ดังกล่าว เราต้องทราบจำนวน CPU cores ที่มีในโหนด เพื่อไม่ให้เกิดการระบุค่าเกินจำนวน core ที่มีอยู่จริง ซึ่งจะส่งผลให้ Job ของเราติด ST (Status) PD หรือ Pending โดยมี NODELIST (REASON) แบบ PartitionConfig ซึ่งแสดงให้เห็นได้จากคำสั่ง squeue

นอกจากนี้เรายังทำการ export ค่า WORLD_SIZE ซึ่งเป็นค่าที่เราคำนวณได้เองจากการที่เราระบุจำนวน Node ที่ต้องการ (-N) และเราสามารถทราบจำนวน GPUs ที่มีต่อโหนดที่เราเรียกใช้จาก partition ดังกล่าวได้ เช่น จาก cluster information page (https://thaisc.io/en/resources/ ) เป็นต้น โดยค่า WORLD_SIZE จะถูกนำไปใช้ในการแจกงานภายใน DDP ต่อไป

Examples

Simple DDP

โดยในตัวอย่างนี้ เราจะใช้งาน SimpleDDP.py คู่กับ SimpleScriptDDP.sh ในการส่งคำสั่งเข้า Slurm ผ่าน sbatch

หากเราโฟกัสที่ SimpleDDP.py เราจะพบว่า เป็นการใช้งาน Linear Regression Model ง่าย ๆ แต่เน้นไปที่การทำความเข้าใจกับ parameter เพิ่มเติมอย่าง rank, world_size, gpus_per_node, local_rank เป็นต้น

SimpleDDP.py

rank, world_size, gpus_per_node, local_rank

SLURM_PROCID เป็น Slurm environment variable ที่มีค่าระหว่าง [0 - N-1], โดยที่ N เป็นค่า the number of tasks ที่รันภายใต้ srun ยกตัวอย่างเช่น การสั่งให้ Slurm config ที่จำนวน nodes = 2 และให้ใช้ ntasks-per-node=4

หมายความว่า Python interpreter จะสร้าง 2x4 = 8 processes ซึ่งจะมีค่า SLURM_PROCID ดังนี้ 0, 1, 2, 3, 4, 5, 6, 7

ซึ่งค่า rank ได้มาจาก SLURM_PROCID ดังได้กล่าวไปข้างต้น

ส่วน world_size แท้จริงแล้วก็คือจำนวน GPUs ที่มีทั้งหมดในการรันครั้งนี้ซึ่งได้มาจาก จำนวน GPU node(s) x จำนวน GPUs หรือ --ntasks-per-node ที่มีใน Node นั้น ๆ เช่น หากเรียกใช้งาน DGX node 1 node ซึ่งประกอบไปด้วย GPUs 8 ตัว ค่า world_size ที่ได้จะเท่ากับ 1x8=8 หรือหากเรียกใช้งาน GPU node 2 node ซึ่งประกอบไปด้วย GPUs 2 ตัวต่อโหนด ค่า world_size ที่ได้จะเท่ากับ 2x2=4 เป็นต้น

ทั้งนี้ค่า gpus_per_node และ local_rank จะถูกใช้ในการระบุ device ให้กับ torch ดังแสดงข้างต้น

อย่างไรก็ดี ค่าเหล่านี้จำต้องได้รับการเซ็ตจาก Slurm environment เอง หากมีการเซ็ตไว้โดยแอดมินของ HPC cluster หรือเราสามารถ export เองได้ใน launch script ดังแสดงในตัวอย่างสคริปต์ SimpleScriptDDP.sh

Output SimpleDDP

SimpleScriptDDP.sh

ใน SimpleScriptDDP.sh ดังแสดงข้างต้นประกอบไปด้วยส่วนหลักอยู่ 4 ส่วนด้วยกันกล่าวคือ

ส่วนของ SBATCH configuration ที่มีสัญลักษณ์ #SBATCH แสดงต้นประโยค ในที่นี้ขออธิบายรายละเอียดที่เกี่ยวข้องกับการทำ distributed training นั่นคือ
1. -p ยกตัวอย่างคลัสเตอร์ TARA จะมีเพียง 3 partitions ที่มี gpu กล่าวคือ partition dgx, dgx-preempt, และ partition gpu ทำให้ตัวเลือก -p เป็นไปได้ทั้งหมด 3 ค่าดังกล่าว
2. ส่วนจำนวน -N ที่เป็นไปได้นั้น สำหรับ partition dgx เป็นไปได้ตั้งแต่ 1 ถึง 3 เท่ากับจำนวนโหนดที่มีมากสุดในแต่ละ partition (-N 2 สำหรับ gpu)
3. --ntask-per-node สำหรับ dgx มีค่ามากสุดเท่ากับ 8 ส่วน gpu มีค่ามากสุดที่ 2 ตามจำนวน gpus ที่มี
ส่วนการ export ค่าต่าง ๆ ที่จำเป็นต่อการรันแบบ distributed pytorch
โดย MASTER_PORT จำเป็นต้องถูกเซ็ต เช่นเดียวกับ WORLD_SIZE ดังได้อธิบายไปก่อนหน้า ส่วน MASTER_ADDR สามารถเรียกได้จาก scontrol show hostnames ดังแสดงใน script ด้านล่างนี้
ส่วนของการจัดการกับ training environment ซึ่งในตัวอย่างนี้ เลือกใช้การจัดการ environment ของ python ด้วย conda เมื่อรันบน HPC เราจึงต้องทำการ activate ทั้ง module ที่เกี่ยวข้องและเรียกใช้ conda ที่เราใช้งานได้พร้อมกัน ซึ่งในกรณีนี้เราต้องเรียกใช้ version cuDNN ที่สามารถใช้งานร่วมกับ environment ของ conda ที่เรา activate ขึ้นมาได้ ซึ่งในตัวอย่างนี้มีการ source ~/.bashrc เนื่องจากมีการเขียนฟังก์ชั่น myconda ในการ activate conda main ไว้นั่นเอง
ส่วนของการสั่ง srun เพื่อรันงาน python simpleDDP.py แบบ distributed parallel srun python simpleDDP.py

Table of Content (in page)

1 เกริ่นนำ
2 Distributed Data-Parallel (DDP) ใน PyTorch
3 Setup หลักเพื่อใช้ DDP PyTorch บน HPC
- 3.1 Setup ส่วนของ PyTorch
- 3.2 Setup ส่วนของ HPC (Slurm configuration)
4 Examples
- 4.1 Simple DDP
- 4.2 MNIST DDP
5 การติดตั้งอื่น ๆ ที่จำเป็น
- 5.1 อย่าลืม Pre-download ข้อมูลที่ frontend-node ก่อนรัน
6 Further Reading
7 Acknowledgment

MNIST DDP

โดยในตัวอย่างนี้เป็นการแปลง python code แต่เดิมของ MNIST ที่เป็น single gpu ให้เป็น DDP

DDP.py

--cpus-per-task หรือ SLURM_CPUS_PER_TASK

ซึ่ง Python script DDP.py ดังกล่าวข้างต้นถูกใช้งานร่วมกับ script.sh ดังแสดงด้านล่างนี้ โดยมี 1 บันทัดที่เพิ่มเติมขึ้นมาสำหรับ Data Loader นั่นคือ

#SBATCH --cpus-per-task=5 # number of cpu per task (5x8=40 cpus)

ซึ่งการกำหนด --cpus-per-task นั้นต้องคำนึงถึงจำนวน cpu ที่มีทั้งหมดของ node หารด้วยจำนวน ntasks-per-node ที่เซ็ตไว้ซึ่งทั้ง partition dgx และ gpu มีจำนวน CPUs เท่ากันที่ 40 CPUs ดังนั้นสำหรับ dgx ค่า cpus-per-task จึงกำหนดได้เป็น 5 และสำหรับ partition gpu ค่า cpus-per-task จึงกำหนดได้เป็น 20

script-N-1-worldsize-8.sh

Output DDP กับ 1 Node DGX-1 (8 GPUs)

ผลลัพธ์ของ script-N-1-worldsize-8.sh และ DDP.py แสดงได้ดังต่อไปนี้

Output DDP กับ 2 Nodes DGX-1 (16 GPUs)

หากเราต้องการเทรนนิ่งด้วย 2 DGX nodes หรือ 16 GPUs V100 เราสามารถปรับแต่งค่าที่เกี่ยวข้องต่าง ๆ ใน Slurm configuration ของ script-N-1-worldsize-8.sh ให้กลายเป็น script-N-2-worldsize-16.sh ได้ง่าย ๆ ดังนี้

ผลลัพธ์ของ script-N-2-worldsize-16.sh และ DDP.py แสดงได้ดังต่อไปนี้

การติดตั้งอื่น ๆ ที่จำเป็น

อย่าลืม Pre-download ข้อมูลที่ frontend-node ก่อนรัน

ซึ่งตอนนี้ข้อมูล MNIST จาก datasets ที่ torch มีให้ก็ถูก download ลงมาที่ /data/MNIST/raw บนเครื่องที่เรารันเรียบร้อย

Acknowledgment

Thank you for the very nice resources from Princeton University, as this article is written based on your work.

https://github.com/PrincetonUniversity/multi_gpu_training/tree/main/02_pytorch_ddp

TARA user guide

Multi-GPUs Training with PyTorch -- Distributed Data-Parallel