UCDataset을 사용하여 Unity Catalog 테이블에서 GPU 학습 가속화

GPU는 일반적으로 모든 학습 실행에서 가장 비용이 많이 드는 부분이지만, 대부분의 팀은 GPU를 효율적으로 활용하지 못합니다. 병목 현상은 대개 모델에 있는 것이 아닙니다. 데이터를 저장된 위치에서 GPU로 충분히 빠르게 가져오는 데 있으며, 칩이 대기하는 매 순간 시간과 비용이 낭비됩니다.

오늘 AI Runtime에서 Unity Catalog Dataset (UCDataset) 베타 버전을 출시합니다. UCDataset은 내보내기, 복사, 커스텀 로더 코드 없이 Unity Catalog 테이블을 학습 작업의 모든 GPU로 직접 스트리밍하는 기본 PyTorch 데이터 로더입니다.

시뮬레이션된 프로덕션 워크로드에서 UCDataset은 동일한 데이터의 내보낸 Parquet 복사본을 읽는 것과 비교하여 학습 시간을 10배 단축했습니다.

문제점: GPU가 데이터를 기다리고 있습니다

모델을 위한 대부분의 학습 데이터는 이미 Unity Catalog 테이블에 있습니다. 이를 분산 학습 작업으로 가져오려면 세 가지 임시 방편 중 하나를 선택해야 하며, 각각 시간과 비용이 소모됩니다.

• 실행 전마다 거치는 느리고 번거로운 내보내기: 가장 일반적인 패턴은 매번 실행하기 전에 볼륨의 Parquet(또는 다른 샤드 형식)로 테이블을 내보낸 다음 데이터셋을 파일에 지정하는 것입니다. 두 번째 복사본을 저장하는 비용을 지불해야 하고, 새로운 데이터가 들어오는 순간 데이터가 오래되며, Unity Catalog 계보(lineage)에서 벗어나게 됩니다.
• 전체 테이블을 메모리로 가져오는 방식은 확장성이 떨어짐: toPandas()를 사용하여 테이블을 pandas로 읽는 것은 작은 테이블에는 작동하지만 큰 테이블에서는 실패합니다. 또한 모든 데이터가 단일 노드에 로드되므로 32개의 GPU에서 학습하는 작업에는 도움이 되지 않습니다.
• 웨어하우스를 통한 스트리밍은 느림: 웨어하우스는 원래 GPU에 데이터를 공급하도록 구축되지 않았으며 최적화가 없으면 많은 GPU가 동시에 읽을 때 속도가 저하됩니다.

로더 자체도 문제입니다. 모든 GPU와 모든 워커는 데이터의 균형 잡힌 슬라이스를 필요로 합니다. 팀은 이를 수동으로 구축하는 데 시간을 허비하며, 미세한 버그로 인해 아무도 모르게 모델이 중복되거나 건너뛴 행으로 학습될 수 있습니다.

결과: 비용이 많이 드는 GPU가 유휴 상태로 방치되고, 학습에 예상보다 오랜 시간이 걸리며, ML 엔지니어는 모델 개선 대신 데이터 인프라를 관리하게 됩니다.

UCDataset 소개

UCDataset은 Unity Catalog 테이블을 분산 PyTorch 학습을 위한 일급 데이터 소스로 만들어 줍니다. 읽을 대상(테이블, 필요한 열, 행 필터)을 설명하면 UCDataset이 학습에 바로 사용할 수 있는 배치를 각 GPU로 스트리밍합니다. 주요 이점은 다음과 같습니다.

• 속도에 최적화됨: UCDataset은 클라우드 스토리지에서 테이블 파일을 직접 읽으므로 데이터와 GPU 사이에 아무것도 가로막지 않습니다. 일치하는 행을 모든 GPU와 데이터 로더 워커에 고르게 분할합니다.
• 복사본이 아닌 레이크에서 직접 학습: 실시간 테이블을 지정하세요. 데이터를 읽을 때 필터와 삭제 벡터(deletion vector)가 적용되며, 다음 실행에서는 새로 추가된 행을 자동으로 가져옵니다.
• 엔드투엔드 거버넌스: 학습 작업은 실행하는 사용자의 권한으로 데이터를 읽으며, Unity Catalog 권한이 Databricks의 다른 모든 곳과 정확히 동일하게 적용됩니다.

벤치마크: 추천 모델 워크로드에서 더 빠른 학습

시뮬레이션된 프로덕션 워크로드에서 UCDataset을 벤치마킹한 결과, 300GB 기능 테이블을 8개의 NVIDIA H100 GPU에 스트리밍하여 추천 모델을 학습시켰을 때 Parquet에서 직접 데이터를 로드하는 것에 비해 학습 시간이 평균 10배 향상되는 것을 확인했습니다.

사용 방법:

팀이 main.recsys.ctr_features 테이블에 저장된 클릭 레이블과 함께 사용자 및 아이템 기능을 유지한다고 가정해 보겠습니다. 새로운 노출 데이터가 매일 테이블에 추가됩니다. 사용자는 테이블에서 직접 두 노드의 16개 GPU에서 클릭률(CTR) 모델을 학습시키고자 합니다.

1. 읽을 내용을 설명하고 로더를 생성합니다. train.py에서 테이블을 선택하고, 레이블이 지정된 학습 행을 유지하며, 모델에 필요한 열만 선택합니다. 테이블 파일을 읽을 때 필터가 적용되므로 일치하지 않는 행은 학습 코드에 도달하지 않습니다. UCDataset은 이미 전체 배치를 생성하므로 batch_size=None으로 설정하여 표준 PyTorch DataLoader로 감싸면 됩니다.

2. 두 노드에서 학습을 시작합니다. 표준 PyTorch DDP 루프로 학습합니다. UCDataset은 16개의 GPU 각각에 일치하는 행의 자체 공유분을 할당하므로 작성해야 할 샘플러나 샤딩 코드가 없습니다. workload.yaml에서 작업을 설명하고, 모든 노드에서 torchrun으로 train.py를 시작한 다음, AI Runtime CLI로 실행합니다:

더 많은 GPU에서 학습하려면 workload.yaml에서 num_accelerators를 변경하세요. 학습 코드는 동일하게 유지됩니다. 사용해 볼 수 있는 샘플 테이블이 포함된 전체 train.py 및 workload.yaml은 AI Runtime 쿡북에서 확인할 수 있습니다: Unity Catalog 테이블의 클릭률 모델.

시작하기

PyPI에서 패키지를 설치합니다. 노트북에서 %pip install databricks-uc-dataset를 실행하세요.

UCDataset은 AI Runtime에서 베타 버전으로 제공됩니다. 사용해 보려면:

• 문서 읽기: UCDataset을 사용하여 Unity Catalog 테이블을 PyTorch로 로드하기
• 작동하는 예제부터 시작하기: 쿡북 레시피
• 계정 팀에 문의해 주세요. 소중한 피드백을 기다립니다.

Unity Catalog 볼륨에 저장된 비정형 데이터의 경우 UCVolumeDataset을 사용하여 로컬 캐싱으로 파일을 스트리밍하세요. UCDataset은 정형 테이블을 처리하는 동시에, FILE 형식을 통한 멀티모달 학습 지원도 확대하여 테이블에서 직접 이미지, 오디오, 문서로 학습할 수 있도록 지원합니다. 베타 버전에서는 Unity Catalog의 Delta 테이블을 읽을 수 있습니다.

자주 묻는 질문

UCDataset이란 무엇인가요?

UCDataset(Unity Catalog Dataset)은 Unity Catalog 테이블을 학습 작업의 GPU로 직접 스트리밍하는 기본 데이터 로더입니다. Databricks AI Runtime의 베타 버전에서 사용할 수 있습니다. 테이블과 필요한 열을 선택하기만 하면, UCDataset이 내보내기 단계나 커스텀 로더 코드 없이 학습에 바로 사용할 수 있는 배치를 전달합니다.

UCDataset은 어떻게 GPU 학습 속도를 높이나요?

UCDataset은 중간에 내보낸 복사본 없이 클라우드 스토리지에서 테이블의 데이터 파일을 직접 읽어 GPU가 계속 작동하도록 유지합니다. 일치하는 행을 모든 GPU와 데이터 로더 워커에 고르게 분할하므로 각 GPU는 자신의 몫만 읽습니다.

학습을 위해 Unity Catalog 테이블을 Parquet으로 내보내야 하나요?

아니요. UCDataset은 Unity Catalog의 라이브 테이블을 읽으므로 내보내기 단계나 저장하고 동기화 상태를 유지해야 하는 두 번째 복사본이 필요 없습니다. 읽을 때 필터와 테이블의 삭제 벡터가 적용되며, 새로운 각 실행은 마지막 실행 이후에 추가된 행을 가져옵니다.

UCDataset이 Unity Catalog 권한을 적용하나요?

네. AI Runtime에서 학습 작업은 실행하는 사용자의 권한으로 데이터를 읽으며, 테이블에 대한 Unity Catalog 권한이 적용됩니다. 작업이 거버넌스가 적용된 테이블을 직접 읽으므로 별도로 보안을 유지해야 하는 내보낸 복사본이 없습니다. 베타 버전에서는 행 필터나 열 마스크가 있는 테이블을 UCDataset으로 읽을 수 없습니다.

UCDataset은 어떤 테이블을 읽을 수 있나요?

UCDataset은 삭제 벡터가 있는 테이블을 포함하여 Unity Catalog에 등록된 Delta 테이블을 읽습니다.

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