arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2603.27296cs.SEcs.AI
收藏

面向深度学习模型从TensorFlow到JAX迁移的多智能体AI系统

A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX

Google(谷歌)

Stoyan Nikolov, Bernhard Konrad, Moritz Gronbach, Niket Kumar, Ann Yan, Varun Singh, Yaning Liang, Parthasarathy Ranganathan

AI总结:

本文提出一种多智能体AI系统,用于自动迁移TensorFlow深度学习模型至JAX,通过AI规划器、协调器和程序员的协作,显著提升迁移效率,减少6.4至8倍的时间消耗。

摘要:

人工智能驱动的产品快速发展的背景下,深度学习框架持续创新。谷歌在数十种产品中推动机器学习应用。维护不同ML框架和版本的模型源代码是重大挑战。目前,迁移和维护工作主要由人工专家完成。本文描述了我们构建的基于AI的多智能体系统,以支持TensorFlow模型到JAX模型的自动迁移。我们有三个主要贡献:首先,展示如何利用静态分析与AI指令混合的AI规划器,为非常复杂的代码组件生成迁移计划,该计划通过协调器和程序员的组合可靠执行,使用AI生成的基于示例的 playbook。其次,定义质量度量和基于AI的评判者,加速开发过程,当评估代码没有测试且必须遵守严格的风格和依赖要求时。第三,展示系统如何在大型超大规模环境中的商业实际用例中加速代码迁移。我们的方法显著减少了深度学习模型迁移时间(6.4x-8x加速),并创建了一个良性循环,其中有效AI支持其自身开发流程。我们预计,本文描述的技术和方法可以推广到其他框架迁移和一般代码转换任务。

英文摘要:

The rapid development of AI-based products and their underlying models has led to constant innovation in deep learning frameworks. Google has been pioneering machine learning usage across dozens of products. Maintaining the multitude of model source codes in different ML frameworks and versions is a significant challenge. So far the maintenance and migration work was done largely manually by human experts. We describe an AI-based multi-agent system that we built to support automatic migration of TensorFlow-based deep learning models into JAX-based ones. We make three main contributions: First, we show how an AI planner that uses a mix of static analysis with AI instructions can create migration plans for very complex code components that are reliably followed by the combination of an orchestrator and coders, using AI-generated example-based playbooks. Second, we define quality metrics and AI-based judges that accelerate development when the code to evaluate has no tests and has to adhere to strict style and dependency requirements. Third, we demonstrate how the system accelerates code migrations in a large hyperscaler environment on commercial real-world use-cases. Our approach dramatically reduces the time (6.4x-8x speedup) for deep learning model migrations and creates a virtuous circle where effectively AI supports its own development workflow. We expect that the techniques and approaches described here can be generalized for other framework migrations and general code transformation tasks.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑