近期免費(fèi)上線的視頻生成平臺(tái)Video Ocean，支持任意角色、任意風(fēng)格，可以文生視頻、圖生視頻、角色生視頻，引起廣泛關(guān)注與肯定，登上Product Hunt全球產(chǎn)品熱度榜單第三。

體驗(yàn)地址：https://video.luchentech.com/zh-CN

Video Ocean如何以極低成本快速完成迭代？開(kāi)源解決現(xiàn)已方案發(fā)布。

二次開(kāi)發(fā)分享至開(kāi)源社區(qū)，還可領(lǐng)取500元GPU算力代金券。

開(kāi)源地址：https://github.com/hpcaitech/Open-Sora

Colossal-AI

在Video Ocean背后，離不開(kāi)AI大模型訓(xùn)練推理系統(tǒng)Colossal-AI的基礎(chǔ)支持，其在GitHub全球AI訓(xùn)推系統(tǒng)開(kāi)源領(lǐng)域指標(biāo)位列世界第一，已獲近4萬(wàn)Stars。

開(kāi)源地址：https://github.com/hpcaitech/ColossalAI

它基于PyTorch，可通過(guò)高效多維并行、異構(gòu)內(nèi)存等，降低AI大模型訓(xùn)練/微調(diào)/推理的開(kāi)發(fā)與應(yīng)用成本，已與多家世界/中國(guó)500強(qiáng)企業(yè)聯(lián)合開(kāi)發(fā)和優(yōu)化AI大模型。

針對(duì)類Sora視頻大模型開(kāi)發(fā)，Colossal-AI做了多方面的優(yōu)化，模型算力利用率MFU相對(duì)現(xiàn)有其他開(kāi)源方案總體最高提升可達(dá)2.61倍，顯著降低成本。

異步Checkpoint

在使用大規(guī)模集群訓(xùn)練時(shí)，因集群規(guī)模擴(kuò)大，故障率會(huì)迅速上升，導(dǎo)致訓(xùn)練極易中斷。在這種情況下快速保存Checkpoint不僅可以加速整體訓(xùn)練效率，也有利于故障容錯(cuò)，快速恢復(fù)訓(xùn)練。

為此，Colossal-AI推出了異步Checkpoint功能。針對(duì)10B量級(jí)的視頻生成模型，可將DiT模型、EMA模型、優(yōu)化器的保存時(shí)間從300s+，降至10s以內(nèi)，節(jié)省了高達(dá)97%的保存時(shí)間。

Checkpoint保存主要分為GPU->CPU（D2H）和硬盤(pán)寫(xiě)兩個(gè)步驟，通過(guò)將這兩個(gè)步驟通過(guò)流水線的形式執(zhí)行，極大的提高了保存的效率。同時(shí)此步驟通過(guò)多線程（C++）的形式在后臺(tái)完成，不會(huì)阻塞訓(xùn)練主進(jìn)程。GPU->CPU通過(guò)單獨(dú)的CUDA Stream完成，也不會(huì)阻塞主計(jì)算Stream。

流水線保存Checkpoint

除此之外，通過(guò)使用safetensors格式，因其安全、零拷貝的特性，也提升了讀取的性能。

Zero內(nèi)存/通信優(yōu)化

常見(jiàn)的ZeRO通信方式

優(yōu)化后的ZeRO通信方式

Colossal-AI在常見(jiàn)的ZeRO通信方式基礎(chǔ)上，通過(guò)進(jìn)一步將參數(shù)的All-gather和下一輪訓(xùn)練的前向計(jì)算重疊，以達(dá)到更高的訓(xùn)練效率。

使用bucket來(lái)進(jìn)行Zero-DP的通信是常見(jiàn)的優(yōu)化手段，但是在使用bucket的過(guò)程中存在大量的內(nèi)存拷貝操作。當(dāng)集群規(guī)模擴(kuò)大時(shí)，內(nèi)存拷貝操作的開(kāi)銷會(huì)逐漸增大。為解決這個(gè)問(wèn)題，通過(guò)將內(nèi)存拷貝操作進(jìn)行融合，降低了內(nèi)存拷貝的開(kāi)銷。

同時(shí)，隨著集群規(guī)模擴(kuò)大，通信算子All-Gather和Reduce-Scatter的速度衰減非常嚴(yán)重。在bucket size較小時(shí)，降速更加明顯（例如Torch DDP默認(rèn)的25MB）。

Colossal-AI從兩個(gè)方面來(lái)解決這個(gè)問(wèn)題。

首先可以通過(guò)增大bucket size來(lái)減緩?fù)ㄐ沤邓?，但是bucket size不能無(wú)限增大，過(guò)大的bucket size會(huì)影響計(jì)算和通信的重疊。通過(guò)以下公式粗略搜索得到較優(yōu)的bucket size，再實(shí)際測(cè)試進(jìn)行精調(diào)。

其中，φ_B為bucket size，φ為模型大小，T_bwd為反向計(jì)算時(shí)間，T_comm為bucket單次通信時(shí)間。

其次，當(dāng)集群規(guī)模很大時(shí)，Colossal-AI引入2d torus方式的通信能夠減緩?fù)ㄐ沤邓賳?wèn)題。

2d torus通信

經(jīng)過(guò)聯(lián)合優(yōu)化后，在視頻模型訓(xùn)練的場(chǎng)景下，集群規(guī)模很大時(shí)也能保證scaling > 95%，在大規(guī)模多機(jī)訓(xùn)練中能達(dá)到~30%的加速。

數(shù)據(jù)加載優(yōu)化

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

PyTorch dataloader提供了自動(dòng)Pin memory的功能，能大幅度提高把數(shù)據(jù)從CPU移動(dòng)到GPU的時(shí)間，是非常實(shí)用的一項(xiàng)功能。其通過(guò)Python多線程來(lái)實(shí)現(xiàn)自動(dòng)Pin memory。

但是由于GIL的存在，Python的多線程并非傳統(tǒng)意義上的多線程。同時(shí)Pin memory操作調(diào)用的cudaMallocHost 可能會(huì)阻塞主進(jìn)程（影響主CUDA Stream）。當(dāng)使用高清/長(zhǎng)視頻進(jìn)行訓(xùn)練時(shí)，pin memory需要申請(qǐng)的內(nèi)存較大，這個(gè)問(wèn)題會(huì)更加明顯。

具體表現(xiàn)為，開(kāi)啟Pin memory之后，某個(gè)進(jìn)程的某一部分操作可能比別的進(jìn)程更慢，從而造成一定的不同步性，而不同步性在大規(guī)模集群訓(xùn)練時(shí)對(duì)整體訓(xùn)練效率影響較大。

為解決這個(gè)問(wèn)題，Colossal-AI將dataloader進(jìn)行了改造，通過(guò)預(yù)分配和緩存pin memory的機(jī)制，盡量避免在訓(xùn)練過(guò)程中調(diào)用cudaMallocHost。如果設(shè)置合理，緩存命中率可以達(dá)到100%，即不會(huì)影響訓(xùn)練速度，并且不會(huì)消耗過(guò)多的RAM cache。

FP8混合精度訓(xùn)練

Colossal-AI支持主流的BF16(O2) + FP8(O1)的新一代混合精度訓(xùn)練方案。

僅需一行代碼，即可對(duì)主流大模型能夠獲得平均30%的加速效果，并保證訓(xùn)練收斂性，降低相應(yīng)大模型開(kāi)發(fā)成本。

使用時(shí)，僅需在初始化plugin時(shí)開(kāi)啟FP8即可：

from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin
...
plugin = LowLevelZeroPlugin(..., use_fp8=True)
plugin = GeminiPlugin(..., use_fp8=True)
plugin = HybridParallelPlugin(..., use_fp8=True)

除此之外，無(wú)需引入額外的手寫(xiě)CUDA算子，避免了較長(zhǎng)的AOT編譯時(shí)間及復(fù)雜的編譯環(huán)境配置。

序列并行優(yōu)化

Colossal-AI針對(duì)VideoOcean模型支持了多種序列并行范式，包括Tensor sequence parallelism，Ring attention （context parallelism）和Sequence parallelism（Ulysses），這幾種范式可以單獨(dú)使用也可以聯(lián)合使用。

同時(shí)根據(jù)視頻數(shù)據(jù)的特征（激活值特別大），進(jìn)一步優(yōu)化了Ring attention的通信，使用ND-ring來(lái)應(yīng)對(duì)復(fù)雜的硬件配置。

當(dāng)視頻模型scale到數(shù)百億參數(shù)量級(jí)，并且使用高清、較長(zhǎng)的視頻訓(xùn)練時(shí)，大規(guī)模多機(jī)訓(xùn)練和混合并行訓(xùn)練幾乎是標(biāo)配。

在這種情況下，Colossal-AI對(duì)序列并行的優(yōu)化可以應(yīng)對(duì)各種情形，尤其是大視頻導(dǎo)致序列需要跨機(jī)的情況加速效果顯著。

卷積層張量并行優(yōu)化

Colossal-AI對(duì)適用于高清視頻、長(zhǎng)視頻的VAE進(jìn)行了針對(duì)性優(yōu)化。

對(duì)這類數(shù)據(jù)，CUDNN的3D卷積會(huì)產(chǎn)生非常大的激活值，為此Colossal-AI實(shí)現(xiàn)了分塊卷積和張量并行。

與Transformer中的張量并行不同，Colossal-AI對(duì)VAE用了一種新的張量并行方式以適配其巨大的激活值，最終在完全不損失精度的情況下完成了加速和內(nèi)存優(yōu)化。