HPC & GPU Engineering/AI Infrastructure Engineer

[HPC/GPU 클러스터 운영 Zero to Hero 39편] 데이터 파이프라인 설계 – 프리페치와 체크포인트 전략

ygtoken 2025. 8. 12. 10:54
728x90

 

왜 데이터 파이프라인 설계가 중요한가

 
HPC와 AI 클러스터에서 GPU의 연산 성능을 100% 활용하는 것은 단순히 GPU 개수나 스펙의 문제가 아닙니다.
많은 경우 병목은 GPU가 아니라 데이터 I/O에서 발생합니다.
 

  • 데이터가 GPU에 늦게 도착 → GPU는 idle 상태로 대기
  • 중간 장애로 학습이 중단 → 처음부터 다시 시작
  • 대규모 데이터셋 → 매 epoch마다 스토리지에서 반복 로드

 
이 문제를 해결하려면 **데이터 프리페치(prefetch)**로 I/O 지연을 줄이고, **체크포인트(checkpoint)**로 장애 시 빠르게 복구하는 전략이 필요합니다.
 


 

프리페치(Prefetch) 전략

 
프리페치는 GPU가 데이터를 필요로 하기 전에 미리 읽어와 메모리에 적재하는 기법입니다.
HPC/AI 환경에서는 다음 세 가지 계층에서 적용할 수 있습니다.
 
 

1. 스토리지 계층 프리페치

 

  • Parallel File System: Lustre, BeeGFS의 read-ahead 설정
  • 오브젝트 스토리지: MinIO, Ceph RGW의 다중 파트 다운로드
# Lustre read-ahead 설정 예시
lctl set_param llite.*.max_read_ahead_mb=64

 

2. CPU 메모리 계층 프리페치

 

  • 데이터 로더에서 멀티스레드·멀티프로세스 활용
  • PyTorch 예시:
DataLoader(dataset, batch_size=256, num_workers=8, pin_memory=True, prefetch_factor=4)

 

3. GPU 메모리 계층 프리페치

 

  • CUDA Stream을 사용해 비동기 데이터 전송
data = data.to(device, non_blocking=True)

 


 

체크포인트(Checkpoint) 전략

 
체크포인트는 작업 진행 상태를 저장해 장애나 재시작 시 이어서 실행할 수 있도록 하는 기능입니다.
대규모 학습이나 시뮬레이션에서는 필수입니다.
 
 

체크포인트 구성 요소

 

  • 모델 가중치: model.state_dict()
  • 옵티마이저 상태: optimizer.state_dict()
  • 학습 진행 상황: epoch, step, loss 등
  • 환경 정보: 랜덤 시드, 하이퍼파라미터

 
 

PyTorch 체크포인트 저장 예시

torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss
}, 'checkpoint.pt')

 

재시작 시 로드

checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

 


 

프리페치 + 체크포인트 통합 설계

 

  1. 프리페치로 GPU 대기 시간 최소화
    • 스토리지 I/O와 GPU 연산을 최대한 겹치도록 구성
  2.  
  3. 체크포인트로 장애 복구 시간 최소화
    • 주기적으로 모델 상태 저장
  4. 분산 환경 동기화
    • Horovod, DeepSpeed, PyTorch DDP 등에서 rank=0 노드만 저장하도록 설정
if hvd.rank() == 0:
    torch.save(...)

 


 

HPC 환경에서의 운영 사례

 
 

사례 1 – AI 모델 학습

  • Lustre read-ahead + PyTorch DataLoader prefetch
  • 매 30분마다 체크포인트 저장 → 72시간 학습에서 장애 발생 시 10분 내 복구

 

사례 2 – 기후 시뮬레이션

  • 입력 데이터 S3 오브젝트 스토리지에서 로컬 SSD로 프리페치
  • 시뮬레이션 상태를 매 100 step마다 저장

 


 

성능 최적화 팁

 

  • 프리페치
    • num_workers 값은 CPU 코어 수와 맞추되 I/O 성능에 따라 조정
    • S3 → 로컬 캐시 디렉토리 전략으로 네트워크 부하 완화
  • 체크포인트
    • 파일 사이즈가 크면 압축 사용 (torch.save(..., _use_new_zipfile_serialization=True))
    • 스토리지에 쓰기 병목이 없도록 비동기 저장 사용

 


 

Slurm Job Script 예시

#SBATCH --job-name=training
#SBATCH --gres=gpu:4
#SBATCH --time=48:00:00
#SBATCH --output=out.log

srun python train.py --prefetch --checkpoint_dir /mnt/lustre/checkpoints

 

  • Lustre 경로에 체크포인트 저장
  • 장애 발생 시 같은 Job Script 재실행

 


 

장점과 단점

 
장점

  • GPU Idle Time 최소화 → 학습 속도 향상
  • 장애 발생 시 재시작 시간 단축
  • 장기 작업의 안정성 확보

 
단점

  • 프리페치 메모리 사용량 증가
  • 체크포인트 저장 빈도가 지나치면 I/O 병목 발생 가능

 


 

정리하며

 
프리페치와 체크포인트 전략은 HPC·GPU 클러스터 운영에서 성능안정성을 동시에 보장하는 핵심입니다.
적절한 프리페치 설정으로 GPU를 최대한 활용하고, 주기적인 체크포인트로 장애 시 빠르게 복구하면, 장기·대규모 작업의 효율성을 크게 높일 수 있습니다.
 
 

728x90