포스트

PyTorch 멀티 GPU가 느린 이유: DataLoader, AMP, DDP 병목 체크리스트

GPU를 여러 장 붙였는데 학습이 빨라지지 않는다면 GPU 개수보다 데이터 공급 속도, 0번 GPU 집중, 프로세스별 batch, sampler 설정을 먼저 확인해야 한다. 멀티 GPU는 자동으로 메모리와 시간을 모두 해결하는 기능이 아니라, 병렬 작업을 나누는 구조다.

한 장짜리 정상 baseline을 만든 뒤 DataLoader, AMP, DDP를 하나씩 켜고 batch 처리량과 validation 결과를 비교해야 한다. 여러 최적화를 동시에 적용하면 어느 변경이 속도와 정확도를 바꿨는지 알 수 없다.

GPU가 기다린다면 DataLoader부터 본다

num_workers는 데이터를 불러오는 데 사용할 프로세스 수다. 원문이 참고한 PyTorch forum 토의에는 GPU 수 × 4가 실험상 좋았다는 경험칙이 소개된다.

1
2
3
import torch

num_workers = 4 * torch.cuda.device_count()

하지만 이 값을 정답처럼 고정하면 안 된다. worker가 너무 적으면 GPU가 데이터를 기다리고, 너무 많으면 프로세스 생성과 메모리 사용의 overhead가 커질 수 있다. 실제 선택은 0부터 몇 단계씩 늘리면서 epoch 시간과 CPU, GPU 사용률을 비교하는 방식이 안전하다.

pin_memory=True는 CPU의 pinned memory를 이용해 host-to-device 전송을 준비한다.

pin_memory

그림의 전송 구조는 NVIDIA 설명에서 가져왔다. 데이터가 작거나 다른 구간이 병목이면 효과가 눈에 띄지 않을 수 있으므로 “True면 무조건 빠르다”가 아니라 실제 loader 대기 시간을 확인한다.

데이터 단계에서 기록할 값은 간단하다.

  • batch를 받는 데 걸린 시간
  • GPU 연산 한 step 시간
  • worker 수를 바꿨을 때 전체 epoch 시간
  • GPU마다 들어가는 실제 batch 크기

AMP는 메모리를 줄이지만 검증이 필요하다

AMP(Automatic Mixed Precision)는 모든 연산을 32bit로 고정하지 않고 16bit 연산을 섞어 메모리와 연산 비용을 줄인다. 원문은 Mixed Precision Training 논문의 아이디어와 PyTorch 1.6부터 제공된 내장 기능을 사용했다.

amp

1
2
3
4
5
6
7
8
9
10
11
12
scaler = torch.cuda.amp.GradScaler()

for inputs, labels in loader:
    optimizer.zero_grad()

    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

이 코드는 model, loader, optimizer, criterion과 device 이동이 준비됐다는 전제의 학습 loop 핵심 조각이다. AMP를 켠 뒤에는 단순히 실행 여부만 보지 말고 다음을 비교한다.

  • 같은 batch에서 loss가 정상적으로 감소하는가?
  • 메모리가 실제로 얼마나 줄었는가?
  • 줄어든 메모리로 batch를 늘렸다면 학습 조건도 함께 기록했는가?
  • baseline과 최종 metric이 비슷한가?

속도가 좋아질 수 있다는 기대와 결과가 안정적이라는 검증은 별개다.

DataParallel보다 DDP가 유리한 이유

nn.DataParallel은 한 프로세스 안에서 batch를 GPU별로 나누고 결과를 0번 GPU에 모은다.

1
model = torch.nn.DataParallel(model)

사용은 간단하지만 큰 모델과 데이터에서는 0번 GPU에 작업이 몰리고 GPU 사이 불균형이 생길 수 있다. DistributedDataParallel(DDP)은 GPU마다 프로세스를 두고 동일한 모델로 각자 batch를 처리한다.

원문 전체 예제는 pytorch_multi_gpu 저장소에 있다. 아래는 구조를 읽기 위한 핵심 조각이며 dataset, model, optimizer, train 함수가 빠져 있어 완전한 실행 코드가 아니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def main():
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(
        main_worker,
        nprocs=world_size,
        args=(world_size,),
    )


def main_worker(gpu, world_size):
    torch.distributed.init_process_group(
        backend='nccl',
        init_method='tcp://127.0.0.1:3456',
        world_size=world_size,
        rank=gpu,
    )

    torch.cuda.set_device(gpu)
    model = model.cuda(gpu)
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[gpu]
    )

    sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset
    )
    loader = torch.utils.data.DataLoader(
        train_dataset,
        shuffle=False,
        sampler=sampler,
    )

DistributedSampler를 쓸 때 shuffle=False로 두는 이유는 DataLoader와 sampler가 동시에 순서를 섞지 않게 하기 위해서다. global batch를 GPU 수로 나눌지, GPU당 batch를 유지해 global batch를 키울지도 명확히 정해야 한다.

checkpoint는 0번 프로세스에서만 저장한다.

1
2
if gpu == 0:
    save()

모든 프로세스가 동시에 같은 파일을 쓰면 손상되거나 서로 덮어쓸 수 있다. 반대로 load는 각 프로세스가 자신의 GPU 위치로 map해 같은 상태를 가져와야 한다. dist.barrier()는 모든 프로세스가 그 지점에 들어올 때까지 기다리므로, 위치를 잘못 잡으면 느림이나 멈춤처럼 보일 수 있다.

“GPU 수만큼 빨라졌다”보다 먼저 확인할 것

멀티 GPU 성능은 GPU만의 숫자로 설명되지 않는다. 다음 순서로 한 단계씩 바꾸면 병목을 찾기 쉽다.

  1. 단일 GPU에서 데이터 로딩 시간과 연산 시간을 분리한다.
  2. AMP만 켜 메모리, 시간, metric을 비교한다.
  3. DDP에서 프로세스별 batch와 sampler가 겹치지 않는지 확인한다.
  4. 0번 GPU 메모리와 사용률이 유독 높은지 본다.
  5. checkpoint와 validation을 모든 프로세스가 중복 실행하는지 확인한다.

가장 큰 함정은 여러 최적화를 한 번에 켜고 빨라진 이유를 모르는 것이다. num_workers, pinned memory, AMP, DDP를 각각 측정해야 느려졌을 때 되돌릴 기준도 생긴다.

한 GPU baseline에서 무엇을 측정하나

고정된 batch 수에 대해 데이터 기다림, forward, backward, optimizer step과 전체 시간을 구분한다. GPU 사용률이 주기적으로 0에 가까워지고 CPU가 바쁘다면 DataLoader 후보를 먼저 본다. Model 계산이 매우 작다면 여러 GPU 통신 비용을 상쇄할 작업이 부족할 수 있다.

Batch 크기, 입력 shape, precision과 gradient accumulation을 기록한다. GPU 수를 늘리며 global batch가 함께 커지면 같은 학습 조건 비교가 아니다. Learning rate와 step 수가 어떤 기준으로 유지됐는지도 명시해야 처리량과 모델 결과를 함께 해석할 수 있다.

Validation과 checkpoint 시간을 학습 step에서 분리한다. 한 epoch이 느린 이유가 모든 process의 중복 평가나 저장이라면 model forward만 최적화해도 전체 시간은 거의 줄지 않는다.

DataLoader 병목을 단계별로 줄이는 법

먼저 dataset의 한 sample 읽기와 변환 시간을 잰다. 디스크 I/O, decode, augmentation 중 어느 부분이 오래 걸리는지 보고 num_workers를 한 단계씩 바꾼다. Worker 수가 많을수록 항상 빠른 것은 아니며 메모리와 CPU 경합이 생길 수 있다.

Batch가 GPU로 이동하는 시간도 확인한다. Pinned memory와 non-blocking transfer를 사용할 때 실제 overlap이 생기는지는 측정해야 한다. Data를 미리 GPU에 모두 올리는 임시 benchmark와 실제 pipeline을 비교하면 입력 공급의 상한을 볼 수 있다.

DDP에서는 각 process가 같은 sample을 반복하지 않도록 sampler와 epoch별 shuffle을 확인한다. Process별 batch 수와 첫 index를 잠깐 출력하면 데이터 중복, 누락을 조기에 찾을 수 있다.

DDP가 올바른지 확인하는 체크포인트

GPU 하나당 process 하나가 의도대로 배치됐는지 rank와 device를 로그에 남긴다. Model과 입력이 같은 local device에 있고, 각 process가 다른 batch를 받는지 확인한다. 한 process만 먼저 끝나거나 step 수가 다르면 동기화 지점에서 멈출 수 있다.

Loss logging은 process별 값과 전체 집계를 구분한다. Checkpoint와 화면 출력은 필요한 rank에서만 실행하도록 해 파일 충돌과 중복 비용을 줄인다. Validation metric을 합칠 때 sample 수와 중복 여부도 확인한다.

한 GPU와 여러 GPU의 첫 몇 step에서 loss가 비슷한 방향인지 비교한다. Global batch와 seed 차이로 완전히 같을 필요는 없지만 gradient가 전혀 달라지거나 한 rank만 NaN이면 데이터, AMP, 동기화 범위를 좁힌다.

Speedup을 어떻게 보고해야 하나

초당 sample, step 시간, epoch 전체 시간과 GPU별 peak memory를 함께 쓴다. GPU 수 대비 speedup과 효율을 계산할 때 warm-up, validation과 checkpoint 포함 여부를 명시한다. 가장 빠른 짧은 구간만 고르지 않는다.

AMP, DDP를 적용한 모델의 validation 결과와 저장한 checkpoint를 다시 불러온 결과도 baseline과 비교한다. 속도가 빨라도 sampler 중복 때문에 일부 데이터만 보거나 checkpoint가 불완전하면 성공한 병렬화가 아니다.

함께 읽으면 이해가 이어지는 글

자주 묻는 질문

GPU를 두 장 쓰면 학습 시간도 절반이 되나요?

보장되지 않습니다. 데이터 로딩, gradient 통신, 작은 batch, validation, checkpoint 중복과 순차 구간이 남아 있어 전체 속도 향상은 제한될 수 있습니다.

DataParallel과 DDP의 가장 중요한 구조 차이는 무엇인가요?

DataParallel은 한 process에서 0번 GPU 중심으로 scatter, gather하고, DDP는 보통 GPU별 process가 각자 모델을 실행하며 gradient를 동기화합니다.

AMP를 켜면 정확도 검증 없이 써도 되나요?

안 됩니다. 메모리와 속도 이점은 실제 모델, 장치에서 측정하고 loss 안정성, validation 지표와 checkpoint 복원을 단정밀도 baseline과 비교해야 합니다.

THE END / OPSOAI

여기까지 읽었습니다

핵심 장면을 한 번 더 떠올려 보세요. 이해가 남았다면 이 책은 제 역할을 다했습니다.

다른 책 고르기
표지 1

키와 좌우 스와이프를 지원합니다. 읽던 페이지는 이 기기에 저장됩니다.

CONTENTS

이 책의 목차

    10개 장 15 분읽는 시간