IT之家 11 月 20 日音讯，蚂蚁集团今天宣告开源万亿参数强化学习高功能权重交流结构 Awex。
据介绍，Awex 是为极致功能打造的练习推理引擎权重同步结构，处理 RL 流程中练习权重参数同步到推理模型的中心难题，可在秒级完结 TB 级大规划参数交流，显着下降 RL 模型练习推迟，主要特征如下：
一致模型适配层：主动处理训推引擎并行战略与引擎间的 Tensor 格局 / 布局差异，兼容多种模型架构；
零冗余 Resharding 传输与原地更新：仅传输必要分片（Shard），推理侧原地更新显存，防止重分配与复制开支；
多形式传输支撑：支撑 NCCL、RDMA、同享内存多种传输形式，充沛的发挥 NVLink / NVSwitch / RDMA 带宽并削减长尾推迟；
异构布置兼容：适配共卡 / 分卡形式，支撑同步和异步 RL 算法练习场景，一起 RDMA 传输形式支撑推理实例动态扩缩容；
灵敏可插拔架构：支撑对不同模型定制化权重 Sharing 和 Layout 行为，一起支撑新的练习和推理引擎接入。
WeightWriter：在每个练习进程内运转，担任当时练习进程的权重 Shard 的元数据搜集上报、权重转化、权重发送 Resharding 方案构建、权重发送等功能；
WeightReader：在每个推理实例的操控进程上运转，其会在推理实例办理的每张 GPU 上面发动一个 WorkerWeightsReader，与练习进程的 WeightWriter 相对应，担任每个推理进程的权重 Shard 的元数据搜集上报、权重转化、权重承受 Resharding 方案构建、权重承受等功能；
MetaServer：Job 等级大局 Server，用于训推引擎的服务发现和权重元数据交流，以及共卡情况下的事情告诉等功能；
训推权重一致转化：担任将不同并行战略和 Tensor 布局的练习引擎和推理引擎的权重转化成一致的格局，用于后续的权重元数据核算和权重传输；
大局权重元数据核算与交流：将训推权重转化成一致的格局后，搜集每个 Worker 的一切权重 Shard 元数据，并上签到 Meta Server，用于接下来的的权重传输方案构建；
P2P 权重传输执行方案：练习和推理引擎拿到大局一切 Worker 的练习和推理权重 Shard 元数据，然后别离各自构建对等的发送和承受确定性传输方案；
NCCL 权重传输：运用 NCCL 的 send / recv API 根据构建的传输方案进行对等的权重发送与接纳；
RDMA 权重传输：运用 NUMA 亲和，面向大局负载均衡传输方案的 RDMA 通讯来进行权重的更新；
一起 Awex 也支撑对权重进行 Tensor 等级的校验，将经过文件体系形式加载的权重跟经过传输形式加载的权重进行 Tensor 等级的细粒度比对，逐一判别差异，确保传输形式的正确性。
在千卡规划的集群上，Awex 运用 NCCL 传输数据能够在一秒内完结 10B 规划的模型权限交流，二十秒内完结 1T 规划的模型权重交流，运用 RDMA 进行传输，1T 模型权重交流耗时能愈加进一步缩短到六秒钟。
未来展望，Awex 是蚂蚁 ASystem 强化学习体系的中心组件之一，而 ASystem 是百灵万亿模型练习的坚实基础。咱们将在未来连续开源 ASystem 的其他中心 RL 组件，加强完善开源强化学习练习生态。现在 Awex 开源版已支撑 Megatron 和 SGLang 引擎。
