用 AI 提高 AI 的效率，它們就能變得像人類大腦一樣高效？

我們的大腦只用了 20 瓦的能量就能完成復(fù)雜思考，而現(xiàn)代 AI 系統(tǒng)卻需要成排的高功率 GPU 和驚人的電力消耗。這種差距如何縮小？

日本 AI 初創(chuàng)公司 Sakana AI 團隊提出了一個大膽的愿景：利用 AI 本身來優(yōu)化 AI。他們開發(fā)的「AI CUDA 工程師」是這一理念的具體實踐。

「AI CUDA 工程師」是第一個用于全自動 CUDA 內(nèi)核發(fā)現(xiàn)和優(yōu)化的綜合智能體框架。這種方法不僅開創(chuàng)性地將進化計算與大型語言模型相結(jié)合，更展示了 AI 自我優(yōu)化的巨大潛力。 

CUDA 是一個 low-level 軟件層，可直接訪問 NVIDIA GPU 用于并行計算的硬件指令集。CUDA 內(nèi)核是用 CUDA 語言編寫的在 GPU 上運行的函數(shù)。通過直接在 CUDA 內(nèi)核層編寫指令，工程師可以為 AI 算法實現(xiàn)更高的性能。然而，使用 CUDA 需要相當(dāng)多的 GPU 知識，實際上，大多數(shù)機器學(xué)習(xí)算法都是在 PyTorch 或 JAX 等更高級別的抽象層中編寫的。

「AI CUDA 工程師」生成的高度優(yōu)化 CUDA 內(nèi)核示例。詳情請參見：https://pub.sakana.ai/ai-cuda-engineer

「AI CUDA 工程師」是一個利用前沿 LLM 的智能體框架，旨在自動將標(biāo)準(zhǔn) PyTorch 代碼轉(zhuǎn)換為高度優(yōu)化的 CUDA 內(nèi)核。通過使用進化優(yōu)化，并利用進化計算中的概念，如「交叉」操作和「創(chuàng)新檔案」來發(fā)現(xiàn)有前途的「踏腳石」內(nèi)核，該團隊提出的框架不僅能夠自動化將 PyTorch 模塊轉(zhuǎn)換為 CUDA 內(nèi)核的過程，而且高度優(yōu)化的 CUDA 內(nèi)核通常能夠?qū)崿F(xiàn)顯著更快的運行時間加速。

該團隊相信這項技術(shù)能夠?qū)崿F(xiàn)加速，從而加快 LLM 或其他生成式 AI 模型等基礎(chǔ)模型的訓(xùn)練和運行（推理），最終使 AI 模型在 NVIDIA 硬件上運行得更快。

「AI CUDA 工程師」能夠生成比常見 PyTorch 操作加速 10-100 倍的 CUDA 內(nèi)核。它還能生成比生產(chǎn)環(huán)境中常用的現(xiàn)有 CUDA 內(nèi)核快得多的高度優(yōu)化的 CUDA 內(nèi)核（加速高達 5 倍）。

AI CUDA 工程師智能體框架的高級概述。

它的運行流程如下：   

第 1 和第 2 階段（轉(zhuǎn)換和翻譯）：「AI CUDA 工程師」首先將 PyTorch 代碼翻譯成可運行的 CUDA 內(nèi)核。即使不明確針對這些目標(biāo)，也能觀察到初始運行時的改進。

第 3 階段（進化優(yōu)化）：受生物進化的啟發(fā)，該框架利用進化優(yōu)化（「適者生存」）來確保只生成最佳的 CUDA 內(nèi)核。此外，該團隊引入了一種新穎的內(nèi)核交叉提示策略，以互補的方式組合多個優(yōu)化的內(nèi)核。

第 4 階段（創(chuàng)新檔案）：正如文化進化如何利用我們祖先幾千年文明的知識來塑造我們的人類智慧一樣，「AI CUDA 工程師」也利用從過去的創(chuàng)新和發(fā)現(xiàn)中學(xué)到的東西（第 4 階段），從已知高性能 CUDA 內(nèi)核的家族中建立創(chuàng)新檔案，利用以前的踏腳石來實現(xiàn)進一步的翻譯和性能提升。

該項目發(fā)布后，不少研究者給予了很高的評價，比如英偉達高級 AI 研究科學(xué)家 Jim Fan 稱這是他最近見過的最酷的自動編程智能體，認(rèn)為用當(dāng)前的計算資源來提高未來計算效率，這是最具回報的投資策略 ，「AutoML is so back!」

不過，也有人發(fā)現(xiàn)了問題。比如 NVIDIA 杰出工程師 Bing Xu 指出「AI CUDA 工程師」的技術(shù)報告中存在幾個誤導(dǎo)性部分：

1. Torch C++ 代碼并不是 CUDA 內(nèi)核，它在底層是調(diào)用 CUDNN 庫。
2. 報告重點強調(diào)的 Conv3D GroupNorm 示例中，卷積代碼根本沒有被生成。如果數(shù)值計算結(jié)果不正確，聲稱的速度提升就沒有意義。
3. 報告中聲稱 WMMA 可以比 PyTorch（CUBLAS）更快，這絕對是錯誤的。很可能是基準(zhǔn)測試出現(xiàn)了問題。

看來，這個「AI CUDA 工程師」的效果還有待驗證。

「AI CUDA 工程師」發(fā)現(xiàn)的內(nèi)核運行時加速

「AI CUDA 工程師」穩(wěn)健地發(fā)現(xiàn)了用于常見機器學(xué)習(xí)操作的 CUDA 內(nèi)核，其速度比 PyTorch 中的原生和編譯內(nèi)核快 10-100 倍。該團隊的方法還能將整個機器學(xué)習(xí)架構(gòu)轉(zhuǎn)換為優(yōu)化的 CUDA 內(nèi)核。下面是幾個完全自主發(fā)現(xiàn)的顯著加速：

這些優(yōu)化 CUDA 內(nèi)核的更多詳情可在交互式網(wǎng)站的排行榜上查看：https://pub.sakana.ai/ai-cuda-engineer/leaderboard

該團隊的方法為矩陣乘法、常見的深度學(xué)習(xí)等操作找到了更高效的 CUDA 內(nèi)核，截至撰寫本文時，它發(fā)現(xiàn)的 CUDA 內(nèi)核在 KernelBench 上實現(xiàn)了 SOTA 的性能。

技術(shù)報告和數(shù)據(jù)集

Sakana AI 發(fā)布了 AI CUDA 的技術(shù)報告，整個技術(shù)報告有 80 多頁。

技術(shù)報告：https://pub.sakana.ai/static/paper.pdf

報告內(nèi)容如下：

• 介紹了一個端到端的智能體工作流，能夠?qū)?PyTorch 代碼翻譯成可工作的 CUDA 內(nèi)核，優(yōu)化 CUDA 運行時性能，并自動融合多個內(nèi)核。
• 構(gòu)建了各種技術(shù)來增強 pipeline 的一致性和性能，包括 LLM 集成、迭代分析反饋循環(huán)、本地內(nèi)核代碼編輯和交叉內(nèi)核優(yōu)化。
• 報告顯示，「AI CUDA 工程師」穩(wěn)健地翻譯了被考慮在內(nèi)的 250 個 torch 操作中的 230 多個，并且對大多數(shù)內(nèi)核實現(xiàn)了強大的運行時性能改進。此外，該團隊的方法能夠有效地融合各種內(nèi)核操作，并且可以超越幾種現(xiàn)有的加速操作。
• 發(fā)布了一個包含超過 17,000 個經(jīng)驗證內(nèi)核的數(shù)據(jù)集，這些內(nèi)核涵蓋了廣泛的 PyTorch 操作。

報告還給出了一些發(fā)現(xiàn)的 CUDA 內(nèi)核的顯著例子，這些內(nèi)核在 AI 模型的關(guān)鍵計算操作上實現(xiàn)了顯著的加速。

AI CUDA Engineer 發(fā)現(xiàn)的優(yōu)質(zhì)內(nèi)核

利用新的 LLM 驅(qū)動的進化內(nèi)核優(yōu)化程序，研究團隊穩(wěn)健地獲得了各種考慮因素的加速。更具體地說，在考慮的 229 個任務(wù)中，81% 的性能優(yōu)于 PyTorch 原生運行時。此外，在所有已發(fā)現(xiàn)的 CUDA 內(nèi)核中，有 20% 的內(nèi)核速度至少是 PyTorch 實現(xiàn)的兩倍。

「AI CUDA 工程師」穩(wěn)健地發(fā)現(xiàn)了優(yōu)于 PyTorch 實現(xiàn)的 CUDA 內(nèi)核。

下面展示了一部分內(nèi)核。它們突顯了「AI CUDA 工程師」可以成功部署的不同操作的多樣性。這包括 normalization 方法、損失函數(shù)、特殊矩陣乘法，甚至整個神經(jīng)網(wǎng)絡(luò)架構(gòu)：

「AI CUDA 工程師」生成的高度優(yōu)化 CUDA 內(nèi)核示例。詳情請參見：https://pub.sakana.ai/ai-cuda-engineer

?

「AI CUDA 工程師檔案」

?

17,000 多個經(jīng)驗證的 CUDA 內(nèi)核數(shù)據(jù)集

「AI CUDA 工程師檔案」的文本嵌入可視化顯示，發(fā)現(xiàn)的內(nèi)核可以按任務(wù)（例如 MatMul、Pooling、Convolution）和實現(xiàn)策略（展開、融合、矢量化）分組。該檔案可公開訪問，可用于 LLM 的下游微調(diào)。

隨論文一起發(fā)布的還有「AI CUDA 工程師檔案」，這是一個由「AI CUDA 工程師」生成的超過 30,000 個 CUDA 內(nèi)核組成的數(shù)據(jù)集。它在 CC-By-4.0 許可下發(fā)布，可通過 HuggingFace 訪問：???https://huggingface.co/datasets/SakanaAI/AI-CUDA-Engineer-Archive??。

該數(shù)據(jù)集包括 torch 參考實現(xiàn)、torch、NCU 和 Clang-tidy 分析數(shù)據(jù)、每個任務(wù)的多個內(nèi)核、錯誤消息以及針對 torch 本地和編譯運行時的加速分?jǐn)?shù)。

「AI CUDA 工程師檔案」的摘要統(tǒng)計數(shù)據(jù)，包含超過 30,000 個內(nèi)核和超過 17,000 個正確驗證的實現(xiàn)。大約 50% 的所有內(nèi)核都優(yōu)于 torch 原生運行時。

研究團隊設(shè)想此數(shù)據(jù)集可以使開源模型的后訓(xùn)練執(zhí)行更好的 CUDA 啟用模塊。這包括離線強化學(xué)習(xí)、偏好優(yōu)化和標(biāo)準(zhǔn)監(jiān)督微調(diào)。

在「AI CUDA 工程師檔案」中探索 17,000 多個內(nèi)核

該團隊還發(fā)布了一個交互式網(wǎng)站，用于交互式檢查超過 17,000 個經(jīng)驗證的內(nèi)核及其配置文件，包括 torch、NCU 和 Clang-Tidy 數(shù)據(jù)：??https://pub.sakana.ai/ai-cuda-engineer??。

該網(wǎng)站允許探索 230 個任務(wù)中的各種高性能內(nèi)核。它帶有一個自定義排行榜，可用于檢查跨實驗和 LLM 的相關(guān)內(nèi)核。

「AI CUDA 工程師」發(fā)現(xiàn)的內(nèi)核排行榜：??https://pub.sakana.ai/ai-cuda-engineer/leaderboard??

此外，你還可以可視化內(nèi)核，檢索相關(guān)內(nèi)核，下載代碼以驗證實現(xiàn)和加速，以及查看獲得的分析數(shù)據(jù)。最后，可以深入了解優(yōu)化實驗。

優(yōu)化的實例 Normalization 內(nèi)核的詳細視圖，包括分析數(shù)據(jù)、評估腳本的下載、相關(guān)內(nèi)核和發(fā)現(xiàn)實驗詳細信息。

局限性和有趣的意外發(fā)現(xiàn) 

雖然將進化優(yōu)化與 LLM 結(jié)合非常強大，但這種組合有時也會找到意想不到的方法繞過驗證系統(tǒng)。比如，Twitter 用戶 @main_horse 幫助測試 CUDA 內(nèi)核時就發(fā)現(xiàn)了一個有趣的情況：「AI CUDA 工程師」竟然找到了一種「投機取巧」的方法。這個 AI 系統(tǒng)在評估代碼中發(fā)現(xiàn)了一個內(nèi)存漏洞，在一小部分情況下成功避開了正確性檢查：

發(fā)現(xiàn)這個問題后，該團隊立即加強了評估框架的安全性，堵住了這類漏洞，并更新了實驗結(jié)果。

有趣的是，這并非該團隊第一次遇到類似情況。在之前的「AI 科學(xué)家」項目中，AI 也曾找到方法修改并運行自己的評估腳本。它沒有努力讓代碼運行得更快，而是直接嘗試修改代碼來延長超時時間！研究文獻中已經(jīng)記錄了這種現(xiàn)象：AI 系統(tǒng)常常能找到創(chuàng)造性的解決方案，這些方案往往出人意料，令開發(fā)者感到驚訝。

此外，該團隊還發(fā)現(xiàn)前沿 LLM 在使用 TensorCore WMMA 方面存在明顯局限。雖然 LLM 能夠生成基礎(chǔ)的 CUDA 代碼，但在實現(xiàn)現(xiàn)代 GPU 架構(gòu)提供的特殊矩陣乘法加速功能時卻常常力不從心。這可能表明 LLM 的訓(xùn)練數(shù)據(jù)中缺乏這方面的信息，或者模型對這類高級硬件優(yōu)化的理解還不夠深入。

隨著前沿 LLM（特別是那些具有強大代碼推理能力的模型）變得越來越強大，該團隊預(yù)計像他們這樣的代碼優(yōu)化系統(tǒng)將繼續(xù)面臨這些挑戰(zhàn)。他們設(shè)想未來的發(fā)展方向是：人類工程師與代碼優(yōu)化 AI 系統(tǒng)協(xié)同工作，共同創(chuàng)造最佳且最可靠的結(jié)果。

「AI CUDA 工程師」的未來影響

AI 革命才剛剛開始，現(xiàn)在只是處于轉(zhuǎn)型周期的最初階段。該團隊認(rèn)為，今天的 LLM 是這一代的「大型主機計算機」?，F(xiàn)在仍處于 AI 的早期階段，由于市場競爭和全球創(chuàng)新（尤其是那些在資源限制下進行創(chuàng)新的國家 / 地區(qū)），這項技術(shù)的效率將提高一百萬倍，這是不可避免的。

目前，AI 系統(tǒng)消耗大量資源，如果技術(shù)繼續(xù)擴展而不考慮效率和能源消耗，結(jié)果將不可持續(xù)。沒有根本原因說明為什么 AI 系統(tǒng)不能像人類智能一樣高效（甚至更高效）。該團隊相信，實現(xiàn)這種更高效率的最佳途徑是利用 AI 使 AI 更加高效。

這是 Sakana AI 正在追求的方向，這個項目是使 AI 快一百萬倍的重要一步。就像早期笨重的大型主機計算機向現(xiàn)代計算發(fā)展一樣，人類使用 AI 的方式在幾年內(nèi)將發(fā)生重大變化。

本文轉(zhuǎn)自 機器之心  ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/uWoySPIqrOWEYAm64w18Kw??