bfloat16训练时段错误: 数组越界的一种方案

问题

在使用半精度 bfloat16 训练lora时有概率报错: 数组越界.

一个快速复现

将 examples/stable_diffusion_xl/model_training/train.py 中的第 30 行中的 torch.dtype=torch.float32 改为 torch.dtype=torch.bfloat16 然后在训练的时候有概率会产生上述报错.

原因(猜测)

报错的具体位置在 diffsynth/diffusion/ddim_scheduler.py 中的 90 行, 猜测是因为使用流匹配的损失函数时间步的范围是 0-999 但是在ddim调度器中可能取到 1000 导致的数组越界, 我在报错的位置输出了当前的时间步确实是 1000, 原因我猜测是因为将 float32 精度的 999.0 转换为 bfloat16 精度时由于浮点数精度导致这个数值变成了 1000.0, 下面的代码展示了这个现象:

import torch

x0 = torch.tensor(999.0, dtype=torch.float32)
x1 = x0.to(dtype=torch.bfloat16)
x2 = x1.to(dtype=torch.float32)

print(x0, x1, x2, sep='\n')

他的输出是:

tensor(999.)
tensor(1000., dtype=torch.bfloat16)
tensor(1000.)

解决

我在 diffsynth/diffusion/ddim_scheduler.py 的 add_noise 函数的一开始加上了 timestep = timestep.to(torch.float32).clamp(0, 999).long() 这样一行代码, 然后就没有再出现上面的问题

Reproduce the issue by changing torch.dtype to torch.bfloat16 in examples/stable_diffusion_xl/model_training/train.py, then inspect diffsynth/diffusion/ddim_scheduler.py and its add_noise function around line 90. Verify how timestep values are converted and indexed, and confirm that training no longer produces an out-of-bounds error while preserving normal scheduler behavior.