新智元报道

编辑:编辑部 HYZ

【新智元导读】随着Sora震撼发布,视频生成技术成为了AI领域新风口。不过,高昂的开发成本是一大瓶颈。国产平台Video Ocean不仅成功登上全球热榜第三,还将视频生成模型开发成本降低50%。而且,模型构建和性能优化方案现已开源,还能免费获得500元GPU算力。

近期免费上线的视频生成平台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用了一种新的张量并行方式以适配其巨大的激活值,最终在完全不损失精度的情况下完成了加速和内存优化。

领取GPU算力代金券

为回馈广大开发者的关注与肯定,基于Colossal-AI或OpenSora

  1. 构建有实际意义的高质量项目,

    如微调、预训练模型、应用、算法论文等开源项目

    奖励:领取潞晨云500元或hpc-ai.com的H200 GPU 100美元算力代金券。

  2. 发布相关开源项目

    奖励:领取潞晨云50元或hpc-ai.com的H200 GPU 10美元算力代金券。

  • 领取详情:https://colossalai.org/zh-Hans/docs/get_started/bonus

  • 开源地址:https://github.com/hpcaitech/Open-Sora