Pytorch训练优化技巧及其在Diffusion上的应用

扩散生成模型 翳明LiuYisong 李英伟达 稳定扩散训练优化 议程 优化思想的概述 性能故障 细节的优化方法 稳定扩散培训管道(v1.4) 文本 嵌入 文本编 码器 图像解 码器 训练数据集 VAE 编码器 图像信息的创造者 (Unet调度器) 培训：文本编码器FP+VAE编码器FP+unetFP&BP “尤达”“钢铁侠” 口袋妖怪的数据集 text_to_image整合结果 稳定扩散培训管道 模块 时间(女士) 比 vae。编码器生产 33.173 7.60% text_encoder录象 13.946 3.20% Unet录象 81.677 18.72% Unetbwd 125.197 28.69% clip_grad_norm 32.124 7.36% 更新 96.457 22.11% ema_unet 37.659 8.63% 其他 16.109 3.69% 总 436.342 100.00% Nsys时间表 FP16: 稳定扩散培训管道 TF32: 8 8 8 8 8 优化思想的概述 batch_size avg时间/iter 吞吐量 内存 加速 评论 基线(tf32) 0.9865 8.1098 69479mib 1.0000 默认配置HuggingFace训练食谱 基线(fp16) 0.8375 9.5520 71695mib 1.1778 默认配置HuggingFace训练食谱 dataloader 0.6983 11.4564 71695mib 1.4127 使用多GPU时，数据加载器将成为瓶颈，DeepSpeed fusedlayerfusedadam 0.6319 12.6600 66949mib 1.5611 中长期规划融合线性的,，LayerNorm和亚当 flash经办人 0.4428 18.0670 25481mib 2.2278 利用FlashAttention针对Pytorch标准关注实现 48 2.0363 23.5727 72979mib 2.9067 多流道融合教育津贴 48 2.0209 23.7511 72929mib 2.9287 我们看到由EMA引起的间歇之间的气泡，这将介绍频繁malloc和free的行为 Zero2 56 2.3302 24.0327 72439mib 2.9634 ZeRO2可以帮助减少GPU内存的使用，增加BS可以提高训练吞吐量 注意:穿孔80gA100*1 优化1:Dataloader 的想法 使pin_mem&num_workers调优dataloader。 对我们来说，检查这两个参数是否正确设置总是一个好主意 我们还可以检查图像解码和预处理是否在Nsys时间内开销（如果是，请尝试DALI） 设置培训DataLoader 时间为一个培训时代 [“num_workers”:1、“pin_memory”:真正的] 2.93秒 [“num_workers”:2,“pin_memory”:真正的] 2.87秒 [“num_workers”:4,“pin_memory”:真正的] 2.85秒 [“num_workers”:8“pin_memory”:真正的] 2.80秒 [“pin_memory”:“num_workers”:16日真] 2.92秒 注意:穿孔80gA100*8,b=56 笔记 使pin_mem&num_workers调优dataloader。 在这种情况下，CPU上的计算开销较小，例如：H2D+图像解码 优化2:内核融合 •利用APEX的构建块进行层融合和优化融合. NVIDIAAPEX（存储库、文档）提供优化的、可重用的构建块。APEX可以从github进行pip安装： https://github.com/NVIDIA/apex#quick-startAPEX预装在PyTorchNGCdocker容器中组件: 熔融层: apex.fused_dense。融合致密apex.fused_dense。FusedDenseGeluDenseapex.normalization.FusedLayerNorm 融合优化: apex.optimizers.FusedAdamapex.optimizers.FusedLAMBapex.optimizers.FusedNovoGradapex.optimizers.FusedSGD 优化2:内核融合 的想法 将text_encoder/unet和Adam模块中的层替换为APEX的优化构建块 细节 •取代神经网络。线性与fused_dense.FusedDense •延时换成fused_dense.FusedDenseGeluDense •替换nn。LayerNormwithapex.normalization.FusedLayerNorm •将torch.optim.AdamW替换为apex.optimizers.FusedAdam 笔记 更换模块可能会导致参数（MLP）的名称更改，从而导致NAN丢失，我们需要匹配密钥 优化2:内核融合 融合亚当优化器 •在融合之前: •125.622毫秒 •6860年vectorized_elementwise_kernel •融合后:之前 •25.022毫秒 •47multi_tensor_apply_kernel 除了减少内核执行时间: •减少内核启动开销 •减少memI/O 后 优化2:内核融合 之前 一次迭代: •在内核融合: •Vectorized_elementwise_kernel，每个线程一个浮动4对齐，提高内存效率 •Unrolled_elementwise_kernel，未float4对齐，在螺纹中循环 •Elementwise_kernel,没有优化 •在内核融合: •所有elementwise内核比例降低下来 后 优化3:Flash关注(HazyResearch) 主要观点: •在A100等现代GPU上，张量核心的速度如此之快，以至于注意力最 终会因从GPU内存读取和写入而受到瓶颈 •FlashAttention通过减少GPU内存的读写来加快注意力 •FlashAttention将注意力计算的MHA融合到一个内核中 •仅存储软最大值的范数系数而不是注意力矩阵（N^2 ） 目前支持: •图灵或安培gpu •fp16和bf16(bf16需要安培gpu) •头部尺寸为8的倍数，最多为128（例如，8、16、24、...、128） 。头部调暗>64向后需要A100 优化3:Flash的注意 本机注意 •元素内核被Flash注意力融合，比例降低 •带闪光灯的注意力,融合尼古拉斯是不需要的了 闪光注意 优化3:Flash的注意 本机注意 内核启动时间/数量大幅减少元素内核被FlashAttention融合比例降低下来 Flash的关注 本机 注意 优化3:Flash的注意 •闪存注意可以减少内存占用：67G->25G，相同 batch_size=8 •我们可以增加时批量大小 使用FlashAttentionFlash的关 注 Batch_size 本机的关注 Flash的关注 加速 8 12.531 17.651 1.408 16 伯父 19.497 1.556 48 伯父 21.162 1.589 教育津贴 优化4:Muti-streamEMA和融合 教育津贴是什么? EMA是一条移动平均线，只能通过知道最后一个EMA值和当前值来计算。EMA不需要保留最后N个数据点，因此非常节省内存。 EMA在稳定扩散 稳定扩散使用模型权重的指数移动平均来提高结果图像的质量，并避免过度拟合到最近训练的图像。 异步教育津贴 EMA与unet训练是独立的，它只保留将在下一次迭代之前使用的权重，因此我们可以使用多流将EMA与其他计算重叠 优化4:Muti-streamEMA和融合教育津贴 创建EMACUDA流 新流中的计算EMA，与模型计算重叠 确保在下一次迭代更新之前已完成EMA 添加最后一个EMA以确保它与原始实现等效 优化4:Muti-streamEMA和融合教育津贴 教育津贴 优化4:Muti-streamEMA和融合 融合的实现 本机实现 优化5:0 的想法 通过ZeRO减少梯度、激活和碎片内存，并尽可能增加BS 细节 笔记 8 8 8 8 8 总结 batch_size avg时间/iter 吞吐量 内存 加速 评论 基线(tf32) 0.9865 8.1098 69479mib 1.0000 默认配置HuggingFace训练食谱 基线(fp16) 0.8375 9.5520 71695mib 1.1778 默认配置HuggingFace训练食谱 dataloader 0.6983 11.4564 71695mib 1.4127 使用多GPU时，数据加载器将成为瓶颈，DeepSpeed fusedlayerfusedadam 0.6319 12.6600 66949mib 1.5611 中长期规划融合线性的,，LayerNorm和亚当 flash经办人 0.4428 18.0670 25481mib 2.2278 利用FlashAttention针对Pytorch标准关注实现 48 2.0363 23.5727 72979mib 2.9067 多流道融合教育津贴 48 2.0209 23.7511 72929mib 2.9287 我们看到由EMA引起的间歇之间的气泡，这将介绍频繁malloc和free的行为 Zero2 56 2.3302 24.0327 72439mib 2.9634 ZeRO2可以帮助减少GPU内存的使用，增加BS可以提高训练吞吐量 谢谢!