梯度检查点 + dropout 导致损失分歧

突出显示可视性: 此仓库中的自定义检查点辅助程序在反向传播期间重新运行前向传播,而不恢复 RNG 状态。在反向传播期间,区块内的每个随机层(如 dropout)都会看到不同的随机掩码,因此梯度与损失不匹配。因此,在启用梯度检查点的情况下,具有非零 dropout 会导致损失分歧。