1. 介绍
Cut Your Losses in Large-Vocabulary Language Models
Cut-Cross-Entropy(CCE) 是ICLR25 投稿的高分paper(10;10;8;6),CCE算法非常优雅,采用与Flash Attention类似的分块思想来优化LLM训练,Flash Attention优化的是注意力层,而CCE优化语言模型的输出头,两者都是memory-efficient类的训练技巧。
CCE是一种提高LLM训练效率的技术,由于语言模型有较大的词表,而LLM模型的输出层是一个巨大的分类器,比如Llama-3
的词表达到128k, 那么其输出的就有128k个类别,那么超大的词表带来的问题是什么?
我们给定LLM解码器输出为数据长度, 维度为 的矩阵, 而分类器投影层参数 ,其中为词表大小。矩阵 经由分类器投影层 后得到矩阵, 经过行softmax后可以得到概率,通常我们会取概率对数 的形式, 与标签 计算Cross Entropy(CE) loss
那么我们可以观察到, 即LLM在训练时当文本长度增加时,概率对数的存储将会非常大,对于Gemma-2
的词表大小为, 当单条数据为长度为 时,对数概率就能塞满一个80GB的显卡。
下图中,分析Gemma2在一个示例训练时的显存占用,其对数概率(绿色区域)的占用占了显存的89%。而CCE能够魔法般的将Gemma2的对数概率占用消失,并且能够跑更大的batch, 从而加速10x训练时间,将训练时的分类头层的显存占用从28GB降低至1GB。
那么CCE是如何做到的?CCE整体前向&后向计算方式
2. 前置推导
2.1 交叉熵计算
交叉熵前向
给定序列输入向量和预测概率,那么我们有交叉熵损失:
而输入向量 经过softmax归一化后的向量。 我们化简如下, 前一项为原值,而后一项我们称为LSE(Log-Sum-Exp), 而我们将log和softmax结合,也是我们常见的log_softmax函数,它的数值更加稳定。
LSE版本的好处在于,可以独立计算损失。
交叉熵反向
我们反向求导
分别求导,而 表示当 时 , 否则为0
而LSE的求导即是softmax
那么交叉熵损失对输入的梯度为:
其中。 这里的梯度与无LSE版本无区别。
2.2 通用矩阵乘优化
计算大型矩阵乘法需要大量的显存,首先我们可以先对矩阵 行切分,对列切分,一次只计算一小块的矩阵乘法。
我们可以将E、C子矩阵通过下图上方的方式行切分成四块,即, C矩阵同理列切分,而乘法计算结果为块状
我们可以进一步在的列向量进一步划分, 那么乘积结果为,那么有:
通用矩阵乘法(GEMM)有较成熟的并行和显存优化技术。 在Flash Attention里,通过计算块状的注意力,减少HBM和SRAM的交换次数,从而提高训练效率。
3. Cut-Cross-Entropy 前向算法
3.1 矩阵乘法
我们有矩阵, 将外循环为E块,可以理解为一条数据中从序列上切分(红色),子序列里来求损失函数,即下图左边所示,但是这里的C矩阵需要完整矩阵,较耗存储。
借助分块GEMM算法,我们将内循环从向量维度d进行切分,那么可以减少C的存储,传入到SRAM进行矩阵乘法计算。
设定
V =128# 词表大小 D =16# 向量大小 BS =8# 批量大小 block_size =4# 块数量 n_b = BS // block_size# 批量块大小 d_b = D // block_size# 向量块大小 v_b = V // block_size# 词表块大小 C = torch.randn(D, V) E = torch.randn(BS, D) Y = torch.randn(BS, V) Y = torch.softmax(Y, dim =1)
实现Memory-efficient indexed matrix multiplication
# CCE : Algorithm 1 O = torch.zeros(bs, V) O_blocks = torch.split(O, n_b, dim =0)# 结果行切割 O_result = [] for e, oin zip(E_blocks, O_blocks):# i e_d_blocks = torch.split(e, d_b, dim =1) c_blocks = torch.split(C, d_b, dim =0) for ed, cdin zip(e_d_blocks, c_blocks): a = ed @ cd o = o + a O = torch.cat(O_blocks, dim =0)
3.2 CCE前向
我们写出LSE版本的矩阵乘法为, 这里的的行独立,右式 表示行求和,而 下标为 矩阵的列向量 , 而。
我们可以先在E矩阵进行切分,计算log§, 由于使用LSE的版本,那么我们需要存储行LSE。
我们将 进行列切分,那么在LSE形式下的交叉熵损失里,我们就更新全局的LSE
上述我们可以再进一步切分EC,即E列切分,C行切分
可以写出代码实现CCE,矩阵乘法能够控制在小块的计算。
# CCE : Algorithm 2 E_blocks = torch.split(E, n_b, dim =0) C_blocks = torch.split(C, v_b, dim =1) LSE = torch.ones(bs,1) *-100000.0# -inf # LSE part for n, ein enumerate(E_blocks):# E loop for v, cin enumerate(C_blocks):# C loop A = torch.zeros(n_b, v_b) e_d_blocks = torch.split(e, d_b, dim =1) c_d_blocks = torch.split(c, d_b, dim =0) for ed, cdin zip(e_d_blocks, c_d_blocks):# Ed loop, Cd loop A += ed @ cd# 块状矩阵乘法 # 计算块行LSE LSE_nv = torch.log( torch.sum( torch.exp(A), dim =1, keepdim=True ) ) # 计算行LSE LSE[n_b * n : n_b * (n+1)] = torch.log( torch.exp(LSE[n_b * n : n_b * (n+1)]) + torch.exp(LSE_nv) ) print(LSE.shape) print(LSE)
4.Cut-Cross-Entropy 反向算法
在2.1时我们计算了LSE形式下Cross Entropy loss的求导,LSE的导数即是 , 计算梯度时,需要重计算. 那么我们可以反算得到:
在反向的过程中,也依然保持小块状的矩阵计算。另外当块状的softmax重新算时值都很小,那么就忽略该块的梯度计算。事实上对于10w量级别的分类概率,大多数类别的概率为0,表现非常稀疏。
以下仅求出
部分的导数```
CCE : Algorithm 3 is_filter =False epsilon =1/4096# 2^{-12} dE = torch.zeros_like(E) dC = torch.zeros_like© dE_blocks = torch.split(dE, n_b, dim =0) dC_blocks = torch.split(dC, v_b, dim =1) print(Y.shape) d_LSE = dY_dLSE print(d_LSE) # LSE part for n, ein enumerate(E_blocks): for v, cin enumerate(C_blocks): A = torch.zeros(n_b, v_b) e_d_blocks = torch.split(e, d_b, dim =1) c_d_blocks = torch.split(c, d_b, dim =0) for ed, cdin zip(e_d_blocks, c_d_blocks): A += ed @ cd# A_nv S = torch.exp(A - LSE[n_b * n : n_b * (n+1)])# softmax # 当softmax所有值都是很小的数字,那么梯度可以忽略,减少计算量 if is_filter: if (S < epsilon).all(): print(‘skip’, n, v) continue # 梯度更新 d_LSE_block = d_LSE[n_b * n : n_b * (n+1), :] d_index =0 for ed, cdin zip(e_d_blocks, c_d_blocks): dE_blocks[n][ :, d_index * d_b: (d_index+1)*d_b ] += (d_LSE_block * S) @ cd.t() dC_blocks[v][ d_index * d_b: (d_index+1)*d_b, : ] += ed.t() @ (d_LSE_block * S) d_index = d_index +1 # dLSE gradient part dE_part_LSE = torch.cat(dE_blocks, dim =0) dC_part_LSE = torch.cat(dC_blocks, dim =1)
另外仍需要对 做混合的求导,详见paper里的Algorithm4,就不再进一步展开了。
5. 总结
-
大语言模型通常有较大的词表,当一次吞吐较多的token时,输出层的显存占用非常大,如果采用分块的思想那么可以极大减少显存占用
-
Cut-Cross-Entropy 借助LSE来将交叉熵线性化,而反向时可以从概率阈值来忽略部分块的梯度计算,并且能保证一定的精度和训练的收敛性
-
CCE除了伪代码,原作用triton进行了实现,有更多的并行化处理。
-
对于小模型<7B 规模的模型,训练时分类头存储占比高,当模型越大时CCE的收益越小。
Reference
Cut Your Losses in Large-Vocabulary Language Models
如何学习AI大模型 ?
“最先掌握AI的人,将会比较晚掌握AI的人有竞争优势”。
这句话,放在计算机、互联网、移动互联网的开局时期,都是一样的道理。
我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。
我意识到有很多经验和知识值得分享给大家,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。【保证100%免费】🆓
优快云粉丝独家福利
这份完整版的 AI 大模型学习资料已经上传优快云,朋友们如果需要可以扫描下方二维码&点击下方优快云官方认证链接免费领取 【保证100%免费】
读者福利: 👉👉优快云大礼包:《最新AI大模型学习资源包》免费分享 👈👈
对于0基础小白入门:
如果你是零基础小白,想快速入门大模型是可以考虑的。
一方面是学习时间相对较短,学习内容更全面更集中。
二方面是可以根据这些资料规划好学习计划和方向。
👉1.大模型入门学习思维导图👈
要学习一门新的技术,作为新手一定要先学习成长路线图,方向不对,努力白费。
对于从来没有接触过AI大模型的同学,我们帮你准备了详细的学习成长路线图&学习规划。可以说是最科学最系统的学习路线,大家跟着这个大的方向学习准没问题。(全套教程文末领取哈)
👉2.AGI大模型配套视频👈
很多朋友都不喜欢晦涩的文字,我也为大家准备了视频教程,每个章节都是当前板块的精华浓缩。
👉3.大模型实际应用报告合集👈
这套包含640份报告的合集,涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师,还是对AI大模型感兴趣的爱好者,这套报告合集都将为您提供宝贵的信息和启示。(全套教程文末领取哈)
👉4.大模型落地应用案例PPT👈
光学理论是没用的,要学会跟着一起做,要动手实操,才能将自己的所学运用到实际当中去,这时候可以搞点实战案例来学习。(全套教程文末领取哈)
👉5.大模型经典学习电子书👈
随着人工智能技术的飞速发展,AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型,如GPT-3、BERT、XLNet等,以其强大的语言理解和生成能力,正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。(全套教程文末领取哈)
👉6.大模型面试题&答案👈
截至目前大模型已经超过200个,在大模型纵横的时代,不仅大模型技术越来越卷,就连大模型相关的岗位和面试也开始越来越卷了。为了让大家更容易上车大模型算法赛道,我总结了大模型常考的面试题。(全套教程文末领取哈)
👉学会后的收获:👈
• 基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;
• 能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;
• 基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;
• 能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习
优快云粉丝独家福利
这份完整版的 AI 大模型学习资料已经上传优快云,朋友们如果需要可以扫描下方二维码&点击下方优快云官方认证链接免费领取 【保证100%免费】
读者福利: 👉👉优快云大礼包:《最新AI大模型学习资源包》免费分享 👈👈