텐서플로우에서 JAX로의 딥러닝 모델 이전을 위한 다중 에이전트 AI 시스템

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

인공지능 기반 제품 및 그 기반 모델의 급속한 발전은 딥러닝 프레임워크의 지속적인 혁신을 야기하고 있습니다. 구글은 수십 개의 제품에 걸쳐 머신러닝 기술을 선도적으로 활용해 왔습니다. 다양한 머신러닝 프레임워크 및 버전에서 관리되는 다수의 모델 소스 코드를 유지하는 것은 상당한 과제입니다. 지금까지 유지 보수 및 이 작업은 주로 인간 전문가에 의해 수동으로 수행되었습니다. 본 연구에서는 텐서플로우 기반 딥러닝 모델을 JAX 기반 모델로 자동 이전을 지원하기 위해 개발된 AI 기반 다중 에이전트 시스템을 설명합니다. 본 연구는 세 가지 주요 기여를 합니다. 첫째, 정적 분석과 AI 지침을 결합한 AI 플래너가 매우 복잡한 코드 구성 요소에 대한 이사 계획을 생성하고, AI가 생성한 예제 기반 플레이북을 사용하여 오케스트레이터와 코더가 안정적으로 이 계획을 따르도록 하는 방법을 보여줍니다. 둘째, 테스트가 없는 코드나 엄격한 스타일 및 종속성 요구 사항을 충족해야 하는 코드를 평가할 때 개발 속도를 가속화하는 품질 지표 및 AI 기반 판단 기준을 정의합니다. 셋째, 본 시스템이 상용 실제 사용 사례를 대상으로 하는 대규모 클라우드 환경에서 코드 이전을 가속화하는 방법을 보여줍니다. 본 연구의 접근 방식은 딥러닝 모델 이전에 소요되는 시간을 6.4배에서 8배까지 단축하며, 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.