Chapter 8: 경쟁 기술 & 학술 연구

대규모 ML 학습의 내결함성(Fault Tolerance)을 위한 업계 솔루션과 학술 연구를 비교 분석합니다. AWS Checkpointless Training의 차별점을 명확히 이해하고 Q&A에 대비합니다.

1. PyTorch Elastic (torchrun) PyTorch

공식 문서: pytorch.org/docs/stable/distributed.elastic.html

1.1 아키텍처 개요

PyTorch Elastic (torchelastic)은 분산 학습에서 동적 스케일링과 내결함성을 제공하는 프레임워크입니다. torchrun 명령어를 통해 실행되며, 기존 torch.distributed.launch를 대체합니다.

┌─────────────────────────────────────────────────────────────┐
│                    PyTorch Elastic Architecture              │
├─────────────────────────────────────────────────────────────┤
│                                                              │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐               │
│  │  Agent   │    │  Agent   │    │  Agent   │               │
│  │ (Node 0) │    │ (Node 1) │    │ (Node 2) │               │
│  └────┬─────┘    └────┬─────┘    └────┬─────┘               │
│       │               │               │                      │
│       └───────────────┼───────────────┘                      │
│                       │                                      │
│              ┌────────▼────────┐                            │
│              │   Rendezvous    │                            │
│              │    Backend      │                            │
│              │  (etcd/c10d)    │                            │
│              └─────────────────┘                            │
│                                                              │
│  Worker Group: [W0, W1, W2, ... Wn]                         │
│  - 동적으로 크기 조절 가능                                    │
│  - min_nodes ~ max_nodes 범위 내 탄력적 운영                  │
└─────────────────────────────────────────────────────────────┘

1.2 Rendezvous Backends

Rendezvous는 분산 학습 참여자들이 서로를 발견하고 동기화하는 메커니즘입니다.

# torchrun 실행 예시
torchrun \
    --nnodes=1:4 \           # min:max 노드 수
    --nproc_per_node=8 \     # 노드당 GPU 수
    --rdzv_backend=etcd \    # Rendezvous 백엔드
    --rdzv_endpoint=etcd-server:2379 \
    --rdzv_id=my_job \
    train.py

1.3 Dynamic Scaling 메커니즘

1.4 Checkpoint-based Recovery

def train():
    # 체크포인트 로드 (존재하는 경우)
    if os.path.exists(CHECKPOINT_PATH):
        checkpoint = torch.load(CHECKPOINT_PATH)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']

    for epoch in range(start_epoch, total_epochs):
        train_one_epoch()

        # 주기적 체크포인트 저장
        if epoch % save_interval == 0:
            torch.save({
                'epoch': epoch,
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
            }, CHECKPOINT_PATH)

1.5 한계점

2. NVIDIA Resiliency Extension (NVRx) NVIDIA

GitHub: github.com/NVIDIA/nvidia-resiliency-ext

상태: Experimental (활발한 개발 중, Breaking changes 예상)

2.1 아키텍처 개요

NVRx는 PyTorch 기반 대규모 학습의 생산성 극대화를 위한 통합 내결함성 플랫폼입니다.

┌─────────────────────────────────────────────────────────────┐
│  Fault Tolerance Layer                                      │
│  ┌──────────────┐  ┌──────────────┐  ┌───────────────┐     │
│  │ Hung Rank    │  │ Straggler    │  │ In-Process    │     │
│  │ Detection    │  │ Detection    │  │ Restart       │     │
│  └──────────────┘  └──────────────┘  └───────────────┘     │
├─────────────────────────────────────────────────────────────┤
│  Checkpointing Layer                                        │
│  ┌──────────────┐  ┌──────────────┐                        │
│  │ Async        │  │ Hierarchical │                        │
│  │ Checkpointing│  │ (Local+Remote)│                        │
│  └──────────────┘  └──────────────┘                        │
├─────────────────────────────────────────────────────────────┤
│  Integration Layer                                          │
│  ┌──────────────┐  ┌──────────────┐                        │
│  │ PyTorch      │  │ NVIDIA NeMo  │                        │
│  │ Lightning    │  │ Framework    │                        │
│  └──────────────┘  └──────────────┘                        │
└─────────────────────────────────────────────────────────────┘

2.2 In-Process Restart 메커니즘

핵심 혁신: SLURM 재할당 없이 동일 Job 내에서 장애 복구

기존 방식 (PyTorch Elastic):
장애 → Job 종료 → SLURM 재할당 → 재시작
[====== 수 분 ~ 수십 분 ======]

NVRx In-Process Restart:
장애 → 프로세스 내 복구 → 학습 재개
[== 수 초 ~ 수십 초 ==]

총 복구 시간: 수 분 → 수 초

2.3 Hung Rank Detection

학습 중 응답하지 않는 Rank를 탐지하는 메커니즘입니다.

# NVRx Hung Rank Detection 개념
class HungRankDetector:
    """각 Rank의 heartbeat를 모니터링하여 hang 상태 탐지"""
    def __init__(self, timeout_seconds=300):
        self.timeout = timeout_seconds
        self.last_heartbeat = {}

    def check_ranks(self):
        for rank, last_time in self.last_heartbeat.items():
            if time.time() - last_time > self.timeout:
                self.handle_hung_rank(rank)

탐지 시나리오:

• NCCL collective 연산 중 무한 대기
• GPU 메모리 오류로 인한 hang
• 네트워크 파티션으로 인한 통신 불가

2.4 Straggler Detection

Straggler: 다른 노드보다 느리게 연산을 수행하여 전체 학습 속도를 저하시키는 노드

Performance Monitoring:
┌─────────┐  ┌─────────┐  ┌─────────┐  ┌─────────┐
│ Rank 0  │  │ Rank 1  │  │ Rank 2  │  │ Rank 3  │
│  100ms  │  │  102ms  │  │   98ms  │  │  250ms ⚠│
└─────────┘  └─────────┘  └─────────┘  └─────────┘
                                            │
                                  Straggler 감지!

대응 옵션:
1. 경고 로깅 및 모니터링
2. 해당 노드 제외 및 재할당
3. 동적 배치 크기 조절

2.5 Hierarchical Checkpointing

2단계 체크포인트 전략으로 I/O 병목 해결:

2.6 Async Checkpointing

체크포인트 저장을 비동기로 수행하여 학습 중단 최소화:

class AsyncCheckpointer:
    def save_async(self, state_dict, path):
        # 메모리에 스냅샷 생성 (빠름)
        snapshot = self._create_snapshot(state_dict)

        # 백그라운드 스레드에서 저장
        thread = threading.Thread(
            target=self._persist_snapshot,
            args=(snapshot, path)
        )
        thread.start()

        # 학습은 즉시 계속
        return

2.7 AWS Checkpointless와의 비교

3. DeepSpeed Microsoft

공식 문서: deepspeed.ai/training/

3.1 아키텍처 개요

DeepSpeed는 Microsoft에서 개발한 대규모 모델 학습 최적화 라이브러리입니다. 메모리 효율성과 성능에 초점을 맞추며, 내결함성은 체크포인트 시스템을 통해 제공됩니다.

기존 Data Parallelism (모든 것을 복제):
GPU 0: [Model] [Optimizer] [Gradients]
GPU 1: [Model] [Optimizer] [Gradients]  ← 중복!
GPU 2: [Model] [Optimizer] [Gradients]  ← 중복!

ZeRO Stage 1 (Optimizer State Partitioning):
- Optimizer 상태를 GPU 간 분할
- 메모리 절약: ~4x

ZeRO Stage 2 (+ Gradient Partitioning):
- Gradient도 분할
- 메모리 절약: ~8x

ZeRO Stage 3 (+ Parameter Partitioning):
- 모델 파라미터까지 분할
- 메모리 절약: GPU 수에 비례

3.2 Checkpointing 설정 옵션

{
  "checkpoint": {
    "tag_validation": "Warn",
    "load_universal": true,
    "use_node_local_storage": false,
    "parallel_write": {
      "pipeline_stage": true
    }
  }
}

3.3 Universal Checkpointing

클러스터 크기 변경 시에도 체크포인트 호환성 유지:

사용 사례:

• 8 GPU → 16 GPU 클러스터 업스케일
• TP=4, PP=2 → TP=8, PP=1 설정 변경
• 다른 하드웨어로 마이그레이션

3.4 ZeRO-Infinity

CPU/NVMe로 오프로드하여 GPU 메모리 한계 극복:

• 단일 GPU에서 수 조 개 파라미터 학습 가능
• 자동 데이터 이동 (prefetch/offload)
• NVMe 대역폭 활용으로 CPU 메모리 한계도 극복

3.5 한계점

4. Google Pathways/OCS Google

4.1 Pathways 아키텍처

논문: "Pathways: Asynchronous Distributed Dataflow for ML" (MLSys 2022)

Pathways는 Google의 대규모 ML 학습을 위한 비동기 분산 데이터플로우 시스템입니다.

┌─────────────────────────────────────────────────────────────┐
│              Single Controller Model                         │
│  • 복잡한 병렬화 패턴을 단일 프로그램으로 표현                 │
│  • 제어 플레인과 데이터 플레인 분리                           │
└─────────────────────────────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────┐
│         Sharded Dataflow Graph                               │
│  • 비동기 연산자들의 샤딩된 그래프                            │
│  • Futures 기반 데이터 의존성 관리                            │
└─────────────────────────────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────────┐
│           Accelerator Clusters                               │
│  ┌───────┐  ┌───────┐  ┌───────┐  ┌───────┐                │
│  │TPU Pod│  │TPU Pod│  │TPU Pod│  │TPU Pod│                │
│  └───────┘  └───────┘  └───────┘  └───────┘                │
│           Gang Scheduling + Data Transfer                    │
└─────────────────────────────────────────────────────────────┘

핵심 성과:

• ~100% Accelerator Utilization: 2,048 TPU에서 SPMD 학습 시 달성
• 유연한 병렬화: 16단계 파이프라인 병렬화, 지리적으로 분산된 가속기 그룹 지원
• 비동기 데이터플로우: Futures 기반으로 데이터 의존성 있어도 제어 플레인 병렬 실행

4.2 Optical Circuit Switches (OCS)

Google은 데이터센터 네트워크에 광학 회로 스위치를 사용하여 동적 네트워크 재구성을 지원합니다.

4.3 TPU Multislice Training

TPU v4/v5에서 여러 Pod Slice를 연결하여 대규모 학습:

TPU v4 Pod (4,096 chips):
┌─────────────────────────────────────────────────────────────┐
│  Slice 0      Slice 1      Slice 2      Slice 3            │
│  ┌─────┐      ┌─────┐      ┌─────┐      ┌─────┐            │
│  │1024 │      │1024 │      │1024 │      │1024 │            │
│  │chips│      │chips│      │chips│      │chips│            │
│  └──┬──┘      └──┬──┘      └──┬──┘      └──┬──┘            │
│     └────────────┴────────────┴────────────┘                │
│              Inter-Slice Network (ICI)                      │
└─────────────────────────────────────────────────────────────┘

Multislice Fault Tolerance:
• Slice 단위 장애 격리
• 장애 Slice 제외하고 나머지로 학습 계속
• 자동 리밸런싱

4.4 하드웨어 접근법 vs AWS 소프트웨어 접근법

5. 종합 비교표

6. CheckFreq (FAST '21) Paper

제목: "CheckFreq: Frequent, Fine-Grained DNN Checkpointing"
저자: Mohan et al.
학회: USENIX FAST '21
Source: usenix.org/conference/fast21/presentation/mohan

6.1 핵심 문제

대규모 DNN 학습에서 체크포인트 주기 설정의 딜레마:

• 자주 저장: I/O 오버헤드로 학습 속도 저하
• 드물게 저장: 장애 시 많은 작업 손실

6.2 Snapshot-and-Persist 분리

학습 스레드                    체크포인트 스레드
┌─────────┐                   ┌─────────────────┐
│ Forward │                   │                 │
│ Backward│ ──────────────→   │ Snapshot (메모리│
│ Update  │     빠름          │   복사, ~ms)   │
└─────────┘                   └────────┬────────┘
    │                                  │
    │                                  ▼
계속 학습                     ┌─────────────────┐
                              │ Persist (디스크 │
                              │   쓰기, ~초)    │
                              └─────────────────┘
                                    비동기

6.3 I/O Latency Hiding

• 학습 연산과 체크포인트 I/O를 오버랩
• GPU 연산 중 백그라운드로 디스크 쓰기

6.4 핵심 결과 & 한계

7. Bamboo (NSDI '23) Paper

제목: "Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs"
저자: Thorpe et al.
학회: USENIX NSDI 2023

7.1 핵심 아이디어: Redundant Computation

파이프라인 버블을 활용한 중복 연산으로 즉시 대체(Instant Takeover) 구현

기존 Pipeline Parallelism의 버블:
Stage 0: [F0][F1][F2][F3][ ][ ][ ][B3][B2][B1][B0]
Stage 1: [ ][F0][F1][F2][F3][ ][B3][B2][B1][B0][ ]
                        ▲              ▲
                       버블           버블

Bamboo의 접근법 - 버블에 중복 연산 삽입:
Stage 0: [F0][F1][F2][F3][F0'][F1'][B3][B2][B1][B0]
Stage 1: [F0'][F0][F1][F2][F3][F1'][B3][B2][B1][B0]
             ▲
             이웃 스테이지의 연산을 중복 실행

7.2 Instant Takeover 메커니즘

정상 상태:
┌────────┐    ┌────────┐    ┌────────┐    ┌────────┐
│Stage 0 │───→│Stage 1 │───→│Stage 2 │───→│Stage 3 │
│        │    │ (중복: │    │ (중복: │    │        │
│        │    │Stage 2)│    │Stage 1)│    │        │
└────────┘    └────────┘    └────────┘    └────────┘

Stage 2 장애 발생:
┌────────┐    ┌────────┐                  ┌────────┐
│Stage 0 │───→│Stage 1 │─────────────────→│Stage 3 │
│        │    │즉시    │                  │        │
│        │    │Stage 2 │                  │        │
│        │    │역할 수행│                  │        │
└────────┘    └────────┘                  └────────┘
                 │
                 ▼
          체크포인트 로드 없이 즉시 계속!

7.3 핵심 결과 & 한계

8. Varuna (EuroSys '22 Best Paper) Paper

제목: "Varuna: Scalable, Low-cost Training of Massive Deep Learning Models"
저자: Microsoft Research
학회: EuroSys '22 (Best Paper Award)
Source: github.com/microsoft/varuna

8.1 Dynamic Model Re-partitioning (Job Morphing)

8 GPU 상태:              4 GPU로 축소:
┌──┬──┬──┬──┐           ┌────┬────┐
│S0│S1│S2│S3│    ──→    │S0+1│S2+3│
│S4│S5│S6│S7│           │S4+5│S6+7│
└──┴──┴──┴──┘           └────┴────┘

CutPoint로 분할 경계 정의
리소스 변화에 따라 자동 재구성

8.2 CutPoint 정의

from varuna import CutPoint

class MyModel(nn.Module):
    def __init__(self):
        self.layer1 = nn.Linear(...)
        self.cutpoint1 = CutPoint()  # 여기서 분할 가능
        self.layer2 = nn.Linear(...)
        self.cutpoint2 = CutPoint()  # 여기서도 분할 가능
        self.layer3 = nn.Linear(...)

    def forward(self, x):
        x = self.layer1(x)
        x = self.cutpoint1(x)  # 파이프라인 스테이지 경계
        x = self.layer2(x)
        x = self.cutpoint2(x)  # 다른 경계
        x = self.layer3(x)
        return x

8.3 Spot VM Training

• Spot/Low-priority VM은 On-demand 대비 5배 저렴
• SIGUSR1 시그널로 graceful shutdown
• 체크포인트 저장 후 종료, 새 리소스에서 재시작

8.4 핵심 결과

• 기존 대비 18배 빠른 학습 시간
• 다른 파이프라인 병렬화 대비 26% 성능 향상
• 200B 파라미터 모델을 Spot VM에서 학습 성공

9. Oobleck (SOSP '23) Paper

제목: "Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates"
학회: ACM SOSP 2023

9.1 Pipeline Template 기반 동적 재구성

Pipeline Template = 미리 정의된 파이프라인 구성으로 노드 장애/추가에 즉시 대응

Template A (8 GPU):  [S0|S1|S2|S3|S4|S5|S6|S7]
Template B (6 GPU):  [S0+1|S2|S3|S4|S5+6+7]
Template C (4 GPU):  [S0+1|S2+3|S4+5|S6+7]
Template D (2 GPU):  [S0+1+2+3|S4+5+6+7]

장애 발생 시:
1. 현재 가용 GPU 수 파악
2. 적합한 Template 선택
3. 모델 상태 재배치
4. 학습 즉시 재개

9.2 Distributed Caching

• 모델 상태를 클러스터 전체에 분산 캐싱
• Template 전환 시 네트워크를 통해 빠르게 상태 재배치

9.3 핵심 결과

• Spot VM에서 안정적인 대규모 모델 학습
• 노드 장애 시 수 초 내 재구성
• 기존 체크포인트 기반 복구 대비 10배 이상 빠른 복구

10. MegaScale (ByteDance 2024) Paper

제목: "MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs"
저자: ByteDance
Source: arxiv.org/abs/2402.15627

10.1 핵심 성과

10.2 Full-Stack Co-Design 접근법

10.3 안정성 및 내결함성

• 진단 도구: 시스템 컴포넌트 모니터링, 장애 원인 식별
• Straggler 완화: 느린 노드 감지 및 대응
• 체크포인트 전략: 장기 학습을 위한 안정적인 체크포인트

11. AWS의 차별점

11.1 기술 발전 세대

1세대: 체크포인트 기반
└─→ 주기적 저장, 장애 시 전체 재시작
    예: PyTorch Elastic, DeepSpeed

2세대: 최적화된 체크포인트
└─→ 비동기/계층적 체크포인트, 빠른 복구
    예: NVRx, CheckFreq

3세대: 중복 연산 기반
└─→ 파이프라인 버블 활용, Instant Takeover
    예: Bamboo

4세대: 인메모리 복제 ★ AWS Checkpointless
└─→ 디스크 I/O 제거, Hot Spare, 즉시 복구
    예: AWS Checkpointless Training

11.2 학술 연구와의 관계

11.3 Production-Ready Managed Service

11.4 선택 가이드