【ICLR高分10/10/8/6】天塌了!LLM训练交叉熵吃这么多显存??

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. 总结

  1. 大语言模型通常有较大的词表,当一次吞吐较多的token时,输出层的显存占用非常大,如果采用分块的思想那么可以极大减少显存占用

  2. Cut-Cross-Entropy 借助LSE来将交叉熵线性化,而反向时可以从概率阈值来忽略部分块的梯度计算,并且能保证一定的精度和训练的收敛性

  3. CCE除了伪代码,原作用triton进行了实现,有更多的并行化处理。

  4. 对于小模型<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籍就是非常不错的学习资源。(全套教程文末领取哈)
img

在这里插入图片描述

👉6.大模型面试题&答案👈

截至目前大模型已经超过200个,在大模型纵横的时代,不仅大模型技术越来越卷,就连大模型相关的岗位和面试也开始越来越卷了。为了让大家更容易上车大模型算法赛道,我总结了大模型常考的面试题。(全套教程文末领取哈)

在这里插入图片描述
👉学会后的收获:👈
基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;

能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;

基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;

能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习

优快云粉丝独家福利

这份完整版的 AI 大模型学习资料已经上传优快云,朋友们如果需要可以扫描下方二维码&点击下方优快云官方认证链接免费领取 【保证100%免费】

读者福利: 👉👉优快云大礼包:《最新AI大模型学习资源包》免费分享 👈👈

(👆👆👆安全链接,放心点击)
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值