AI-Hypercomputer/accelerator-agents

来源：https://github.com/AI-Hypercomputer/accelerator-agents

加速器智能体

加速器智能体是一套基于人工智能的工具集合，旨在加速在Google Cloud TPU上的机器学习开发。本仓库托管的智能体可协助完成代码迁移、内核优化和性能调优，帮助开发者更高效地释放TPU的全部潜力。

免责声明： 这不是Google官方支持的产品。本项目不符合Google开源软件漏洞奖励计划（https://bughunters.google.com/open-source-security）的资格。

概述

随着机器学习模型复杂度的增长，为TPU等特定硬件加速器优化模型变得日益困难。本项目旨在提供一套由Gemini驱动的“智能体”——专业AI工具，以自动化方式协助处理这些复杂任务。

项目包含两个主要智能体：

1. MaxCode

MaxCode智能体可简化现有PyTorch模型和代码库到JAX的转换过程。其设计目的是帮助用户将工作负载迁移至TPU高效运行，并充分利用MaxText（https://github.com/AI-Hypercomputer/maxtext）等高性能框架。

注：MaxCode正处于积极开发阶段，我们持续优化迁移质量并扩展模型支持范围。

功能特点：

• 自动化转换： 将PyTorch中的函数代码块和模型层转换为JAX格式。
• MaxText集成： 生成与MaxText框架兼容的JAX代码，可在TPU上直接进行训练和推理。
• 人机协作： 旨在生成初步实现方案，供开发者审阅与优化。

2. MaxKernel

MaxKernel智能体是专为TPU高性能内核开发设计的工具。它协助工程师编写、优化和调试自定义内核，重点关注Pallas（JAX的内核语言）。

功能特点：

• 内核编写： 从零开始或基于JAX参考实现草拟Pallas内核。
• CUDA到Pallas转换： 协助将自定义CUDA/GPU内核移植到TPU上实现最优运行。
• 优化与分析： 提供性能剖析洞察和优化建议，提升内核性能（MFU）。
• 测试框架生成： 自动生成用于正确性测试和编译检查的样板代码。

快速开始

前置条件

• Google Cloud虚拟机。运行MaxCode仅需CPU虚拟机，而MaxKernel建议使用TPU虚拟机。
• Python 3.11+
• Gemini API访问权限（用于智能体推理功能）。

安装步骤

克隆仓库：bash git clone https://github.com/AI-Hypercomputer/accelerator-agents.git cd accelerator-agents

（注：每个智能体的具体安装说明可在其对应子目录中找到。）

贡献指南

欢迎参与贡献！请查阅CONTRIBUTING.md了解提交拉取请求、报告问题及参与项目的详细信息。

许可证

本项目采用Apache License 2.0许可证。完整许可文本请参见LICENSE。