Top 10 Performance Tuning Practices for Pytorch

Wonyoung Shin
네이버 쇼핑 개발 블로그
13 min readApr 26, 2021

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

목차

General optimizations

  1. Use async data loading
  2. Pin memory, transfer data asynchronously
  3. Efficiently zero-out gradients
  4. Increase batch size

GPU specific optimizations

5. Use 16-bit precision

6. Enable cuDNN autotuner

7. Avoid unnecessary CPU-GPU synchronization

8. Construct tensors directly on GPUs

Distributed optimizations

9. Use DistributedDataParallel not DataParallel

10. Balance workload on multiple GPUs

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단계로 이루어집니다.

  1. FP16으로 casting.
  2. 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 코드

  1. 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.

References

--

--