VeRL-Omni: Diffusion 및 Omni-Modality 모델을 위한 RL 학습 프레임워크
VeRL-Omni: Diffusion 및 Omni-Modality 모델을 위한 RL 학습 프레임워크
vLLM은 멀티모달 생성 모델을 위해 설계된 범용 강화 학습(RL) 사후 학습 프레임워크인 VeRL-Omni의 프리릴리스(pre-release)를 발표했습니다. verl 및 vllm-omni를 기반으로 구축된 VeRL-Omni는 diffusion transformer 및 통합 이해 및 생성 아키텍처와 같은 비자기회귀(non-autoregressive) 및 omni-modal 모델의 RL 학습 시 발생하는 고유한 과제들을 해결합니다.
멀티모달 RL의 기술적 과제
VeRL-Omni는 멀티모달 생성 RL과 관련된 세 가지 주요 기술적 난제를 해결하도록 설계되었습니다:
- 아키텍처 확장: 이 프레임워크는 RL 기능을 diffusion transformer 백본(예: Qwen-Image), 혼합 AR-DiT 아키텍처(예: Qwen-Omni), 그리고 통합 모델(예: BAGEL, HunyuanImage3.0)로 확장합니다.
- 이질적 Rollout 파이프라인: 텍스트 기반 RL와 달리, 멀티모달 rollout은 연속적인 잠재 공간(latent space)에서의 디노이징 궤적(denoising trajectories)입니다. 이러한 파이프라인은 텍스트 인코더, Diffusion Transformers (DiT), Variational Autoencoders (VAE)와 같은 여러 구성 요소를 포함하는 경우가 많습니다.
- 워크로드 스케줄링: 멀티모달 RL은 보상 함수 자체가 멀티모달 모델(VLM judge 또는 OCR scorer 등)인 복잡한 워크플로우를 조율해야 하며, 생성 rollout은 일반적으로 텍스트 생성보다 더 높은 메모리 피크를 나타냅니다.
주요 프레임워크 기능
VeRL-Omni는 멀티모달 RL 학습을 위한 모듈식이고 효율적인 스택을 제공합니다:
- 효율적인 멀티모달 Rollout: vLLM-Omni를 통합함으로써, 이 프레임워크는 단계별 연속 배칭(step-wise continuous batching)과 임베딩 캐싱을 사용하여 효율성을 최적화하고 멀티모달 생성을 위한 고처리량 비동기 서빙을 활용합니다.
- 유연한 보상 엔진: 시스템은 규칙 기반 및 모델 기반 보상(예: VLM-as-judge)을 모두 지원합니다. vLLM은 효율적인 보상 모델 추론에 사용되며, 보상 계산은 지연 시간을 줄이기 위해 rollout 및 학습 프로세스와 중첩(overlap)됩니다.
- 모듈식 학습 백엔드: 프레임워크에는 DiffusersFSDP, Megatron, VeOmni와 같은 다양한 트레이너가 포함되어 있으며, diffusion 및 omni-modal 모델에 대한 내장 최적화와 FSDP, USP, TP를 포함한 병렬화 전략을 지원합니다.
- 하드웨어 호환성: VeRL-Omni는 NVIDIA GPU와 Ascend NPU를 모두 지원합니다.
지원 모델 및 알고리즘
VeRL-Omni는 다양한 아키텍처와 RL 알고리즘을 지원하며, 현재 여러 개가 출시되었거나 개발 중입니다:
| 모델 | 아키텍처 | 모달리티 | 알고리즘 | 상태 |
|---|---|---|---|---|
| Qwen-Image | DiT | Text → Image | FlowGRPO, MixGRPO, GRPO-Guard | 출시됨 |
| BAGEL | Unified understand + gen | Text + Image | FlowGRPO | PR 준비됨 |
| Qwen3-Omni-Thinker | AR | Text / Image / Video / Audio | GSPO | PR 준비됨 |
| Wan2.2 | DiT | Text → Video | DanceGRPO | WIP |
| SD3.5 | DiT | Text → Image | DPO | WIP |
| HunyuanImage-3.0 | Unified understand + gen | Text + Image | MixGRPO, SRPO | 계획됨 |
FlowGRPO 구현 및 성능
VeRL-Omni는 flow-matching 모델을 위한 온라인 정책 방법인 FlowGRPO를 구현합니다. 학습 워크플로우는 4단계로 구성됩니다: rollout 생성(궤적 및 이미지 수집), 보상 모델 점수 산정, CLIP 스타일 손실을 사용한 정책 최적화, 그리고 트레이너에서 rollout 워커로의 주기적인 가중치 동기화.
성능 벤치마크
OCR 보상 작업에 대해 Qwen-Image의 LoRA 미세 조정을 위해 NVIDIA H800 GPU를 사용한 테스트에서, 프레임워크는 다음과 같은 처리량을 보여주었습니다:
- Colocated Training: 단계당 시간 420s, 0.305 images/GPU/s.
- Async Reward: 단계당 시간 360s, 0.280 images/GPU/s. 보상 모델을 전용 GPU로 이동하면 단계당 wall-clock 시간이 약 14% 감소합니다.
4 × NVIDIA H200 GPU에서 Qwen-Image(non-CFG)의 전체 모델 미세 조정을 위해, 프레임워크는 약 250 s/step에서 0.510 images/GPU/s를 달성했습니다.
향후 로드맵
VeRL-Omni는 현재 프리릴리스 단계입니다. 향후 개발 로드맵은 다음과 같습니다:
- 모델 지원 확대: 이미지, 비디오 및 오디오 생성을 위한 더 많은 오픈 소스 diffusion 및 omni-modal 모델 추가.
- 고급 알고리즘: DiffusionNFT와 같은 새로운 RL 알고리즘 통합.
- 완전 비동기 RL: GPU/NPU 활용도를 높이기 위해 actor, rollout, reward 간의 엔드 투 엔드 비동기 파이프라인 개발.
- vLLM-Omni 공동 최적화: 고급 병렬화, 양자화 및 배칭을 통해 rollout 가속화.
- 최적화된 트레이너: Megatron-core 및 VeOmni를 기반으로 한 더 많은 트레이너 엔진 출시.
- 하드웨어 확장: Ascend NPU 경로 강화 및 더 많은 하드웨어 백엔드 추가.