摘要
苹果公司发布并迅速撤回了一篇关于强化学习框架RLAX的论文。该框架专为在谷歌TPU上进行大规模分布式语言模型训练而设计，实验中使用了亚马逊的AWS Lambda服务进行验证，并以阿里巴巴的QwQ-32B模型为测试对象。结果显示，RLAX在1024个TPU v5p上运行约13小时，成功将QwQ-32B的代码生成准确率（pass@8）提升了12.8%。论文的众多核心作者近期已从苹果离职，加入了Meta、OpenAI、Anthropic及xAI等竞争对手公司。
线索
此次事件揭示了两个关键的投资与风险信号。风险方面，苹果公司正面临严峻的AI人才流失问题。论文的核心作者，包括前AI基础设施主管和多位资深工程师，已集体流向Meta、OpenAI等头部AI企业。这种顶尖人才的集中流失可能削弱苹果在AI基础设施领域的长期竞争力，影响其未来AI产品的迭代速度和技术深度。机会方面，该事件印证了AI基础设施领域的实用主义趋势。苹果在自家拥有M系列芯片的情况下，仍选择谷歌TPU、亚马逊云服务和中国的开源模型进行组合，表明最佳性能和效率是技术选型的首要标准。这利好于在特定领域（如TPU、Serverless计算）提供顶尖服务的云厂商，也证明了中国开源大模型（如Qwen、DeepSeek）在代码推理等细分赛道已具备全球领先的竞争力，成为行业公认的“基准测试工具”，围绕其生态的创业和投资机会值得关注。
正文
一篇题为《RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs》的论文于2025年12月11日在arXiv上公开后，被迅速撤回。该论文的v1版本已被网络存档。
论文揭示了苹果公司开发的一个名为RLAX的可扩展强化学习框架，该框架专为在谷歌TPU上高效运行大规模分布式RL算法而设计。
作者信息
RLAX论文的四名核心作者为Runlong Zhou、Lefan Zhang、Shang-Chen Wu和Kelvin Zou。通讯作者为Kelvin Zou和Cheng Leong。Kelvin Zou曾任苹果首席工程师，现已入职Meta担任AI研究科学家。Cheng Leong是苹果AI基础设施（AI Infra）主管，已在苹果工作超过13年。
此外，作者名单中还包括庞若鸣等六位前苹果员工，论文注明他们已在受雇于苹果期间为该工作做出贡献，但现已离职。这些作者的当前去向包括：Kelvin Zou和Chong Wang加入了Meta，Hanzhi Zhou入职OpenAI，Ye Ke加入了Anthropic，Floris Weers成为一家隐身模式创业公司的创始工程师，Yi Zhang则在xAI研究模型推理。
技术框架
RLAX框架采用参数-服务器架构。主训练器定期将更新的模型权重推送至参数服务器，而一组推理工作器则拉取最新权重以生成新的采样数据。
该框架通过系统级技术实现了训练器、推理工作器和验证器在逻辑上的分离，允许为各组件灵活且独立地分配计算资源。RLAX完全支持抢占式调度，当有更高优先级任务（如在线推理）需要资源时，系统可立即回收TPU资源而不中断训练。
在策略支持上，RLAX通过可编程配置选项，允许用户设置“陈旧度界限”，控制推理工作器拉取新权重的频率和训练器能容忍的最大数据陈旧度，从而在On-policy（在线策略）和Off-policy（离线策略）RL之间进行灵活选择。
验证环境
为了对训练语料库中的代码进行执行验证，RLAX的验证器将标准Python依赖项容器化。该系统调用了亚马逊的AWS Lambda服务来执行大规模代码测试，并将其内部命名为“Oubliette”。这个命名旨在描述其无状态验证环境：代码和测试数据在执行完毕后，整个环境随即销毁。
实验结果
实验环节的技术栈组合包括：
* 算力底座：谷歌的TPU v5p，使用了1024张芯片。
* 验证环境：亚马逊的AWS Lambda服务。
* 基础模型：阿里巴巴团队开源的QwQ-32B。
实验结果显示，RLAX在1024个TPU v5p上，用时12小时48分钟，将QwQ-32B模型的pass@8准确率提升了12.8%，同时在训练期间保持了对任务抢占的鲁棒性。
技术挑战与解决方案
论文披露了一个在TPU训练中遇到的关键技术挑战。在On-policy强化学习中，重要性采样比率理论上应恒等于1.0。但在实际TPU训练中，由于bfloat16浮点数格式的非结合律特性，推理端和训练端的计算结果存在微小差异，导致该比率不等于1.0，进而引发训练崩溃。
问题根源在于，推理时JAX编译器为追求速度会进行大量算子融合，而训练时为计算梯度需保留中间值，导致算子融合策略不同，计算顺序的差异在bfloat16精度下被放大。
苹果的解决方案是在训练器中强制进行“重算”，禁用大部分激活值的保存，使训练端的计算图模仿推理端的计算顺序。此方法以少量速度为代价，消除了该数值问题。
发布时间
2025-12-15 15:46:58
评论 ( 0 )