【大模型学习】快速下载大模型权重和评估大模型占用显存

### 大模型训练时显存占用量的计算方法 在大模型训练期间，显存主要被以下几个方面占据： - **模型参数**：这是指网络权重和其他可学习变量所需的内存空间。对于不同的数据类型，每种类型的字节数有所不同[^4]。 对于浮点数表示法而言，在单精度（`fp32`）下每个数值需4个字节；而在半精度（`fp16` 或 `bf16`）则只需2个字节；整型量化至8位(`int8`)仅需1个字节来表达一个值。因此，当采用较低的数据精度时能够有效减少显存量的需求。 - **激活函数产生的中间结果**：这些是在前向传播过程中由各层运算所产生的临时张量，它们会在反向传播阶段用于梯度计算并最终释放掉。这部分显存消耗取决于批次大小(batch size)以及输入序列长度等因素的影响。 - **梯度信息**：针对全量参数训练模式下的情形，整个神经网络的所有参数都会参与到梯度更新之中，这意味着除了原本就存在的模型参数外还需要额外相同体积的空间用来保存对应的梯度值。然而，在应用低秩自适应(LoRA)[^2]或其他特定策略的情况下，则无需为原生模型分配相应的梯度存储区域，因为只有新增加的小规模模块才会经历实际意义上的调整过程。 - **优化器状态**：像Adam这样的高级优化算法会维护一些辅助性的统计量比如一阶矩估计(mean of gradients) 和二阶矩估计(uncentered variance of the gradients)，这同样增加了总的GPU RAM负担。不过值得注意的是，在某些简化版或者定制化的场景里可能不会涉及到此类开销项。 综上所述，为了更精确地评估一次完整的迭代周期内的总显存需求，可以按照以下公式进行粗略估算: \[ \text{Total GPU Memory} = (\text{Model Parameters Size}) + (Batch\,Size × Sequence\,Length × Hidden\,Dimension × Precision) + (\text{Gradient Storage for Full Training or Only LoRA Part}) + (\text{Optimizer States if any}) \] 其中，“隐藏维度”指的是每一时刻处理单元内部特征的数量；而“精度”则是指上述提到的不同数值格式所带来的差异性影响因素之一。 ```python def estimate_gpu_memory_usage( model_params_size: int, batch_size: int, seq_length: int, hidden_dim: int, precision_bytes: float, full_training=True): activation_mem = batch_size * seq_length * hidden_dim * precision_bytes gradient_mem = model_params_size * precision_bytes \ if full_training else 0 optimizer_state_mem = model_params_size * 1.2 * precision_bytes \ if not isinstance(optimizer_states_multiplier, type(None)) else 0 total_memory = sum([ model_params_size * precision_bytes, # Model parameters memory activation_mem, # Activation memory during forward pass gradient_mem, # Gradient storage optimizer_state_mem # Optimizer states memory ]) return total_memory / (1024 ** 3) # Convert bytes to GBs ```