NVIDIA/Megatron-Energon

Megatron Energon – 대규모 멀티모달 모델 학습용 데이터 로더

정의 – NVIDIA의 Megatron‑LM 학습 프레임워크를 위해 대규모 멀티모달 데이터셋(이미지, 비디오, 텍스트 등)을 준비, 인덱싱, 스트리밍하는 Python 패키지(및 작은 CLI)입니다. 단독으로도 사용할 수 있습니다.

중요성 – 오늘날의 대규모 비전-언어 또는 비디오-언어 모델 학습에는 여러 스토리지 시스템에 분산된 페타바이트 규모의 데이터가 필요합니다. Megatron Energon은 이러한 무거운 작업을 처리합니다:

주요 기능 (README에 설명됨)

기능	설명
CLI (`energon`)	WebDataset 인덱스 준비, 미디어 메타데이터 수집, 데이터셋을 읽기 전용 가상 파일 시스템으로 마운트 등.
Python API	`get_train_dataset` 및 `get_loader`와 같은 함수가 학습 루프에서 반복 가능한 PyTorch 호환 데이터셋/로더를 반환합니다.
분산 지원	`WorkerConfig`를 사용하여 순위, 월드 크기, 작업자 수를 지정할 수 있어 동일한 코드가 단일 GPU 또는 전체 멀티노드 클러스터에서 실행됩니다.
멀티모달 친화성	이미지, 비디오, 텍스트 등을 지원하며, `TaskEncoder`에 등록된 “쿠커”를 사용하여 원시 샘플을 모델이 기대하는 텐서로 변환합니다.
재현성 및 재개 가능성	인덱스는 `.nv-meta` 폴더에 저장됩니다. 로더는 전체 데이터셋을 다시 처리하지 않고 재시작할 수 있습니다.
설치	`pip install megatron-energon`(안정 버전) 또는 리포지토리에서 직접 설치합니다.

일반적인 워크플로

데이터 준비 – energon prepare /path/to/dataset(또는 JSONL 파일에 대해) 실행하여 .nv-meta 인덱스와 선택적 분할을 생성합니다.
(선택 사항) 미디어 통계 수집 – energon prepare-media …로 이미지 크기, 비디오 길이 등을 수집하여 필터링 또는 버케팅에 사용할 수 있습니다.
코드에서 로드 – Python API를 사용하여 train_ds와 train_loader를 얻은 다음 학습 스크립트에서 배치를 반복합니다.
Megatron‑LM과 결합 – 배치를 Megatron의 모델 병렬 학습 루프에 공급하여 대규모 LLM 또는 멀티모달 모델 학습을 수행합니다.

대상 – 확장 가능하고 재현 가능한 데이터 파이프라인이 필요한 매우 큰 비전-언어, 비디오-언어 또는 기타 멀티모달 모델을 구축하는 연구원 및 엔지니어. 연구용 코드이므로 API는 변경될 수 있습니다.

자세히 알아보기 – 전체 문서는 https://nvidia.github.io/Megatron-Energon/ 및 GitHub 리포지토리에서 문제, 테스트, 예제를 확인하세요.