为AIGC客户降本增效。
作者 | 赵健
AIGC是当前AI领域最热门的话题之一。
尤其是伴随着Stable Diffusion、Midjourney、NovelAI、DALL-E等为代表的文本生成图像的跨模态应用涌现,AIGC更是火爆出圈。
最常见的AIGC应用莫过于AI绘画了。抖音在11月底上线了AI绘画特效,根据照片生成二次元头像,目前已经有近2600万人使用。
抖音AI绘画特效背后的技术支撑,来自于字节跳动的云服务平台火山引擎,抖音AI绘画的上线也意味着火山引擎正式进入AIGC领域,为抖音以及全行业提供底层的技术支撑、工程优化等能力。
AIGC由人工智能生成声音、图像、视频等内容,被认为是继PGC(专业生产内容)、UGC(用户生产内容)之后,新一代的内容生产方式。
但AIGC的大规模普及,仍然尚需时日,还需要克服高昂的成本等问题。
火山引擎的机器学习平台,正是为了解决这些问题,以此来推动AIGC产业的发展。「甲子光年」采访了火山引擎机器学习平台相关负责人,来探究火山引擎对AIGC的投入与创新。
火山引擎相对PyTorch原生基线训练提升40%、推理提升3.47倍。对此,火山引擎是如何做到的?
1.降低AIGC成本,需要云的助力
AIGC火了之后,也催生了新兴的AI独角兽,最具代表性的就是主打AI文字生成的Jasper.ai,以及主打AI生成绘画的Stability AI。
科技巨头也不甘落后,纷纷推出各自的AI工具,比如英伟达推出了全新的Magic3D,能够基于文本描述自动生成3D模型;Meta推出了高性能推理引擎AITemplate,大幅简化多GPU后端部署。
今年7月Stability AI这家公司将其AI文本-图像生成模型Stable Diffusion开源,后续开发者只需要在此开源模型的基础上做一定的调整开发工作,就可以得到属于自己的AI绘画创作。
但是,AIGC仍然存在一定的门槛,尤其是对于国内的开发者来说。
现有Diffusion模型主要使用英文数据进行训练。由于中英文在文化和表达上的的差异性,如果直接将英文数据翻译成中文进行模型训练,产出的模型通常无法建模中文特有的现象。此外,通用的Stable Diffusion模型由于数据源的限制,很难用于生成特定领域、特定场景下的高清图片。
同时,高昂的成本也一定程度上限制了AIGC产业的大规模普及。
以Stable Diffusion为例,最初版本的Stable Diffusion动用了256块英伟达A100 GPU,花费了15 万小时的总训练时长,仅这一项成本就是60万美元。即使是Fine Tune(微调),也需要消耗几天时间。换算成钱来看,一次完整的预训练可能需要上百万的成本,Fine Tune阶段也需要二十多万/次。目前AIGC市场火爆,各种画风和产品形态频频出现且快速演进,这样的预训练和Fine Tune的次数肯定不会少。
对于创业公司来说,开发这样体量的大模型困难重重,不仅仅是写一套算法那么简单,还包括数据准备、模型开发、训练调优、模型部署等多个环节,甚至在模型的训练与模型的推理阶段还需使用不同的芯片,每一个环节工作量都不小。
其次,在模型推理方面,随着AIGC火爆全网,各个产品的QPS(Queries Per Second,每秒查询率)压力都非常大,原先预计的几百上千资源,可能需要短时间扩容至上千甚至是万规模,才能保证前端的服务质量。而目前的开源加速方案,仅适用于标准的Stable Diffusion加速,需要用户有较强的工程适配能力。
火山引擎就观察到,火山上的很多用户,由于自己的模型有些许改动、且时间有限,都是用原始框架直接部署。
因此,不少AI初创公司选择把模型开发这项“AI炼丹”的工作通过云来完成,以期相对低成本、又能无限算力的方式获取云端资源。
Stability AI选择了AWS的云服务,目前拥有一个在 AWS 中运行的由 4000 多个英伟达A100 GPU 组成的集群,用于训练包括 Stable Diffusion在内的 AI 系统。在今年的亚马逊re:Invent大会上宣布双方达成战略合作。
云厂商基本也是基于这个逻辑布局AIGC,即为AIGC产业提供基础设施。
作为最年轻的一朵云,火山引擎在AIGC领域的表现却一点也不像个新手,其模型推理速度达到了66it/s,是目前已公开的模型优化速度最快的平台。
火山引擎凭什么做到?
2.火山引擎,极致优化
火山引擎在性能、显存消耗以及接入方式这3个方面,都做到了极致。
首先,是训练和推理的GPU性能。火山引擎计算加速团队对Stable Diffusion的整个模型做了详尽的分析,然后按照如下思路逐步优化训练和推理的性能。
第一步,确定Kernel Fusion优化范围。
众所周知,Kernel Fusion是GPU优化最常用有效的技巧,通过将多个连续的或者可以并行执行的操作,融合为尽可能少的GPU kernels,可以大量减少重复的显存读写、降低Kernel launch的开销,从而实现更优的性能。
然而,由于Kernel Fusion后的算子粒度往往较原始实现大很多,有时候大量的Kernel Fusion会极大地限制优化适用的灵活性。
对此,火山引擎优化团队充分考虑模型内在的计算逻辑、不同GPU硬件的硬件参数及特性、以及现有的GPU基础库提供的基础能力,通过设计多个层级不同粒度的优化算子,同时兼顾优化性能以及模型未来迭代的发展趋势。
比如,通过扩展Cutlass,火山引擎在卷积操作前后可以额外添加丰富灵活的前后处理,比如add bias或者residual等,这就打破了以往以卷积为界限的Kernel Fusion断点,在Stable Diffusion这样卷积操作很重且Tensor很大的网络中就可以取得很好的性能收益。再比如,依托Flash-Attention算法,火山引擎可以突破片上shared memory的限制,将整个Attention操作融合为一个操作。
经过Kernel Fusion后,Stable Diffusion v1 训练和推理的CUDA Kernels数量相比原始的PyTorch实现,分别减少了50%和60%。
第二步,并行算法重设计。
对于每个算子,均会从问题的本质出发,并结合现有的GPU硬件参数,设计更优的并行算法实现。
比如,LayerNorm是深度学习里非常典型常用的一个操作,其反向操作均涉及到大量的按列求和操作,PyTorch使用比较基础的实现来组装,而火山引擎通过更好的并行算法实现,仅用2个kernel就完成了全部的反向操作。再比如,GroupNorm是StableDiffusion中耗时较高的一个操作,其性能非常受Tensor Layout排布的影响。火山引擎对两种Layout下的GroupNorm操作,做了不同的并行算法设计,最终该操作前后向分别加速50%和30%。
第三步,SOL评估。对每个操作,都会仔细评估其需要的最少数据读写、以及浮点操作数量。过程中会借助一系列Profiling工具,深入SASS Code层,争取每个CUDA Kernel都能达到SOL的80-90%的水平。
其次,降低显存消耗。
性能优化的过程中,通过大量的Kernel Fusion,就可以大幅减少中间Tensor的读写。除此之外,火山引擎进一步,在训练优化过程中,将部分轻量级操作以re-compute的方式来进行反向推导,可以进一步减少显存消耗。最终,可以使单张上可以同时计算更大的batch_size,进一步提高整体吞吐。
最后,火山引擎在接入方式上也做了极致简化。
因为火山引擎提供的加速算子,同时支持高效训练和推理,从而真正实现了训推一体,因此可以从模型训练到推理部署做到无缝衔接。接入方式上提供由Pytorch Custom op搭建出的相同模型结构,完全兼容原始Stable Diffusion模型的Pipeline, 接入方便且无需改动框架。
在推理部署时,高性能算子库可以支持不同GPU硬件,包括英伟达V100、A100、A10、A30等不同系列的GPU卡,而且不需要做额外的模型转换工作。
未经优化时,绝大多数模型只能在最高端的A100 80GB训练卡上跑起来。而火山引擎提供的训练加速,能够摆脱特定计算卡的限制,让客户在更广泛的范围内寻找适合自己需求的解决方案。
火山引擎目前提供40%的训练加速,这样在同类型的云产品上训练,无论是迭代速度,还是单次的训练成本,都有了明显的提升和节省。
3.经过抖音内部实践的考验
今年7月,Stable Diffusion的V1版本开源时,火山引擎已经做了初步的尝试。11月25日Stable Diffusion的V2版本发布,火山引擎在此基础上又花了不到2周的时间,做了基础工程的优化。这是抖音AI绘画特效上线的技术准备工作。
除了做好技术的基本功,还需要有一个好的练兵场。对于火山引擎而言,抖音就在字节跳动内部为火山引擎提供了一个绝佳的内部实践。
火山引擎机器学习总监吴迪告诉「甲子光年」:“要把包含了AIGC以及机器翻译、人机交互、语音识别等多媒体模型的训练和推理做好,是有一定门槛的,需要长时间的技术和经验积累。而抖音的产品形态中包含了丰富的多模态内容,抖音的高标准要求和多样化的产品形态在不停地打磨火山引擎机器学习平台的工程团队。在这样严苛的长期锤炼下, 火山机器学习团队有信心,把最有竞争力的机器学习优化能力带给客户。帮助客户用更低的成本去敏捷迭代。”
在内部与抖音的合作过程中,火山引擎的机器学习平台团队主要做工程优化方面的工作。有时候如果仅仅依靠工程团队仍然没有办法满足上线需求,就会与算法团队一起做更透彻的优化,比如进一步的量化,或者是对模型进行一些蒸馏、裁剪等。
抖音对于火山引擎所提出的诉求,往往也代表了所有市场用户的基本诉求,首当其冲的就是“降低成本”。
作为一款国民级短视频应用,抖音所持有的机器总量很大,因此即使是细微的性能节省,都将带来巨大的资金回报。因此,当火山引擎的AIGC模型,优化速度提升40%,硬件成本降低百分之二三十,训练成本降低80%,相比目前最快的TensorRT部署方案,推理还快10%,将会带来巨大的成本优势。
除了更省钱之外,还必须要稳定,不能出错。
近期OpenAI推出的聊天机器人模型chatGPT,就因为太过火爆而一度服务器崩溃。作为一款现象级产品,这并没有影响到用户对chatGPT的热情,但用户体验总归不够友好,同样的事情如果发生在抖音,可能就是一次产品事故。
为了避免这种情况发生,火山引擎必须做到足够的稳定性,能够抗住激增流量的压力。为此,火山引擎不仅要做模型的优化,还要在服务框架层面做一些稳定性的优化。
第三是要保证充足的资源供应。
火山引擎管理了大量计算资源,来支撑字节跳动的整体业务。但这些卡片不会是永远忙碌的,一定有的处在业务峰值,有的处于相对闲置的状态。此时,就需要有一个自动化的资源调动系统,能够“集中力量干大事”。比如,当AIGC的应用在短时间内需要大量资源的时候,火山引擎可以在小时级或者分钟级的时间内,从其他相对闲置的场景资源去调取,形成一种灵活、弹性的资源供应。
最后就是满足差异化的场景。
AIGC的火爆也带火了机器学习在多媒体方面的应用。其实在抖音内部,多媒体的机器学习应用早就层出不穷了,火山引擎的模型优化要去理解每一帧图片、每一个视频、文字与语音之间的关联等等,其丰富程度通常远大于市面上常见的一些形态。
同时,因为抖音也有强大的算法团队,意味着模型的变化幅度与速度都是非常快的,这也是市场的的一些优化方案仅仅支持好一些基础版本的模型就可以,而火山引擎一定要考虑算法整体设计方案的灵活性,满足“变来变去”的需求。
如今,抖音的AI绘画特效已经有接近2600万人使用。
AIGC目前仍然处在大规模爆发的前夜,而诸如AI绘画这样的早期应用,正是打开AIGC市场的一把钥匙。而随着火山引擎这样的云平台不断把降本、增效这件事做到极致,AIGC产业大规模爆发的那一刻,有望在不久的将来快速实现。
END.
亲爱的朋友们,不知不觉2022即将过去,感谢大家又一年的陪伴,在2022的最后几周,我们邀请所有的朋友们来分享你在2022难忘的回忆,我们将挑选其中的一些故事放入甲子光年2022年终回顾文章中,扫码下方二维码,分享属于你的2022故事吧!
继续阅读
阅读原文