first take care of checkpointing for pairformer stack

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -10,7 +10,7 @@
 from torch import nn
 from torch import Tensor
 import torch.nn.functional as F
-from loguru import logger
+from torch.utils.checkpoint import checkpoint, checkpoint_sequential
 
 from torch.nn import (
     Module,
@@ -42,22 +42,20 @@
 from alphafold3_pytorch.inputs import (
     IS_MOLECULE_TYPES,
     IS_PROTEIN_INDEX,
+    IS_DNA_INDEX,
+    IS_RNA_INDEX,
     IS_LIGAND_INDEX,
     IS_METAL_ION_INDEX,
     IS_BIOMOLECULE_INDICES,
+    IS_PROTEIN,
+    IS_DNA,
+    IS_RNA,
+    IS_LIGAND,
+    IS_METAL_ION,
     NUM_MOLECULE_IDS,
     ADDITIONAL_MOLECULE_FEATS
 )
 
-
-IS_DNA_INDEX = 1
-IS_RNA_INDEX = 2
-
-IS_PROTEIN, IS_DNA, IS_RNA, IS_LIGAND, IS_METAL_ION = map(
-    lambda x: IS_MOLECULE_TYPES - x if x < 0 else x, [
-        IS_PROTEIN_INDEX, IS_DNA_INDEX, IS_RNA_INDEX, IS_LIGAND_INDEX, IS_METAL_ION_INDEX])
-
-
 from frame_averaging_pytorch import FrameAverage
 
 from taylor_series_linear_attention import TaylorSeriesLinearAttn
@@ -70,6 +68,8 @@
 
 from tqdm import tqdm
 
+from loguru import logger
+
 from importlib.metadata import version
 
 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
@@ -169,6 +169,21 @@ def unpack_one(to_unpack, unpack_pattern = None):
 def exclusive_cumsum(t, dim = -1):
     return t.cumsum(dim = dim) - t
 
+# checkpointing utils
+
+@typecheck
+def should_checkpoint(
+    self: Module,
+    inputs: Tuple[Tensor, ...],
+    check_instance_variable: str | None = 'checkpoint'
+) -> bool:
+
+    return (
+        self.training and
+        any([i.requires_grad for i in inputs]) and
+        (not exists(check_instance_variable) or getattr(self, check_instance_variable, False))
+    )
+
 # decorators
 
 def maybe(fn):
@@ -350,8 +365,7 @@ def repeat_consecutive_with_lens(
 
     # final mask
 
-    if mask_value is None:
-        mask_value = False if dtype == torch.bool else 0
+    mask_value = default(mask_value, False if dtype == torch.bool else 0)
 
     output = einx.where(
         'b n, b n ..., -> b n ...',
@@ -1101,6 +1115,8 @@ def __init__(
         pair_bias_attn_heads = 16,
         dropout_row_prob = 0.25,
         num_register_tokens = 0,
+        checkpoint = False,
+        checkpoint_segments = 1,
         pairwise_block_kwargs: dict = dict(),
         pair_bias_attn_kwargs: dict = dict()
     ):
@@ -1136,6 +1152,11 @@ def __init__(
 
         self.layers = layers
 
+        # checkpointing
+
+        self.checkpoint = checkpoint
+        self.checkpoint_segments = checkpoint_segments
+
         # https://arxiv.org/abs/2405.16039 and https://arxiv.org/abs/2405.15071
         # although possibly recycling already takes care of this
 
@@ -1150,6 +1171,80 @@ def __init__(
             self.pairwise_row_registers = nn.Parameter(torch.zeros(num_register_tokens, dim_pairwise))
             self.pairwise_col_registers = nn.Parameter(torch.zeros(num_register_tokens, dim_pairwise))
 
+    @typecheck
+    def to_layers(
+        self,
+        *,
+        single_repr: Float['b n ds'],
+        pairwise_repr: Float['b n n dp'],
+        mask: Bool['b n'] | None = None
+
+    ) -> Tuple[Float['b n ds'], Float['b n n dp']]:
+
+        for _ in range(self.recurrent_depth):
+            for (
+                pairwise_block,
+                pair_bias_attn,
+                single_transition
+            ) in self.layers:
+
+                pairwise_repr = pairwise_block(pairwise_repr = pairwise_repr, mask = mask)
+
+                single_repr = pair_bias_attn(single_repr, pairwise_repr = pairwise_repr, mask = mask) + single_repr
+                single_repr = single_transition(single_repr) + single_repr
+
+        return single_repr, pairwise_repr
+
+    @typecheck
+    def to_checkpointed_layers(
+        self,
+        *,
+        single_repr: Float['b n ds'],
+        pairwise_repr: Float['b n n dp'],
+        mask: Bool['b n'] | None = None
+
+    ) -> Tuple[Float['b n ds'], Float['b n n dp']]:
+
+        inputs = (single_repr, pairwise_repr, mask)
+
+        def pairwise_block_wrapper(layer):
+            def inner(inputs, *args, **kwargs):
+                single_repr, pairwise_repr, mask = inputs
+                pairwise_repr = layer(pairwise_repr = pairwise_repr, mask = mask)
+                return single_repr, pairwise_repr, mask
+            return inner
+
+        def pair_bias_attn_wrapper(layer):
+            def inner(inputs, *args, **kwargs):
+                single_repr, pairwise_repr, mask = inputs
+                single_repr = layer(single_repr, pairwise_repr = pairwise_repr, mask = mask) + single_repr
+                return single_repr, pairwise_repr, mask
+            return inner
+
+        def single_transition_wrapper(layer):
+            def inner(inputs, *args, **kwargs):
+                single_repr, pairwise_repr, mask = inputs
+                single_repr = layer(single_repr) + single_repr
+                return single_repr, pairwise_repr, mask
+            return inner
+
+        wrapped_layers = []
+
+        for _ in range(self.recurrent_depth):
+            for (
+                pairwise_block,
+                pair_bias_attn,
+                single_transition
+            ) in self.layers:
+
+                wrapped_layers.append(pairwise_block_wrapper(pairwise_block))
+                wrapped_layers.append(pair_bias_attn_wrapper(pair_bias_attn))
+                wrapped_layers.append(single_transition_wrapper(single_transition))
+
+        single_repr, pairwise_repr, _ = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs)
+
+        return single_repr, pairwise_repr
+
     @typecheck
     def forward(
         self,
@@ -1175,19 +1270,20 @@ def forward(
             if exists(mask):
                 mask = F.pad(mask, (num_registers, 0), value = True)
 
-        # main transformer block layers
+        # maybe checkpoint
 
-        for _ in range(self.recurrent_depth):
-            for (
-                pairwise_block,
-                pair_bias_attn,
-                single_transition
-            ) in self.layers:
+        if should_checkpoint(self, (single_repr, pairwise_repr)):
+            to_layers_fn = self.to_checkpointed_layers
+        else:
+            to_layers_fn = self.to_layers
 
-                pairwise_repr = pairwise_block(pairwise_repr = pairwise_repr, mask = mask)
+        # main transformer block layers
 
-                single_repr = pair_bias_attn(single_repr, pairwise_repr = pairwise_repr, mask = mask) + single_repr
-                single_repr = single_transition(single_repr) + single_repr
+        single_repr, pairwise_repr = to_layers_fn(
+            single_repr = single_repr,
+            pairwise_repr = pairwise_repr,
+            mask = mask
+        )
 
         # splice out registers
 
@@ -3644,6 +3740,9 @@ def compute_lddt(
         is_rna: boolean tensor indicating RNA atoms
         pairwise_mask: boolean tensor indicating atompair for which LDDT is computed
         """
+
+        atom_seq_len, device = pred_coords.shape[1], pred_coords.device
+
         # Compute distances between all pairs of atoms
         pred_dists = torch.cdist(pred_coords, pred_coords)
         true_dists = torch.cdist(true_coords, true_coords)
@@ -3669,7 +3768,7 @@ def compute_lddt(
         )
 
         # Compute mean, avoiding self term
-        mask = inclusion_radius & ~torch.eye(pred_coords.shape[1], dtype=torch.bool, device=pred_coords.device)
+        mask = inclusion_radius & ~torch.eye(atom_seq_len, dtype=torch.bool, device=device)
 
         # Take into account variable lengthed atoms in batch
         if exists(coords_mask):
@@ -3700,7 +3799,7 @@ def compute_chain_pair_lddt(
         plddt between atoms maked by asym_mask_a and asym_mask_b
         """
 
-        if coords_mask is None:
+        if not exists(coords_mask):
             coords_mask = torch.ones_like(asym_mask_a)
 
         if asym_mask_a.ndim == 1:

alphafold3_pytorch/inputs.py

@@ -64,10 +64,17 @@
 
 IS_MOLECULE_TYPES = 5
 IS_PROTEIN_INDEX = 0
+IS_DNA_INDEX = 1
+IS_RNA_INDEX = 2
 IS_LIGAND_INDEX = -2
 IS_METAL_ION_INDEX = -1
 IS_BIOMOLECULE_INDICES = slice(0, 3)
 
+IS_PROTEIN, IS_DNA, IS_RNA, IS_LIGAND, IS_METAL_ION = tuple(
+    (IS_MOLECULE_TYPES - i if i < 0 else i)
+    for i in [IS_PROTEIN_INDEX, IS_DNA_INDEX, IS_RNA_INDEX, IS_LIGAND_INDEX, IS_METAL_ION_INDEX]
+)
+
 MOLECULE_GAP_ID = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES)
 MOLECULE_METAL_ION_ID = MOLECULE_GAP_ID + 1
 NUM_MOLECULE_IDS = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES) + 2

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.62"
+version = "0.2.63"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_af3.py

@@ -173,9 +173,11 @@ def test_centre_random_augmentation():
     assert augmented_coords.shape == coords.shape
 
 
+@pytest.mark.parametrize('checkpoint', (True, False))
 @pytest.mark.parametrize('recurrent_depth', (1, 2))
 @pytest.mark.parametrize('enable_attn_softclamp', (True, False))
 def test_pairformer(
+    checkpoint,
     recurrent_depth,
     enable_attn_softclamp
 ):
@@ -187,6 +189,7 @@ def test_pairformer(
         depth = 4,
         num_register_tokens = 4,
         recurrent_depth = recurrent_depth,
+        checkpoint = checkpoint,
         pair_bias_attn_kwargs = dict(
             enable_attn_softclamp = enable_attn_softclamp
         )
@@ -201,6 +204,10 @@ def test_pairformer(
     assert single.shape == single_out.shape
     assert pairwise.shape == pairwise_out.shape
 
+    if checkpoint:
+        loss = single_out.sum() + pairwise_out.sum()
+        loss.backward()
+
 def test_msa_module():
 
     single = torch.randn(2, 16, 384)