1. 显存估算的核心逻辑与常见误区
1.1 为什么LoRA微调的显存总是不够用
很多人第一次跑LoRA微调,看到网上说“7B模型LoRA只要16GB显存”就兴冲冲上手,结果OOM(Out of Memory)报错糊脸。问题出在哪?显存占用从来不是一个单一数字,它是一条链路上多个环节叠加的结果。你看到的“16GB能跑”往往是在特定配置下的极限值,换一个序列长度、换一个batch size、换一个优化器,数字就完全不一样了。
先把这条链路拆开。LoRA微调时的显存消耗主要来自以下几块:基座模型权重、LoRA适配器参数、优化器状态、梯度、前向传播的中间激活值、CUDA上下文与碎片开销。这六块里面,前两块相对固定,后面几块跟你的训练配置强相关,波动非常大。
拿一个7B参数的模型举例。如果基座模型用FP16加载,权重占14GB左右。LoRA适配器通常只占几十MB到几百MB,这部分可以忽略不计。但优化器状态就不一样了——如果你用AdamW,它需要为每个可训练参数维护一阶矩和二阶矩,每个都是FP32精度。LoRA的可训练参数虽然少,但如果你不小心把modules_to_save配多了,或者用了全参数微调的思路去配LoRA,优化器状态就会膨胀。
真正的大头往往在中间激活值上。序列长度512和2048,激活值占用可能差4倍以上。batch size从1调到4,激活值线性增长。这就是为什么同样的模型,别人能跑你不能跑——配置不同,显存需求天差地别。
1.2 显存估算的实用公式与快速心算
我不喜欢给一个“万能公式”然后让大家往里套,因为实际场景里变量太多。但可以给一个快速估算框架,帮你在动手之前心里有个底。
基座模型权重占用(GB)≈ 参数量(B)× 精度字节数 ÷ 1024
- FP16/BF16:每参数2字节
- FP32:每参数4字节
- INT8:每参数1字节
- INT4:每参数0.5字节
7B模型FP16加载:7 × 2 = 14GB。13B模型FP16:26GB。这样一算,32GB GPU跑13B FP16基座就已经很紧张了,因为还要留空间给激活值和优化器。
LoRA可训练参数占用 ≈ 基座参数量 × LoRA rank × 2 × 目标模块数 ÷ 模型隐藏维度
这个公式看起来复杂,实际心算可以简化:对于7B模型,rank=8、target_modules为q_proj和v_proj时,可训练参数大约在400万到800万之间。FP32的优化器状态就是每个参数8字节(一阶矩4字节+二阶矩4字节),800万参数也就64MB,很小。
激活值占用才是真正的变量。一个粗略的经验公式:
激活值(GB)≈ batch_size × seq_len × hidden_size × num_layers × 精度字节数 × 系数
这个系数通常在2到6之间,取决于是否使用gradient checkpointing、flash attention等优化。不用任何优化时系数偏高,用了gradient checkpointing可以降到1左右,但会牺牲约20%到30%的训练速度。
实际心算时,我习惯用这个经验值:7B模型、seq_len=512、batch_size=1、FP16、开启gradient checkpointing,激活值大约2到3GB。如果seq_len拉到2048,激活值可能到8到10GB。如果batch_size=4,再乘4。
把这几块加起来:14GB(权重)+ 0.1GB(LoRA)+ 0.1GB(优化器)+ 2到10GB(激活值)+ 1到2GB(CUDA上下文和碎片)=17到26GB。这就是为什么32GB GPU跑7B LoRA比较从容,但跑13B就需要精打细算。
1.3 那些容易被忽略的显存黑洞
有几个地方特别容易吃显存,但很多人排查时想不到。
第一个是数据加载和预处理。如果你用HuggingFace的datasets库,默认会把整个数据集缓存到内存里,但某些操作(比如tokenize时的map函数)如果num_proc设得太大,会开多个进程,每个进程都占一份内存和显存映射。更隐蔽的是,如果dataloader_num_workers设得过高,每个worker都会复制一份数据,显存占用可能翻倍。
第二个是评估阶段。很多人训练时显存刚好够,一跑evaluation就OOM。因为评估时虽然不存梯度,但如果你没设eval_accumulation_steps,模型会把所有预测结果堆在显存里,等评估完再一起算指标。序列一长、样本一多,直接爆掉。
第三个是日志和回调。像wandb、tensorboard这些工具,默认会记录很多中间张量。如果你开了log_gradients或者log_weights,它们会把梯度或权重复制到CPU或显存里做直方图,显存占用会突然飙升。
第四个是CUDA上下文本身。一个空的CUDA上下文在PyTorch里大约占300到500MB。如果你用了多个进程(比如DDP),每个进程都有自己的上下文。8卡训练时,光上下文就吃掉2到4GB。
注意:排查OOM时,先用
torch.cuda.memory_summary()看当前分配和保留的显存,再用nvidia-smi看整体占用。两者对不上时,往往是碎片或缓存问题。
2. 32GB GPU上的LoRA训练配置实战
2.1 模型加载精度与量化策略选择
32GB显存说大不大,说小不小。跑7B模型LoRA很舒服,跑13B需要动脑筋,跑34B基本没戏(除非上量化)。所以第一步是决定基座模型用什么精度加载。
FP16/BF16全精度加载是最省心的,兼容性最好,训练速度也快。7B模型占14GB,13B占26GB。13B的情况下,剩下6GB给激活值和优化器,只够跑seq_len=512、batch_size=1的配置,而且必须开gradient checkpointing。
INT8量化加载可以把权重占用减半。7B模型降到7GB,13B降到13GB。但INT8加载需要bitsandbytes库,而且训练时只有LoRA部分是可训练的,基座权重是冻结的。这里有个坑:INT8加载的模型,前向传播时会把权重反量化回FP16做计算,所以激活值占用并不会减少,只是权重占用少了。另外INT8训练的速度通常比FP16慢20%到40%,因为多了量化/反量化的开销。
INT4量化加载(比如QLoRA的思路)可以把7B模型压到3.5GB,13B压到6.5GB。这样32GB GPU跑13B LoRA就非常宽裕了,甚至能跑34B。但INT4的精度损失比INT8明显,特别是对于需要精细理解的任务(比如代码生成、数学推理),效果下降可能比较明显。而且INT4训练速度更慢,通常比FP16慢50%以上。
我的建议是:7B模型直接用FP16,13B模型优先考虑INT8,34B模型才上INT4。如果你对训练速度不敏感,或者显存实在紧张,再往下调。
| 加载精度 | 7B权重占用 | 13B权重占用 | 34B权重占用 | 训练速度 | 精度损失 |
|---|---|---|---|---|---|
| FP16/BF16 | 14GB | 26GB | 68GB | 基准 | 无 |
| INT8 | 7GB | 13GB | 34GB | 慢20-40% | 轻微 |
| INT4 | 3.5GB | 6.5GB | 17GB | 慢50%+ | 中等 |
2.2 LoRA rank与target_modules的取舍
LoRA的rank决定了适配器的表达能力。rank越大,可训练参数越多,模型能学到的信息越丰富,但显存和计算开销也越大。常见的rank取值是4、8、16、32、64。
rank=8是默认起点,适合大多数指令微调任务。rank=16到32适合更复杂的任务,比如多轮对话、长文本理解。rank=64以上通常只在数据量很大、任务很复杂时才需要,而且收益递减明显。
target_modules的选择更关键。最保守的做法是只对q_proj和v_proj加LoRA,这是原始LoRA论文的配置。但实践发现,对k_proj、o_proj、gate_proj、up_proj、down_proj也加LoRA,效果通常更好,特别是对于需要模型改变行为模式的任务。
但每多一个target module,可训练参数就多一份。以7B模型为例:
- 只加q_proj和v_proj:约400万可训练参数
- 加上k_proj和o_proj:约800万
- 全部7个模块都加:约2000万
2000万参数在FP32下的优化器状态是160MB,其实也不算大。真正的影响在计算图:更多的LoRA层意味着前向和反向传播时要多算很多矩阵乘法,激活值也会相应增加。实测下来,全模块LoRA比只加q/v的激活值占用高15%到25%。
我的经验是:如果显存充裕,直接上全模块LoRA,rank=16。如果显存紧张,先保q_proj和v_proj,rank=8。不要为了省显存把rank降到4以下,效果损失往往得不偿失。
2.3 batch size与gradient accumulation的配合
batch size是显存占用的线性放大器。batch_size=1和batch_size=8,激活值占用差8倍。但batch size太小会导致训练不稳定,梯度噪声大。这时候就需要gradient accumulation来救场。
gradient accumulation的思路是:用小的micro-batch跑多次前向和反向,把梯度累加起来,等累加到一定步数再更新一次参数。这样等效于更大的batch size,但显存占用只跟micro-batch大小有关。
配置示例:
per_device_train_batch_size = 2 gradient_accumulation_steps = 8 # 等效batch size = 2 × 8 × GPU数量在32GB GPU上跑7B LoRA,我通常这样配:
- seq_len=1024:micro-batch=2,grad_accum=8,等效batch=16
- seq_len=2048:micro-batch=1,grad_accum=16,等效batch=16
- seq_len=4096:micro-batch=1,grad_accum=32,等效batch=32
注意gradient accumulation不会减少激活值占用,因为每次前向传播的中间结果还是要存着做反向。它只是让你能用小batch跑出大batch的效果。真正减少激活值要靠gradient checkpointing。
2.4 gradient checkpointing的代价与收益
gradient checkpointing(也叫activation checkpointing)的原理是:前向传播时不保存中间激活值,只保存几个检查点。反向传播时,从最近的检查点重新计算需要的激活值。这样显存占用大幅降低,但计算量增加,训练速度变慢。
在HuggingFace Trainer里,开启方式很简单:
training_args = TrainingArguments( gradient_checkpointing=True, ... )实测数据:7B模型、seq_len=2048、batch_size=2,不开checkpointing激活值约12GB,开了之后降到3GB左右。但训练速度从每秒2.5个step降到1.8个step,慢了约28%。
所以这是一个显存换时间的权衡。32GB GPU跑7B模型时,如果seq_len不超过1024,其实可以不开checkpointing,速度更快。跑13B或者seq_len超过2048时,checkpointing基本是必开的。
提示:开启gradient checkpointing后,记得把
model.config.use_cache设为False,否则会报错或显存异常。因为checkpointing和KV cache不兼容。
3. 完整训练流程与关键参数配置
3.1 环境准备与依赖安装
先把环境搭好。我习惯用conda建一个独立环境,避免跟系统Python冲突。
conda create -n lora_train python=3.10 -y conda activate lora_train pip install torch==2.1.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install transformers==4.36.0 peft==0.7.0 accelerate==0.25.0 bitsandbytes==0.41.3 pip install datasets==2.15.0 trl==0.7.4 wandb==0.16.0版本很重要。peft和transformers的版本不匹配会导致LoRA层注入失败,或者训练时loss不下降。我踩过好几次坑,最后锁定这套组合比较稳定。
bitsandbytes在Windows上安装比较麻烦,建议用WSL2或者Linux。如果非要在Windows上跑,可以找预编译的wheel包,但版本兼容性要自己试。
3.2 模型加载与LoRA配置代码
以Qwen2.5-7B为例,完整加载和配置代码如下:
import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training model_name = "Qwen/Qwen2.5-7B-Instruct" # 量化配置(可选,FP16加载时去掉) bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, ) tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=bnb_config, # FP16加载时删掉这行 device_map="auto", trust_remote_code=True, torch_dtype=torch.bfloat16, ) # 准备k-bit训练 model = prepare_model_for_kbit_training(model) model.gradient_checkpointing_enable() model.config.use_cache = False # LoRA配置 lora_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", ) model = get_peft_model(model, lora_config) model.print_trainable_parameters()lora_alpha通常设为rank的2倍。rank=16时alpha=32,rank=8时alpha=16。这个比例不是绝对的,但大多数情况下效果不错。lora_dropout设0.05到0.1可以防止过拟合,数据量少时尤其有用。
prepare_model_for_kbit_training这个函数做了几件事:把LayerNorm层转成FP32、启用gradient checkpointing、关闭cache。如果你用FP16加载,不需要这个函数,但需要手动把LayerNorm转FP32:
for name, module in model.named_modules(): if isinstance(module, torch.nn.LayerNorm): module = module.to(torch.float32)3.3 训练参数与显存监控
训练参数配置直接决定显存占用和训练效果。以下是我在32GB GPU上跑7B模型LoRA的常用配置:
from transformers import TrainingArguments training_args = TrainingArguments( output_dir="./lora_output", per_device_train_batch_size=2, gradient_accumulation_steps=8, learning_rate=2e-4, lr_scheduler_type="cosine", warmup_ratio=0.03, num_train_epochs=3, logging_steps=10, save_strategy="epoch", evaluation_strategy="no", fp16=False, bf16=True, optim="adamw_8bit", gradient_checkpointing=True, dataloader_num_workers=2, report_to="wandb", max_grad_norm=0.3, group_by_length=True, )几个关键点:
optim="adamw_8bit"用bitsandbytes的8-bit AdamW,优化器状态占用从FP32的8字节/参数降到2字节/参数。对于LoRA这种可训练参数很少的场景,节省的显存有限,但积少成多。
bf16=True比fp16更稳定,不容易出现梯度溢出。但需要GPU支持BF16,RTX 30系和40系都支持。如果用的是V100或更老的卡,只能用fp16,并且要配fp16_opt_level="O2"。
group_by_length=True把长度相近的样本分到同一个batch,减少padding浪费。这个对显存优化很有帮助,特别是数据长度分布不均匀时。
max_grad_norm=0.3比默认的1.0更严格,LoRA训练时梯度通常比较小,用0.3可以防止个别batch梯度爆炸。
训练过程中要实时监控显存。我习惯开两个终端,一个跑训练,一个跑:
watch -n 1 nvidia-smi如果看到显存占用持续上涨不回落,大概率是内存泄漏。常见原因是dataloader_num_workers设太大,或者数据集里有超长样本导致padding爆炸。
3.4 数据集准备与tokenize技巧
数据集格式对显存影响很大。如果所有样本都padding到最大长度,显存浪费严重。推荐用动态padding:
from datasets import load_dataset dataset = load_dataset("json", data_files="train.json") def tokenize_function(examples): return tokenizer( examples["text"], truncation=True, max_length=2048, padding=False, # 不在这里padding ) tokenized_dataset = dataset.map( tokenize_function, batched=True, remove_columns=dataset["train"].column_names, num_proc=4, )然后在DataCollator里做动态padding:
from transformers import DataCollatorForLanguageModeling data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=False, pad_to_multiple_of=8, # 对齐到8的倍数,提升GPU利用率 )pad_to_multiple_of=8是个小技巧。GPU的显存分配和计算都是以8或16为对齐单位的,padding到8的倍数可以减少碎片,提升约5%到10%的吞吐。
如果数据里有超长样本(比如超过4096 token),建议先截断或过滤掉。一个10000 token的样本,padding后可能占掉几GB显存,直接把batch撑爆。
4. 常见问题排查与避坑指南
4.1 OOM报错的系统排查流程
OOM是LoRA训练最常见的问题。排查时按以下顺序来:
第一步:确认是训练时OOM还是加载时OOM。加载时OOM说明模型权重就放不下,需要换量化精度或换更小的模型。训练时OOM说明激活值或优化器状态超了,需要调batch size或开checkpointing。
第二步:看报错信息里的具体数字。PyTorch的OOM报错会显示“Tried to allocate X GiB”和“GPU has Y GiB total capacity”。X就是当前操作需要的显存,如果X特别大(比如好几GB),说明是某个中间张量爆炸,通常是序列太长或batch太大。
第三步:用torch.cuda.memory_summary()看详细分配。这个命令会输出当前显存的分配情况,包括已分配、已保留、碎片大小。如果“reserved”远大于“allocated”,说明碎片严重,可以试torch.cuda.empty_cache(),但效果有限。
第四步:逐步缩小配置。先把batch_size降到1,seq_len降到512,关掉所有优化,看能不能跑。如果能跑,再逐步往上加,找到临界点。
常见OOM原因速查表:
| 报错特征 | 可能原因 | 解决方法 |
|---|---|---|
| 加载时OOM | 模型权重太大 | 换INT8/INT4量化 |
| 第一个step就OOM | 序列太长或batch太大 | 降seq_len或batch_size |
| 训练中途OOM | 显存泄漏或数据异常 | 检查dataloader和数据集 |
| 评估时OOM | 预测结果堆积 | 设eval_accumulation_steps=1 |
| 保存时OOM | 模型复制到CPU失败 | 设save_on_each_node=False |
4.2 loss不下降或训练不收敛
LoRA训练loss不降,通常不是显存问题,而是配置问题。按以下顺序排查:
学习率是否合适。LoRA的推荐学习率是1e-4到3e-4,比全参数微调高一个数量级。因为LoRA只训练少量参数,需要更大的学习率才能有效更新。如果用的是1e-5,loss基本不动。
target_modules是否覆盖了关键层。只加q_proj和v_proj有时不够,特别是当任务需要模型改变输出风格或知识时。试试加上gate_proj和up_proj。
lora_alpha是否匹配。alpha太小(比如等于rank)会导致LoRA更新幅度不够。通常alpha设为rank的2倍。
数据格式是否正确。指令微调时,prompt和response的拼接方式很重要。如果没加正确的special token或者attention mask,模型学不到东西。
是否冻结了基座模型。用get_peft_model后,基座参数自动冻结,但如果你手动改了requires_grad,可能把基座也解冻了,导致训练不稳定。
4.3 训练速度慢的优化手段
32GB GPU跑LoRA,速度慢通常有几个原因:
gradient checkpointing开销。前面说过,checkpointing会慢20%到30%。如果显存够,关掉它。
dataloader瓶颈。如果dataloader_num_workers=0,数据加载在主进程里做,会阻塞训练。设成2到4,让数据预取和训练并行。
flash attention没开。如果模型支持flash attention(比如Qwen2、Llama3),开启后训练速度能提升30%以上,显存也能省一些。开启方式:
model = AutoModelForCausalLM.from_pretrained( model_name, attn_implementation="flash_attention_2", ... )需要安装flash-attn库,而且对GPU架构有要求(Ampere及以上)。
优化器选择。adamw_8bit比标准AdamW省显存,但速度可能稍慢。如果显存够,用adamw_torch更快。
batch size太小。batch_size=1时GPU利用率很低,大部分时间在等数据。用gradient accumulation配合稍大的micro-batch,能提升吞吐。
4.4 模型保存与合并的注意事项
LoRA训练完后,保存的只是适配器权重,通常几十MB。要得到完整的模型,需要把LoRA权重合并回基座:
from peft import PeftModel base_model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.bfloat16, device_map="auto", ) lora_model = PeftModel.from_pretrained(base_model, "./lora_output") merged_model = lora_model.merge_and_unload() merged_model.save_pretrained("./merged_model") tokenizer.save_pretrained("./merged_model")合并时要注意:如果基座是INT4/INT8加载的,合并前要先反量化回FP16。否则合并后的模型精度会有问题。merge_and_unload()会自动处理这个,但前提是基座模型加载时没用load_in_4bit。如果用了,需要先model.dequantize()再合并。
另一个坑是tokenizer的special token。LoRA训练时如果加了新的special token,保存时要确保tokenizer也保存了这些token,否则推理时会出现token不匹配。
注意:合并后的模型显存占用跟基座FP16一样。7B合并后是14GB,13B是26GB。如果推理时显存不够,可以再用INT8/INT4量化加载。
5. 显存优化的进阶技巧与实战心得
5.1 序列打包与动态batch的妙用
序列打包(sequence packing)是把多个短样本拼成一个长序列,减少padding浪费。比如你有10个长度为100的样本,padding到512的话,每个样本浪费412个token的位置。打包后,10个样本拼成1000 token,只需要2个512的序列,浪费大幅减少。
HuggingFace的trl库提供了ConstantLengthDataset来做打包:
from trl import ConstantLengthDataset train_dataset = ConstantLengthDataset( tokenizer, dataset, dataset_text_field="text", seq_length=2048, chars_per_token=3.6, )chars_per_token是估算值,英文通常3.5到4,中文1.5到2。设得不准会导致打包后的序列长度不对。
打包的代价是样本边界被打破,模型可能看到跨样本的注意力。对于指令微调,这通常不是大问题,因为模型主要学的是指令和回复的对应关系。但对于需要严格样本独立的任务,要谨慎使用。
动态batch是另一个技巧:根据序列长度动态调整batch size,长序列用小batch,短序列用大batch。HuggingFace Trainer不直接支持,但可以通过自定义Sampler实现。这个比较复杂,收益也不如打包明显,一般不建议新手折腾。
5.2 CPU offload的取舍与配置
CPU offload是把部分显存占用转移到内存里。比如优化器状态可以放在CPU上,需要时再拷到GPU。accelerate库支持这个:
from accelerate import Accelerator accelerator = Accelerator( cpu_offload=True, ... )或者在TrainingArguments里:
training_args = TrainingArguments( ... optim="adamw_torch", optim_args="offload", )CPU offload能省显存,但代价是训练速度大幅下降。因为每次更新参数都要在CPU和GPU之间传输数据,PCIe带宽成为瓶颈。实测下来,offload优化器状态会让训练慢2到3倍。
我的建议是:32GB GPU上跑7B LoRA,完全不需要offload。跑13B时,如果显存差一点,可以offload优化器状态,但要做好速度慢的心理准备。更好的选择是降精度或减rank,而不是offload。
5.3 多卡训练时的显存分配
如果你有两张32GB GPU,可以用DDP(DistributedDataParallel)做数据并行。每张卡跑一份完整的模型副本,各自处理不同的batch,梯度通过all-reduce同步。
DDP的显存占用跟单卡一样,因为每张卡都有完整的模型和优化器状态。但通信开销会占一些显存,通常每张卡多占500MB到1GB。
启动方式:
accelerate launch --num_processes=2 train.py或者用torchrun:
torchrun --nproc_per_node=2 train.py多卡训练时要注意batch size的换算。per_device_train_batch_size=2、2张卡、gradient_accumulation_steps=4,等效batch size是2×2×4=16。
如果模型太大,单卡放不下,可以用模型并行(比如device_map="auto"把不同层放到不同卡上)。但模型并行对LoRA训练不太友好,因为LoRA层可能跨卡,通信开销大。32GB GPU跑7B/13B LoRA,单卡足够,不需要模型并行。
5.4 我踩过的那些坑与独家建议
坑一:device_map="auto"和DDP冲突。如果你用device_map="auto"加载模型,再用DDP训练,会报错。因为device_map会把模型分散到多卡,而DDP要求每张卡有完整模型。解决方法是单卡训练时用device_map,多卡时用torch.cuda.set_device(local_rank)手动指定设备。
坑二:padding_side设错导致loss异常。对于decoder-only模型,tokenizer的padding_side应该是"right"。如果设成"left",padding token会出现在序列开头,模型会学到错误的注意力模式。检查方法:
print(tokenizer.padding_side) # 应该是 "right"坑三:eos_token和pad_token相同导致梯度问题。很多模型默认pad_token等于eos_token。训练时如果padding位置也算loss,模型会学着预测eos,导致推理时过早结束。解决方法是在DataCollator里设mlm=False并确保labels中padding位置是-100。
坑四:学习率调度器选错。LoRA训练推荐用cosine或linear,不要用constant。cosine在训练后期学习率降下来,loss更稳定。warmup_ratio设0.03到0.1,让模型先热身再全力学。
坑五:保存checkpoint时OOM。训练时显存刚好够,一保存就OOM。因为保存时要把模型状态字典复制到CPU,如果CPU内存不够或者复制过程中显存峰值超标,就会崩。解决方法是设save_on_each_node=False,或者用save_strategy="steps"配合save_total_limit=2,减少保存频率。
独家建议:用torch.cuda.memory_snapshot()做详细分析。这个命令会输出显存分配的完整快照,包括每个张量的大小、类型、分配位置。对于排查显存泄漏特别有用。用法:
import torch torch.cuda.memory_snapshot()输出是一大串JSON,可以保存到文件里慢慢看。重点找那些“allocated”很大但“active”很小的张量,它们可能是泄漏点。
另一个建议:训练前先跑一个dry run。用1个batch的数据跑完整的前向、反向、优化器更新,看显存峰值。这样能在正式训练前发现问题,避免跑了几小时才OOM。
# Dry run for batch in train_dataloader: outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() print(torch.cuda.max_memory_allocated() / 1024**3, "GB") break这个数字就是你的显存峰值。如果接近32GB,正式训练时就要小心了,因为实际训练中还会有波动。
最后说一个心态问题:显存优化没有银弹。每个配置都是权衡,降精度损失效果,开checkpointing损失速度,减batch size损失稳定性。我的原则是先保效果,再保速度,最后才抠显存。32GB GPU已经能跑大多数7B和13B的LoRA微调了,没必要为了省那几GB把效果搞崩。实在跑不动,换更小的模型或者用云GPU,比在本地硬扛划算得多。