HF模型权重拆解
- 模型配置文件
config.json
: 存储模型架构的元数据,包括网络层数、隐藏层维度、注意力头数、词表大小、激活函数类型等参数。 - 模型权重文件:torch和hf权重文件
- Pytorch格式
pytorch_model.pth
:用于保存模型权重参数的二进制文件,以字典形式保存,键为层的名称,值为对应的权重张量数据。 - SafeTensors格式
model.safetensor
:和pth权重文件功能相同,区别是采用安全序列化格式,避免反序列化漏洞风险。 - 分片权重文件:对于超大型模型权重,会被分割为多个文件,解决单文件体积限制问题。
- Pytorch格式
- 分词相关文件:
tokenizer_config.json
: 分词器类型(Qwen2Tokenizer
)、特殊标记([CLS]、[SEP])定义。tokenizer.json
:将分词器的所有配置信息(包括词汇表vocab、合并规则merges、归一化、预分词、后处理等)打包在一起,目的是让训练和推理过程中的分词行为完全一致。- 文本生成参数配置文件
generation_config.json
: 定义文本生成策略,包括温度(temperature)、Top-P采样等超参数。
Tokenizer.json 解析及特殊tokens
- added_tokens: 原有词汇表基础上额外添加的特殊token([PAD]、[CLS])
- post_processor: 后处理模块,设置单句、双句对输入时自动添加[CLS]和[SEP]等特殊标记。
- decoder:解码器类型。
- model:底层分词模型的类型及其参数。如qwen2.5-1.5B 模型采用BPE分词算法。
- vocab:定义了词汇表,其中key为token字符串,value为对应的id
- merges:定义了BPE的子tokens合并成更长token的规则
Special_tokens_map.json 定义与文本序列处理相关的特殊token
- bos_token:序列开始标记,文本序列的起始位置。
- eos_token:序列结束标记,表示文本序列的结束位置。
- eop_token: 段落的结束标记,表示段落结束的特殊标记。
- pad_token:填充标记,它用于将batch形式文本序列填充到相同长度时使用的特殊token。
LLM参数量计算量分析
[(masked)multi_headed_attention --> layer_normalization --> MLP --> layer_normalization] * N --> Linear --> softmax --> output probs
- LLM参数量和输入序列长度无关。对于有n层decode block块的LLM参数量为n(12h^2 + 13h)+Vh。当h较大时,可以忽略一次项,近视为12nh^2。
- 推理总需要内存约等于 1.2 x 模型参数内存(20%时经验,不同框架可能不一样)
推理优化技术
张量并行
-
集合通信是机器学习集群中实现分布式训练推理系统的基础,其作为并行计算的一个重要概念,集合通信经常被用来构建高性能的单程序流/多数据流(Single Program-Multiple Data, SPMD)程序。其本质上是一个进程组的所有进程都参与的全局通信操作。
-
常见的集合通信算子包括:Broadcast、Reduce、AllGather、Scatter和 AllReduce。
-
单点到多点:广播(broadcast)、发散(scatter);
-
多点到单点:收集(gather)、规约(reduce);
-
多点到多点:全收集(all-gather)、全规约(all-reduce)、全交换(all-to-all)、规约发散(reduce-scatter)。
-
-
Scatter是数据的1对多的分发,它将一张GPU卡上的数据进行分片再分发到其他所有的GPU卡上,反向操作对应Gather。
AllGather和AllReduce是多对多的集合通信算子,图中AllReduce算子作用是将Reduce规约函数-加法的结果存至所有设备上。
AllReduce 算子原理
通信算子 AllReduce 的目标是将所有进程的数据通过特定操作(如求和、取最大值)聚合后,将结果同步到所有进程。常见实现方式包括 Ring AllReduce 和 Tree AllReduce。
AllReduce 的实现主要分为两个过程 Reduce-Scatter 和 All-Gather。
-
分块传输:将待传输的数据切分为 N 块(N 为进程数也是设备数),形成逻辑环状结构。
-
Reduce-Scatter 阶段:每个进程依次向右邻进程发送下一块数据,并进行局部聚合(如累加、累减等)。经过 N−1 次传递后,每个设备上都有一块数据拥有了对应位置完整的聚合。
-
AllGather 阶段:每个进程聚合后的数据块都在环中传播,按照“相邻设备对应位置进行通讯”的原则,最终所有进程获得完整结果。
AllReduce 通信成本分析
- 假设有 N 个设备,输入(也是 Reduce)总数据大小为 K,一次 AllReduce 过程中,会进行 N-1 次 Scatter-Reduce 操作和 N-1 次 Allgather 操作,每一次操作所需要传递的数据大小为 K/N,所以整个 Allreduce 过程所传输的数据大小为 2(N-1) * K/N,随着 N 的增大,Ring AllReduce 通信算子的通信量可以近似为 2K。
- 分布式系统 AllReduce 的通信速度只受限于逻辑环中最慢的两个 GPU 的连接;(每次需要通信的数据大小仅为 K/N,随着 N 增大,通信量减少);
- 总结:Ring Allreduce 的通信速度恒定,和设备数量无关,完全由系统中GPU 之间最慢的连接决定。
MLP层的张量并行
将第一个线性层权重矩阵 A 沿列方向切分为 A = [A1, A2]。这种切分方式使得每个 GPU 设备能独立完成部分 GEMM 运算,并应用 GeLU 激活函数。列切分第一个线性层权重方法的优势在于避免了前向传播中一次全局同步通信(行切分方案需先同步合并分块结果,因为激活函数是非线性操作)。
MLP 的张量并行过程的形状变换公式拆解如下:
-
GPU0 上计算 Z1: [b, s, h] * [h, h/2] -> [b, s, h/2] * [h/2, h] -> [b, s, h];
-
GPU1 上计算 Z2: [b, s, h] * [h, h/2] -> [b, s, h/2] * [h/2, h] -> [b, s, h];
-
Z1 + Z2 = Z: [b, s, h] + [b, s, h] -> [b, s, h]
MLP 模块中两个线性层权重先列后行的切分方法,其张量并行的前向传播过程拆解如下:
- 第一个 GEMM(如 XA_1 和 XA_2):每个 GPU 独立计算,无需通信。
- 第二个 GEMM(如 Y_1B_1 和 Y_2B_2 之后,需要一次 All-Reduce 操作合并结果,再将结果输入 Dropout 层。
-
之所以需要 All-Reduce 归约操作,是因为第二个线性层的权重是按行切分得到 Z_1 和 Z_2,所以需要执行加法操作得到最终的 Z。
- 具体来说,第二个 GEMM(如 Y_1B_1 和 Y_2B_2) 之后,需要一次 All-Reduce 操作合并结果(对 Y_1B_1 和 Y_2B_2 求和)。
- 激活函数(如 GeLU):本地计算,无需通信。
总结:
-
为什么对 A 采用列切割,对 B 采用行切割。根本原因是尽量保证各 GPU 上的计算相互独立,从而减少通讯量;如果对 A 采用行切分,那么在执行非线性函数 GeLU 激活之前,必须做一次 AllReduce 操作,这会产生额外的通讯量。
-
通过将 MLP 模块中的两个 GEMM 拆分到多个 GPU 上独立并行计算,在正向传播中仅需一次 All-Reduce 操作(对应 g 运算符),在反向传播中也仅需一次 All-Reduce 操作(对应 f 运算符)。
-
f 与 g 运算符的前向和反向过程互为对称关系(例如,f 在前向传递中为恒等运算,反向传递中执行 All-Reduce;g 的行为则相反),这种设计可通过 PyTorch 的自动微分机制高效实现。
-
All-Reduce 的过程分为两个阶段,Reduce-Scatter 和 All-Gather,每个阶段的通讯量是相等的。如果输入张量大小为 [b, s, h],则每个阶段的通讯量为 bsh,每次 All-Reduce 操作通讯量自然为 2bsh。MLP 层在 forward(前向推理时) 做一次 All-Reduce 操作,在 backward(前向推理时) 做一次 All-Reduce 操作。
-
模型训练和推理阶段,MLP 层的张量并行通信量如下所示:
-
模型训练时,包含前向传播和反向传播两个过程,即两次 All-Reduce 操作,所以 MLP 层训练时的总通讯量为:4bsh。
- 模型推理时,只有前向传播过程,即一次 All-Reduce 操作,所以 MLP 层推理时的总通讯量为:2bsh。
-
MHA层的张量并行
在设计上,MHA 层对于每个头(head),就有都有独立的 q/k/v 三个线性变换层以及对应的 self-attention 计算结构,然后将每个 head 输出的结果做拼接 concat,最后将拼接得到结果做线性变换得到最终的注意力层输出张量。
从多头注意力结构看,其计算机制真的是天然适合模型的张量并行计算,即每个头上都可以在每个设备上独立计算,这也意味着可以把每个头(也可以是 n 个头)的参数放到一块 GPU 上,最后将子结果 concat 后得到最终的张量。
MHA层张量并行-过程
-
多头注意力结构的张量并行计算过程拆解如下:
-
对三个权重参数矩阵 W_Q、W_K、W_V,按照“列切割”,每个头放到一块 GPU 上,即每个注意力头的 Q/K/V 线性变换的矩阵乘法可以在单个 GPU 上独立完成。
-
对注意力输出线性层 B,按照“行切分”。
-
-
MHA 层的张量并行通信量如下所示:
-
模型训练时,包含前向传播和反相传播两个过程,即两次 All-Reduce 操作,所以 MLP 层的总通讯量为:4bsh。
-
模型推理时,只有前向传播过程,即一次 All-Reduce 操作,所以 MLP 层的总通讯量为:2bsh。
-