Molt: 에이전트형 강화학습을 위한 확장 가능한 PyTorch 네이티브 훈련 프레임워크

왜 중요한가

에이전트형 강화학습 연구는 알고리즘 수정과 실험 반복이 잦아 프레임워크 복잡성이 연구 속도의 병목으로 작용한다. Molt는 코드 가독성을 우선 설계해 연구자가 아이디어를 한 파일 범위에서 추적하고 수정할 수 있도록 하며 AI 코딩 어시스턴트가 전체 흐름을 이해할 수 있게 한다. 동시에 Molt는 vLLM과 FSDP2 같은 고성능 구성 요소를 그대로 활용해 경량화가 성능 저하를 초래하지 않음을 보장한다.

핵심 기여

가독성 우선의 프레임워크 설계

코드베이스를 한 번에 읽을 수 있도록 불필요한 추상화를 제거하고 단일 진입점과 국지적 메커니즘을 유지해 연구자와 AI 코딩 어시스턴트가 제어흐름과 데이터흐름을 추적할 수 있게 만들었다. 이 설계는 변경이 알고리즘 구성 요소 하나에만 국한되도록 하여 실험 반복 비용을 낮춘다. 성능 손실을 허용하지 않는 제약을 병행해 실제 훈련 경로에서 가독성과 성능을 동시에 확보했다.

토큰 우선(agent boundary) 경계와 일관성 보장

SDK 기반 에이전트 트래픽을 토큰 ID와 로그확률로 캡처해 별도의 통합 코드 없이 원시 토큰으로부터 학습 데이터가 생성되도록 했다. 토큰 정체성, 정책 버전 의미론, 전방 일관성이라는 불변식을 정의해 생성과 훈련 간 수치적 불일치를 탐지하고 차단하는 메커니즘을 도입했다. 이로 인해 학습에 투입되는 그라디언트가 실제 생성 행동을 정확히 반영하도록 보장한다.

비동기 PyTorch 네이티브 경로와 고성능 조합

Ray 큐와 vLLM 롤아웃 엔진, NeMo AutoModel의 단일 FSDP2 액터를 결합한 비동기 루프를 구현해 프롬프트 스트리밍·일시정지·refit·재개를 지원한다. 직접 NCCL을 통한 가중치 refit과 FSDP2의 텐서·전문가·컨텍스트 병렬을 활용해 다중모달 MoE 워크로드를 밀집 모델 프로그래밍 모델을 바꾸지 않고 처리했다. 실험에서 Megatron 기반 스택과 매칭된 프로토콜 하에서 처리량이 통계적으로 유사함이 확인되었다.

오픈 구현과 매칭 프로토콜 평가

코드와 레시피를 GitHub에 공개하고, 35B 멀티모달 MoE 워크로드와 700B MoE(전문가 병렬 256)까지 동일한 비동기 루프를 유지하며 평가를 수행했다. 평가에서 Molt는 코드베이스 규모를 크게 줄인 상태에서도 Megatron 기반 스택과 처리량 동등성을 보였다. 모든 비대칭성을 명시한 매칭 프로토콜을 사용해 성능 비교의 공정성을 확보했다.

핵심 아이디어 이해하기

에이전트형 강화학습 연구는 자주 알고리즘을 바꾸는 반복적 실험을 요구하는데 기존 대규모 스택은 다중 백엔드와 플러그인 레지스트리 등으로 복잡도가 높아 연구 편의성을 저해했다. 복잡한 인프라가 필요한 이유는 초거대 규모의 상용 훈련에 최적화되었기 때문이며, 연구 목적에서는 그 모든 계층이 불필요한 비용으로 작용한다. Molt는 이 점에서 출발해 '단일 백엔드 채택'과 '코드 최소화'로 복잡도를 줄이는 것을 기본 원리로 삼았다.

토큰 일관성이 핵심 동작 원리다. Molt는 생성 시 샘플된 토큰 ID와 로그확률을 그대로 캡처해 학습으로 흐르게 함으로써 생성-훈련 간의 수치적 미스매치를 제거한다. 이 방식은 토큰화·샘플링·멀티모달 확장·MoE 라우팅 등 여러 단계에서 일어날 수 있는 비일관성을 단일 데이터 표현으로 고정시켜 디버깅과 검증을 단순화한다.

연구자가 프레임워크를 한 번에 이해하도록 설계한 결과로서의 구조적 선택이 성능 제약을 전제로 한다는 점이 핵심이다. Molt는 vLLM, NeMo AutoModel, FSDP2 같은 고성능 구성 요소를 포기하지 않고 단일 비동기 루프에 결합해 처리량 손실 없이 경량화를 달성했다. 따라서 연구자는 아이디어를 코드 한 곳에서 수정하고 즉시 재실험할 수 있으며, 성능은 Megatron 기반 스택과 통계적으로 동등한 수준을 유지한다.

방법론

전체 접근은 세 가지 결합 요소와 하나의 비동기 루프를 중심으로 한다. Molt는 사용자 에이전트를 순수 Python으로 표현하고 vLLM을 롤아웃 엔진으로 두며 NeMo AutoModel 기반의 단일 FSDP2 정책 액터를 Ray 비동기 큐로 연결한다. 이 구성은 각 구성요소를 포크하지 않으므로 upstream 개선이 바로 반영되는 특성을 가진다.

핵심 메커니즘은 토큰-퍼스트 경계와 지속적 프롬프트 그룹 스트리밍이다. 에이전트의 API 트래픽은 토큰 ID와 대응 로그확률로 캡처되어 채팅 포맷의 단일 데이터 경로로 흘러가며, 컨텍스트 압축이 일어나면 자동으로 트래젝토리가 세그먼트된다. 이 과정에서 토큰 일치성, 정책 버전 의미론, 전방 일관성이라는 세 가지 불변식이 유지되어 생성과 학습의 수치적 동형성을 확보한다.

학습 경로는 비동기성을 활용해 엔진 자원을 효율적으로 운영하도록 설계되었다. 지속적 프롬프트 그룹을 스트리밍하면서 pause/refit/resume 패턴으로 가중치 재적합을 수행하고, 가중치 동기화는 NCCL을 통해 엔진으로 직접 반환해 라우터를 우회한다. FSDP2의 텐서·전문가·컨텍스트 병렬과 옵티마이저 오프로드를 조합해 대형 MoE 및 밀집 모델을 동일한 코드 경로로 처리했다.

주요 결과

주요 벤치마크에서 Molt는 Megatron 기반의 스택과 매칭된, 완전 비동기 프로토콜 하에서 처리량이 통계적으로 유사한 성능을 보였다. 논문은 35B 멀티모달 MoE 워크로드에 대한 서빙 가속 및 메모리-처리량 트레이드오프를 정량적으로 제시했으며, 동일한 비동기 루프를 700B MoE와 전문가 병렬 256 환경까지 확장해 적용한 결과를 보고했다. 이들 실험은 Molt가 경량 코드베이스로도 초대형 모델 워크로드를 처리할 수 있음을 보여준다.

비교 실험은 모든 비대칭성을 명시한 매칭 프로토콜을 사용해 수행되었다. 이 프로토콜은 롤아웃-학습 경로의 비대칭과 엔진 수준 최적화가 비교에 미치는 영향을 최소화하기 위한 절차를 포함하며, 그 결과 처리량 차이가 통계적으로 유의미하지 않음을 확인했다. 또한 Molt 코드베이스는 Megatron 기반 스택보다 코드량이 여러 배 적다는 정성적·정량적 근거를 함께 제시했다.

추가로 논문은 성능 외에도 가독성과 수정 편의성 측면의 실험 사례를 제공해 연구자가 알고리즘을 손쉽게 수정하고 재현할 수 있음을 보였다. 이 사례들은 코드 경로가 단일 모듈에 집중되어 있어 알고리즘 변경이 최소한의 영향 범위에서 가능한 점을 입증했다.

기술 상세

전체 아키텍처는 세 구성요소와 하나의 비동기 루프로 요약된다. 사용자 에이전트를 순수 Python으로 구현하고 vLLM을 롤아웃 엔진, NeMo AutoModel 기반의 단일 FSDP2 정책 액터를 사용해 Ray 비동기 큐로 연결함으로써 생성에서 학습까지 토큰이 정확히 전달되는 파이프라인을 구성했다. 이 구성은 upstream 컴포넌트를 포크하지 않으므로 외부 개선이 즉시 수용된다.

토큰 일관성의 수학적/시스템적 기반은 샘플된 토큰 ID와 대응 로그확률을 생성 시점에 캡처해 동일한 데이터 표현을 훈련 경로로 유지하는 것이다. 입력은 엔진에서 샘플된 토큰 ID와 로그확률, 컨텍스트 압축 단계에서 수정된 프리픽스 정보이며 처리 흐름은 Ray 큐를 통해 트래젝토리로 조립되고 FSDP2 액터가 이를 소비해 손실과 그라디언트를 계산한 뒤 NCCL로 가중치 refit을 엔진에 반영하는 과정이다. 이 흐름은 전방 일관성 검사를 포함해 라우팅과 토크나이저 차이에 따른 silent divergence를 탐지하고 차단한다.

MoE 관련 구현은 라우팅 일관성 보장이 핵심이다. Molt는 MoE 라우팅 결정이 추론 엔진과 학습 액터에서 동일하게 재현되도록 라우팅 상태를 명시적으로 관리하고 일치하지 않을 경우 실패하도록 설계했다. 또한 FSDP2의 전문가 병렬과 텐서 병렬, 컨텍스트 병렬을 조합해 700B MoE에서 전문가 병렬 256의 설정까지 동일한 코드 경로로 확장할 수 있게 했다.

성능·메모리 최적화는 엔진 차원의 최적화와의 조합으로 달성되었다. Molt는 엔진의 최적화(예: vLLM의 서빙 가속)를 유지하면서도 학습 경로에 오버헤드를 추가하지 않도록 NCCL 직접 동기화, 옵티마이저 오프로드, persistent prompt-group 스트리밍 등을 도입했다. 이러한 설계는 경량 코드베이스에서도 처리량을 유지하도록 만든다.

한계점

논문에서 명시한 제한 중 하나는 단일 백엔드를 고집하는 설계의 의도된 트레이드오프이다. 한 가지 백엔드만 지원함으로써 코드 복잡도는 줄었지만 배포 옵션과 백엔드 다양성은 의도적으로 제한되었으며 특정 인프라 요건에서는 제약이 될 수 있다. 논문에서는 이 설계 선택이 연구 반복과 가독성이라는 목표를 위해 합리적이라고 결론지었다.

실무 활용

Molt는 오픈소스로 공개되어 있어 연구실 환경에서 새로운 에이전트형 RL 알고리즘을 빠르게 프로토타이핑하고 대형 모델 워크로드로 확장할 수 있다. 제공되는 레시피와 컨테이너는 vLLM과 NeMo AutoModel 기반 구성에서 재사용 가능한 실행 단위를 제공해 실무 적용 진입 장벽을 낮춘다. 단일 백엔드 선택으로 인프라 복잡도는 낮아지지만 배포 선택지는 좁아지는 트레이드오프가 존재한다.

연구자가 새로운 advantage estimator나 트래젝토리 필터링 단계를 하루 단위로 실험하고자 할 때 빠른 프로토타이핑 환경으로 활용할 수 있다.
멀티모달 MoE나 대형 모델의 비동기 온라인 학습 파이프라인을 구성해 생성-훈련 일관성 검증이 필요한 실험에 적용할 수 있다.
vLLM과 NeMo AutoModel을 사용하는 환경에서 지속적 프롬프트 그룹 스트리밍과 pause/refit/resume 패턴을 도입해 서빙-훈련 통합 워크플로를 구현할 수 있다.

코드 공개 여부: 공개

코드 저장소 보기

키워드

Agentic Reinforcement Learning(에이전트형 강화학습)vLLM(vLLM 인퍼런스 엔진)FSDP2(FSDP2 분산 학습)Mixture-of-Experts(MoE, 전문가 혼합)Token-first trajectory(토큰 우선 트래젝토리)

용어 해설

전문가 혼합(MoE)(Mixture of Experts): — 모델 내부에 여러 전문가(expert) 서브네트워크를 두고 입력에 따라 일부 전문가만 활성화해 연산을 줄이는 구조이다. 라우팅 메커니즘이 어떤 전문가가 활성화될지 결정하고, 활성된 전문가들만 순전파와 역전파에 참여하므로 대형 모델의 용량을 늘리면서도 계산량을 제어할 수 있다. Molt에서는 MoE 라우팅과 훈련/추론 간 일관성 유지가 핵심 정확성 요구사항으로 다루어진다.
FSDP2(완전분산 데이터병렬 v2)(FSDP2): — 파라미터와 옵티마이저 상태를 여러 장비에 샤딩해 메모리 사용을 줄이고 대형 모델을 병렬로 학습하는 기술이다. FSDP2는 텐서·컨텍스트·전문가 수준 병렬화와 조합되어 Molt의 단일 액터 환경에서 대형 MoE와 밀집 모델을 동일한 경로로 학습할 수 있게 한다. Molt는 FSDP2 네이티브 지원을 통해 구현 복잡성을 낮추고 성능을 유지한다.
토큰 일치성(token identity)(Token Identity): — 생성 시 샘플된 토큰 ID와 학습에 사용되는 토큰 ID가 정확히 일치하도록 추적하는 원리이다. 토큰화·샘플링 변환·로그확률 보관 등 모든 단계에서 동일한 토큰 표현을 유지해 훈련된 그라디언트가 실제 생성 행동을 반영하게 한다. Molt는 이 원칙을 핵심 불변식으로 삼아 생성-훈련 간 미미한 불일치도 버그로 처리한다.
비동기 훈련 루프(Asynchronous Loop): — 생성(rollout)과 학습(update)이 서로 다른 타이밍으로 중첩 실행되는 구조로, 엔진에서 토큰이 스트리밍되는 즉시 큐에 삽입하고 별도 학습 루틴이 이를 소비한다. Molt는 Ray 큐와 단일 FSDP2 액터를 결합해 지속적 프롬프트 그룹 스트리밍·일시정지·refit·재개를 지원하며 엔진 자원을 비효율적으로 대기시키지 않고 처리량을 유지한다.