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

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

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

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

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

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

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

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

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

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

它的運(yùn)行流程如下：   

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

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

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

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

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

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

看來，這個(gè)「AI CUDA 工程師」的效果還有待驗(yàn)證。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

?

「AI CUDA 工程師檔案」

?

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

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

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

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

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

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

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

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

該網(wǎng)站允許探索 230 個(gè)任務(wù)中的各種高性能內(nèi)核。它帶有一個(gè)自定義排行榜，可用于檢查跨實(shí)驗(yàn)和 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)核，下載代碼以驗(yàn)證實(shí)現(xiàn)和加速，以及查看獲得的分析數(shù)據(jù)。最后，可以深入了解優(yōu)化實(shí)驗(yàn)。

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

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

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

發(fā)現(xiàn)這個(gè)問題后，該團(tuán)隊(duì)立即加強(qiáng)了評(píng)估框架的安全性，堵住了這類漏洞，并更新了實(shí)驗(yàn)結(jié)果。

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

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

隨著前沿 LLM（特別是那些具有強(qiáng)大代碼推理能力的模型）變得越來越強(qiáng)大，該團(tuán)隊(duì)預(yù)計(jì)像他們這樣的代碼優(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)型周期的最初階段。該團(tuán)隊(duì)認(rèn)為，今天的 LLM 是這一代的「大型主機(jī)計(jì)算機(jī)」。現(xiàn)在仍處于 AI 的早期階段，由于市場競爭和全球創(chuàng)新（尤其是那些在資源限制下進(jìn)行創(chuàng)新的國家 / 地區(qū)），這項(xiàng)技術(shù)的效率將提高一百萬倍，這是不可避免的。

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

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

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

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