DataParallel을 활용한 다중 GPU 학습

최신 모델과 배치는 단일 GPU의 메모리와 연산 능력을 넘어섭니다. 여러 GPU를 사용하면 더 큰 모델을 학습하거나 더 큰 배치를 실제 경과 시간이 더 짧은 동안 처리할 수 있습니다. PyTorch는 이를 위한 여러 방법을 제공하며, 가장 간단한 방법은 DataParallel입니다.

데이터 병렬 처리는 모든 GPU에 모델을 복제하고 각 입력 배치를 여러 GPU에 나눕니다. 각 GPU는 할당된 데이터 조각에 대해 계산한 다음, 모든 복제본이 동기화된 상태를 유지하도록 그래디언트를 결합합니다.

nn.DataParallel은 한 줄로 모델을 감쌀 수 있게 해 줍니다. 사용할 GPU ID를 전달하면 PyTorch가 입력을 자동으로 분산하고 출력을 수집합니다.

순전파 중에 DataParallel은 나열된 GPU 전체에 0번 차원을 따라 배치를 분할합니다. 두 개의 GPU에서 배치 64개는 32개씩인 두 개의 하위 배치가 됩니다. 각 GPU는 자신의 하위 배치에 동일한 모델을 병렬로 실행합니다.

역전파 중에 GPU별 그래디언트가 주 장치로 수집되고 평균화됩니다(실제로는 합산한 후 스케일 조정). 따라서 모델 업데이트가 전체 배치를 반영하게 됩니다. 그런 다음 다음 step을 위해 복제본이 다시 동기화됩니다.

가장 좋은 점은 학습 반복문이 거의 바뀌지 않는다는 것입니다. 데이터를 주 GPU로 옮기고 평소처럼 래핑된 모델을 호출하면 됩니다. DataParallel이 내부에서 분산 처리를 수행합니다.

DataParallel에는 잘 알려진 결함이 있습니다. 주 GPU(대개 GPU 0)가 모든 출력을 수집하고 손실을 계산하므로 메모리와 연산을 더 많이 부담합니다. GPU가 많아지면 GPU 0이 병목이 되고 다른 GPU보다 먼저 메모리가 부족해질 수 있습니다.

DataParallel은 단일 파이썬 프로세스에서 실행되며 스레드를 사용해 GPU를 구동합니다. 파이썬 전역 인터프리터 잠금과 분산 및 수집 오버헤드 때문에 특히 GPU가 2~4개를 넘으면 확장 효율이 제한됩니다.

이러한 이유로 본격적인 다중 GPU 작업에는 DistributedDataParallel (DDP)이 권장됩니다. DDP는 GPU마다 하나의 프로세스를 실행하고, 효율적인 전체 집계를 사용해 그래디언트를 동기화하며, GPU 0 병목을 없애 거의 선형적인 확장을 제공합니다.

DataParallel은 GPU 2개를 사용하는 단일 머신에서 빠르게 실험할 때는 사용할 수 있습니다. 이 경우 한 줄로 작성할 수 있다는 간결함이 비효율성을 상쇄합니다. 여러 노드에서 학습하거나 GPU가 많다면 대신 DDP를 사용하십시오.

래핑된 모델의 state_dict에는 module. 접두사가 붙습니다. 정리된 체크포인트를 저장하려면 model.module.state_dict()를 저장하십시오. 그러면 나중에 래핑되지 않은 모델에 올바르게 불러올 수 있습니다.

DataParallel에 대해 배운 내용을 확인해 보십시오.

DataParallel을 사용한 다중 GPU 학습을 배웠습니다:

• nn.DataParallel(model, device_ids=[0, 1])은 모델을 복제하고 배치를 분할합니다
• 각 step마다 GPU 전체에서 그래디언트가 평균화됩니다
• GPU 0 불균형과 단일 프로세스 설계가 확장을 제한합니다
• GPU가 많거나 여러 노드에서 학습할 때는 DDP를 선호하십시오