Distributing Training

TRL을 사용한 멀티 GPU 학습

TRL의 트레이너는 PyTorch의 분산 런처인 torchrun으로 실행됩니다. GPU당 하나의 프로세스를 시작합니다:

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

trl CLI도 내부적으로 동일한 방식으로 동작하므로, trl sft ...는 해당 머신의 모든 GPU에서 실행됩니다. GPU 수를 줄여 사용하려면 --nproc_per_node를 전달하세요.

내부적으로 🤗 Accelerate는 GPU마다 하나의 모델을 생성합니다. 각 프로세스는:

• 자체 데이터 배치 처리
• 해당 배치에 대한 손실과 기울기를 계산합니다
• 모든 GPU에 걸쳐 그라디언트 업데이트를 공유합니다

유효 배치 크기는 다음과 같이 계산됩니다: $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

여러 GPU로 확장할 때 일관된 배치 크기를 유지하려면 per_device_train_batch_size와 gradient_accumulation_steps를 그에 맞게 업데이트해야 합니다.

예를 들어, 이러한 구성은 동등하며 동일한 결과를 생성해야 합니다:

GPU당 하나의 모델을 배치하면 메모리 사용량이 높아져 대형 모델이나 저메모리 GPU에서는 실행이 불가능할 수 있습니다. 이러한 경우, 모델 샤딩, Zero Redundancy Optimizer, 혼합 정밀도 학습, CPU 또는 NVMe 오프로딩과 같은 최적화를 제공하는 DeepSpeed를 활용할 수 있습니다. 자세한 내용은 DeepSpeed Integration 가이드를 참고하세요.

매우 긴 시퀀스에 대한 학습은 별도의 가이드가 있습니다: 1M 토큰 이상 학습.

멀티 노드 학습

단일 머신에 GPU가 충분하지 않을 경우, torchrun은 여러 머신(노드)에 걸쳐 학습을 확장할 수 있습니다. 각 노드는 전체 노드 수, 자신의 랭크, 그리고 메인 노드의 rendezvous 엔드포인트와 함께 동일한 명령을 실행합니다.

옵션 1: 수동 실행 (Non-HPC)

각 노드에서 다음을 수동으로 실행합니다:

# Node 0 (main node, IP 10.0.0.1)
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 0 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

# Node 1
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 1 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

10.0.0.1를 rank 0(메인) 노드의 실제 IP 주소로 교체하세요.

옵션 2: SLURM 실행 (HPC 클러스터)

SLURM 작업 스케줄러를 사용하는 클러스터의 경우, 작업 스크립트(예: slurm_job.sh)를 생성합니다:

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

MAIN_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
srun torchrun --nnodes 2 --nproc_per_node 8 --node_rank $SLURM_NODEID --rdzv_backend c10d --rdzv_endpoint $MAIN_NODE:29500 train.py

그런 다음 작업을 제출합니다:

sbatch slurm_job.sh

srun은 노드당 하나의 torchrun을 시작하고, 각각의 torchrun은 GPU당 하나의 프로세스를 시작합니다.

주요 SLURM 지시문:

• --nodes=2: 2개의 컴퓨트 노드를 요청합니다
• --gpus-per-node=8: 노드당 8개의 GPU를 할당합니다(총 16개)
• --job-name: 작업 큐에서 추적을 위한 라벨

트레이닝 설정에서 deepspeed를 설정하여 멀티노드와 DeepSpeed를 조합할 수 있습니다. DeepSpeed 통합 가이드를 참조하세요.

추가 자료

• torchrun (탄력적 시작)
• SLURM Workload Manager 문서 - 클러스터 작업 스케줄링을 위한