在单卡或本地跑 DMAD 少步蒸馏，显存爆掉时的处理顺序建议是：先把批大小压到 1，确认整条链路能走通；再在 batch=1 的条件下，把分辨率加到训练真正需要的最小值；最后用梯度累积把等效批大小补回来。这个顺序之所以成立，是因为先确认能不能跑通，再讨论跑得快不快；分辨率通常比批大小更容易把显存吃掉，先解决它更省试错。前提是手上要有三组能复现的测量：不训练时的显存底数、只改批大小后的显存与单步耗时、同一批样本换分辨率后的峰值。没有这三组数，先降分辨率还是先降批大小都只能靠猜。

判断与做法：单卡上优先保分辨率、先降批大小，再用梯度累积补等效批大小，是相对省事的顺序；每一步都用 nvidia-smi 峰值和单步耗时复核，并给显存留出余量。适用场景是训练与推理共用一张卡、显存只够跑小 batch 的情况。边界提示：注意力实现、精度、是否开梯度检查点在不同仓库差别很大，别人机器上的数字不能直接搬过来用。

先量基线：不蒸馏直接推理时的显存峰值

这一步只为拿到底数：一张卡上跑推理（不反向、不做蒸馏）时占多少显存，把它从训练预算里减掉，后面才知道多少显存是训练本身花掉的。记录时机建议至少三个：进程起来加载完权重之后、首次前向之后、以及（如果能跑）首次反向之后。首次前向往往包含编译、缓存分配和临时张量，峰值比稳态更高，因此 max_memory_allocated 比瞬时 memory_allocated 更有参考价值。

外部视角用 nvidia-smi 轮询，能看到进程整体占用、CUDA context 和其他进程的干扰：

nvidia-smi `--query-gpu`=index,memory.used,memory.total `--format`=csv -l 1
watch -n 1 nvidia-smi

进程内视角用 PyTorch 的分配器计数，两者要对照看：nvidia-smi 里通常包含 context、缓存块和别的进程，Torch 看到的只是本进程申请的部分，所以训练预算要以更保守的那个数为准。

import torch

def mem(tag):
    a = torch.cuda.memory_allocated() / 1024**3
    p = torch.cuda.max_memory_allocated() / 1024**3
    print(f"{tag} alloc={a:.2f}GiB peak={p:.2f}GiB")

torch.cuda.reset_peak_memory_stats()
load_weights()
mem("after_load")
out = model(x)          # 首次前向
mem("after_first_forward")

记下这两三个数就够了，不必在这一步追求精度；关键是把“推理也要占掉的那部分”和“训练新增的部分”分清楚，否则后面很容易把推理占用误算成批大小的开销。

只改批大小时记录显存与单步耗时

这一节只动一个变量：批大小。取 2 到 3 档，例如 1 / 2 / 4（能不能上 4 由实测决定），其余全部钉死：分辨率、精度（fp32/fp16/bf16 保持同一个）、是否开梯度检查点、num_workers、随机种子、样本列表都必须一致。每档单独起一次进程，避免上一轮的缓存块影响峰值。

每次记录这些字段，数值由自己实测填写，这里不给预设结论：

• 批大小（实际生效值，不是配置文件里写的值）；
• 显存峰值（max_memory_allocated 与 nvidia-smi 各记一个）；
• 单步耗时：先跑若干步 warmup，再对固定步数取平均，把数据加载时间和计算时间分开记；
• 是否溢出（OOM 发生在哪一步：第一次前向、第一次反向，还是第 N 步）。

import time, torch
for i in range(warmup + steps):
    torch.cuda.synchronize(); t0 = time.perf_counter()
    loss = train_step(batch)
    torch.cuda.synchronize(); dt = time.perf_counter() - t0
    if i >= warmup:
        print(i, f"{dt*1000:.0f}ms", f"peak={torch.cuda.max_memory_allocated()/1024**3:.2f}GiB")

如果 OOM 出现在第一次反向，说明激活占用是主要压力，这时降批大小通常见效很直接；如果批大小已经降到 1 还在首次前向就溢出，问题多半不在批大小，应该去看分辨率、精度和是否开启了梯度检查点。

提高分辨率后重测同一批样本

分辨率要单独测一轮，条件尽量和上一节对齐，只改分辨率。配置里通常涉及这些位置，名称以实际仓库为准：数据集配置中的 image_size / resolution、训练分辨率与验证分辨率、裁剪尺寸（crop size）、以及数据增强里的 resize。

需要同步检查的字段：模型侧的输入尺寸或位置编码是否需要插值、预计算好的 latent 缓存是否必须删掉重建（分辨率变了缓存就失效了）、以及任何写死尺寸的 collate 逻辑。少步蒸馏还要留意教师模型和学生模型的输入尺寸是否一致，否则比较的不是同一件事。

重测时要保持一致的其他变量：同一批样本（固定文件列表或固定 sampler 与种子）、同一个批大小（建议先固定在 1）、同一精度、同一个训练步数。记录字段和上一节相同，只是把批大小换成分辨率。观察两个数：峰值随分辨率的增长幅度，以及单步耗时的增长幅度。通常 token 数变多后注意力部分的显存增长会快于线性，但具体陡到什么程度必须由这台机器上的数说话。

用梯度累积替代大 batch 的可行性验证

梯度累积的目标是在显存不变的前提下把等效批大小提上去。通用配置骨架大致如下，字段名各仓库不同，先看自己仓库的配置类和训练脚本，找不到对应字段就以官方文档为准，不要凭猜测写键名：

train:
  batch_size: 1
  gradient_accumulation_steps: 4   # 等效批大小 = batch_size × 累积步数
  max_steps: 2000
optimizer:
  name: adamw
scheduler:
  name: cosine
  # 注意按 optimizer step 还是按 micro step 计数

生效验证方式建议做三件事。第一，启动日志里确认实际打印的 batch size 与累积步数，并自己算出等效批大小。第二，确认 optimizer.step() 只在累积边界调用、梯度在边界处清零；看到日志里 step 数减少到原来的 1/累积步数，基本说明生效。第三，比较梯度量级：用真实批大小跑一次、再用累积跑一次，看梯度范数是否接近；如果差得明显，检查 loss 是否被多次平均（累积时通常要把 loss 除以累积步数，或用 sum 再统一除）。

还有两处容易漏：学习率调度器和 EMA 的更新频率是按 micro step 还是按 optimizer step 推进的；以及累积期间是否有多余的日志、验证或对齐损失在每次前向都计算。若仓库完全没有累积字段，可以看是否支持 accelerate 一类训练器的累积参数，或者手动在脚本里包一层，但这时步数统计和调度器需要自己对齐。

把三组数据整理成取舍表

把前面三节的记录合成一张表，只保留能指导决策的列，填完这张表，取舍就不用再靠感觉：

填写方式：每行对应一次单独启动的进程，分辨率填配置里实际生效的输入尺寸，批大小填日志打印的实际值，显存峰值优先填 max_memory_allocated，单步耗时填去掉 warmup 后的平均值，“是否溢出”填“否 / 前向 OOM / 反向 OOM / 第 N 步 OOM”。可以再加一列备注，记下当次是否开梯度检查点、用的什么精度。

用法上比较稳的几条：显存峰值逼近卡上限时，先动批大小，因为它的收益通常更线性；单步耗时成为瓶颈而显存还有余量时，再考虑动分辨率；峰值离上限还有明显空间、但 batch 想加又不敢加时，用梯度累积补等效批大小，接受单步变慢的代价。这张表只在同机、同精度、同仓库版本下可比，换卡或换精度要重测。