大模型微调中显存占用和训练时间的影响因素

BatchSize

显存占用：与batch_size呈线性关系，可理解为 $M_{total}=M_{fixed}+BatchSize*M_{per-sample}$ ，其中 $M_{fixed}$ 指的是模型本身固定占用的显存（由参数数量决定）和优化器状态（也由参数数量决定）

总训练时间：理论上与BatchSize无关（总数不变，单步训练时间增加，总步数减少），但实际中随BatchSize越大，总时间可能减少（硬件并行效率提升），直到显存或硬件并行能力达到瓶颈。

截断长度（输入序列分词后的最大长度，即每条样本被大模型读取的最大长度）

1. 显存占用

在大型语言模型（如 Transformer）中，显存占用主要与模型的激活值（Activations）有关，而激活值的大小受到输入序列长度（即截断长度）的直接影响。以下是逐步分析：

激活值的定义

激活值是指模型在正向传播过程中每一层计算出的中间结果，通常存储在显存中，以便反向传播时计算梯度。对于 Transformer 模型，激活值主要与注意力机制（Self-Attention）和前馈网络（Feed-Forward Network, FFN）的计算相关。

显存占用的组成

显存占用主要包括：

模型参数（权重和偏置）：与模型规模（层数、隐藏维度）相关，与截断长度无关。
激活值：与输入序列长度（截断长度 $L$ ）、批次大小（batch size $B$ ）、隐藏维度（hidden size $H$ ）和层数（ $N$ ）成正比。
梯度（训练时）：与参数量和激活值大小相关。

对于激活值部分，显存占用主要来源于：

注意力机制：计算 $\cdot K^T$ 的注意力分数矩阵，尺寸为 $(B, L, L)$ ，每层需要存储。
中间张量：如 $V$ 的加权和、前馈层的输出等。

数学表达式

假设： $L$ ：截断长度（序列长度）， $B$ ：批次大小， $H$ ：隐藏维度， $N$ ：模型层数， $P$ ：浮点数精度（如 FP32 为 4 字节，FP16 为 2 字节）

激活值的显存占用近似为：
$\text{显存}_{\text{激活值}} \approx N \cdot B \cdot L \cdot H \cdot P + N \cdot B \cdot L^2 \cdot P$

第一项 $\cdot B \cdot L \cdot H \cdot P$ ：表示每层的线性张量（如 $Q, K, V$ 或 FFN 输出）的显存占用。
第二项 $\cdot B \cdot L^2 \cdot P$ ：表示注意力分数矩阵的显存占用（仅在标准注意力机制中显著，若使用优化如 FlashAttention，则可能减少）。

结论：显存占用与截断长度 $L$ 呈线性（ $O (L)$ ）到二次方（ $O(L^2)$ ）的关系，具体取决于注意力机制的实现方式。

2. 训练时间

训练时间主要与计算量（FLOPs，浮点运算次数）和硬件并行能力有关，而截断长度会影响计算量。

计算量的组成

注意力机制：每层的计算量与 $L^2$ 相关，因为需要计算 $\times L$ 的注意力矩阵。
前馈网络：每层的计算量与 $L$ 线性相关，因为对每个 token 独立计算。

总计算量（FLOPs）近似为：
$\text{FLOPs} \approx N \cdot B \cdot (2 \cdot L^2 \cdot H + 4 \cdot L \cdot H^2)$

$\cdot L^2 \cdot H$ ：注意力机制的矩阵乘法（如 $\cdot K^T$ 和 $\text{softmax} \cdot V$ ），
$\cdot L \cdot H^2$ ：前馈网络的计算（假设 FFN 隐藏层维度为 $4 H$ ）。

训练时间

训练时间与 FLOPs 成正比，同时受硬件并行能力（如 GPU 的计算核心数）影响。假设每秒浮点运算能力为 $F_{\text{GPU}}$ （单位：FLOPs/s），则单次前向+反向传播的训练时间为：
$\text{时间} \approx \frac{\text{FLOPs}}{F_{\text{GPU}}} \approx \frac{N \cdot B \cdot (2 \cdot L^2 \cdot H + 4 \cdot L \cdot H^2)}{F_{\text{GPU}}}$