Clamp output when rounding weight to prevent Nan.

comfyanonymous · comfyanonymous · commit 73e3a9e67654 · 2024-10-19T19:07:10.000-04:00
diff --git a/comfy/float.py b/comfy/float.py
@@ -41,6 +41,8 @@ def manual_stochastic_round_to_float8(x, dtype, generator=None):
         (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x
     )
 
+    inf = torch.finfo(dtype)
+    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
     return sign
 
 

Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,8 @@ def manual_stochastic_round_to_float8(x, dtype, generator=None):`
`41`	`41`	`(2.0 ** (-EXPONENT_BIAS + 1)) * abs_x`
`42`	`42`	`)`
`43`	`43`
	`44`	`+ inf = torch.finfo(dtype)`
	`45`	`+ torch.clamp(sign, min=inf.min, max=inf.max, out=sign)`
`44`	`46`	`return sign`
`45`	`47`
`46`	`48`