정렬 튜닝: 대형 언어 모델 사후 정렬을 위한 모듈형 툴킷

초록

AlignTune은 TRL과 Unsloth 두 백엔드를 통합한 단일 인터페이스를 제공해, SFT와 RLHF(특히 DPO, PPO, GRPO 등) 훈련을 일관된 방식으로 수행한다. 백엔드 간 간섭을 차단하는 환경 변수 기반 격리, 43개의 내장 보상 함수와 보상 모델 학습 파이프라인, 그리고 Hugging Face Hub, CSV·Parquet·JSON 등 다양한 데이터 소스를 지원하는 데이터 레이어를 통해 재현 가능한 정렬 실험 환경을 구축한다.

상세 분석

본 논문은 사후 정렬(post‑training alignment) 연구에서 가장 흔히 마주치는 세 가지 문제—백엔드 간 간섭, 보상 함수의 파편화, 그리고 파이프라인 재현성 부족—를 명확히 규정하고, 이를 해결하기 위한 툴킷인 AlignTune을 제시한다. AlignTune의 핵심 설계는 “백엔드 팩토리”라는 단일 진입점으로, 사용자는 모델명, 데이터셋, 백엔드(TRL 혹은 Unsloth), 알고리즘(DPO, PPO 등)만 지정하면 동일한 API를 통해 훈련 객체를 얻을 수 있다. 이는 백엔드별 파라미터 차이를 숨기고, 동일 실험을 서로 다른 백엔드에서 비교할 때 설정 오류를 최소화한다.

백엔드 격리 메커니즘은 네 가지 요소로 구성된다. 첫째, 환경 변수(PURE_TRL_MODE, DISABLE_UNSLOTH_FOR_TRL 등)를 자동 설정해 Unsloth이 트랜스포머 스택을 전역 패치하는 것을 방지한다. 둘째, lazy import를 통해 Unsloth 모듈을 필요 시에만 로드한다. 셋째, 문자열 기반 백엔드 선택으로 import 순서에 의한 부작용을 회피한다. 넷째, Unsloth이 호환되지 않을 경우 자동으로 TRL로 폴백한다. 실험 섹션(5.3)에서는 동일 모델·데이터·하이퍼파라미터 조건에서 격리 여부가 학습 수렴 속도와 최종 성능에 미치는 영향을 정량적으로 보여준다.

보상 시스템은 추상 베이스 클래스 RewardFunction을 중심으로, 30여 개의 RewardType 열거형과 RewardConfig 데이터클래스로 구성된다. 사용자는 RewardRegistry에 새로운 보상 함수를 등록하고, CompositeReward를 통해 가중합 형태의 복합 보상을 정의할 수 있다. 특히 코드 실행 보상(CodeExecutionReward)이나 수학 정답 채점(MathCorrectnessReward)처럼 복잡한 평가 로직을 별도 모듈로 캡슐화함으로써, 훈련 로직과 보상 로직의 결합도를 크게 낮춘다.

보상 모델 학습 파이프라인은 Rule‑based 보상으로 라벨링된 데이터를 자동 생성하고, 이를 기반으로 Transformer 기반 보상 모델을 학습한다. RewardModelTrainer, RewardModelDataset, RewardModelValidator, RewardModelLoader 네 클래스가 각각 데이터 준비, 학습, 검증, 추론을 담당한다. 이렇게 학습된 보상 모델은 TRL의 PPO 트레이너와 직접 연동돼, 실제 RLHF 실험에서 정책 업데이트에 활용된다.

데이터 레이어는 Hugging Face Hub, 로컬 JSON/CSV/Parquet, 디렉터리 기반 데이터셋을 일관된 Dataset 객체로 추상화한다. 이는 다양한 형식의 데이터 소스를 동일한 파이프라인에 투입할 수 있게 해, 연구 재현성을 크게 향상시킨다.

실험에서는 TRL과 Unsloth 각각에 대해 throughput, 메모리 사용량, 평가 지표를 비교하였다. Unsloth은 4‑bit 양자화와 커스텀 커널 덕분에 동일 하드웨어에서 최대 2.3배 빠른 학습 속도와 30% 이하의 메모리 절감을 보였으며, 성능(예: Alpaca 테스트 정확도)에서는 차이가 미미했다. 또한, 동일 설정에서 백엔드 교체 시 결과가 일관된 것을 확인해, 팩토리 기반 설계의 유효성을 입증했다.

한계점으로는 현재 지원되는 알고리즘이 기존 논문에 정의된 것에 국한되고, 백엔드 간 기능 격차(예: GSPO, GBMPO는 Unsloth 미지원)가 존재한다는 점을 언급한다. 향후 계획으로는 새로운 RLHF 변형(예: KL‑regularized DPO)과 더 많은 도메인‑특화 보상 함수를 플러그인 형태로 추가하고, 배포 단계에서 모델 서빙까지 연결하는 파이프라인을 확장할 예정이다.

전반적으로 AlignTune은 사후 정렬 연구에 필요한 인프라를 모듈화·표준화함으로써, 실험 설계·재현·비교를 크게 단순화하고, 백엔드 선택에 따른 성능·자원 효율성을 객관적으로 평가할 수 있는 기반을 제공한다.