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

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

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

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

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

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

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

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

為了解決上述問題，MuseAI團隊從以下幾個方面進行了優(yōu)化：

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

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

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

1. Pinned memory內(nèi)存池：避免重復(fù)malloc內(nèi)存，減少內(nèi)存分配和釋放操作。
2. 直接讀取到pinned memory：消除一次內(nèi)存拷貝，提高效率。
3. 兩級內(nèi)存池設(shè)計：根據(jù)模型大小分配內(nèi)存，實現(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計算能力提升推理速度。

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

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

3. 實驗結(jié)果與分析

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

4. 結(jié)論

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

聯(lián)系作者

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