过去几年，强化学习已成为推动大模型边界扩张的核心技术。从 ChatGPT 的 RLHF，到 DeepSeek、Claude、Llama 的后训练体系，无不依赖强化学习让模型更符合人类偏好、具备更强的推理能力。
今年 10 月，蚂蚁百灵正式开源了两款业界领先的万亿参数旗舰模型：非思考模型Ling-1T和思考模型Ring-1T。其中，Ring-1T如此大规模思考模型的训练，对系统工程提出了极高的要求，它不仅是算法层面的突破，更是一项融合了极致技术深度与匠心精神的系统工程挑战。
为了让大家更全面、更深入地理解蚂蚁百灵如何驾驭万亿模型训练，我们精心策划了为期 5-6 周的「ASystem系统开源」核心技术解析系列分享。我们将逐一解密 ASystem 如何实现高性能计算，以及我们为克服万亿规模工程挑战所做出的系统级创新。
请锁定我们的公众号，每周四不见不散，共同进入高性能计算的「心流之境」！
本周，我们首先为大家带来系列分享的首个重磅组件——蚂蚁自研的高性能权重交换框架 Awex（Asystem Weight Exchange Framework）。
Awex 是专为极致性能打造的训练推理引擎权重同步框架，旨在解决 RL 流程中，训练权重参数同步到推理模型这一核心延迟难题。
Awex 的能力体现在如下两个维度：
-
极致性能：可在秒级完成 TB 级大规模参数交换，显著降低 RL 模型训练延迟。 -
实战验证：Ring-1T 正是基于 Awex，在数千卡 GPU 集群上，实现了 5~10 秒内完成万亿参数级别的权重同步，性能业界领先。
代码已正式开源：https://github.com/inclusionAI/asystem-awex
本文将深入为您介绍 Awex 框架的设计理念、我们面对的核心挑战、创新的解决方案及详尽的性能数据。
问题与挑战
在大模型强化学习训练中，为提升推理效率，Rollout 阶段通常使用独立的高性能推理引擎（如 SGLang/vLLM）。权重同步是指将训练引擎每完成一个 step 后的模型权重传输给推理引擎，确保推理使用最新参数生成样本，形成完整的训练闭环。
随着 RL 模型参数、训练集群规模的扩大，以及不同RL算法的演进与变化，权重同步面临以下技术挑战：
训练推理引擎的异构并行策略适配对齐：训练引擎一般采用 DP + PP/VPP + TP + EP + CP (*)，而推理引擎采用 DP + TP + EP 居多，不同并行策略导致权重分片不一致，需要高效的 Resharding 机制；
训练推理引擎的差异化权重精度与格式：推理引擎常用低精度（FP8/INT4）提升性能，训练引擎使用高精度（BF16/FP32）保证收敛，不同引擎对权重 Layout 和 Fusion 有不同要求；
权重参数的增长带来的交换效率：随着 Scaling Law 的持续演进，大模型参数规模逐步增长到万亿级别，在千卡规模的集群上需传输 100TB+ 数据，传输耗时成为训练瓶颈；
共卡&分卡部署模式：同步 RL 算法多采用共卡模式（Co-locate），训推分时共用 GPU，最大化资源利用率；异步 RL 算法则采用分卡模式（Separate），训推使用不同 GPU，并行运行；
推理实例的动态扩容：训练迭代时模型输出逐步变长，Rollout 耗时增加，需要动态扩容推理实例，权重同步需支持新实例接入。
（*注：异构并行策略常见的为如下几种：Data Parallelism DP 数据并行；Pipeline Parallelism PP 流水线并行；Tensor Parallelism TP 张量并行；Expert Parallelism EP 专家并行；Context Parallelism CP 上下文并行）
现有方案的局限性
现有的权重传输方案主要分为两类，在模型参数规模较大的场景下仍然存在性能瓶颈：
基于分布式文件系统来交换模型权重：训练引擎在完成单个 step 的模型训练后，将模型权重逐层 Gather 并写入共享文件系统，之后推理引擎从共享文件系统加载模型权重。受到共享文件系统的性能限制，10B 参数模型需要数分钟，万亿参数模型换则耗时小时级以上；
基于 NCCL + Full Tensor 逐层同步权重：这是部分开源框架采用的权重交换方式，训练引擎逐层 Gather Full Tensor 到每张卡上，然后通过 CUDA IPC 方式映射到推理引擎进行权重加载。这种方式逐层串行执行，无法充分利用带宽，在 TP 较大时存在较多冗余传输， 模型参数较大时仍需要分钟级以上的延迟。
Awex 简介
Awex 是为极致性能打造的训练推理引擎权重同步框架，解决 RL 流程中训练权重参数同步到推理模型的核心难题，可在秒级完成 TB 级大规模参数交换，显著降低 RL 模型训练延迟，主要特点如下：
⚡ 极速同步性能：千卡集群万亿参数模型 6 秒内全量同步，性能领先；
🔄 统一模型适配层：自动处理训推引擎并行策略与引擎间的 Tensor 格式/布局差异，兼容多种模型架构；
💾 零冗余 Resharding 传输与原地更新：仅传输必要分片（Shard），推理侧原地更新显存，避免重分配与拷贝开销；
🚀 多模式传输支持：支持 NCCL、RDMA、共享内存多种传输模式，充分发挥 NVLink/NVSwitch/RDMA 带宽并减少长尾延迟；
🔌 异构部署兼容：适配共卡/分卡模式，支持同步和异步RL算法训练场景，同时 RDMA 传输模式支持推理实例动态扩缩容；
🧩 灵活可插拔架构：支持对不同模型定制化权重Sharing和Layout行为，同时支持新的训练和推理引擎接入。
GitHub Repo: https://github.com/inclusionAI/asystem-awex
Awex 总体设计
Awex 权重交换框架整体主要由三个组件组成：
WeightWriter：在每个训练进程内运行，负责当前训练进程的权重 Shard 的元数据收集上报、权重转换、权重发送 Resharding 计划构建、权重发送等功能；
WeightReader：在每个推理实例的控制进程上运行，其会在推理实例管理的每张GPU上面启动一个 WorkerWeightsReader，与训练进程的 WeightWriter 相对应，负责每个推理进程的权重 Shard 的元数据收集上报、权重转换、权重接受 Resharding 计划构建、权重接受等功能；
MetaServer：Job级别全局 Server，用于训推引擎的服务发现和权重元数据交换，以及共卡情况下的事件通知等功能。
权重交换的核心功能模块主要由 5 个部分组成，我们将逐一展开介绍：
训推权重统一转换：负责将不同并行策略和 Tensor 布局的训练引擎和推理引擎的权重转换成统一的格式，用于后续的权重元数据计算和权重传输；
全局权重元数据计算与交换：将训推权重转换成统一的格式后，收集每个 Worker 的所有权重 Shard 元数据，并上报到 Meta Server，用于接下来的的权重传输计划构建；
P2P 权重传输执行计划：训练和推理引擎拿到全局所有 Worker 的训练和推理权重 Shard 元数据，然后分别各自构建对等的发送和接受确定性传输计划；
NCCL 权重传输：使用 NCCL 的 send/recv API基于构建的传输计划进行对等的权重发送与接收；
RDMA 权重传输：使用 NUMA 亲和，面向全局负载均衡传输计划的 RDMA 通信来进行权重的更新。
同时 Awex 也支持对权重进行 Tensor 级别的校验，将通过文件系统模式加载的权重跟通过传输模式加载的权重进行 Tensor 级别的细粒度比对，逐个判断差异，保证传输模式的正确性，限于篇幅，这里将不做展开。
（1）训推权重统一转换
训推引擎由于计算负载不同，一般采取不同的并行策略，Megatron 训练引擎采用 5D 并行策略，DeepSpeed/FSDP 采用 Zero + DP 数据并行，而 SGLang 和 VLLM 推理引擎则采用 DP + TP + EP 居多。同时不同引擎在加载完权重后，还会对权重做合并转置量化等优化，以适配高性能算子。
为了能够消除不同引擎的差异，用于后续的权重交换，Awex 构建了统一的权重转换层，主要做以下转换：
权重的拆分：将合并的权重（如 FFN 的 gate/up）拆分为独立权重，支持跨 TP 的 Resharding；
权重名的统一：将所有引擎的内部权重都统一转换到相同的命名空间，从而建立训推引擎之间的权重映射关系；
Attention 权重的ReGroup：在训练引擎侧对 QKV 权重沿着推理引擎的 TP/DPAttention 并行策略进行重分组和对齐，避免细粒度拆分带来的 Shard 膨胀；
量化、精度与格式转换：在训练侧自动将权重按照推理侧的精度与格式进行转换，同时这种低精度转换也可以减少传输的数据量。
整个权重转换适配层在实现上是一个可插拔的结构，可以在引擎层、模型权重转换、Sharding 层进行全面的自定义适配，满足复杂场景的定制化需求。
（2）全局权重元数据管理
每个训练和推理进程需要感知全局所有训练与推理进程的的权重元数据，用于构建后续的权重传输计划，同时 Awex 也会在这一步对训推的权重元数据进行一致性校验，主要流程如下：
训练引擎的每个进程执行权重转换，并获取转换后的 Shard 的元数据
通过 all_gather_object 使每个 Rank 获取全局的训练 Shard 元数据
训练侧的 Rank0 将全局元数据序列化并上报到 Meta Server
推理侧的推理实例 0 做类似的工作，其它推理实例元数据相同，不需要额外计算
所有训推进程从 MetaServer 获取全局元数据
训推引擎会各自基于进行 Shard 级别的 Meta 一致性和兼容性校验
（3）P2P 权重传输执行计划
在获得全局权重元数据后，Awex 在每个训推进程内构建确定性的点对点传输计划。
核心策略（NCCL 模式）：
对相同 Tensor Shard 的每个副本，通过 Round Robin 进行将训练的 Shard 分配给推理 Shard，保证均匀拉取；
对于 Overlap 的 Shard 区间，若完全对齐则直接映射，否则采用两次发送给不同 Shard；
提前筛选与当前进程相关的 Shard ，避免构建全局计划（万亿参数下 Shard 可达千万级）；
保证 NCCL send/recv 的严格顺序一致性；
RDMA 相比 NCCL 更加灵活，使用了一套单独的传输计划，我们会在后续文章进一步展开。
（4）NCCL权重传输
Awex 支持 NCCL（NVIDIA Collective Communications Library）和 RDMA（Remote Direct Memory Access）两种传输模式，NCCL 模式更加易用，RDMA 模式更加灵活且具备更高的性能。
NCCL 传输模式主要使用 NCCL 的 send/recv 接口来进行权重的传输，在分卡和共卡模式下，Awex 的实现也有一些差异，这里我们将详细展开。
NCCL 分卡权重传输
在分卡传输模式下，Awex 首先构建训推联合的 NCCL Process Group，用于后续从训练到推理的权重传输。接下来便是根据初始化阶段创建的 NCCL 发送和传输计划，在训练和推理引擎的每个进程上进行 P2P 的对等有序的权重发送和接收，整体流程如下：
CUDA IPC 共卡零拷贝权重映射
由于 NCCL 单个通信组的 Rank 数量只能等于 GPU 卡数，而在训推共卡的情况下，Rank 数量是 GPU 卡数的两倍，无法直接建立训推联合通信组。
在这种情况下，Awex 使用 CUDA IPC 将训练进程的显存零拷贝地映射到推理进程，并对所有推理进程建立全局的通信组，然后使用该通信组进行 NCCL send/recv，完成从训练到推理引擎的权重交换：
在实现上我们也做了一些性能优化：
问题：每个 CUDA IPC Handle 的 Open/Close 有较大开销，MOE 等模型单卡可能有数千上万个权重Tensor 需要进行 IPC 序列化；
解决：IPC 序列化前按 shape 和 dtype 合并 Tensor，将数量降至数十个，大幅降低 CUDA IPC 开销。
（注意：CUDA IPC 不支持 CUDA 虚拟内存。未来计划在训练引擎开启虚拟显存时，额外分配物理显存空间，用于权重合并和传输）
（5）RDMA 权重传输
尽管 NCCL 传输模式已经能够大幅提升权重交换的性能，但是 NCCL 模式存在两个主要的限制：
训练和推理侧的 NCCL 版本需要保持兼容，不然 NCCL 传输会出现 hang 的情况，导致训推引擎无法独立更新迭代与升级；
NCCL 静态拓扑对通信域的扩缩容不友好，RL 训练的持续会使得推理输出逐步增长，workload 逐步加大，需要对推理实例进行扩容，NCCL 需要销毁整个通信组并重建；
考虑到这两个原因，我们也开发了基于 RDMA 的传输实现，通过单个配置参数可一键切换传输模式。
RDMA 模式优势：
解除 NCCL 版本绑定，支持训推引擎独立迭代
更灵活的传输计划优化空间
支持推理实例动态扩缩容
性能进一步提升（1T 模型从 20 秒降至 6 秒）
开源计划：RDMA 模式实现将在近期开源，敬请关注。
性能数据
在千卡规模的集群上，Awex 使用 NCCL 传输数据可以在一秒内完成 10B 规模的模型权限交换，二十秒内完成 1T规模的模型权重交换，使用 RDMA 进行传输，1T 模型权重交换耗时可以进一步缩短到六秒钟。
|
权重参数规模 |
权重数据量大小 |
Verl 耗时 |
Awex NCCL 传输耗时 |
Awex RDMA 传输耗时 |
|
10B |
31GB |
3.5S |
0.8S |
0.5S |
|
100B |
191GB |
35S |
9S |
3.2S |
|
1000B |
1000GB (FP8) |
/ |
20S |
6S |
快速开始 Get Started
Awex 是一个纯 Python 库，可以一键安装使用，支持 Python 3.8 及以上版本。
pip install awexMegatron 训练引擎权重发送示例：
from awex import NCCLWeightsWriterfrom awex.engine.mcore import MegatronEngine# inittrain_engine = MegatronEngine(awex_config, hf_config, mcore_model)writer = NCCLWeightsWriter(train_engine)writer.initialize()# write weightswriter.write_weights(step_id=1)
SGLang 推理引擎权重更新示例：
from awex import WeightsReaderfrom awex.engine.sglang import SGLangEngineimport sglang as sglsgl_engine = sgl.Engine(model_path="xxx", tp_size=2, random_seed=42)inference_engine = SGLangEngine(awex_confg, sgl_engine)reader = WeightsReader(inference_engine)reader.initialize()# update weightsreader.update_weights(step_id=1)
未来展望
Awex 是蚂蚁 ASystem 强化学习系统的核心组件之一，而 ASystem 是百灵万亿模型训练的坚实基础。我们将在未来陆续开源 ASystem 的其他核心 RL 组件，进一步完善开源强化学习训练生态。敬请关注！
目前 Awex 开源版已支持 Megatron 和 SGLang 引擎，还有很多待完成的工作，这方面我们期待与社区携手向前，结合社区的力量，逐步扩展对更多主流引擎的支持：
🔧 期待支持的训练引擎：DeepSpeed、FSDP 等
🚀 期待支持的推理引擎：vLLM、TensorRT-LLM 等
🤖 模型适配：内置更多主流模型的开箱即用支持
欢迎每一位对强化学习和大语言模型感兴趣的开发者使用 Awex，并提供宝贵的反馈与建议，一起推动强化学习系统持续创新！
📦 GitHub 仓库：https://github.com/inclusionAI/awex ⭐ 欢迎 Star 和 Fork，也期待您的 PR！
也欢迎持续关注蚂蚁百灵模型的新发布：
🤗 Hugging Face：https://huggingface.co/inclusionAI
🤖 魔搭社区：https://www.modelscope.cn/models/inclusionAI
同时蚂蚁 ASystem 团队持续招聘业界优秀人才，欢迎对强化学习、训练和推理引擎感兴趣的同学加入我们团队，在变化的世界里，持续探索前沿系统的边界。
感兴趣的同学可以通过网站投递：
蚂蚁集团-训推系统研发专家-杭州/北京/上海 https://talent.antgroup.com/off-campus-position?positionId=25052904956438&tid=0b442eeb17633881544247991e1cc0
或邮件发送简历到 ASystem@service.alipay.com
