Use torch.log1p (deepspeedai#6930)

This function provides greater precision than `log(1 + x)` for small values of `x`. Found with TorchFix https://github.com/pytorch-labs/torchfix/ Signed-off-by: siqi <siqi@tecorigin.com>
siqi654321 · Feb 7, 2025 · f74b1e5 · f74b1e5
1 parent dbbfd01
commit f74b1e5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/sequence/fpdt_layer.py b/deepspeed/sequence/fpdt_layer.py
@@ -47,7 +47,7 @@ def _update_out_and_lse(
     block_out = block_out.to(torch.float32)
     block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
 
-    new_lse = lse + torch.log(1 + torch.exp(block_lse - lse))
+    new_lse = lse + torch.log1p(torch.exp(block_lse - lse))
 
     out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out