Top 10 Performance Tuning Practices for Pytorch

Published in

네이버 쇼핑 개발 블로그

13 min readApr 26, 2021

Pytorch 모델의 학습 및 추론을 가속화 할 수 있는 10가지 팁을 공유드립니다. 코드 몇 줄만 바꿈으로써 속도를 개선하고 모델의 품질 또한 유지할 수 있습니다.

1. Use async data loading

torch.utils.data.DataLoader(dataset, num_workers=num_workers)

num_workers = 0: 메인 프로세스가 데이터를 디스크에서 동기식으로 로딩합니다.
num_workers > 0: 여러 프로세스를 사용하여 디스크에서 데이터를 비동기식으로 읽고, 학습과 데이터로딩이 overlapping될 수 있도록 허용합니다. CPU의 데이터 로딩을 빠르게 처리하는 용도로 사용합니다.
(주의) num_workers 증가 → 메모리 사용량 증가, IO 사용량 증가.

num_workers 튜닝하는 방법

CPU 코어 수, GPU 사용량, 배치 크기, IO 사용량 및 모델에 따라 값을 튜닝해야 하는데, 고정된 배치 크기로 num_workers를 늘리고 학습 속도가 더 이상 향상되지 않을 때의 값으로 설정해도 좋습니다.
Guidelines for assigning num_workers 참고

2. Pin memory, transfer data asynchronously

torch.utils.data.DataLoader(dataset, pin_memory=True)
batch.to(device, non_blocking=True)

GPU가 pageable host 메모리에서 곧바로 데이터를 가져올 수 없기 때문에, pinned (page-locked) 메모리를 활용합니다.
: pageable host memory -> temporary pinned host memory -> GPU device memory
자세한 설명은 CUDA 의 pageable memory 와 pinned memory를 참고하시길 바랍니다.

출처: NVIDA Developer Blog “How to Optimize Data Transfers in CUDA C/C++”

pin_memory=True

데이터 텐서를 자동으로 pinned 메모리로 가져오기 때문에, 데이터 전송이 빠릅니다.

pin_memory=True, non_blocking=True

pinned 메모리에 있는 데이터에 한해서 GPU로 비동기식으로 데이터를 전송합니다.
GPU 데이터 전송 이후의 연산이 GPU 데이터를 필요로 하지 않는 경우, 속도 개선 효과를 볼 수 있습니다. 데이터 전송이 모두 완료되기 전에, 기다리지 않고 즉시 연산을 실행하기 때문입니다.
(주의) page-locked memory은 다른 작업에 의해 memory deallocation 되지 않기 때문에, 너무 많은 메모리를 점유하게 될 경우, 다른 데이터가 메모리에 못 올라오는 문제가 생길 수 있습니다.

3. Efficiently zero-out gradients

Don’t

model.zero_grad()

Do

for param in model.parameters():
    param.grad = None

param.grad = None 권장하는 이유는 아래와 같습니다.

모든 파라미터마다 memset을 실행하지 않습니다.
Gradient를 업데이트할 때, “+=” (read+write)이 아닌 “=” (write)를 사용합니다.
Pytorch 백엔드에서 더 효율적으로 gradient을 0으로 만듭니다 (zero out).

(참고) pytorch >= 1.7인 경우, model.zero_grad(set_to_none=True) 또는 optimizer.zero_grad(set_to_none=True) 으로 대체해도 됩니다.

4. Increase batch size

배치 크기를 키워서 GPU 메모리를 최대한 활용하는 것이 학습 시간을 단축하는데 큰 도움이 됩니다. 배치 크기가 크면, 수렴이 느려질 수 있기 때문에 아래와 같은 방법을 사용해서 보완할 수 있습니다.

Tune learning rate, tune weight decay
Add learning rate warm-ups & decay
큰 배치 크기로 학습할 때, 최적화된 optimizer 사용하는 것도 좋습니다.
(LARS, LAMB, NVLAMB, NovoGrad 등)

5. Use 16-bit precision

scaler = torch.cuda.amp.GradScaler()with torch.cuda.amp.autocast(enabled=use_fp16):
    output = model(input)
    loss = loss_fn(output, target)if use_fp16:
    scaler.scale(loss).backward()    if max_norm is not None:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)    scaler.step(optimizer)
    scaler.update()

Mixed precision training은 FP16, FP32를 같이 사용해서 학습하는 방법입니다. 일반적으로 2단계로 이루어집니다.

FP16으로 casting.
FP16 숫자가 0으로 되지 않도록 loss / gradient scaling.
: FP16이 나타낼 수 있는 수의 최소 범위 (2²⁴) 보다 숫자가 작아서 0으로 강제 변환하는 문제를 scaling으로 해결.

Mixed precision training을 사용했을 때 다음과 같은 이점이 있습니다.

FP32으로만 학습할 때와 비슷한 정확도.
필요한 메모리 사이즈 감소.
학습 시간 감소.
Tensor core 를 지원하는 gpu는(e.g. V100) mixed precision training 을 위한 하드웨어 가속을 제공하기 때문에 효과 극대화. V100 기준, 1.5~5배 speedup.

PyTorch 블로그에서는 아래와 같은 장점 때문에apex.amp 보다 torch.cuda.amp (torch >= 1.6) 사용을 권장하고 있습니다. (apex.amp 를 maintenance 모드로 전환했으며 deprecate할 예정이라고 합니다.)

PyTorch의 일부이므로 PyTorch 버전 호환성 보장
Extension을 구축 할 필요가 없음
모델 체크 포인트를 bit 단위로 정확한 저장 / 복원
DataParallel 과 intra-process 모델 병렬화 (가장 성능이 좋은 접근 방식으로 torch.nn.DistributedDataParallel과 프로세스 당 하나의 GPU를 사용하는 것을 권장.)
Gradient penalty (double backward)
Sparse gradient support
torch.cuda.amp.autocast() 는 활성화된 구역에만 영향을 미치기 때문에, 이전apex.amp.initialize()을 여러 번 호출할 때 어려움을 겪었던 사례를 쉽게 처리 가능. 동일한 스크립트에서 여러 convergence run을 실행할 때 각각 새로운 GradScaler 인스턴스를 사용해야하지만, GradScaler는 가볍고 독립적이므로 문제가 되지 않음.

Multi-gpu per process 사용할 때는 (e.g. torch.nn.DataParallel) 사용법을 참고하시길 바랍니다.

6. Enable cuDNN autotuner

torch.backends.cudnn.benchmark = True

Nvidia cuDNN은 convolution (CNN)을 계산하기 위해 다양한 알고리즘을 지원하고 있습니다.
Autotuner는 짧은 benchmark 실행하고, 하드웨어와 input 크기에 최적화된 알고리즘을 선택합니다.
(주의) 고정된 input 크기일 때만 효과적이고, input 크기가 동적으로 변하면 매번 최적화된 알고리즘을 찾게 되어 시간이 더 오래 걸릴 수도 있습니다.
(참고) Batch size, input / output size가 최소 64, 이상적으로는 256으로 나뉘어지는 수로 선택하기를 권장합니다.

7. Avoid unnecessary CPU-GPU synchronization

Don’t

.item()
.cuda()
.cpu()
.to(device)
.nonzero()
print(tensor)

불필요하게 GPU, CPU간 데이터를 전송하는 경우, 성능이 크게 저하됩니다.
cuda tensor의 operation에 의존하는 경우, 성능이 저하됩니다.
e.g. (cuda_tensor != 0).all()

8. Construct tensors directly on GPUs

Don’t

t = tensor.rand(2,2).cuda()

Do

t = tensor.rand(2,2, device=torch.device('cuda:0'))

Don’t: CPU에 tensor를 생성한 후에 GPU로 전송하기 때문에, 시간이 오래 걸립니다.
Do: Tensor를 device에 곧바로 생성하는 것을 권장합니다.

9. Use DistributedDataParallel not DataParallel

DataParallel vs DistributedDataParallel

DistributedDataParallel 코드

Spawn processes

nprocs: 현재 머신에서 생성(spawn) 할 프로세스 수
(GPU당 하나의 프로세스를 생성하는 경우, nprocs=gpu_num)

torch.multiprocessing.spawn(main_worker, nprocs=args.gpu_num, args=(args,))

2. Environment variable initialization

MASTER_ADDR: rank 0 머신의 주소.
하나의 머신으로 학습할 경우, “127.0.0.1”로 설정.
MASTER_PORT: rank 0 머신의 free port.
WORLD_SIZE: init_process_group에서 세팅 가능.
RANK: init_process_group에서 세팅 가능.
rank 0 머신이 모든 connection을 setup.

os.environ['MASTER_ADDR'] = master_address
os.environ['MASTER_PORT'] = str(master_port)

3. Initialize process group

backend: nccl (GPU용. backend 가장 빨라서 권장) / gloo (CPU용)
init_method: peer process를 어디서/어떻게 찾을 수 있는 지 설정. [참고]
환경변수로 MASTER_ADDR, MASTER_PORT 세팅했으면, 'env://'로 설정 가능.
world_size: 동시에 실행되는 총 애플리케이션 프로세스 수.
rank: 모든 프로세스 중 global rank

torch.distributed.init_process_group(backend='nccl',
                                     init_method='env://',
                                     world_size=world_size,
                                     rank=rank)

4. Distributed Data Parallel

device_ids: 코드가 작동할 GPU device id. 일반적으로 프로세스의 local rank.

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

10. Balance workload on multiple GPUs

시퀀스 데이터를 사용할 때 load imbalance 발생할 수 있어 해결 방법으로,

ex) 일정한 수의 토큰들과 가변적인 수의 시퀀스로 배치를 생성.
ex) 비슷한 시퀀스 길이를 가진 샘플들을 bucketing.
ex) 시퀀스 길이로 데이터셋 sorting.

Top 10 Performance Tuning Practices for Pytorch

목차

1. Use async data loading

2. Pin memory, transfer data asynchronously

3. Efficiently zero-out gradients

Don’t

Do

4. Increase batch size

5. Use 16-bit precision

6. Enable cuDNN autotuner

7. Avoid unnecessary CPU-GPU synchronization

Don’t

8. Construct tensors directly on GPUs

Don’t

Do

9. Use DistributedDataParallel not DataParallel

DataParallel vs DistributedDataParallel

DistributedDataParallel 코드

10. Balance workload on multiple GPUs

References

Written by Wonyoung Shin