跳到正文
北京时间
原文
蚂蚁 inclusionAI：GitHub 新仓库· inclusionAI·· 2025-12-03精选AI 评分63

蚂蚁集团开源AState：面向强化学习的高性能状态管理系统

inclusionAI/asystem-astate

AI 导读

蚂蚁集团开源了AState，这是一个专为强化学习设计的高性能通用状态数据管理系统。它旨在解决RL训练与推理中的I/O效率低下、权重同步性能不足及状态恢复不鲁棒等核心挑战。系统采用三层架构：提供张量原生接口的API层、支持多种部署模式的服务层以及具备高效可扩展传输能力的基础层。其关键特性包括统一的张量级API、高性能权重同步和拓扑感知设计。在万亿参数规模下，AState能在约6秒内完成权重同步，远低于业界常见的分钟级延迟，目前已作为ASystem的关键组件在蚂蚁内部生产环境部署。

推荐理由

蚂蚁把万亿参数 RL 训练的权重同步从分钟级压到 6 秒，这套 AState 系统是真刀真枪的工程解法，做大规模 RL infra 的团队值得拆一拆它的 RDMA P2P 架构。

正文 · AI 翻译

AState：面向强化学习的高性能状态管理

AState

AState 是一个面向强化学习工作负载的通用状态数据管理系统。它旨在解决强化学习中的几个核心挑战：

  • 训练和推理中的 I/O 效率低下；
  • 权重同步性能不足；
  • 缺乏针对多轮对话和外部环境调用的稳健状态恢复能力。

在蚂蚁集团内部，AState 已作为 ASystem 的关键组件部署，并已在生产环境中支撑大规模强化学习训练的高性能权重交换。在万亿参数规模下，AState 可在约 6 秒内完成权重同步，而业界许多方案的延迟为分钟级。

系统架构

Figure 1

为应对大规模强化学习权重同步的挑战，AState 为强化学习工作负载提供了统一的权重管理 API。它可支持任意模型架构以及多种部署和流水线模式，无需对强化学习框架本身进行侵入式修改或额外适配。如图 1 所示，AState 的架构分为三层：

  • API 层

    提供张量原生接口，可快速与不同的训练/推理框架集成。这些 API 暴露单边读写语义，天然支持更复杂或异步的计算与数据交换流程。

  • 服务层

    为了用单一 API 支持多种部署模式和流水线模式，AState 引入了一个中间服务层，提供不同的权重同步服务和协议，将上层与交换和调度的复杂性隔离开来。例如：

    • 在共置训练/推理中，推理节点可以按需拉取权重。
    • 在离策略（off-policy）设置中，训练和推理可以执行完全异步的权重更新。该服务层还提供零冗余的张量分片管理，以及能够感知 RL 拓扑和集群布局的同步计划。

  • 传输层

    这一基础层提供高效、可扩展的数据传输能力，包括：

    • NUMA 拓扑与亲和性感知；
    • 多种传输后端（PCIe / NVLink / RoCE / InfiniBand 等），以充分释放底层硬件带宽和拓扑的潜力。

功能特性

  • 统一的张量级 API

    • 单向读/写语义（put/get/multi_put/multi_get/complete）
    • 与具体训练/推理框架解耦

  • 高性能权重同步

    • 面向行并行与列并行张量的零冗余传输
    • 基于 RDMA 的 P2P 数据通路，配合 DMA 零拷贝
    • 推理侧原地权重更新（无需额外 GPU 显存拷贝）

  • 拓扑感知、可扩展的设计

    • NUMA 感知调度与 CPU/GPU/RDMA 亲和性
    • 多通道传输：RDMA/RoCE/IB、共享内存、NCCL 等
    • 全局执行规划，避免热点与长尾延迟

  • RL 原生状态管理（进行中）

    • 面向长上下文 RL 的 KV cache 持久化与恢复
    • 激活值缓存与内存卸载
    • 智能体状态（多轮对话、工具调用）缓存与恢复

快速开始

git clone https://github.com/inclusionAI/asystem-astate.git
cd asystem-astate
make deps install      # installs as a Python library

# development build:
make develop      # build in develop mode + install as Python library
make test

简单用法示例

该内部库目前仅支持 RDMA 网卡，不支持使用 soft-RoCE（TCP）。基于 TCP 的演示将在集成 UCX 后进行测试。

# start train process
cd python/example
bash setup_trainer_env.sh

# start infer process
bash setup_infer_env.sh

RL 中的典型用法：

  • 训练侧每 K 步或每次迭代将更新后的权重写入 AState。
  • 推理侧通过 get/multi_get 定期或异步拉取最新权重。
  • AState 为你处理分片、重分片、传输和原地更新。

性能

方法 支持的部署模式 支持的流水线模式 数据冗余 所需外部服务 端到端同步（100B¹） 端到端同步（1T¹）
从 FS 加载数据 全部 同策略（On-policy） 是 是 约 10 分钟 不适用²
逐层 gather → NCCL P2P 全部？ 全部？ 是 是 约 30 秒 约 1 分钟
AState：同时进行 RDMA P2P 全部 全部 否 否 约 4 秒 ~6 s

$^?$ 这可能不是最优解。例如，共置模式引入了额外的 GPU 显存卸载流水线；此外，由于 NCCL 的双边语义，训练和推理在数据传输期间必须同步，这实际上可能导致 off-policy 方案退化为 on-policy。

$^1$ 100B(FP16) [训推并行配置：训练：tp=1, etp=1；推理：tp=4]；

$^1$ 1T(FP8) [训推并行配置：训练：tp=1, etp=1；推理：tp=16]。

$^2$ 时间占比过长，几乎无法使用。

路线图（高层概览）

  • 支持使用 UCX 作为底层 RDMA 依赖，适配 NCCL 及行业生态。
  • 面向长上下文 RL 的 KV cache 持久化与快速故障切换。
  • 激活值缓存与 GPU 显存扩展。
  • 弹性训练运行时支持。
  • 统一智能体状态（对话 + 工具）缓存与恢复。

社区与链接

非常欢迎通过 Issues / PRs / Discussions 贡献。如果 AState 对你的 RL 系统有帮助，欢迎点一个 ⭐ Star！

来源：蚂蚁 inclusionAI：GitHub 新仓库 · github.com

#Agent#开源/仓库#部署/工程