add silut impl

Author: caic99

deepmd/pt/cxx_op.py

@@ -87,7 +87,7 @@ def load_library(module_name: str) -> bool:
     return False
 
 
-ENABLE_CUSTOMIZED_OP = load_library("deepmd_op_pt")
+ENABLE_CUSTOMIZED_OP = None#load_library("deepmd_op_pt")
 
 __all__ = [
     "ENABLE_CUSTOMIZED_OP",

deepmd/pt/train/training.py

@@ -976,6 +976,7 @@ def log_loss_valid(_task_key="Default"):
 
             elapsed_batch = self.num_steps - self.start_step
             if self.timing_in_training and elapsed_batch // self.disp_freq > 0:
+                log.info(msg=f"reserved memory: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
                 if self.start_step >= 2 * self.disp_freq:
                     log.info(
                         "average training time: %.4f s/batch (exclude first %d batches)",

deepmd/pt/utils/dataloader.py

@@ -189,6 +189,7 @@ def print_summary(
         name: str,
         prob: list[float],
     ) -> None:
+        return
         rank = dist.get_rank() if dist.is_initialized() else 0
         if rank == 0:
             print_summary(

deepmd/pt/utils/utils.py

@@ -17,11 +17,44 @@
 )
 from .env import PRECISION_DICT as PT_PRECISION_DICT
 
+class SiLUT(torch.nn.Module):
+    def __init__(self, threshold=3.0):
+        super().__init__()
+
+        def sigmoid(x):
+            return 1 / (1 + np.exp(-x))
+
+        def silu(x):
+            return x * sigmoid(x)
+
+        def silu_grad(x):
+            sig = sigmoid(x)
+            return sig + x * sig * (1 - sig)
+
+        self.threshold = threshold
+        self.slope = float(silu_grad(threshold))
+        self.const = float(silu(threshold))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        silu_part = F.silu(x)
+        mask = x >= self.threshold
+        if torch.any(mask):
+            tanh_part = torch.tanh(self.slope * (x - self.threshold)) + self.const
+            return torch.where(x < self.threshold, silu_part, tanh_part)
+        else:
+            return silu_part
 
 class ActivationFn(torch.nn.Module):
     def __init__(self, activation: Optional[str]) -> None:
         super().__init__()
         self.activation: str = activation if activation is not None else "linear"
+        if self.activation.lower().startswith(
+            "silut"
+        ) or self.activation.lower().startswith("custom_silu"):
+            threshold = (
+                float(self.activation.split(":")[-1]) if ":" in self.activation else 3.0
+            )
+            self.silut = SiLUT(threshold=threshold)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Returns the tensor after applying activation function corresponding to `activation`."""
@@ -43,6 +76,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             return F.silu(x)
         elif self.activation.lower() == "linear" or self.activation.lower() == "none":
             return x
+        elif self.activation.lower().startswith("custom_silu"):
+            assert self.silut is not None
+            return self.silut(x)
         else:
             raise RuntimeError(f"activation function {self.activation} not supported")