AI 속도 향상: Google Colossus를 GCSFS와 Rapid Bucket을 통해 PyTorch에 도입

Source: Google Developers Blog

도전 과제: GPU에 데이터 공급 유지

모델 규모가 커짐에 따라 데이터 로딩과 체크포인팅이 훈련의 주요 병목 현상이 되는 경우가 많습니다. 대규모 모델을 위한 데이터를 준비하려면 원격 스토리지 시스템(예: 객체 스토어)에서 테라바이트에서 페타바이트에 이르는 데이터를 가져오고 처리해야 합니다.

표준 REST 기반 스토리지 접근 방식은 현대 분산 훈련이 요구하는 극한 처리량 및 저지연 요구 사항을 충족하기 어려워 GPU 사이클이 낭비되는 결과를 초래할 수 있습니다.

Rapid Bucket: 양방향 gRPC를 통한 빠른 스토리지

우리의 새로운 Rapid Bucket 솔루션은 전용 존 버킷에서 고성능 객체 스토리지를 제공합니다. 레거시 REST API를 우회하고 지속적인 gRPC 양방향 스트림을 활용함으로써 YouTube와 Google Search를 구동하는 Colossus—파일시스템‑상태 프로토콜의 힘을 직접 PyTorch 생태계에 전달합니다.

핵심 성능 지표

모든 수치는 전용 존 버킷에서 수행된 내부 벤치마크 테스트를 기반으로 합니다.

Fsspec – PyTorch의 파이썬식 파일 인터페이스

fsspec은 PyTorch 생태계에서 파일 시스템을 위한 광범위한 파이썬식 인터페이스입니다. 이미 다음 용도로 사용되고 있습니다:

• Data preparation: Dask, Pandas, Hugging Face Datasets, Ray Data
• Checkpoints: PyTorch Lightning, torch.distributed, Weights & Biases
• Inference: vLLM

다양한 스토리지 시스템에 대한 fsspec의 백엔드 구현이 존재하며, 이들은 모두 하나의 레이어 아래에서 통합될 수 있습니다. 이를 통해 각 백엔드마다 별도의 코드를 작성할 필요가 사라집니다. Rapid Storage를 gcsfs(Google Cloud Storage용 fsspec 구현)와 통합하면, 개발자는 복잡한 코드 재작성 없이도 간단한 fsspec.open() 호출만으로 Rapid이 제공하는 속도 향상을 활용할 수 있습니다.

Under the Hood: Leveraging Colossus

Rapid Buckets 로 성능을 향상시키기 위해 전체 데이터 경로를 최적화했습니다:

• Stateful gRPC‑based streaming – 양방향 gRPC가 연결을 지속시켜 연결 설정, 인증, 메타데이터 교환과 같은 작업당 오버헤드를 제거합니다. 이를 통해 단일 객체 내에서 여러 읽기 또는 추가 작업에 대해 효율적인 상태 유지 데이터 전송이 가능합니다.
• Direct path – Rapid Buckets는 direct connectivity를 사용하여 gRPC 양방향 스트리밍 API(BidiReadObject, BidiWriteObject)에 직접 연결합니다. 클라이언트는 기본 Colossus 파일에 바로 연결되며, 일반 GCS 트래픽에서 발생하는 추가 네트워크 홉을 우회합니다. 그 결과 읽기/쓰기 지연 시간이 크게 감소합니다. 자세한 내용은 블로그 게시물 How the Colossus stateful protocol benefits Rapid Storage 를 참고하세요.
• Zonal co‑location – 데이터를 컴퓨팅 리소스와 동일한 영역(예: us-central1-a)에 저장하면 영역 간 지연이 사라집니다. Rapid Buckets 이전에는 리전 버킷과 컴퓨트 가속기가 서로 다른 영역에 있어 불필요한 지연이 발생했습니다.
• No‑op user migration – 기존 fsspec API는 그대로 유지하면서 내부 트래픽을 HTTP에서 양방향 gRPC로 전환했습니다. gcsfs에 버킷 유형 자동 감지를 추가함으로써 PyTorch 및 기타 fsspec 클라이언트가 별도 설정 없이 자동으로 Rapid을 활용할 수 있습니다.

결과

134 M 행(≈ 451 GB) 규모의 데이터셋을 16개의 GKE 노드에 로드했으며, 각 노드에는 8개의 A4 GPU가 장착되었습니다. 학습은 100 스텝 동안 진행되었고, PyTorch Lightning을 사용해 매 25 스텝마다 체크포인트를 저장했습니다. 데이터 로드 시간을 포함한 전체 학습 시간을 벤치마크한 결과, 표준 리전 버킷 대비 Rapid Bucket을 사용할 때 **23 %**의 성능 향상이 나타났습니다.

마이크로벤치마킹

마이크로벤치마크(입출력 및 리소스 사용량 측정)를 통해 이러한 향상이 확인되었습니다:

테스트 구성: 48개의 프로세스에 걸친 16 MB I/O 크기.

자세한 내용은 GCSFS 성능 벤치마크를 참고하세요.

일러스트레이션

시작하기

Rapid Bucket에서 GCSFS 를 시작하는 것은 쉽습니다. 기존 코드와 스크립트는 그대로 유지되며, Rapid Bucket 을 지정하기만 하면 성능 향상을 누릴 수 있습니다.

설치

Rapid Bucket 통합은 2026.3.0 버전부터 사용할 수 있습니다.

pip install gcsfs

코드 예시 – GCS Rapid에서 읽기/쓰기

import gcsfs

# 파일 시스템 초기화
fs = gcsfs.GCSFileSystem()

# Rapid 버킷에 쓰기
with fs.open('my-zonal-rapid-bucket/data/checkpoint.pt', 'wb') as f:
    f.write(b"model data...")

# 기존 객체에 추가 쓰기 (Native Rapid 기능)
with fs.open('my-zonal-rapid-bucket/data/checkpoint.pt', 'ab') as f:
    f.write(b"appended data...")