自回归Transformer解码非教师强制训练显存瓶颈问题咨询

1. 计算图的累积存储 #

训练阶段PyTorch会保留每一步前向传播的所有中间张量，用于后续反向传播计算梯度。你写的非教师强制逻辑里，循环迭代20次，每次都独立调用一次Transformer解码器：

• 第1次推理序列长度为1的输入，对应1套解码器各层的中间激活、注意力矩阵等张量需要留存
• 第2次推理序列长度为2的输入，再新增1套更长的中间张量留存
• ...一直到第20次推理序列长度为20的输入，总共会留存(1+2+...+20) = 210单位长度的解码器中间张量，是教师强制模式（仅1次长度20的前向）的10.5倍，直接导致显存占用翻倍溢出。

而推理阶段用with torch.no_grad()时，PyTorch不会保留任何计算图和中间张量，每一步生成后旧的中间张量可以直接释放，所以不会出现显存问题。

2. 冗余张量的重复占用 #

每次循环中你调用torch.cat生成新的predictions张量，所有中间版本的predictions因为和计算图存在依赖关系，在整个前向过程结束前都无法被GC释放，进一步增加了无效显存占用。

3. 注意力计算的累积开销 #

Transformer自注意力模块的中间权重矩阵大小和序列长度的平方成正比，20次迭代下来累积的注意力中间矩阵总大小是(1²+2²+...+20²) = 2870，是教师强制模式（仅20²=400）的7倍多，进一步放大了显存压力。

• 采用*计划采样（Scheduled Sampling）*策略，不需要完全模拟自回归迭代过程，只随机替换部分输入token为上一步的预测结果，既可以混合教师强制和非教师强制的训练效果，也不会产生多倍显存占用。
• 必须模拟逐步生成训练的话，可以用PyTorch的torch.utils.checkpoint做梯度检查点，用计算速度换显存，手动释放不需要的中间张量，仅保留必要的节点用于反向传播。

内容的提问来源于stack exchange，提问作者Shayan Jawed