基于计图框架的 AI 编程优化小技巧 | 新程序员

AI科技大本营 2024-09-02 10:31:02 阅读 84

【编者按】计图框架通过元算子和统一计算图的两大创新,有效提升了开发效率和执行性能。作者基于计图框架分享了 AI 辅助开发技术的新思路,以及 Fitten Code 在 AI 编程领域丰富的实践小技巧。

本文出自 2024 全球软件研发技术大会中的演讲,同时收录于《新程序员 008》。《新程序员 008》聚焦于大模型对软件开发的全面支撑,囊括 Daniel Jackson 和 Daniel Povey 等研发专家的真知灼见与“AGI 技术 50 人”栏目的深度访谈内容,欢迎大家订阅年卡。

作者 | 刘政宁  非十科技CTO,Fitten Code负责人

责编 | 何苗

出品丨AI 科技大本营(ID:rgznai100)

35527093c89640c68cc0a6cf333b5214.png

近年来代码模型领域迎来爆发式增长,在学术界与工业界彰显出双重价值。

在学术界,代码生成是自然语言模型非常重要的应用领域,吸引了大量组织的研究。在工业界,超过四十款 AI 编程助手的陆续面世,证明了代码生成作为大模型最具代表性的应用落地场景,极具潜力。

作为一个创业公司,非十科技主要将力量集中在产品研发和技术上,尽管很难在搜索引擎或者 AI 编程的文章推荐里面看到它的身影,但我们支持的 IDE 插件,包括 VSCode、Visual Studio 以及 JetBrains 全系列产品等在全平台累计下载量突破了 30 万次,在国内 AI 编程产品下载量排名中跻身前三。 

这一点非常让人惊讶,我想这一定是因为做对了一些事情。 

我将聚焦于非十科技所擅长的产品底层技术以及生成准确率、时延等关键技术的优化策略,分享一些实践经验。 

55babfcb279c93982f67fcc44f5352dc.png

AI 编程助手开发的四大挑战 

以 GPT 为代表的自然语言大模型兴起,为代码大模型奠定了坚实的基础。自 2020 年起,微软便基于 GPT 架构探索代码生成技术,“Fitten Code”则站在巨人肩膀上,不断优化、创新,力求为用户提供更小体积却更强大的代码生成工具。 

尽管代码大模型与自然语言模型在基础架构和规模上有着相似之处,但代码的独特属性也带来了特殊挑战:

1、代码的语法严格、逻辑严密,与自然语言的文本描述形式截然不同

在上下文长度方面,代码文件可能长达数千行,项目级别的代码量甚至达到数千万行,远超出普通文本的范畴。直接将通用自然语言大模型应用于代码任务的操作显得不切实际,因此业界普遍倾向于添加更多代码语料或专门训练代码大模型,以适应代码生成的特定需求。

2、自然语言大模型可以与知识图谱结合,而程序没有统一的知识库 

不同于自然语言中的知识图谱,代码的抽象语法树(Abstract Syntax Tree, AST)扮演着核心角色。通过诸如 GACLAVVM 等工具,可以解析代码结构,获取函数与变量间的关联,基于此,可开展代码分析、优化等一系列高级操作,极大地丰富了代码大模型的应用场景。 

3、代码生成对正确性要求远高于自然语言生成,需要被正确执行 

在效果追求上,代码生成对正确性的要求远超自然语言生成。代码不仅需通过静态编译,还要确保动态执行的准确性。面对编程语言的多样性,确保不同语言下的规范性和效果成为一项重大挑战。此外,代码质量的高低对程序的可靠性至关重要,但这在自然语言处理中并非主要考量因素。 

编程语言的进化速度飞快,这就要求代码大模型必须具备持续学习能力,以适应新知识的快速迭代。如华为推出的仓颉语言,展现了编程语言的持续创新,强调了代码大模型需具备动态更新机制,以应对编程语言的快速发展。 

4、代码大模型训练对准确度要求高 

在训练代码大模型时,为满足其对精度的苛刻要求,往往需要庞大的参数量和海量数据支持。例如 Llama 3 在训练 70B 参数的大模型时,使用了 15T 数据,而在处理 8B 参数场景时,通过扩充数据集,代码量增加了四倍,显著提升了代码生成的质量。但这背后隐藏着资源需求与训练时间的矛盾,一般需要更多资源投入才能保证性能达标。 

在 AI 编程插件的实际运用中,补全任务因其高频触发特点,可容忍一定错误率,适用于较小模型。而在问答等深度交互场景中,为了提供更好的用户体验,必须采用更大模型,确保生成代码的准确性和可靠性。 

面对上述挑战,我将从底层优化技术和框架层面深入探讨,分享见解与实践经验,推动代码大模型领域的发展,解决实际应用中的关键问题。

80c57dad0fcebf3d9d8ba8ea9d0d5a30.png

元算子与统一计算图的双重创新 

清华大学推出的计图深度学习框架于 2020 年初春正式亮相,是非十科技创业团队的核心技术,定位为 Patchwork 的全面替代方案,可提供底层硬件支持和多种模型及库的兼容性,通过元算子和统一计算图两大创新,增强对国产芯片和操作系统的支持,以达到简化模型开发与优化过程的目标。

e04328da38885ca376ac3123c786b587.png

计图深度学习框架

计图深度学习框架的架构自底层硬件兼容性着手,覆盖 CPU、GPU 及各类 AI 加速硬件,向上支持广泛模型与模型库,为开发者提供创新实验的理想平台。框架的两大创新亮点在于元算子与统一计算图。

元算子的引入,使得计图能够灵活适应国产芯片,包括华为、曙光在内的多种芯片类型,累计适配数量超过六种,同时兼容多款国产操作系统。元算子作为一种基础运算单元,允许开发者自由组合,构建复杂的模型结构,不仅增强了框架的灵活性,还促进了模型效率的提升。 

统一计算图则提供了全局视角,使得框架能够自动优化计算流程,减少冗余计算,提高整体执行效率。双管齐下,使得计图在支持国产芯片的同时,保持了优异的运行性能。 

元算子融合策略 

在框架开发初期,开发者面对的是深度学习算子库庞大的挑战。TensorFlow 初期包含两千多个算子,即便 PyTorch 将其精简至七百个,每项算子仍需深度优化,这对开发与维护构成了巨大压力。

a631533f96f280db1254f8c29fb90759.png

元算子融合方式

为破解这一难题,我们引入了元算子(Primitive Operator)的理念。元算子是一种底层抽象,非实际算子,而是将深度学习运算归纳为 18 个,细分为三类:

原数据算子,涵盖相同形状向量的基本运算,如加减乘除; 

索引化简算子,用于数据的简化处理,如求最大值、最小值、平均值;

索引操作算子,将低维向量映射至高维空间,用于向量的扩展操作。 

通过这种抽象,我们构建了中间表示,再经后续代码生成阶段,灵活组合成各类深度学习所需的复杂算子,如卷积、注意力机制等。这一设计借鉴了 TVM 等框架,但侧重于替代 Patchwork,通过维护基础元算子与中间代码优化阶段,成功规避了维护庞杂算子库的困扰,大幅简化了开发与优化流程。 

统一计算图思想 

另一项创新是统一计算图。早期深度学习框架如 TensorFlow 采用静态执行模式,需预定义整个网络结构,再将完整计算图发送至计算设备执行,这一模式在调试与性能瓶颈定位上存在局限性。相比之下,PyTorch 的动态执行模式,算子逐个发送,虽便于调试与算法改进,但效率受限,需后续优化。

我们采用了动态切分策略,将用户需要获取中间结果或动态调整网络结构的操作进行分割,形成静态子图进行专门优化。这一设计不仅保留了动态执行的灵活性,还显著提升了执行效率。同时,我们维持了与 PyTorch 一致的接口,便于用户无缝迁移。

去年,PyTorch 2.0 推出的 Eager Compile 功能与我们的设计理念不谋而合,这证实了我们在元算子融合与统一计算图设计上的前瞻性和先进性。 

值得一提的是,我们还具备跨迭代融合的独特机制。这一机制跨越多次迭代,实现更长远的优化策略,适用于需要多次迭代优化的深度学习场景,如循环神经网络(RNN)的训练,显著提升了模型训练的效率与效果。让我们能够更有效地管理计算资源,实现深度学习任务的高效执行。 

563b718350447a7ca3541f1f78f8de38.png

统一计算图的计算策略

计图深度学习框架凭借其创新特性,在支持国产芯片、提升模型效率与简化开发流程方面展现出巨大潜力。“Fitten Code”代码助手作为计图框架在实际应用中的体现,聚焦于代码生成与优化,具备了智能代码补全、错误检测与修复等功能,为开发者提供高效、精准的编程辅助。

007b505353591ec7c3fb5e3ef92e93c8.png

通用大模型的两大优化难点 

在底层技术的坚实基础上,我们构建了丰富的生态体系,覆盖了从生成对抗网络(GANs)、图像分割,到三维点云处理与神经辐射场(NeRF)等前沿应用领域,尤其在二维与三维视觉任务中,取得了卓越的性能表现。

806c7514e71604b58875d23f91a36186.png

计图的大量模型

与 PyTorch 等主流框架相比,我们的模型库在性能上实现了显著提升,在生成对抗网络领域,部分模型的效率提升超过一倍,最高甚至达到 PyTorch 的两倍多,平均效率提升达到了 2.26 倍。这一成就源于 GANs 的双网络结构与复杂优化过程。 

我们逐渐将这些技术沉淀应用到大模型的训练与推理优化中,特别是在代码大模型领域。 

1、显存优化

直面大模型训练中最核心的挑战——显存管理。以 7B 参数模型为例,全量训练至少需要 112GB 显存,而 GPT-3 级别的百亿参数模型,则需高达 2800GB 显存,显然,单卡环境下难以实现。零冗余优化器技术可以很好解决这一难题,核心思想是将模型参数与优化器状态切片分配至多张 GPU 上,利用 GPU 间高速互联网络进行参数交换与同步,有效减少了单卡的显存需求。不仅适用于单机多卡环境,也可扩展至多机分布式场景,为大模型训练提供了有力支持。 

fcfbed867ad74858b1774110d78b354e.jpeg

优化器示意

2、分布式训练 

分布式训练是另一重要议题,主要包括三种策略。

数据并行:将数据集切分为多个子集,分别在不同机器上进行训练,最后汇总梯度进行统一计算,适用于模型较小、数据量较大的场景;

模型并行:将大型模型分割至多设备,尤其适用于如 MOE 架构的超大规模模型,适用于模型较大、数据量适中的场景。 

流水并行:按模型层排列,实现并行计算,适用于模型层间依赖较弱的场景。在实际训练中,通常结合多种并行模式,以达到最佳训练效果,实现资源的高效利用。 

e4c7b43cc557b5a952f00ab9aa31f00f.jpeg

分布式训练计算流水排布

完成上述基础工作后,我们即可实现对大模型的有效训练与微调。

去年年初,我们便投身于 Llama 系列大模型及其衍生微调方法的深入探索,包括 Apache 及其变体。通过广泛的微调实验我们发现,我们的方法在精度上与 DeepSpeed 等主流技术并驾齐驱,基于先前的优化经验,训练与微调速度提升了 20%,内存消耗降低了 30%,这意味着我们能够支持的模型规模可扩大 30%以上,为大模型训练注入新的活力。

30493c96ef8cab6784c32372b4077462.png

代码模型的实用训练优化策略 

探讨通用模型训练方法之余,在代码模型的训练与优化方面,我们也在实践中发现的一些更为实用的策略。

1、海量训练数据采样优化 

这对于大模型训练尤为关键,正确的数据配比与采样策略能够显著影响模型性能。在代码数据预处理阶段,我们遭遇了存储资源的严峻挑战。例如,代码数据集预处理前可能达到近 10TB,处理后仍需数 TB 空间。起初,我们采用 Hugging Face 的均匀采样方案,但其预处理生成的索引文件体积庞大,占用了近十倍于原始数据的空间,对存储造成了巨大压力,并在运行时消耗大量内存,影响了训练效率。为解决这一难题,我们自行设计了一种索引算法,先对训练数据建立简易索引,以此快速定位数据在存储中的位置,实现在不占用大量内存的前提下快速读取整个数据集,显著提升了数据处理效率。 

2、逐层梯度裁剪保证训练稳定 

梯度裁剪用于控制梯度爆炸风险,尤其是在使用较低精度如 FP16 或 BF16 训练时,以确保模型稳定收敛。经典裁剪算法需要计算整个网络各层梯度的平方和,然后统一裁剪,但这导致所有层的梯度计算必须等待最后一层结果,延迟了梯度更新,增加了显存占用与回程次数。我们尝试了逐层裁剪的简化方法,即为每层单独计算裁剪因子,直接对当前层梯度进行规约。虽然理论上依据不足,但实践表明,该方法对模型效果影响甚微,同时减少了 6% 的显存占用,加速了 10%以上,显著提升了训练效率与资源利用率。 

3、长上下文良好支持 

为了实现良好的问答与补全效果,模型通常需要处理较长上下文窗口,这不仅涉及当前文件与函数块,还需考虑相关代码片段与推荐代码,以捕捉更全面的上下文信息。直接扩大窗口长度会导致算力需求呈平方级增长,成本激增。我们采用了基于 Llama 训练方案的策略,先以较小窗口训练 80%数据,再用剩余 20%数据训练所需上下文窗口,这种方法与直接训练长上下文窗口的性能相近,但总计算量减少 40%,有效平衡了训练效率与模型性能。 

通过这些技术创新与优化,我们在训练效率与资源利用上取得了显著进展,为代码大模型训练开辟了新的道路。 

876a8c5bb6d48d9715fb3c5f78cc2607.png

代码补全场景的推理与优化

接下来,让我们转向推理领域,聚焦代码补全场景。 

在代码补全场景中,请求频率高,响应速度要求快,以 GitHub Copilot 为例,虽然能提升编程效率,但其速度仍较慢,跟不上用户输入速度,体验较差。推理系统需要做到高吞吐量,低延迟,这对大模型推理系统提出了挑战。 

对 AI 插件而言,速度与准确性同等重要,代码补全不仅要正确无误,还须迅速响应,确保不打断程序员的创作流程。Fitten Code 在提供高吞吐量的同时兼顾低延迟的推理体验。作为 AI 编程的深度使用者,我们深知快速响应是提升工作效率的关键。我们在去年 3 月推出了计图大语言模型推理库,该库的最大亮点在于大幅降低硬件配置需求,内存与显存使用最多可减少 80%,同时支持 CPU 计算,兼容多种模型。

动态 swap 机制 

大模型在推理过程中,常常碰到参数文件过大,模型加载效率低下等问题。我们通过动态 swap 机制优化数据访问,将参数卸载至容量更大但访问速度较慢的存储介质,从而使低端设备亦能高效运行代码模型。相比 PyTorch,模型加载效率提升 40% 

38bbe7450de25bc9e54c2fd0417b562d.jpeg

动态 swap 机制效果提升示意

统一内存管理 

模型在推理过程中,常常显存和内存消耗过大,无法运行。Jittor 通过基于元算子的内存切分技术与统一内存管理,使得数据可以在显存,内存和硬盘之间快速切换,精度不变的情况下,大幅降低硬件配置要求(减少 80%),普通笔记本纯 CPU 也能跑大模型。 

c60df12dd6bd8ac0cddf98996b612ccb.png

统一内存管理效果提升示意 

前向推理加速 

模型计算量大,推理延迟高,速度越快体验越好。计图基于在服务端,集成了一系列优化技术,如 PagedAttention、FlashAttention、FlashDecoder,这些技术显著减少了显存占用与 GPU 计算,实现低延迟、高吞吐。 

1854ce4b8148c498bed90441a59d37d6.jpeg

效果提升原理示意 

代码场景长上下文支持 

上下文扩展,提供变量、文件提问等友好交互可减少用户描述负担,同时优化跨文件补全与分析,几乎无显著延迟增加,确保了流畅的编程体验。用户可实时进行互联网搜索,获取最新文档信息,减少大模型的幻觉问题,确保获取的信息及时、准确。 

通过技术创新与优化,我们在训练效率与资源利用上取得了显著进展。Fitten Code 基于自训练部署的代码大模型服务,以速度与准确率为卖点,集成 AI 问答功能,支持 80 多种语言。在一月份的统计中,IDE 插件生成速度比 Copilot 提升 30%,准确率提升 10%。 

a8f2b82d6f53161ac427aa8b0641a406.png

AI 程序员之路仍关卡重重 

打造 AI 程序员是整个行业的共同追求,但要实现这一愿景的路径众多。 

我们尤其关注 AI 程序员在理解和生成代码方面的基本能力。这不仅涉及代码补全功能,更重要的是如何让 AI 深入理解项目的整体结构,并在此基础上进行更高级别的抽象思考。例如,除了使用 RAG 技术获取代码片段外,还需考虑如何更好地解析项目架构,从而实现更为精准的功能实现。 

关于未来,数据集的质量和多样性至关重要。目前,我们主要依赖于公共数据集,难以获得企业内部的数据资源。因而如何建立一个既能保护隐私又能促进共享的机制成为了一个重要的议题。另一个关键问题是确保代码质量,先进的代码检测技术和应对复杂场景的方法,才能确保 AI 生成的代码能够满足高标准严要求。 

这一路还有诸多挑战,但我们相信目标并不遥远。

b8a57d2ff0c2bc0c38f2b18b41c52d57.gif

模型刷新一切,让我们有着诸多的迷茫,AI 这股热潮究竟会推着我们走向何方?面对时不时一夜变天,焦虑感油然而生,开发者怎么能够更快、更系统地拥抱大模型?《新程序员 007》以「大模型时代,开发者的成长指南」为核心,希望拨开层层迷雾,让开发者定下心地看到及拥抱未来。

读过本书的开发者这样感慨道:“让我惊喜的是,中国还有这种高质量、贴近开发者的杂志,我感到非常激动。最吸引我的是里面有很多人对 AI 的看法和经验和一些采访的内容,这些内容既真实又有价值。”

能学习到新知识、产生共鸣,解答久困于心的困惑,这是《新程序员》的核心价值。欢迎扫描下方二维码订阅纸书和电子书。

226fb4c089f84f638b2eb6e71fc8e4b1.jpeg



声明

本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。