Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning

Agentic reinforcement learning research is constant algorithm modification, new estimators, new pipeline stages, new rollout schemes, and in mainstream frameworks each change threads through layers of trainer, distributed backend, and rollout glue: the cost lands on the researcher at every iteration. Molt is a PyTorch-native training framework built to keep that cost small: a codebase compact and clean enough for a researcher to hold in their head, and for an AI coding assistant to read and reason about in its entirety, so the algorithm flow can be traced and changed end to end. The agent is an ordinary program, and one asynchronous loop trains multimodal and mixture-of-experts policies while never training on a token it did not generate, consistent in tokens, policy versions, and model semantics. Leanness does not cost performance: under a matched, fully asynchronous protocol, Molt is statistically comparable to a state-of-the-art Megatron-based stack. Molt is open source and provides recipes and containers at https://github.com/NVIDIA-NeMo/labs-molt.

한 줄 요약

Molt는 PyTorch 기반의 가벼운 강화학습 학습 프레임워크로, 연구자와 AI 코드 어시스턴트가 알고리즘 흐름을 직관적으로 추적하고 수정할 수 있도록 설계되었다.

핵심 기여도

**가독성 중심의 프레임워크 설계**: Molt는 인간과 AI 코드 어시스턴트 모두가 알고리즘 흐름을 전체적으로 이해할 수 있도록 설계되었다.
**단일 백엔드 기반의 최소 코드**: Megatron 기반의 대규모 프레임워크보다 코드 크기가 훨씬 작으며, 단일 엔트리 포인트와 모듈화된 구조를 채택.
**토큰 중심의 에이전트 경계**: OpenAI 또는 Anthropic SDK와 직접 호환되며, 토큰 ID와 로그 확률을 기반으로 통합 코드 없이 학습 가능.
**고성능 비동기식 PyTorch 경로**: FSDP2, NCCL, vLLM, Ray를 활용한 비동기 학습 경로를 통해 멀티모달 MoE 학습을 지원하며, Megatron 기반 스택과 유사한 성능을 보인다.

핵심 아이디어

Molt는 기존 강화학습 연구에서 알고리즘 변경이 여러 계층(트레이너, 분산 백엔드, 롤아웃 엔진)을 거쳐야 하는 복잡성을 줄이기 위해 설계되었다. 연구자는 알고리즘 흐름을 단일 루프에서 추적하고 수정할 수 있어야 한다는 통찰에서 출발했다. Molt는 PyTorch 기반의 단일 루프를 통해 비동기식 학습을 수행하며, 토큰 ID, 정책 버전, 모델 의미론을 일관되게 유지한다. 이는 실험의 신뢰성을 높이고, AI 코드 어시스턴트가 코드를 전체적으로 추적할 수 있도록 한다.

기술적 접근법

**모듈 구성**: Ray, vLLM, NeMo AutoModel을 활용한 단일 비동기 루프 구조.
**토큰 일관성**: 생성된 토큰 ID가 아닌 재토큰화된 전사본이 아닌 토큰 ID로 트레이젝토리 정의.
**정책 버전 의미론**: 학습 가능한 토큰은 동일한 행동 정책 로그 확률을 유지하며, 비동기 사용 시 정확히 보정.
**전방 일관성**: 롤아웃과 액터 실행이 모델 의미론(멀티모달 확장, MoE 라우팅)을 동일하게 유지.
**성능 최적화**: FSDP2, NCCL, 옵티마이저 오프로드, 퍼시스턴트 프롬프트 그룹 스트리밍 등을 통해 성능 유지.

주요 결과

Molt는 Megatron 기반 스택과 **통계적으로 유사한 성능**을 보이며, **하이퍼스케일 복잡성 없이도 높은 성능 유지**.
**OpenAI/Anthropic SDK와의 통합**: 코드 통합 없이 토큰 ID와 로그 확률을 기반으로 학습 가능.
**비동기 학습 성능**: 단일 루프 기반의 비동기 학습 경로를 통해 멀티모달 MoE 학습을 지원하며, **Dense 모델 프로그래밍 모델 변경 없이 수행**.

의의 및 한계

Molt는 연구자와 AI 코드 어시스턴트가 알고리즘 흐름을 전체적으로 이해하고 수정할 수 있도록 설계된 **가독성 중심의 강화학습 인프라**로, 기존의 복잡한 분산 학습 스택 대비 **코드 크기와 유지 관리 복잡도를 크게 줄였다**. 그러나 Molt는 **하이퍼스케일 학습 환경에 최적화된 기존 스택과 동일한 성능을 보장**하는 데 초점을 맞추고 있어, **초대규모 모델 학습에 대한 확장성 검증은 추가 연구가 필요**하다. 또한, **복잡한 환경에서의 실험 안정성 검증**도 향후 연구 주제로 제시된다.

실용적 활용

Molt는 **강화학습 알고리즘 연구**, **멀티모달 및 MoE 정책 학습**, **AI 코드 어시스턴트와 협업하는 연구 환경**에 적합하다. 특히, **PyTorch 기반의 연구자와 엔지니어가 빠르게 알고리즘을 프로토타이핑하고 실험할 수 있는 도구**로 활용될 수 있다.