菜叶子芯酸笔记4：大模型训练、分布式训练、显存估算
大模型训练任务主要分为以下三种模型训练过程。
预训练pretrain
监督微调 supervised finetune training
奖励模型 reward model
RLHF
它们之间的顺序联系用RLHF (reinforcement learning with human feedback) 过程来阐释。
首先预训练pretrain得到一个base模型。
到微调阶段，用SFT训练得到一个Actor模型，同时拷贝得到pef_actor模型，为Actor的拷贝，pef_actor不做训练参数优化，用于防止信息遗忘，利用KL散度计算对输出logits做分布近似。
到奖励模型训练阶段，用RM训练方法得到奖励模型，同时拷贝得到critic模型，为RM的拷贝，critic不做训练参数优化，用于提升模型对数据偏好，这里会用到PPO算法。也可以直接用DPO，用偏好数据直接优化。
policy/Actor模型结构。
框图就是RLHF过程。prompt进入tokenize之后。
prompt送入old_policy/ Actor模型，得到response和old_log_prob。这里的输出是迭代反复的过程，因为输出预测的token是连续的，串行的。同时old_values也在更新。至今还是不懂这怎么更新的。
这个博客是这么说的。持有怀疑态度。
prompt 输入到Ref_policy / ref_actor模型里，得到ref_log_probs。【非必需】
prompt + response 输入到critic模型里，得到score。
计算advatanges= reward - old_values。
ref：从0开始实现LLM：7.1、Reward/PPO/DPO/KTO/SimPO详解_dpo 训练、orpo 训练和 simpo 训练-CSDN博客
计算loss。 这里a就是advantages。
更新old_policy模型。
RM训练过程：RLHF之前。或者换个说法 old_policy模型的预训练
偏好数据：
在训练时会把conversations+chosen作为chosen_messages，conversations+rejected作为rejected_messages。如果存在system，则将system拼接到conversations前。将两个messages分别转换为token_ids作为输入（此时数据格式为[bs=2n, max_seq_len]，n是原本设置的bs）。
而label值是将messages的system+conversations部分的token_id置为-100。
将token_ids输入模型，得到output_values，将其拆成chosen_rewards和rejected_rewards。
对于每个输入对，获取system+prompt长度，截取response+padding部分的reward_value。
数据输入：在训练时会把conversations+chosen作为chosen_messages，conversations+rejected作为rejected_messages。如果存在system，则将system拼接到conversations前。将两个messages分别转换为token_ids作为输入（此时数据格式为[bs=2n, max_seq_len]，n是原本设置的bs）。
而label值是将messages的system+conversations部分的token_id置为-100。
将token_ids输入模型，得到output_values，将其拆成chosen_rewards和rejected_rewards。
对于每个输入对，获取system+prompt长度，截取response+padding部分的reward_value。
其中输入最后一个token得到的reward_value结果，作为Reward模型的最终打分。计算损失。
输入输出数据示意图
考虑到每一个chosen_data的token打分都应该比rejected_data的token打分更高。因此在LLM训练阶段，loss使用了全部的response token的loss平均。
Loss
对于chosen和rejected的数据对形式来说，自然是chosen的得分越高越好，rejected的得分越低越好，那么自然是loss=chosen_logits - rejected_logits，此时loss越大越好。
通常来说loss需要在[-1,1]或着[0,1]之间，那么可以变成loss=sigmod(chosen-rejected)。由于在数据非常大或者非常小时，sigmod函数会存在数值溢出问题，且在结果0时容易存在梯度消失问题，可以将sigmod换成logsigmod函数。又考虑到训练时需要loss越低越好，那么可以对结果取负。最终loss公式如下
不同任务的数据格式
pretrain
pretrain 的目的是对模型预训练，使得模型具备基础的知识。这里我们可以把知识理解为记忆面包，统统一股脑喂给模型即可。
1.数据样式
以 wiki_demo.txt 为例，其中每一行信息就是我们上面说到的记忆面包的知识。PT 也是一种自学习的方法，就像之前版本给出的那样，PT 的样本 source 和 target 是一样的。
实现就是设定长度，组队source和target 一样长。
SFT
SFT 有监督微调的数据以 QA 的形式提供，其中 instruction 可以作为 prompt 使用即 Q，如果 instrution 和 input 都有值，则逻辑会把二者拼接在一起作为统一个 Source。output 为 A ，在有监督微调的情况下，前面的 Prompt 会以 mask 的形式进行遮蔽，从而构造 label_ids。
RM
和前面 sft 的逻辑比较相似，sft 只有 source 和 target，这里 prompt 相当于 source，chosen_ids 相当于 positive_target 的 token ids，rejected_ids 相当于 negative_target 的 token ids，最终全部添加至 model_inputs 即可。原始样本的格式我们虽然未给出，但是大家可以自行构建，只需要在 sft 的基础上增加 bad case 即可，不过这一步需要有正负情感的数据标注。
ref：LLM - 数据处理之 Process Dataset For LLM With PT、SFT、RM_llm sft-CSDN博客
分布式计算
显存需求取决于多个因素，包括但不限于以下几点：
1. 模型参数的数量：这是显存需求的主要部分。
2. 激活（中间层输出）：模型在前向传播和反向传播过程中产生的中间结果也需要存储在显存中。
3. 优化器状态：如果使用 Adam 或其他需要额外状态的优化器，这些状态也会占用显存。
4. 批量大小（Batch Size）：更大的批量大小会增加显存需求。
5. 序列长度（Sequence Length）：更长的输入序列会增加显存需求。
估算显存需求的方法 llamaXB 需要3XGB
1. 模型参数
每个参数通常占用 4 字节（对于 float32 类型）。如果是混合精度训练（如 float16），每个参数占用 2 字节。
* Llama 2-7B：
* 参数数量：7,000,000,000
* float32 占用：7,000,000,000 * 4 字节 = 28 GB
* float16 占用：7,000,000,000 * 2 字节 = 14 GB
* Llama 2-13B：
* 参数数量：13,000,000,000
* float32 占用：13,000,000,000 * 4 字节 = 52 GB
* float16 占用：13,000,000,000 * 2 字节 = 26 GB
* Llama 2-30B：
* 参数数量：30,000,000,000
* float32 占用：30,000,000,000 * 4 字节 = 120 GB
* float16 占用：30,000,000,000 * 2 字节 = 60 GB
2. 激活
激活的显存需求取决于模型的结构和层数。假设每层激活的显存需求为 S，则总的激活显存需求可以近似为 S * 层数 * 批量大小 * 序列长度。
3. 优化器状态
对于 Adam 优化器，每个参数需要额外的两个状态变量（梯度的平方和梯度的一阶矩），因此每个参数需要额外的 8 字节（float32）或 4 字节（float16）。
* Llama 2-7B：
* float32 占用：7,000,000,000 * 8 字节 = 56 GB
* float16 占用：7,000,000,000 * 4 字节 = 28 GB
* Llama 2-13B：
* float32 占用：13,000,000,000 * 8 字节 = 104 GB
* float16 占用：13,000,000,000 * 4 字节 = 52 GB
* Llama 2-30B：
* float32 占用：30,000,000,000 * 8 字节 = 240 GB
* float16 占用：30,000,000,000 * 4 字节 = 120 GB
4. 总显存需求
总显存需求是上述各项的总和。假设批量大小为 1，序列长度为 2048，且不考虑其他额外开销，我们可以得到一个大致的估计。
* Llama 2-7B：
* 参数：14 GB (float16)
* 激活：假设每层激活约 1 MB，则总激活约为 1 MB * 32 层 * 1 * 2048 ≈ 64 MB
* 优化器状态：28 GB (float16)
* 总显存：14 GB + 64 MB + 28 GB ≈ 42 GB
* Llama 2-13B：
* 参数：26 GB (float16)
* 激活：假设每层激活约 1 MB，则总激活约为 1 MB * 32 层 * 1 * 2048 ≈ 64 MB
* 优化器状态：52 GB (float16)
* 总显存：26 GB + 64 MB + 52 GB ≈ 78 GB
* Llama 2-30B：
* 参数：60 GB (float16)
* 激活：假设每层激活约 1 MB，则总激活约为 1 MB * 32 层 * 1 * 2048 ≈ 64 MB
* 优化器状态：120 GB (float16)
* 总显存：60 GB + 64 MB + 120 GB ≈ 180 GB
并行策略
数据并行空间复杂度
ADAM 优化器 + 混合精度训练情况下模型状态的显存占用
综上所述，训练过程中的显存占用可分为两大部分：
模型状态：记模型本身参数量为 Φ ，在 Adam + 混合精度训练的情况下，模型状态包括 fp16 的模型参数 2 Φ 和参数梯度 2 Φ 和 fp32 的模型参数备份 4 Φ ，momentum 4 Φ 和 variance 4 Φ ，即总共 2 Φ + 2 Φ + 4 Φ + 4 Φ + 4 Φ = 16Φ 。（注意 fp16 占两个字节，fp32 占四个字节）
剩余状态：即训练中的激活值、临时缓冲区和显存碎片等。
以 GPT-2 为例，GPT-2 模型含有 1.5B 个参数，如果用 fp16 格式，模型本身只占 3GB 显存，但是实际训练过程中的模型状态需要耗费 24GB【1.5*16】！可以看到。模型状态是成倍于模型本身的大小，是显存消耗的大头。并且，对于剩余状态中的激活值等，已经有 activation checkpointing 等以时间换空间的优化方式，可以有效减小这部分显存消耗。因此，优化模型状态的显存占用是重点。
ZeRO 由 ZeRO-DP 和 ZeRO-R 组成，分别是对模型状态和剩余状态的显存优化。
ref : 分布式训练数据并行极致优化：ZeRO_zero分布式训练-CSDN博客
#TODO：
#smooth quant
#deepspeed megatron
#mixtral MOE
#vllm tensorrt-llm
原文地址：https://blog.csdn.net/qq_64955200/article/details/143794300
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！