fix: report gradient snr in db

This commit is contained in:
2026-08-09 13:40:27 +08:00
parent 47b3ed4e44
commit 596c35fd71
2 changed files with 22 additions and 1 deletions
+5 -1
View File
@@ -1,3 +1,4 @@
import math
from typing import Dict
import torch
@@ -27,6 +28,8 @@ class GradSNRTracker:
SNR = E[g]^2 / Var(g) = E[g]^2 / (E[g^2] - E[g]^2)
The reported value is the power ratio in decibels: ``10 * log10(SNR)``.
The tracker accumulates per-parameter EMA moments across optimizer steps.
Call ``update`` after backward (before ``optimizer.step``) and read
``snr`` to get the aggregate SNR across all parameters.
@@ -64,7 +67,8 @@ class GradSNRTracker:
noise = (v - m.pow(2)).clamp(min=0).sum().item()
total_signal += signal
total_noise += noise
return total_signal / (total_noise + self.eps)
snr = total_signal / (total_noise + self.eps)
return 10.0 * math.log10(max(snr, self.eps))
def ctx_get_loss(ctx):