AI Runtime에서 빠르고 내결함성이 뛰어난 PyTorch 학습

대규모 환경에서 학습 효율성은 단 하나의 메트릭, 즉 GPU가 대기하거나 장애에서 복구하는 대신 생산적인 연산에 소비하는 시간의 비율인 "goodput"에 의해 결정됩니다. 대규모 환경에서는 GPU 장애가 발생하기 마련이므로, 장애로부터 신속하고 자동으로 복구하는 능력만이 높은 goodput을 유지하고 전체 GPU 비용을 관리할 수 있는 유일한 방법입니다.

두 개의 하위 시스템이 이러한 복구의 성패를 좌우하지만, 두 시스템 모두 사후 고려 사항으로 취급되는 경우가 많습니다. 바로 가속기에 데이터를 공급하는 데이터 파이프라인과 작업이 재개될 수 있도록 상태를 스냅샷으로 저장하는 체크포인팅 메커니즘입니다. 둘 중 하나라도 잘못되면 장애가 발생할 때마다 필요 이상으로 훨씬 더 많은 GPU 대기 시간이 낭비됩니다. 장애 상황이 아니더라도 가속기의 속도를 따라가지 못하는 데이터 파이프라인은 시스템 충돌과 마찬가지로 GPU를 유휴 상태로 만들고 goodput을 저하시킵니다. 이 글에서는 두 시스템의 메커니즘과 절충안(trade-off), 그리고 각각이 어떻게 goodput과 전체 GPU 비용에 영향을 미치는지 살펴보겠습니다. 코드 포인터와 예제는 관련 학습 성능 및 복구력 가이드를 참조하세요.

동일한 문제의 인프라 측면, 즉 플릿이 작업을 중단시키기 전에 비정상 GPU를 감지하고 격리하는 방법에 대해서는 관련 게시물인 Databricks AI 전반에서 GPU 안정성을 유지하는 방법을 참조하세요.

대규모 환경에서 장애가 발생하기 마련인 이유

작업의 GPU 수가 늘어남에 따라 중단 없이 전체 실행 시간을 버텨낼 확률은 급격히 떨어집니다. 관련 Databricks 게시물의 유용한 대략적인 모델에 따르면 각 GPU의 연간 장애율은 약 1%로 가정합니다. 이 가정하에 해당 게시물에서는 "256개의 GPU 작업이 30일 동안 실행될 때 장애가 발생할 확률은 약 19%입니다. GPU가 1,024개로 늘어나면 이 확률은 57%로 치솟습니다"라고 설명하며, 이는 단지 인프라 수준의 문제일 뿐입니다.

이 추정치를 현실에 대입해 보면, 608개의 H100 GPU로 구성된 delta 슈퍼컴퓨터에서는 1.9시간마다 장애가 발생했습니다. 즉, 32개의 GPU 작업의 경우 평균 장애 발생 시간은 36시간이 됩니다. 여기서 핵심은 학습 작업은 결국 어느 시점에 실패할 가능성이 높으며, 올바른 결정을 내리면 모델의 복구력을 높이고 장애 발생 시 손실되는 총 시간을 줄일 수 있다는 점입니다.

영향 1: 체크포인트 포맷이 저장 빈도를 결정합니다

체크포인팅은 복구력의 성패를 가르는 핵심이며, 선택한 메커니즘은 저장 빈도에 직접적인 영향을 미칩니다. 이는 goodput에 가장 큰 영향을 미치는 요소입니다. 하루에 한 번만 체크포인트를 저장하는 경우, 장애가 발생하면 장애 발생 전 상태로 되돌리기 위해 평균 12시간의 중복 작업을 다시 실행해야 합니다.

단일 구조의 torch.save 병목 현상

대부분의 팀이 처음 작성하는 체크포인트는 rank 0에서의 단순한 torch.save입니다. 모델 학습 방식에 따라 다음과 같은 두 가지 문제가 발생할 수 있습니다.

1. 분산 학습의 경우, 모든 상태를 rank 0으로 모아서 단일 파일을 작성합니다.
2. 단일 프로세스가 전체 체크포인트를 동기식으로 작성합니다. Unity Catalog(UC)와 같은 원격 오브젝트 스토리지에 저장할 때 네트워크 전송 등으로 인해 차단(block)될 수 있습니다.

이러한 차단 동작은 GPU를 유휴 상태로 만들어 goodput을 저하시킵니다. 하지만 GPU가 체크포인팅에 소비하는 시간을 줄일 수 있는 방법이 있습니다. 바로 Torch의 분산 체크포인트 API입니다.

분산 체크포인트(DCP): 각 rank가 자체 샤드를 작성

PyTorch의 분산 체크포인트는 이 설계를 뒤집습니다. 각 rank는 샤드가 전체 텐서로 구성되는 방식을 설명하는 작은 .metadata 파일과 함께 자체의 고유한 샤드를 병렬로 작성합니다.

저장 시간은 rank 수에 따라 대략 1/N로 감소하며, .metadata 파일이 글로벌 레이아웃을 기록하므로 동일한 체크포인트를 다른 수의 GPU에 다시 로드할 수 있습니다. DCP는 각 새로운 rank에 필요한 바이트를 재계획하므로, 노드를 유실한 후 축소된 용량의 클러스터로 복구하는 작업이 원활하게 진행됩니다.

일반 데이터 병렬 작업에서도 DCP는 유용합니다

흔히 DCP는 샤딩된 모델에만 유용하며, 모든 rank가 동일한 가중치 복제본을 갖는 데이터 병렬(DDP) 작업에서는 얻을 게 없다고 생각합니다. 하지만 그렇지 않습니다. DCP는 DDP 학습 작업에서도 모델 상태를 샤딩하고 각 워커에 병렬로 작성합니다.

또한 FSDP나 텐서 병렬 처리로 전환할 때 필요한 것과 동일한 API이므로, 조기에 도입하면 가장 안 좋은 타이밍에 복구 코드를 다시 작성해야 하는 번거로움을 피할 수 있습니다.

비동기 저장을 통해 저장 빈도에 따른 비용이 거의 들지 않습니다

병렬 쓰기를 사용하더라도 동기식 저장은 바이트가 스토리지에 완전히 저장될 때까지 학습을 차단하므로, 원격 볼륨에 대한 대규모 체크포인트의 경우 수십 초의 가속기 대기 시간이 발생합니다. async_save은 이 작업을 분할합니다. 스테이징 버퍼로 빠르게 복사한 다음, 계속 진행되는 학습과 겹치도록 백그라운드 업로드를 수행합니다.

학습 루프는 업로드가 아닌 스테이징 복사에 대한 비용만 지불합니다. 이전에는 수십 초의 대기 시간이 소요되던 체크포인트 비용이 이제는 거의 들지 않으므로, 다음 섹션에서 설명할 빈번한 체크포인팅을 부담 없이 수행할 수 있습니다.

AI Runtime에서 UCVolumeWriter 및 UCVolumeReader은 UC 볼륨에 대해 DCP를 구현하여 로컬 NVMe를 통해 I/O를 스테이징하고 데이터가 완전히 저장된 후에만 체크포인트를 완료된 것으로 표시합니다. 자세한 내용과 코드 예제는 성능 및 복구력 가이드를 참조하세요.

위 수치는 torch.save에 대한 네트워크 스토리지 시간을 제외한 것입니다.

영향 2: 체크포인트 빈도가 복구 비용을 결정합니다

여기서 여러 요소가 복합적으로 작용합니다. 작업이 실패하면 마지막 유효한 체크포인트 이후의 모든 데이터가 손실되므로 이를 다시 계산해야 합니다. 따라서 장애당 예상되는 낭비 작업은 체크포인트 간격의 약 절반이며, 비용이 적게 드는 비동기 저장을 사용하면 이 간격을 좁힐 수 있습니다.

간격을 10분의 1로 줄이면 예상 복구 시간도 10분의 1로 줄어듭니다. 하루에 약 8.6회의 중단이 발생하는 Llama 3의 수치를 떠올려 보세요. 이 장애율에서 2시간마다 체크포인트를 저장하면 재학습에 하루 평균 8.6시간을 낭비하게 되어 goodput은 64%가 됩니다. 반면 30분마다 체크포인트를 저장하면 2.15시간만 소요되므로 goodput이 91%에 달합니다.

또한 복구는 자동으로 이루어져야 합니다. 재시작 시 작업은 충돌로 인해 불완전하게 작성된 체크포인트를 건너뛰고, 작성이 완료된 가장 최근의 체크포인트를 찾아 사람의 개입 없이 자동으로 재개해야 합니다. DCP는 이를 안정적으로 수행합니다. 모든 샤드가 저장된 후에만 .metadata 파일이 작성되므로, 이 파일의 존재 여부는 선택할 수 있는 신뢰할 수 있는 "저장 완료" 마커가 됩니다.

영향 3: 데이터 로딩이 GPU의 유휴 여부를 결정합니다

학습 작업은 가장 느린 입력 속도에 맞춰 진행됩니다. 가속기가 다음 배치를 대기할 때 GPU가 유휴 상태가 되므로 goodput이 감소합니다. 이 문제를 해결하는 유일한 방법은 아래 그림과 같이 입력 파이프라인이 현재 단계의 연산과 다음 단계의 데이터 준비를 중첩하여 처리하도록 하는 것입니다.

데이터 로딩과 연산을 중첩하는 방식으로 전환한 고객의 경우 실제 소요 시간이 20~50% 감소하는 것을 자주 확인합니다.

원격 스토리지에서 직접 읽을 때 발생하는 비용

거버넌스가 적용된 플랫폼에서 학습 데이터는 원격 오브젝트 스토리지에 저장됩니다. AI Runtime에서 Unity Catalog(UC) 볼륨은 네트워크 마운트로 표시됩니다.

액세스할 때마다 해당 마운트에서 직접 파일을 읽으면 단계별 소요 시간이 네트워크 지연 시간에 묶이게 되고 매 에포크(epoch)마다 동일한 파일을 다시 다운로드하게 됩니다. 이에 대한 해결책은 첫 번째 액세스 시 각 파일을 빠른 로컬 스토리지에 복사하고, 이후의 읽기 작업은 해당 로컬 캐시에서 처리하며, GPU가 연산하는 동안 병렬로 다음 파일을 가져오는 데이터로더를 사용하는 것입니다.

AI Runtime을 사용하면 UCVolumeDataset 및 DataLoader이 정확히 이 작업을 수행합니다(코드 예시는 가이드를 참조하세요). UCVolumeDataset은 UC 볼륨에서 파일을 스트리밍하고, 첫 액세스 시 각 파일을 로컬 NVMe에 캐싱하며, 모든 가속기가 서로 겹치지 않는 독립적인 슬라이스를 받도록 랭크(rank)와 워커(worker) 전체에 파일을 분할합니다. Databricks의 DataLoader은 PyTorch DataLoader의 드롭인(drop-in) 하위 클래스로, 기본값이 이 경로에 맞게 튜닝되어 있습니다. 따라서 GPU가 계산하는 동안 파일이 동시에 가져와지고 캐싱되므로, 학습 스레드에서 한 번에 하나씩 처리할 필요가 없습니다.

예시: UC 파일로 이미지 모델 학습시키기

간단한 이미지 분류 워크로드를 생각해 보겠습니다. UC 볼륨에서 JPEG 파일을 디코딩하고, 증강(augment)한 다음, 비전 모델을 학습시키는 과정입니다. 동일한 GPU, 모델, 배치 크기에서 이 작업을 수행하는 두 가지 방법을 살펴보겠습니다. 하나는 UC 볼륨에서 읽는 기본 PyTorch Dataset를 사용하는 것이고, 다른 하나는 UCVolumeDataset에 Databricks DataLoader 기본값을 조합하여 사용하는 것입니다.

시간이 어디에 소요되는지 추측할 필요가 없습니다

DataLoader을 엔지니어링하면서 메트릭이 MLflow에 기록되도록 하여, 데이터 파이프라인이 학습을 방해하고 있는지 한눈에 쉽게 파악할 수 있도록 했습니다.

fetch_seconds 메트릭은 데이터로더가 배치를 생성하는 데 걸리는 시간을 명시적으로 측정하며, 이 시간 동안 GPU는 대기 상태(idle)로 있게 됩니다.

영향 4: 데이터 파이프라인을 간과하면 모델이 소리 없이 손상됩니다

에러 메시지도 없고, 크래시도 발생하지 않으며, 작업 실패도 없지만, 모델 성능이 예상보다 미묘하게 떨어지는 마지막 회복 탄력성(resilience) 버그가 하나 있습니다. 이는 모델, 옵티마이저, 스텝은 체크포인트로 저장하지만, 데이터셋 내에서 데이터 파이프라인의 위치는 저장하지 않을 때 발생합니다.

에포크 중간에 작업이 중단된 경우를 생각해 보겠습니다. 모델은 올바르게 복원되고 학습 루프가 재개되지만, 데이터로더는 데이터셋의 처음부터 다시 시작합니다.

재개된 작업은 이번 에포크에서 이미 확인한 예제를 다시 학습하고, 아직 도달하지 못한 예제는 건너뛸 가능성이 있습니다. 대규모 환경에서 흔히 발생하는 수많은 재시작 과정에서, 이는 데이터 분포에 소리 없이 편향을 유발합니다. 모델은 여전히 학습되지만, 데이터의 잘못된 샘플링을 바탕으로 학습되는 것입니다. 이는 작업이 완료된 후 메트릭 결과가 실망스럽게 나올 때까지 아무도 문제를 발견하지 못하기 때문에 가장 비용이 많이 드는 소리 없는 실패 유형입니다.

해결책은 데이터 위치를 체크포인트의 일부로 취급하는 것입니다. 파이프라인에 따라 샘플 또는 샤드(shard) 오프셋을 추적하고 재개 시 앞으로 건너뛰거나, 커스텀 데이터셋이 자체 위치를 직렬화하도록 하거나, 에포크 경계에서 체크포인트를 생성하는 것을 의미합니다. 이 모든 것은 한 가지 전제 조건인 결정론(determinism)에 기반합니다. 셔플링과 증강은 난수 생성기(RNG)를 사용하므로, 이러한 시드(seed)와 RNG 상태도 체크포인트의 일부여야 합니다. 그렇지 않으면 재시작 후의 데이터 순서가 재시작 전의 순서와 일치하지 않아 저장된 위치가 잘못된 샘플을 가리키게 됩니다.

시드, 재현 가능한 순서, 재개 가능한 데이터 파이프라인은 하나의 개념을 세 가지로 표현한 것입니다. 가이드에서 코드와 함께 각 전략을 다룹니다.

요약

빠르고 결함 허용적인(fault-tolerant) 학습은 다음과 같은 몇 가지 결정들이 결합되어 이루어집니다:

1. DDP의 경우에도 torch.save 대신 Distributed Checkpoint를 사용하세요. 그래야 직렬 병목 현상이 발생하지 않고 병렬로 저렴하게 저장할 수 있습니다.
2. 비동기식으로 저장하여 체크포인트 비용을 거의 제로로 만드세요. 이를 통해 자주 저장할 수 있습니다.
3. 가장 최근의 유효한 체크포인트로 자동 복구하세요. 실패 시 재계산에 몇 시간이 아닌 몇 분만 소요됩니다.
4. 원격 스토리지에서 캐싱 및 프리페칭(prefetching)을 수행하여 데이터 로딩과 계산을 중첩(overlap)시키세요. 가속기가 입력을 기다리며 대기하는 일이 없어집니다. 이는 매 스텝마다 GPU 시간을 지속적으로 절약해 줍니다.
5. 데이터 파이프라인과 RNG 상태를 체크포인트로 저장하세요. 재개된 작업이 모델을 소리 없이 손상시키는 대신 올바른 데이터에서 계속 진행됩니다.

통합된 원칙은 다음과 같습니다. 빈번하고 저렴하며 완전한 체크포인트는 하드웨어 오류를 작업을 종료시키는 치명적인 이벤트에서 단순한 단수 차이(rounding error) 수준으로 바꾸어 주며, 중첩된 입력 파이프라인은 그 사이에 가속기가 계속 작동하도록 유지합니다. 저렴한(비동기) 저장은 빈번한 저장을 가능하게 하고, 완전한 저장(모델, 데이터, RNG)은 올바른 복구를 보장합니다. 이 두 가지가 모두 갖춰지고, 오류가 발생한 하드웨어를 감지하고 격리하는 플릿(fleet)이 작동하면, 그 아래의 클러스터가 아무리 불안정하더라도 실제 유효 학습 시간은 하드웨어가 허용하는 한계치에 가까워집니다.

참고 자료

• Meta, The Llama 3 Herd of Models (2024): §3.3.1, 최대 16,384개의 H100 GPU에서의 학습 안정성 및 중단 분석.
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): Delta 시스템에 대한 2.5년간의 현장 연구, 작업 크기별 MTBE 및 실패 확률.
• Databricks, How we keep GPUs reliable across Databricks AI: 플릿 수준의 GPU 상태 점검 및 확률적 실패 모델.

시작할 준비가 되셨나요? 전체 코드는 Databricks AI Runtime 문서의 학습 성능 및 회복 탄력성 가이드를 참조하시고, 인프라 측면의 이야기는 How we keep GPUs reliable across Databricks AI를 읽어보세요.

(이 글은 AI의 도움을 받아 번역되었습니다. 원문이 궁금하시다면 여기를 클릭해 주세요)