texttron/tevatron

해결하는 문제

Tevatron은 대규모 뉴럴 리트리버를 구축하고 학습시키기 위한 통합 툴킷을 제공합니다. 계산 및 메모리 효율성을 유지하면서 다양한 언어와 모달리티(텍스트 및 이미지)에 걸쳐 수십억 규모의 LLM 기반 리트리버를 학습시키는 과제를 해결합니다.

작동 방식

이 툴킷을 사용하면 사용자는 LoRA를 통한 매개변수 효율적 미세 조정(PEFT)을 사용하여 사전 학습된 모델(Mistral-7B, BGE-Embedding, Instruct-E5 등)을 미세 조정할 수 있습니다. PyTorch (GPU) 및 JAX (TPU/GPU) 백엔드를 모두 지원합니다. 성능 최적화를 위해 vLLM, DeepSpeed, FlashAttention, GradCache와 같은 고효율 라이브러리와 통합되어 있습니다. 일반적인 워크플로우는 쿼리-패시지 쌍으로 리트리버를 학습시키고, 코퍼스를 임베딩으로 인코딩한 다음, 유사도 검색을 수행하여 관련 문서를 검색하는 과정을 포함합니다.

대상

모델을 수십억 개의 문서 규모로 확장해야 하는 밀집 검색(dense retrieval), 오픈 도메인 질의응답 및 멀티모달 검색 시스템을 연구하는 AI 연구자 및 엔지니어.

주요 특징

• 멀티 백엔드 지원: GPU용 PyTorch 및 TPU/GPU용 JAX와 호환됩니다.
• 규모 및 모달리티: 다국어에 걸친 수십억 규모의 검색과 텍스트 및 이미지 모달리티를 모두 지원합니다.
• 효율성: 매개변수 효율적 튜닝을 위한 LoRA와 빠른 학습을 위한 DeepSpeed/FlashAttention을 통합했습니다.
• 즉시 사용 가능한 데이터: 멀티모달 및 다국어 검색 작업을 위한 독립형 HuggingFace 데이터셋을 포함합니다.