토스가 다양한 ML 모델을 만드는 법: Feature Store & Trainkit

토스는 여러 팀이 각자 Feature를 만들고 학습 파이프라인을 구현하면서 생기는 중복, 재현성 저하, 학습·서빙 간 불일치 문제를 해결하기 위해 Feature Store와 Trainkit을 구축했다. Feature Store는 Feature 메타데이터와 오프라인·온라인 데이터를 관리하며, 온라인 서빙에는 저지연 Key-Value 저장소인 Aerospike를 사용한다. Trainkit은 Feature Store의 메타정보를 바탕으로 여러 Feature Package를 결합하고, 전처리와 데이터셋 생성을 표준화하며 시간 Shift와 런타임 PIT 조인을 제공한다. 이를 통해 Feature 재사용과 학습 파이프라인의 일관성을 높이고, 데이터 누수와 조인 과정의 메모리 낭비를 줄이는 설계를 제시한다.

• Feature Store는 Feature의 저장소를 넘어 학습·서빙 간 일관성, 재사용, 품질 모니터링과 접근 권한을 함께 관리하는 시스템이다.
• Feature Table은 물리 데이터 테이블을 감싸고, Feature Service는 모델 학습과 추론에 쓰이는 논리적 Feature 묶음으로 구분한다.
• Trainkit은 PyTorch Lightning의 LightningDataModule을 확장해 데이터 로딩부터 전처리와 배치 생성까지 공통 인터페이스로 정형화한다.
• 시간 파티션 Shift로 학습 시점의 데이터 누수를 줄이고, PIT 조인은 런타임에 수행해 반복 Feature의 메모리 중복을 줄인다.