stabilize fp16 training

Author: yxlllc

modules/commons/common_layers.py

@@ -114,7 +114,15 @@ def forward(self, x):
         # out, gate = x.chunk(2, dim=self.dim)
         # Using torch.split instead of chunk for ONNX export compatibility.
         out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
-        return out * F.silu(gate)
+        gate = F.silu(gate)
+        if x.dtype == torch.float16:
+            out_min, out_max = torch.aminmax(out.detach())
+            gate_min, gate_max = torch.aminmax(gate.detach())
+            max_abs_out = torch.max(-out_min, out_max).float()
+            max_abs_gate = torch.max(-gate_min, gate_max).float()
+            if max_abs_out * max_abs_gate > 65504:
+                return (out.float() * gate.float()).clamp(-65504, 65504).half()             
+        return out * gate
 
 
 class Conv1d(torch.nn.Conv1d):