顯卡在偷懶？阿里大模型創(chuàng)作平臺(tái) MuseAI 極速模型切換技術(shù)提升 AI 創(chuàng)作效率

MuseAI 是一款專為設(shè)計(jì)專業(yè)人士量身定制的先進(jìn) AI 繪圖工具，旨在提供卓越的繪畫(huà)體驗(yàn)，并為設(shè)計(jì)團(tuán)隊(duì)打造一個(gè)既穩(wěn)定又易于管理的創(chuàng)作平臺(tái)。

原標(biāo)題：顯卡在偷懶？阿里大模型創(chuàng)作平臺(tái) MuseAI 極速模型切換技術(shù)提升 AI 創(chuàng)作效率
文章來(lái)源：AI前線
內(nèi)容字?jǐn)?shù)：32246字

阿里MuseAI平臺(tái)AIGC創(chuàng)作工作臺(tái)性能優(yōu)化實(shí)踐

本文介紹了阿里巴巴愛(ài)橙科技研發(fā)的面向阿里內(nèi)部的AIGC創(chuàng)作工作臺(tái)MuseAI（以及其對(duì)公眾開(kāi)放的魔搭社區(qū)AIGC專區(qū)）的性能優(yōu)化實(shí)踐。文章重點(diǎn)分析了平臺(tái)因頻繁切換Diffusion Pipeline導(dǎo)致的用戶體驗(yàn)和資源浪費(fèi)問(wèn)題，并從網(wǎng)絡(luò)傳輸、內(nèi)存管理、Host-to-Device、模型量化等多個(gè)方面進(jìn)行了優(yōu)化。

1. 性能問(wèn)題與挑戰(zhàn)

MuseAI集成了大量的模型（基礎(chǔ)模型、LoRA微調(diào)模型、ControlNet控制模型及輔助性模型），導(dǎo)致頻繁的模型切換。這帶來(lái)了以下性能問(wèn)題：

1. 端到端生成時(shí)間過(guò)長(zhǎng)：模型下載、加載和切換時(shí)間占據(jù)了大部分時(shí)間。
2. 緩存未命中率高：大量模型難以全部緩存到磁盤或內(nèi)存。
3. GPU資源浪費(fèi)：長(zhǎng)時(shí)間的模型加載導(dǎo)致GPU空閑。

2. 優(yōu)化策略與方法

為了解決上述問(wèn)題，MuseAI團(tuán)隊(duì)從以下幾個(gè)方面進(jìn)行了優(yōu)化：

2.1 模型加載優(yōu)化

1. 存儲(chǔ)介質(zhì)選擇：根據(jù)業(yè)務(wù)特性，公司內(nèi)部使用高性能分布式存儲(chǔ)“盤古”+fsfuse，公有云使用NAS。fsfuse通過(guò)緩存機(jī)制和Direct I/O技術(shù)提升讀取效率。
2. NAS最佳實(shí)踐：調(diào)整Linux內(nèi)核參數(shù)（nconnect），增加NAS與客戶端之間的連接數(shù)，提高帶寬利用率。采用多線程并發(fā)讀取模型文件。
3. 盤古+fsfuse最佳實(shí)踐：統(tǒng)一掛載模型父目錄，采用順序讀取模式，充分利用fsfuse的預(yù)讀取和緩存機(jī)制，使用Direct I/O技術(shù)。

2.2 模型切換優(yōu)化

1. 執(zhí)行順序優(yōu)化：先加載state dict，再將模型遷移到GPU，避免不必要的CPU-GPU數(shù)據(jù)拷貝。
2. H2D傳輸性能優(yōu)化：使用內(nèi)存池管理pinned memory，減少內(nèi)存分配和拷貝次數(shù)，并使用多線程并發(fā)傳輸。
3. Skip init技術(shù)：跳過(guò)nn.Module構(gòu)造時(shí)冗余的初始化過(guò)程，減少加載時(shí)間。

2.3 內(nèi)存管理與復(fù)用

1. Pinned memory內(nèi)存池：避免重復(fù)malloc內(nèi)存，減少內(nèi)存分配和釋放操作。
2. 直接讀取到pinned memory：消除一次內(nèi)存拷貝，提高效率。
3. 兩級(jí)內(nèi)存池設(shè)計(jì)：根據(jù)模型大小分配內(nèi)存，實(shí)現(xiàn)內(nèi)存復(fù)用。
4. 在pinned memory上構(gòu)造state_dict：直接在預(yù)分配的pinned memory中構(gòu)造tensor，減少內(nèi)存拷貝。

2.4 模型量化

將模型轉(zhuǎn)換為FP8精度，減少模型大小和顯存占用，利用新一代GPU架構(gòu)的FP8計(jì)算能力提升推理速度。

2.5 T5化部署 (技術(shù)儲(chǔ)備)

將T5模型部署為Embedding Server，通過(guò)RPC調(diào)用獲取文本嵌入向量，減少模型加載和卸載操作。該方案因工程復(fù)雜度和穩(wěn)定性風(fēng)險(xiǎn)，最終未應(yīng)用到生產(chǎn)環(huán)境。

3. 實(shí)驗(yàn)結(jié)果與分析

實(shí)驗(yàn)結(jié)果顯示，優(yōu)化后的MuseAI在冷啟動(dòng)和模型切換性能方面均優(yōu)于Diffusers和WebUI-forge，尤其是在NAS環(huán)境下，性能提升顯著。

4. 結(jié)論

通過(guò)一系列優(yōu)化措施，MuseAI平臺(tái)顯著提升了模型加載和切換速度，改善了用戶體驗(yàn)，并有效降低了資源浪費(fèi)。這項(xiàng)工作為AIGC平臺(tái)的性能優(yōu)化提供了寶貴的經(jīng)驗(yàn)和參考。

聯(lián)系作者

文章來(lái)源：AI前線
作者微信：
作者簡(jiǎn)介：面向AI愛(ài)好者、開(kāi)發(fā)者和科學(xué)家，提供大模型最新資訊、AI技術(shù)分享干貨、一線業(yè)界實(shí)踐案例，助你全面擁抱AIGC。