Relax：小红书开源的全模态大模型强化学习引擎，高效异步训练框架

一、Relax是什么

Relax（Reinforcement Engine Leveraging Agentic X-modality）是由小红书AI Infra团队开源的面向全模态与智能体场景的大规模大模型强化学习（RL）后训练引擎。它是专为解决当前多模态大模型强化学习训练中存在的异构数据流、大规模运维鲁棒性、陈旧度-吞吐量权衡三大核心痛点而设计的新一代基础设施。

Relax定位为生产级、全异步、服务化的强化学习训练框架，底层依托Ray Serve实现服务化架构、Megatron-LM提供分布式训练能力、SGLang支撑高吞吐推理，将全模态原生支持、服务化容错架构、端到端异步流水线三大技术维度深度协同设计，形成完整闭环解决方案。与传统单模态、耦合式、同步训练框架不同，Relax从底层架构层面原生支持文本、图像、音频、视频全模态数据的统一RL训练，并实现训练与推理模块的完全解耦、弹性扩展与故障隔离，在Qwen3-Omni-30B等全模态模型上已验证训练稳定收敛，16×H800多机环境下全异步训练相比同步基线提速76%。

二、功能特色

2.1 全模态原生支持，覆盖多类型数据训练

• 全模态统一处理：原生支持文本、图像、音频、视频四种核心模态，无需针对单模态改造适配，通过--multimodal-keys配置即可接入多模态数据，统一处理流程与训练逻辑。

• 模态感知并行：智能识别不同模态数据类型，将图像、文本、音视频分配至对应计算单元并行处理，最后融合结果，避免多模态数据处理瓶颈。

• 精细化预处理：内置CPU预处理池（Processor Pool），通过独立进程+零拷贝共享内存实现图像resize、视频抽帧、音频重采样等CPU密集操作，与GPU推理流水线重叠，完全隐藏预处理延迟。

• ViT高效并行：通过Megatron Bridge兼容HF原生ViT实现，TP维度数据并行编码，各rank独立处理像素切片后AllReduce合并，ViT参数冗余开销仅1-5%，无格式转换风险。

2.2 服务化六层架构，生产级高可用

Relax采用服务化六层架构，所有功能角色封装为独立Ray Serve部署单元，支持弹性调度、故障恢复与热升级：

1. Entrypoints：训练入口（train.py），负责信号处理、CLI解析、Ray集群连接与控制器启动。

2. Orchestration：训练编排层，Controller管理训练循环与全局重启，Service管理部署组与生命周期。

3. Components：核心组件层，包含Actor、Rollout、Critic、GenRM等独立Ray Serve部署单元。

4. Engine：引擎层，集成SGLang推理引擎、奖励函数、请求路由与数据过滤模块。

5. Backends：后端层，训练依托Megatron-LM（支持TP/PP/CP/EP全维度分布式），推理依托SGLang（高吞吐低延迟）。

6. Distributed：分布式层，Ray Actor组+DCS分布式Checkpoint服务，实现权重快速同步与故障恢复。

2.3 全异步训练流水线，极致吞吐量提升

• 训练推理解耦：Rollout（推理生成）与Train（梯度更新）拆分为独立服务，通过TransferQueue数据总线连接，读写完全异步、互不阻塞。

• Micro Batch细粒度流水：摒弃全局Batch同步，推进至Micro Batch级别流水线，数据预处理、模型计算、结果反馈多阶段重叠执行，最大化GPU利用率。

• 双模式灵活切换

• Colocate（同步）：Actor/Rollout分时复用GPU，内存高效，适合资源受限、严格on-policy场景。

• Fully Async（全异步）：各模块跨集群并行，异步权重同步，吞吐量最优，16×H800下比同步基线提速76%。

• 弹性Rollout扩展：支持训练中通过HTTP REST API动态扩缩推理引擎，兼容集群内（ray_native）与跨集群（external）联邦部署。

2.4 智能体化RL能力，适配复杂交互场景

• 多轮交互闭环：支持Agentic多轮对话训练，实现“执行→观察→决策”完整闭环，适配工具调用、逻辑推理、长文本生成等复杂场景。

• 灵活终止策略：支持自定义对话终止条件，适配不同任务的交互长度需求。

• 损失掩码支持：针对多轮交互中无效token、填充token实现损失掩码，提升训练精度。

2.5 丰富算法与工程化特性

• 主流算法全覆盖：内置GRPO、GSPO、SAPO策略优化算法，支持On-Policy Distillation知识蒸馏，新增算法可通过注册ALGOS快速扩展。

• 分布式Checkpoint服务（DCS）：拓扑感知的分布式权重传输系统，支持快速权重同步与故障恢复，解决大模型权重传输瓶颈。

• 全链路监控告警：集成HealthManager自动故障恢复，支持WandB、TensorBoard、ClearML中心化指标监控，Apprise实时告警，满足生产运维需求。

• 跨平台兼容：支持Python 3.12，兼容主流GPU集群（NVIDIA H100/H800/A100等），提供Docker镜像快速部署。

三、技术细节

3.1 核心技术栈

3.2 异步训练核心原理

传统RL训练采用耦合式同步架构，Rollout推理与Train训练强绑定，必须等待前一批数据生成完成才能开始训练，GPU大量时间处于空闲状态。

Relax采用全异步解耦架构，核心原理如下：

1. 服务拆分：将Rollout（SGLang）、Train（Megatron-LM）、Critic、Reward等模块拆分为独立Ray服务，部署在不同GPU节点/集群。

2. 数据总线：通过TransferQueue实现模块间异步数据交互，Rollout生成数据后异步写入队列，Train异步读取数据训练，无阻塞等待。

3. 权重异步同步：通过DCS服务实现训练权重的异步广播，Rollout服务定期拉取最新权重，支持可配置“陈旧度”，平衡训练效率与策略一致性。

4. Micro Batch流水：将全局Batch拆分为多个Micro Batch，实现“上一批训练、下一批生成、再下一批预处理”的三级流水线，GPU利用率接近100%。

3.3 多模态技术实现

• 数据层：统一多模态数据格式，支持文本（token）、图像（pixel_values）、音频（waveform）、视频（frame_features）的联合输入与标注。

• 预处理层：Processor Pool独立进程处理多模态预处理，asyncio实现三阶段流水线，CPU预处理完全隐藏在GPU推理延迟后。

• 模型层：通过Megatron Bridge兼容HF多模态模型（如Qwen3-Omni），ViT采用TP数据并行，文本模型采用TP/PP混合并行，跨模态特征自动对齐融合。

• 训练层：统一策略优化算法，支持多模态联合奖励计算，针对不同模态设计差异化损失权重，确保全模态训练稳定收敛。

3.4 分布式与容错机制

• 弹性分布式：基于Ray集群，支持多机多卡分布式部署，各模块可独立扩缩容，Rollout服务支持动态增加推理节点提升吞吐量。

• 故障自动恢复：HealthManager实时监控各服务状态，模块故障后自动重启、重新加入集群，训练状态通过DCS快速恢复，无全局中断。

• Checkpoint持久化：支持定时Checkpoint与异常退出自动保存，DCS分布式存储Checkpoint，支持断点续训。

四、应用场景

4.1 全模态大模型强化学习后训练

• 文本-图像多模态模型：如Qwen-VL、LLaVA、InternVL等，通过RL优化图文理解、跨模态生成、视觉推理能力。

• 文本-音频多模态模型：如AudioLM、SpeechLLM，优化语音理解、语音生成、语音-文本交互质量。

• 全模态大模型：如Qwen3-Omni、GPT-4V，支持文本、图像、音频、视频联合训练，提升多模态理解、生成、推理综合性能。

4.2 大模型智能体（Agent）训练

• 工具调用智能体：训练模型精准调用API、代码执行、搜索引擎等工具，提升工具使用成功率与准确率。

• 多轮对话智能体：优化长文本对话、上下文理解、逻辑推理、情感交互能力，适配客服、助手、陪伴等场景。

• 自主决策智能体：训练模型在复杂环境中自主感知、决策、执行，适配游戏AI、自动驾驶、工业控制等场景。

4.3 大规模生产级模型调优

• 千亿/万亿参数大模型：依托Megatron-LM分布式能力，支持超大规模模型的RL训练，解决显存与算力瓶颈。

• 商业化模型迭代：为企业级大模型提供高效、稳定、可扩展的RL调优方案，快速迭代模型性能，缩短上线周期。

• 多模态内容生成优化：优化图文笔记、短视频解说、音频内容生成质量，提升内容相关性、原创性、用户满意度。

4.4 科研与算法创新

• 强化学习算法验证：为GRPO、GSPO、PPO等算法提供高效实验平台，快速验证算法在多模态场景的有效性。

• 新模态训练探索：支持3D点云、传感器数据等新兴模态的RL训练研究，拓展大模型应用边界。

• 训练效率优化研究：为异步训练、分布式训练、模型压缩等方向提供底层基础设施支撑。

五、使用方法

5.1 环境准备

5.1.1 硬件要求

• 最低配置：单节点≥8×A100（80GB）GPU，CPU≥32核，内存≥256GB，存储≥1TB

• 推荐配置：多节点≥16×H100/H800 GPU，分布式Ray集群，高速RDMA网络

5.1.2 软件依赖

# 基础依赖
Python 3.12+
CUDA 12.1+
cuDNN 8.9+
NCCL 2.18+

# 核心框架
Ray >= 2.30.0
Megatron-LM >= 0.8.0
SGLang >= 0.4.0
Transformers >= 4.40.0
Torch >= 2.3.0

5.2 安装部署

5.2.1 源码安装

# 克隆仓库
git clone https://github.com/redai-infra/Relax.git
cd Relax

# 安装依赖
pip install -e .

# 安装Megatron-LM与SGLang
bash scripts/install_megatron.sh
bash scripts/install_sglang.sh

5.2.2 Docker部署（推荐）

# 拉取镜像
docker pull redaiinfra/relax:latest

# 启动容器（GPU映射）
docker run -it --gpus all --shm-size 128g \
 -v /path/to/data:/data \
 -v /path/to/checkpoint:/checkpoint \
 redaiinfra/relax:latest

5.3 快速开始

5.3.1 单节点全模态训练（GRPO算法）

# 启动Ray集群
ray start --head

# 启动Relax全模态训练（Qwen3-Omni-30B示例）
python train.py \
 --model_name_or_path Qwen/Qwen3-Omni-30B \
 --dataset_name ./data/multimodal_dataset \
 --multimodal_keys text,image,audio,video \
 --algorithm GRPO \
 --train_mode fully_async \
 --num_rollouts 8 \
 --micro_batch_size 4 \
 --epochs 10 \
 --output_dir ./checkpoints/relax_qwen3_omni \
 --report_to wandb

5.3.2 分布式集群训练

# 节点1（Head）
ray start --head --node-ip 192.168.1.100

# 节点2-16（Worker）
ray start --address 192.168.1.100:6379

# 启动分布式训练
python train.py \
 --model_name_or_path Qwen/Qwen3-Omni-70B \
 --dataset_name ./data/large_multimodal_dataset \
 --multimodal_keys text,image,video \
 --algorithm GSPO \
 --train_mode fully_async \
 --num_nodes 16 \
 --tp_size 8 \
 --pp_size 4 \
 --rollout_nodes 8 \
 --train_nodes 8

5.4 核心配置说明

• --model_name_or_path：预训练模型路径或Hugging Face模型名

• --multimodal_keys：启用的模态类型，可选text/image/audio/video

• --algorithm：训练算法，可选GRPO/GSPO/SAPO/Distillation

• --train_mode：训练模式，colocate/fully_async

• --tp_size/pp_size：张量并行/流水线并行维度

• --num_rollouts：Rollout服务数量，控制推理吞吐量

• --output_dir：Checkpoint与日志输出目录

六、竞品对比

6.1 主流框架对比（Relax vs OpenRLHF vs veRL）

6.2 核心差异总结

1. Relax优势：全模态原生支持、极致全异步、服务化高可用、跨集群弹性，是目前唯一专为多模态场景设计的生产级RL框架，效率与稳定性领先。

2. OpenRLHF优势：轻量易用、生态完善、快速上手，适合中小模型科研与小规模生产，多模态能力较弱。

3. veRL优势：大规模模型优化、算法全面、字节内部验证，文本模型训练高效，多模态支持有限。

七、常见问题解答

Q：Relax支持哪些预训练模型？

A：Relax原生支持Hugging Face生态的主流大模型，包括Qwen系列（Qwen3-Omni/Qwen-VL）、LLaMA系列、LLaVA、InternVL、Mistral等全模态与文本大模型；通过Megatron Bridge可兼容自定义Megatron-LM模型，支持7B~70B+参数规模。

Q：全异步模式与Colocate模式如何选择？

A：资源有限、追求策略一致性（严格on-policy）时选Colocate模式，GPU内存利用率高、无需跨集群；追求极致吞吐量、多机分布式部署时选Fully Async模式，训练效率提升50%+，支持Rollout与Train分离部署。

Q：多模态训练需要额外配置哪些参数？

A：核心配置--multimodal_keys指定启用模态（如text,image,video）；通过relax/utils/multimodal/下配置文件调整图像尺寸、视频帧采样率、音频采样率；ViT并行通过--vit_tp_size配置，默认与文本模型TP一致。

Q：训练过程中Rollout服务如何动态扩缩？

A：Relax提供HTTP REST API接口，训练中通过curl -X POST http://controller:8000/scale_rollout -d '{"num_rollouts": 16}'动态增加Rollout数量；减少时同理，系统自动下线空闲服务，无训练中断。

Q：DCS分布式Checkpoint服务有什么作用？

A：DCS是Relax核心分布式组件，解决大模型权重快速同步与故障恢复问题：训练时异步广播最新权重至所有Rollout节点，故障时快速从DCS拉取最新Checkpoint恢复训练，相比传统文件存储同步速度提升10倍+。

Q：Relax支持单GPU训练吗？

A：支持但不推荐。单GPU可运行Colocate模式，仅适合小模型（7B以下）测试与调试；全模态与大规模模型必须多GPU分布式部署，否则显存不足、训练极慢。

Q：如何监控训练进度与性能指标？

A：默认支持TensorBoard本地监控；通过--report_to wandb启用WandB云端监控，实时查看loss、奖励值、吞吐量、GPU利用率等指标；HealthManager提供服务状态与故障日志，可通过Ray Dashboard查看集群状态。

Q：Relax相比TRL有哪些优势？

A：TRL是Hugging Face推出的轻量级RL框架，适合中小模型科研；Relax定位生产级大规模训练，优势：全模态原生支持、全异步架构、服务化高可用、分布式弹性扩展、吞吐量提升2-3倍、支持70B+超大规模模型。

八、相关链接

• GitHub仓库：https://github.com/redai-infra/Relax

• 官方技术论文：https://arxiv.org/pdf/2604.11554

九、总结

Relax作为小红书AI Infra团队开源的新一代全模态大模型强化学习后训练框架，通过全模态原生支持、服务化六层架构、全异步训练流水线三大核心技术创新，系统性解决了当前多模态大模型RL训练中异构数据流处理、大规模运维鲁棒性、效率-一致性权衡的行业痛点，依托Ray Serve、Megatron-LM与SGLang的技术栈优势，实现了训练与推理的完全解耦、弹性扩展、故障隔离与极致吞吐量，在Qwen3-Omni等全模态模型上验证了稳定收敛与显著效率提升，相比同类框架在多模态适配、异步性能、生产级可用性方面具备明显领先优势，适用于从科研实验到商业化大规模部署的全场景需求，为多模态大模型的强化学习调优提供了完整、高效、可靠的基础设施解决方案。