近期,在我们的内容技术灵媒智算平台(MVAP)上部署的服务数量提升迅猛,部分业务场景对SD系列模型的推理速度有着一定的要求。因此,我们对当前较为流行的SD加速方式进行了调研与测试,并以AI试衣业务场景为例,尝试了多种加速方案。下面是对调研结果与实际落地效果的一些总结与分享。
加速方案介绍
目前主流的加速思路包含算子优化、模型编译、模型缓存、模型蒸馏等,下面将简要介绍一下几种测试中用到的有一定代表性的开源方案。
▐ 算子优化:FlashAttention2
主要针对transformer中的attention模块进行优化。FlashAttention改变了常规的attention计算方式,将完整的attention计算融合到单个cuda kernel中,并结合了前向tiling与反向recompute技巧;FlashAttention2在此基础上进一步减少了冗余计算,并加强了前向与反向过程中的并行计算,从而降低显存占用和计算耗时。由于效果很好且受到广泛使用,FlashAttention2已被集成到pytorch2.2之后的版本中。
▐ 模型编译:oneflow / stable-fast
oneflow通过将模型编译为静态图,结合oneflow.nn.Graph内置的算子融合等加速策略,实现对模型推理的加速。优势在于对于基础SD模型只需一行编译代码即可完成加速、加速效果明显、生成效果差异小、可以结合其他加速方案(如deepcache)使用以及官方更新频率高等。缺点放到后面再说。
stable-fast也是基于模型编译的加速库,并结合了一系列算子融合的加速方法,但它的性能优化依赖于xformer、triton、torch.jit等工具。
▐ 模型缓存:deepcache
SD模型的常规推理过程包含很多步unet计算(我们在基础模型测试中设置的步数为50),而deepcache的作者发现相邻步数的unet在深层特征上只有极其微小的变化,为了避免重复计算这些冗余的深层特征,deepcache只在部分步骤执行完整计算,并缓存深层特征的计算结果;对于剩余的步骤,则仅计算浅层特征,并结合最近的缓存特征结果。这一过程相当于跳过了unet网络中的大部分层,因此能够显著降低计算量。
deepcache在diffusers框架中的使用非常方便,并且可以根据实际情况设置使用缓存的最深层及使用缓存的步数间隔(这里也可以使用非均匀分布的策略,根据不同步数的冗余程度设置不同的间隔),自行在推理速度与生成质量之间进行权衡。
▐ 模型蒸馏:lcm-lora
结合了lcm(Latent Consistency Model)与lora,lcm会对整个sd模型进行蒸馏,从而实现少步数推理,而lcm-lora借助了lora的形式,只对lora部分进行优化,这样既能够实现加速,也可以直接与常规的lora使用相结合。