公众号关注 "GitHubDaily"
设为 "星标",每天带你逛 GitHub!
近期免费上线的视频生成平台 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 。
它基于 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 会影响计算和通信的重叠。通过公式
其次,当集群规模很大时,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
构建有实际意义的高质量项目
如微调、预训练模型、应用、算法论文等开源项目
奖励:领取潞晨云 500 元或 hpc-ai.com 的 H200 GPU 100美元 算力代金券。
发布相关开源项目
奖励:领取潞晨云 50 元或 hpc-ai.com 的 H200 GPU 10 美元算力代金券。
领取详情,请点击阅读原文或下方链接。
领取详情:https://colossalai.org/zh-Hans/docs/get_started/bonus
开源地址:https://github.com/hpcaitech/Open-Sora
没有评论:
发表评论