EMA shadow weights diverge — `sub_` changed to `add_` in commit 4a41f14
Commit 4a41f14 ("Fix feedback when doing an ema with feedback so it uses its own scale") flipped the sign of the EMA shadow update in ExponentialMovingAverage.update.
Before:
tmp = (s_param_float - param_float)
tmp.mul_(one_minus_decay)
s_param_float.sub_(tmp)After:
gap = (s_param_float - param_float)
s_param_float.add_(gap * one_minus_decay)Since gap = s - p, the correct update is s -= (1 - decay) * gap. With add_, the shadow becomes:
s_new = s + (1 - decay) * (s - p) = 1.01 * s - 0.01 * p (at decay=0.99)so the shadow is pushed away from the parameters instead of toward them, multiplying by 1.01 every step.
Impact: any run with use_ema: true saves diverging shadow weights. Training itself is unaffected — the live parameters and the loss curve look completely normal — so there is no visible signal that anything is wrong. At decay=0.99 the growth is ~12x per 250 steps, which exceeds bf16's usable range within a few hundred steps and produces LoRAs that render as noise at any strength multiplier.
Observed, LoRA training on Krea 2, decay 0.99, bf16 save, max |weight| per checkpoint:
| step | max abs weight | ratio vs previous |
|---|---|---|
| 1500 | 1.082e4 | — |
| 1750 | 1.306e5 | 12.07 |
| 2000 | 1.573e6 | 12.04 |
| 2019 | 1.876e6 | 1.19 (19 steps) |
1.01^250 = 12.03 and 1.01^19 = 1.21, matching the observed ratios. An identical config run on b1bf3e4 (before this commit) stays at ~0.04 as expected.
Fix — one line:
gap = (s_param_float - param_float)
s_param_float.sub_(gap * one_minus_decay)The use_feedback branch below it is correct as written; param_float.add_(gap * self.feedback_rate) pulls the parameter toward the shadow, which is the intended direction.
Source: ostris/ai-toolkit