近期免費上線的視頻生成平臺Video Ocean，支持任意角色、任意風格，可以文生視頻、圖生視頻、角色生視頻，引起廣泛關注與肯定，登上Product Hunt全球產品熱度榜單第三。

體驗地址：https://video.luchentech.com/zh-CN

Video Ocean如何以極低成本快速完成迭代？開源解決現已方案發布。

二次開發分享至開源社區，還可領取500元GPU算力代金券。

開源地址：https://github.com/hpcaitech/Open-Sora

Colossal-AI

在Video Ocean背后，離不開AI大模型訓練推理系統Colossal-AI的基礎支持，其在GitHub全球AI訓推系統開源領域指標位列世界第一，已獲近4萬Stars。

開源地址：https://github.com/hpcaitech/ColossalAI

它基于PyTorch，可通過高效多維并行、異構內存等，降低AI大模型訓練/微調/推理的開發與應用成本，已與多家世界/中國500強企業聯合開發和優化AI大模型。

針對類Sora視頻大模型開發，Colossal-AI做了多方面的優化，模型算力利用率MFU相對現有其他開源方案總體最高提升可達2.61倍，顯著降低成本。

異步Checkpoint

在使用大規模集群訓練時，因集群規模擴大，故障率會迅速上升，導致訓練極易中斷。在這種情況下快速保存Checkpoint不僅可以加速整體訓練效率，也有利于故障容錯，快速恢復訓練。

為此，Colossal-AI推出了異步Checkpoint功能。針對10B量級的視頻生成模型，可將DiT模型、EMA模型、優化器的保存時間從300s+，降至10s以內，節省了高達97%的保存時間。

Checkpoint保存主要分為GPU->CPU（D2H）和硬盤寫兩個步驟，通過將這兩個步驟通過流水線的形式執行，極大的提高了保存的效率。同時此步驟通過多線程（C++）的形式在后臺完成，不會阻塞訓練主進程。GPU->CPU通過單獨的CUDA Stream完成，也不會阻塞主計算Stream。

流水線保存Checkpoint

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

Zero內存/通信優化

常見的ZeRO通信方式

優化后的ZeRO通信方式

Colossal-AI在常見的ZeRO通信方式基礎上，通過進一步將參數的All-gather和下一輪訓練的前向計算重疊，以達到更高的訓練效率。

使用bucket來進行Zero-DP的通信是常見的優化手段，但是在使用bucket的過程中存在大量的內存拷貝操作。當集群規模擴大時，內存拷貝操作的開銷會逐漸增大。為解決這個問題，通過將內存拷貝操作進行融合，降低了內存拷貝的開銷。

同時，隨著集群規模擴大，通信算子All-Gather和Reduce-Scatter的速度衰減非常嚴重。在bucket size較小時，降速更加明顯（例如Torch DDP默認的25MB）。

Colossal-AI從兩個方面來解決這個問題。

首先可以通過增大bucket size來減緩通信降速，但是bucket size不能無限增大，過大的bucket size會影響計算和通信的重疊。通過以下公式粗略搜索得到較優的bucket size，再實際測試進行精調。

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

其次，當集群規模很大時，Colossal-AI引入2d torus方式的通信能夠減緩通信降速問題。

2d torus通信

經過聯合優化后，在視頻模型訓練的場景下，集群規模很大時也能保證scaling > 95%，在大規模多機訓練中能達到~30%的加速。

數據加載優化

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

PyTorch dataloader提供了自動Pin memory的功能，能大幅度提高把數據從CPU移動到GPU的時間，是非常實用的一項功能。其通過Python多線程來實現自動Pin memory。

但是由于GIL的存在，Python的多線程并非傳統意義上的多線程。同時Pin memory操作調用的cudaMallocHost 可能會阻塞主進程（影響主CUDA Stream）。當使用高清/長視頻進行訓練時，pin memory需要申請的內存較大，這個問題會更加明顯。

具體表現為，開啟Pin memory之后，某個進程的某一部分操作可能比別的進程更慢，從而造成一定的不同步性，而不同步性在大規模集群訓練時對整體訓練效率影響較大。

為解決這個問題，Colossal-AI將dataloader進行了改造，通過預分配和緩存pin memory的機制，盡量避免在訓練過程中調用cudaMallocHost。如果設置合理，緩存命中率可以達到100%，即不會影響訓練速度，并且不會消耗過多的RAM cache。

FP8混合精度訓練

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

僅需一行代碼，即可對主流大模型能夠獲得平均30%的加速效果，并保證訓練收斂性，降低相應大模型開發成本。

使用時，僅需在初始化plugin時開啟FP8即可：

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

除此之外，無需引入額外的手寫CUDA算子，避免了較長的AOT編譯時間及復雜的編譯環境配置。

序列并行優化

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

同時根據視頻數據的特征（激活值特別大），進一步優化了Ring attention的通信，使用ND-ring來應對復雜的硬件配置。

當視頻模型scale到數百億參數量級，并且使用高清、較長的視頻訓練時，大規模多機訓練和混合并行訓練幾乎是標配。

在這種情況下，Colossal-AI對序列并行的優化可以應對各種情形，尤其是大視頻導致序列需要跨機的情況加速效果顯著。

卷積層張量并行優化

Colossal-AI對適用于高清視頻、長視頻的VAE進行了針對性優化。

對這類數據，CUDNN的3D卷積會產生非常大的激活值，為此Colossal-AI實現了分塊卷積和張量并行。

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