요약(TL;DR) — 우리는 에이전트 우선(agentic-first) 방식의 PyTorch 네이티브 RL 사후 학습(post-training) 프레임워크인 Molt를 구축했습니다. 핵심 RL 코드는 약 9,000라인으로 구성되어 있어 처음부터 끝까지 실제로 읽을 수 있으면서도, Qwen3.5-397B부터 GLM-5.2 753B 규모에 이르기까지 MoE 모델을 학습할 수 있습니다. 전체 스택은 Megatron 없이 오직 Ray + vLLM + NVIDIA AutoModel/FSDP2로만 이루어져 있습니다. 코드: https://github.com/NVIDIA-NeMo/labs-molt 기술 보고서 (DOI): 10.13140/RG.2.2.23375.65447 라이선스: Apache-2.0

문제점
대규모 모델에 대해 에이전트 RL(agentic RL)을 수행해 보았다면, 다음과 같은 과정을 잘 알고 계실 것입니다. 하나의 아이디어를 시도하기 위해—보상을 교체하거나, 도구 호출(tool call)을 추가하거나, 오프-폴리시(off-policy) 토큰 처리 방식을 변경하는 등의 작업—프레임워크를 열어보면, Megatron에 긴밀하게 결합된 5만 라인 이상의 코드를 마주하게 됩니다. 한 가지만 바꿔도 일주일 내내 코드를 읽어야 하며, 자신이 어떤 숨겨진 제어 흐름(control flow)을 건드렸는지조차 확실히 알 수 없습니다. 이러한 트레이드오프(trade-off)는 프로덕션 환경에서는 타당할 수 있습니다. 하지만 빠른 반복(iteration)이 핵심인 연구 분야에서는 매우 힘든 일입니다. 우리는 규모(scale)를 포기하지 않으면서도 몇 분 안에 읽고 수정할 수 있는 프레임워크를 원했습니다.

Molt란 무엇인가
설계는 의도적으로 작게 구성되었습니다: 에이전트는 당신이 작성하는 프로그램인 Gymnasium 스타일의 Env / ChatAgent입니다. 트레이너(trainer)는 단일 학습 가능한 액터(actor)인 PyTorch + AutoModel/FSDP2입니다. 보상(reward)은 규칙, 휴리스틱(heuristics), LLM-as-judge, 도구 반환값, 시뮬레이터 등 임의의 Python 코드로 구현됩니다. 나머지 부분은 세 가지 구성 요소가 담당합니다: 배치(placement) 및 비동기 큐(async queue)를 위한 Ray, 롤아웃(rollout)을 위한 vLLM, 그리고 학습을 위한 NVIDIA AutoModel + FSDP2입니다. 이것이 전체 런타임(runtime)입니다.

아키텍처

Molt architecture
이는 네 개의 박스를 거치는 완전 비동기 루프(fully-async loop)입니다:
vLLM Router/Rollout (고처리량 생성, 연속 배칭(continuous batching), EP=256까지의 전문가 병렬화(expert parallel)) ↔ Agent (당신의 Python: 멀티 턴, 도구, VLM, 채점자, API 서버) → Ray queue (궤적(trajectories) 버퍼링, 백프레셔(backpressure), 결함 허용(fault tolerance)) → RL Trainer (FSDP2/AutoModel 상에서의 advantage + 정책 업데이트). 데이터 경로는 토큰 우선(token-first) 방식입니다. 즉, 토큰 ID, 로그 확률(logprobs), 액션 범위(action ranges), 보상, 멀티모달 텐서가 롤아웃부터 학습까지 정렬된 상태를 유지하며, 라우터로의 가중치 동기화(weight sync)는 NCCL을 통해 비동기적으로 이루어집니다.

단순히 덧붙여진(bolted-on) 것이 아니라, 에이전트 우선(Agentic-first)이라고 부를 만한 특징들이 있습니다. 멀티턴 도구 사용(Multi-turn tool use), 멀티모달 (VLM) 환경, 그리고 LLM-as-judge가 일급 시민(first-class)으로 지원됩니다. 판사(judge) 모델조차 롤아웃(rollout)을 구동하는 것과 동일한 vLLM 엔진을 통해 경로를 호출하므로, 별도의 두 번째 추론 스택을 구축할 필요가 없습니다. Megatron을 사용하지 않는 순수 PyTorch 기반입니다. 학습이 AutoModel + FSDP2로 이루어지기 때문에, 학습 경로를 블랙박스로 취급하는 대신 실제로 읽고 변경할 수 있습니다. 확장성(Scale) 측면에서도 통신/연산 중첩(comm/compute overlap)을 지원하는 TP · EP · CP (· DP)를 제공합니다. Qwen3.5-397B에서 검증되었으며, GLM-5.2 753B 규모의 MoE까지 확장이 가능합니다. 극한의 경우, 가장 큰 액터(actor)들에 대해 optimizer CPU offload를 적용한 --fsdp.ep_size 256 설정의 DeepSeek-V3에서도 작동합니다. 단일 GPU 실험부터 시작된 약 9,000라인 규모의 코드베이스를 그대로 유지합니다. 학습/추론 일관성(Train/inference consistency) — 이는 비동기 RL(async RL)에서 실제로 문제가 되는 부분입니다. 비동기 및 부분적 롤아웃(partial rollout)은 FSDP 액터의 재계산된 로그 확률(recomputed logprobs)을 vLLM의 생성 시점 로그 확률(gen-time logprobs)과 다르게 만듭니다 (서로 다른 커널 사용, 그리고 HTTP 라우터가 관찰할 수 없는 요청 중간의 가중치 스왑 때문). Molt는 이를 vllm_kl = mean(rollout_logprob − actor_logprob)라는 액션 토큰에 대한 일급 지표로 드러내며, 결과적으로 발생하는 오프-폴리시(off-policy) 업데이트를 토큰별 중요도 비율(per-token importance ratio) pi_train / pi_rollout을 통해 교정합니다. 이 비율은 두 개의 직교하는 노브(knob)로 나뉩니다:

• --algo.advantage.is_correction_level {off, token, seq, geo} — 게이트 비율(gated ratio)의 세밀도(granularity)
• --algo.advantage.is_correction_mode {mask, clip, trunc} — 범위를 벗어난 단위(out-of-band units)를 처리하는 방식

geo × mask (기하 평균 비율에 따라 전체 시퀀스를 거부하고, 살아남은 시퀀스에 대해서는 토큰별 IS를 유지)가 기본값이며, token × trunc/mask는 TIS / IcePop 방식을 다룹니다. 또한 R3(router replay) — vLLM의 토큰별 전문가 선택(expert selection)을 캡처하여 학습 포워드(training forward)에서 재현함으로써 MoE 라우팅을 일치시키는 방식 — 및 라우터 동결(router-freeze) 옵션도 있습니다. 이러한 메커니즘은 Molt가 장난감 수준의 설정이 아니라 실제 대규모 모델 실행 과정에서 요구되었기 때문에 존재합니다. 완전 비동기 롤아웃(Fully-async rollout).

--train.async_queue_size는 롤아웃 (rollout)과 학습 (training)을 분리하며, --train.partial_rollout_enable은 vLLM의 pause/resume 기능을 통해 생성 (generation)과 가중치 동기화 (weight sync)를 중첩시킵니다. 하나의 플래그로 엄격한 온-폴리시 (on-policy) 방식으로 전환할 수 있습니다. 지원 범위는 SFT와 전체 RL 스택 (REINFORCE / RLOO / GRPO / DR-GRPO / PPO(GAE) / on-policy distillation), 멀티모달 (multimodal) + 멀티턴 (multi-turn), TP/EP/CP, 옵티마이저 (optimizer)/전체 CPU 오프로드 (full CPU offload), IS 보정 (IS correction), R3 라우팅 안정성 (R3 routing stability)을 포함합니다. Qwen3.x, Nemotron Omni, Kimi, GLM, Gemma, 그리고 DeepSeek 채팅 템플릿을 통해 검증되었습니다. 빠른 시작 (Quick start) bash 명령: python3 -m molt.cli.train_rl_ray \ --actor.model_name_or_path /path/to/automodel \ --data.prompt_dataset /path/to/prompts.jsonl \ --train.agent_path examples/python/agents/math.py \ --vllm.num_engines 2 --vllm.tensor_parallel_size 2 \ --rollout.batch_size 128 --train.batch_size 128 \ --algo.advantage.estimator reinforce \ --fsdp.attn_implementation te \ --ckpt.output_dir ./ckpt/rl examples/scripts/ 에는 Qwen3-4B부터 397B/753B 규모까지 바로 사용할 수 있는 레시피가 준비되어 있습니다. 솔직한 주의 사항을 말씀드리자면, 저는 저자 중 한 명입니다. 비판은 기꺼이 수용하며, PR(Pull Request)이나 이슈(issue) 제기도 환영합니다. Qwen3.5-397B는 엔드 투 엔드 (end to end)로 검증되었습니다. GLM-5.2 753B는 저희가 확장해 온 상한선입니다. 기술 보고서의 수치가 나오기 전까지는 확정된 벤치마크라기보다 "확장을 향해 나아가는 단계"로 간주해 주세요. 이 포스트는 의도적으로 정성적인 내용을 담고 있습니다. 처리량 (throughput) 및 수렴 (convergence) 수치는 기술 보고서를 참조하십시오. 링크: 코드 (Code): https://github.com/NVIDIA-NeMo/labs-molt 기술 보고서 DOI: 10.13140/RG.2.2.23375.65447 DeepWiki (자동 인덱싱된 Q&A): https://deepwiki.com/NVIDIA-NeMo/labs-molt 라이선스 (License): Apache-2.0 여러분의 설정과 비교한 피드백도 매우 환영합니다. /u/seventh_day123에 의해 제출됨 [link] [comments]