AI多代理系统自动将深度学习模型从TensorFlow迁移到JAX,提速6.4至8倍。
A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX
- 用AI规划+示例剧本,自动制定复杂代码迁移方案。
- 在无测试情况下仍能保证代码质量,满足严格风格和依赖要求。
- 已在超大规模商用场景验证,适合框架迁移与代码重构团队使用。
人工智能产品及其底层模型的快速发展推动了深度学习框架的持续创新。谷歌在数十个产品中广泛使用机器学习,但维护不同框架和版本的大量模型源码是一项重大挑战。以往的维护与迁移工作主要依赖人工专家。本文描述了一个基于AI的多代理系统,用于自动将基于TensorFlow的深度学习模型迁移至JAX。我们提出三项主要贡献:第一,展示了一种结合静态分析与AI指令的AI规划器,可生成可靠迁移计划,并由协调器与编码代理执行,基于AI生成的示例剧本;第二,定义了质量评估指标与AI判别器,在缺乏测试的情况下加速开发,确保代码符合严格的风格和依赖规范;第三,证明该系统在大型云服务商环境中对真实商业用例的迁移具有显著加速效果。相比传统方法,迁移速度提升6.4至8倍,并形成AI支持自身开发的良性循环。我们预期这些技术可推广至其他框架迁移与通用代码转换任务。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。