transformers GPU 训练内存解剖从权重到激活的显存构成与运算密度分析【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers本文以 Hugging Face Transformers 官方文档《训练过程解剖》docs/source/ar/model_memory_anatomy.md为主体完整复现用 pynvml 监控 GPU 显存 → 加载 BERT-large → 基线训练 → 拆解显存构成的实验链路并结合仓库中 TrainingArguments 与 Trainer 的源码实现解释训练时每参数 18 字节显存开销混合精度 AdamW从何而来帮助读者建立哪里吃显存、哪里可优化的工程判断力。一、为什么需要解剖训练时的 GPU 显存用较小的批量大小训练模型显存往往就能被填满但模型本身只占一部分显存这一事实常让人困惑把 BERT-large 加载到 GPU 上只占约 1.3GB基线训练却要占满 16GB 的 V100。要回答训练究竟多占了什么显存必须先理解两件事训练过程中 GPU 上执行了哪些运算以及这些运算各自对显存的需求。官方文档给出的方法论就是先实测各阶段显存占用再逐项拆解占用来源。二、实验环境搭建用 pynvml 从 Python 内部监控显存文档第一步是安装监控与训练所需的库pip install transformers datasets accelerate nvidia-ml-py3其中nvidia-ml-py3即 pynvml允许在 Python 进程内直接读取 NVIDIA 设备的显存信息等价于在终端运行nvidia-smi。文档定义了核心训练所需的两个辅助函数 import numpy as np from datasets import Dataset seq_len, dataset_size 512, 512 dummy_data { ... input_ids: np.random.randint(100, 30000, (dataset_size, seq_len)), ... labels: np.random.randint(0, 1, (dataset_size)), ... } ds Dataset.from_dict(dummy_data) ds.set_format(pt)这里构造了 512 条长度 512 的序列token ID 取自 10030000 的随机整数标签为二分类。set_format(pt)使数据集直接以 PyTorch 张量格式喂给 Trainer与 Trainer.train 的标准数据接口兼容。 from pynvml import * def print_gpu_utilization(): ... nvmlInit() ... handle nvmlDeviceGetHandleByIndex(0) ... info nvmlDeviceGetMemoryInfo(handle) ... print(fGPU memory occupied: {info.used//1024**2} MB.) def print_summary(result): ... print(fTime: {result.metrics[train_runtime]:.2f}) ... print(fSamples/second: {result.metrics[train_samples_per_second]:.2f}) ... print_gpu_utilization()print_gpu_utilization读取 0 号设备的used显存print_summary则从trainer.train()返回的result.metrics中取出train_runtime与train_samples_per_second两项指标这两个键是 Trainer 训练结束回调中写回的标准指标名。文档同时给出一个重要实验卫生准则计划跑多个实验时应在实验之间重启 Python 内核以确保显存被正确释放。三、基线测量空载显存与 CUDA 上下文开销实验首先确认起点是干净的 print_gpu_utilization() GPU memory occupied: 0 MB.如果本地不是 0MB需先关闭所有占用该 GPU 的进程。接着文档揭示了一个容易被忽略的开销——CUDA 上下文context。哪怕只把一个1×1的小矩阵搬到 GPU 上 import torch torch.ones((1, 1)).to(cuda) print_gpu_utilization() GPU memory occupied: 1343 MB.可以看到仅初始化 CUDA 上下文就消耗约 1.3GB 显存。也就是说用户并不能使用一块空 GPU的全部显存这解释了为什么框架报告的显存占用总是比模型参数量算出来的理论值大一截。四、加载模型权重本身占多少接下来把google-bert/bert-large-uncased直接加载到 GPU只测权重占用 from transformers import AutoModelForSequenceClassification model AutoModelForSequenceClassification.from_pretrained(google-bert/bert-large-uncased).to(cuda) print_gpu_utilization() GPU memory occupied: 2631 MB.2631MB 减去上下文开销 1343MB约 1.3GB 即 BERT-large约 3.35 亿参数fp32 权重的占用。文档提醒具体数值随 GPU 型号而变且较新架构上权重可能因加载时做了针对加速的内存布局优化而占更多空间。可用命令行nvidia-smi交叉验证同一进程的真实占用。文档记录到的 V10016GB输出节选| 0 Tesla V100-SXM2... On | 00000000:00:04.0 Off | 0 | | N/A 37C P0 39W / 300W | 2631MiB / 16160MiB | 0% Default |Python 侧pynvml读到的 2631MB 与nvidia-smi完全一致验证了监控方法的有效性。五、基线训练小批量为何填满 16GB 显存设置一组标准训练参数后用 [Trainer] 以批量大小 4、不使用任何 GPU 优化技术进行训练default_args { output_dir: tmp, eval_strategy: steps, num_train_epochs: 1, log_level: error, report_to: none, }对照源码可以核对这些参数的含义training_args.py 中per_device_train_batch_size的默认值是 8field(default8, ...)文档显式传 4 正是为了得到可复现的小批量基线per_device_train_batch_size * num_devices * gradient_accumulation_steps共同决定有效批量。 from transformers import TrainingArguments, Trainer, logging logging.set_verbosity_error() training_args TrainingArguments(per_device_train_batch_size4, **default_args) trainer Trainer(modelmodel, argstraining_args, train_datasetds) result trainer.train() print_summary(result)文档记录的实测结果Time: 57.82 Samples/second: 8.86 GPU memory occupied: 14949 MB.关键观察批量大小仅 4 的小训练就占用了 14.9GB接近 V100 16GB 的满容量远超 1.3GB 的模型权重。而更大的批量通常带来更快收敛或更好的最终效果——我们希望按模型需求而非 GPU 限制来调批量。要多占的这部分显存到底去了哪需要解剖训练过程的运算与显存构成。六、解剖模型运算按计算密度分三层官方文档将 Transformer 的运算按算术强度arithmetic intensity分为三类这一分层是分析性能瓶颈的起点矩阵乘法matmul线性层与多头注意力的 QKV/输出投影、前馈层全部是对批量做矩阵-矩阵乘法。这是 Transformer 训练中计算密度最高的部分最适合 GPU 的大规模并行。统计归约运算reductionSoftmax 与 LayerNorm 等。它们先跨一个维度做归约算出统计量如均值、方差再map回原张量计算密度低于 matmul且需要多次访问显存。逐元素运算element-wise偏置、Dropout、激活函数、残差连接等是计算密度最低的一类。该分层源自 2020 年论文《Data Movement Is All You Need: A Case Study on Optimizing Transformers》arXiv:2007.00072其核心论点是优化 Transformer 性能的关键在于减少数据搬运而非增加运算。七、解剖显存构成六个组成部分逐项量化训练时 GPU 显存上共存放六类内容文档给出了每类的字节成本1. 模型权重Model weightsfp32 训练4 字节 × 参数量混合精度训练6 字节 × 参数量——因为内存中同时保留一份 fp32 主权重用于稳定的权重更新和一份 fp16 副本用于前向/反向传播。2. 优化器状态Optimizer states标准 AdamW8 字节 × 参数量分别保存一阶动量与二阶方差两个 fp32 张量8 位量化 AdamW如 bitsandbytes 实现2 字节 × 参数量。这一点在仓库中有对应的集成实现quantizer_bnb_8bit.py 与quantizer_bnb_4bit.py等文件即为 BitsAndBytes 8/4 位量化的接入层说明用量化优化器压显存在 Transformers 中是可落地的路径SGD momentum 一类只保存一个状态的优化器4 字节 × 参数量。3. 梯度Gradients无论 fp32 还是混合精度训练梯度始终按 fp32 保存4 字节 × 参数量。4. 前向激活Forward activations大小取决于序列长度、隐藏维度、批量大小与模型深度是最难估算的变量。这是为什么模型放得下训练却 OOM长序列或大批量会把激活撑爆显存。5. 临时缓存Temporary cachessoftmax、矩阵乘法的中间 scratch 张量等。它们在运算结束后即被释放但峰值时刻可能叠加出额外显存尖峰直接触发 OOM。文档强调写代码时要策略性地考虑这些临时变量必要时在不再需要时显式del释放。6. 功能特定内存Feature-specific memory例如 beam search 生成文本时要同时保留多份输入/输出大词表的 embedding 表本身也可能是显存大户。每参数字节账本文档给出的关键结论用混合精度 AdamW 训练一个典型模型每参数需要 18 字节6 权重 8 优化器状态 4 梯度外加激活显存推理时没有优化器状态与梯度每参数只需 6 字节混合精度外加激活。英文同主题文档docs/source/en/model_memory_anatomy.md给出了一个规模化示例在批量大小 16 下混合精度训练一个 4B 参数模型约需 85GB 显存——权重 6B/param、Adam 状态 8B/param、梯度 4B/param、激活随 batch×seq_len×depth×hidden 浮动量化 Adam 可用 bitsandbytes 把优化器状态压到 2B/param与上述账本完全吻合。组成部分训练混合精度推理混合精度权重6 字节/参数6 字节/参数AdamW 状态8 字节/参数—梯度4 字节/参数—激活随批量/序列长度/深度变化随输入变化临时张量/功能特定内存峰值尖峰可致 OOMbeam search 等可致多份副本八、前向 vs 反向为什么 backward 大约慢 2 倍文档最后指出运算速度的不对称性对卷积与线性层反向传播的浮点运算量是前向的 2 倍通常体现为约 2 倍的耗时有时更多因为反向涉及的张量尺寸分布对硬件更不友好激活类运算多为带宽受限bandwidth-bound。例如一个激活函数的前向只需读一次、写一次反向却要读gradOutput与前向输出两份、再写gradInput一份读放大使反向更慢。这两点共同决定了优化空间显存上可压缩的是优化器状态与激活速度上可优化的是归约与逐元素运算的访存——对应到仓库中就是梯度累积、梯度检查点、混合精度与自定义融合内核这几类现成手段。九、下一步仓库中可用的显存优化手段理解了显存解剖之后文档指引读者结合训练优化专题继续深入。当前仓库英文文档中对应且仍然存在的指南包括梯度累积grad_accumulation在不增加峰值显存的前提下模拟更大的有效批量梯度检查点grad_checkpointing以重算换显存只保留部分中间激活混合精度训练mixed_precision_training把前向/反向放到低精度把权重主副本留在 fp32正是前文6 字节/参数的来源自定义融合内核kernels加速归约与逐元素类运算。需要说明的是阿拉伯语文档末尾引用的perf_train_gpu_one单卡高效训练页面在当前 阿拉伯语目录树 中已被注释停用实际查阅时以上述英文指南与仓库源码training_args.py、trainer.py为准。小结本文完整继承了官方训练过程解剖文档的实验主线pynvml 显存监控 → 空载/上下文/权重三级基线0MB → 1343MB → 2631MB→ 批量 4 基线训练实测 14949MB并落到两张账本上——运算按计算密度分 matmul / 归约 / 逐元素三层显存按权重6B、AdamW 状态8B、梯度4B、激活、临时缓存与功能特定内存六项拆解得出混合精度 AdamW 训练 18 字节/参数、推理 6 字节/参数的核心结论。掌握这套解剖方法后任何一次训练 OOM 都可以沿着是权重、状态、梯度还是激活的路径定位并选择梯度累积、检查点、低精度或量化优化器这类对症下药的优化。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
