feat: use bfloat16 with torch.autocast on training

Author: caic99

deepmd/pt/model/descriptor/repflow_layer.py

@@ -424,7 +424,7 @@ def optim_angle_update(
         sub_edge_update_ik = torch.matmul(edge_ebd, sub_edge_ik)
 
         result_update = (
-            bias
+            bias.to(sub_angle_update.dtype)
             + sub_node_update.unsqueeze(2).unsqueeze(3)
             + sub_edge_update_ij.unsqueeze(2)
             + sub_edge_update_ik.unsqueeze(3)
@@ -463,7 +463,10 @@ def optim_edge_update(
         sub_edge_update = torch.matmul(edge_ebd, edge)
 
         result_update = (
-            bias + sub_node_update.unsqueeze(2) + sub_edge_update + sub_node_ext_update
+            bias.to(sub_node_update.dtype)
+            + sub_node_update.unsqueeze(2)
+            + sub_edge_update
+            + sub_node_ext_update
         )
         return result_update
 
@@ -679,6 +682,7 @@ def forward(
                     )
                 )
 
+            a_sw.to(edge_angle_update.dtype)
             # nb x nloc x a_nnei x a_nnei x e_dim
             weighted_edge_angle_update = (
                 a_sw[..., None, None] * a_sw[..., None, :, None] * edge_angle_update

deepmd/pt/train/wrapper.py

@@ -136,6 +136,7 @@ def share_params(self, shared_links, resume=False) -> None:
                             f"Shared params of {model_key_base}.{class_type_base} and {model_key_link}.{class_type_link}!"
                         )
 
+    @torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True)
     def forward(
         self,
         coord,

deepmd/pt/utils/utils.py

@@ -93,7 +93,7 @@ def forward(ctx, x, threshold, slope, const_val):
                 ctx.threshold = threshold
                 ctx.slope = slope
                 ctx.const_val = const_val
-                return silut_forward_script(x, threshold, slope, const_val)
+                return silut_forward_script(x, threshold, slope, const_val).to(x.dtype)
 
             @staticmethod
             def backward(ctx, grad_output):