separate serial from parallel (probably error in paper) in diffusion transformers and add checkpointing for serial

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1241,7 +1241,7 @@ def inner(inputs, *args, **kwargs):
                 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)
+        single_repr, pairwise_repr, _ = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs, use_reentrant = False)
 
         return single_repr, pairwise_repr
 
@@ -1615,6 +1615,8 @@ def __init__(
         serial = False,
         add_residual = True,
         use_linear_attn = False,
+        checkpoint = False,
+        checkpoint_segments = 1,
         linear_attn_kwargs = dict(
             heads = 8,
             dim_head = 16
@@ -1689,6 +1691,9 @@ def __init__(
                 conditionable_transition
             ]))
 
+        self.checkpoint = checkpoint
+        self.checkpoint_segments = checkpoint_segments
+
         self.layers = layers
 
         self.serial = serial
@@ -1703,7 +1708,7 @@ def __init__(
             self.registers = nn.Parameter(torch.zeros(num_register_tokens, dim))
 
     @typecheck
-    def forward(
+    def to_checkpointed_serial_layers(
         self,
         noised_repr: Float['b n d'],
         *,
@@ -1712,32 +1717,92 @@ def forward(
         mask: Bool['b n'] | None = None,
         windowed_mask: Bool['b nw w (w*2)'] | None = None
     ):
-        w = self.attn_window_size
-        has_windows = exists(w)
 
-        serial = self.serial
+        inputs = (noised_repr, single_repr, pairwise_repr, mask, windowed_mask)
 
-        # handle windowing
+        wrapped_layers = []
 
-        pairwise_is_windowed = pairwise_repr.ndim == 5
+        def efficient_attn_wrapper(fn):
+            def inner(inputs):
+                noised_repr, single_repr, pairwise_repr, mask, windowed_mask = inputs
+                noised_repr = fn(noised_repr, mask = mask) + noised_repr
+                return noised_repr, single_repr, pairwise_repr, mask, windowed_mask
+            return inner
 
-        if has_windows and not pairwise_is_windowed:
-            pairwise_repr = full_pairwise_repr_to_windowed(pairwise_repr, window_size = w)
+        def attn_wrapper(fn):
+            def inner(inputs):
+                noised_repr, single_repr, pairwise_repr, mask, windowed_mask = inputs
+                noised_repr = fn(noised_repr, cond = single_repr, pairwise_repr = pairwise_repr, mask = mask, windowed_mask = windowed_mask) + noised_repr
+                return noised_repr, single_repr, pairwise_repr, mask, windowed_mask
+            return inner
 
-        # register tokens
+        def transition_wrapper(fn):
+            def inner(inputs):
+                noised_repr, single_repr, pairwise_repr, mask, windowed_mask = inputs
+                noised_repr = fn(noised_repr, cond = single_repr) + noised_repr
+                return noised_repr, single_repr, pairwise_repr, mask, windowed_mask
+            return inner
 
-        if self.has_registers:
-            num_registers = self.num_registers
-            registers = repeat(self.registers, 'r d -> b r d', b = noised_repr.shape[0])
-            noised_repr, registers_ps = pack((registers, noised_repr), 'b * d')
+        for linear_attn, colt5_attn, attn, transition in self.layers:
 
-            single_repr = F.pad(single_repr, (0, 0, num_registers, 0), value = 0.)
-            pairwise_repr = F.pad(pairwise_repr, (0, 0, num_registers, 0, num_registers, 0), value = 0.)
+            if exists(linear_attn):
+                wrapped_layers.append(efficient_attn_wrapper(linear_attn))
 
-            if exists(mask):
-                mask = F.pad(mask, (num_registers, 0), value = True)
+            if exists(colt5_attn):
+                wrapped_layers.append(efficient_attn_wrapper(colt5_attn))
 
-        # main transformer
+            wrapped_layers.append(attn_wrapper(attn))
+            wrapped_layers.append(transition_wrapper(transition))
+
+        out = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs, use_reentrant = False)
+
+        noised_repr, *_ = out
+        return noised_repr
+
+    @typecheck
+    def to_serial_layers(
+        self,
+        noised_repr: Float['b n d'],
+        *,
+        single_repr: Float['b n ds'],
+        pairwise_repr: Float['b n n dp'] | Float['b nw w (w*2) dp'],
+        mask: Bool['b n'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None
+    ):
+
+        for linear_attn, colt5_attn, attn, transition in self.layers:
+
+            if exists(linear_attn):
+                noised_repr = linear_attn(noised_repr, mask = mask) + noised_repr
+
+            if exists(colt5_attn):
+                noised_repr = colt5_attn(noised_repr, mask = mask) + noised_repr
+
+            noised_repr = attn(
+                noised_repr,
+                cond = single_repr,
+                pairwise_repr = pairwise_repr,
+                mask = mask,
+                windowed_mask = windowed_mask
+            ) + noised_repr
+
+            noised_repr = transition(
+                noised_repr,
+                cond = single_repr
+            ) + noised_repr
+
+        return noised_repr
+
+    @typecheck
+    def to_parallel_layers(
+        self,
+        noised_repr: Float['b n d'],
+        *,
+        single_repr: Float['b n ds'],
+        pairwise_repr: Float['b n n dp'] | Float['b nw w (w*2) dp'],
+        mask: Bool['b n'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None
+    ):
 
         for linear_attn, colt5_attn, attn, transition in self.layers:
 
@@ -1755,25 +1820,72 @@ def forward(
                 windowed_mask = windowed_mask
             )
 
-            if serial:
-                noised_repr = attn_out + noised_repr
-
             ff_out = transition(
                 noised_repr,
                 cond = single_repr
             )
 
-            if serial:
-                noised_repr = ff_out + noised_repr
-
             # in the algorithm, they omitted the residual, but it could be an error
             # attn + ff + residual was used in GPT-J and PaLM, but later found to be unstable configuration, so it seems unlikely attn + ff would work
             # but in the case they figured out something we have not, you can use their exact formulation by setting `serial = False` and `add_residual = False`
 
             residual = noised_repr if self.add_residual else 0.
 
-            if not serial:
-                noised_repr = ff_out + attn_out + residual
+            noised_repr = ff_out + attn_out + residual
+
+        return noised_repr
+
+    @typecheck
+    def forward(
+        self,
+        noised_repr: Float['b n d'],
+        *,
+        single_repr: Float['b n ds'],
+        pairwise_repr: Float['b n n dp'] | Float['b nw w (w*2) dp'],
+        mask: Bool['b n'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None
+    ):
+        w = self.attn_window_size
+        has_windows = exists(w)
+
+        serial = self.serial
+
+        # handle windowing
+
+        pairwise_is_windowed = pairwise_repr.ndim == 5
+
+        if has_windows and not pairwise_is_windowed:
+            pairwise_repr = full_pairwise_repr_to_windowed(pairwise_repr, window_size = w)
+
+        # register tokens
+
+        if self.has_registers:
+            num_registers = self.num_registers
+            registers = repeat(self.registers, 'r d -> b r d', b = noised_repr.shape[0])
+            noised_repr, registers_ps = pack((registers, noised_repr), 'b * d')
+
+            single_repr = F.pad(single_repr, (0, 0, num_registers, 0), value = 0.)
+            pairwise_repr = F.pad(pairwise_repr, (0, 0, num_registers, 0, num_registers, 0), value = 0.)
+
+            if exists(mask):
+                mask = F.pad(mask, (num_registers, 0), value = True)
+
+        # main transformer
+
+        if self.serial and should_checkpoint(self, (noised_repr, single_repr, pairwise_repr)):
+            to_layers_fn = self.to_checkpointed_serial_layers
+        elif self.serial:
+            to_layers_fn = self.to_serial_layers
+        else:
+            to_layers_fn = self.to_parallel_layers
+
+        noised_repr = to_layers_fn(
+            noised_repr,
+            single_repr = single_repr,
+            pairwise_repr = pairwise_repr,
+            mask = mask,
+            windowed_mask = windowed_mask,
+        )
 
         # splice out registers
 

pyproject.toml

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

tests/test_af3.py

@@ -181,8 +181,8 @@ def test_pairformer(
     recurrent_depth,
     enable_attn_softclamp
 ):
-    single = torch.randn(2, 16, 384)
-    pairwise = torch.randn(2, 16, 16, 128)
+    single = torch.randn(2, 16, 384).requires_grad_()
+    pairwise = torch.randn(2, 16, 16, 128).requires_grad_()
     mask = torch.randint(0, 2, (2, 16)).bool()
 
     pairformer = PairformerStack(
@@ -228,17 +228,26 @@ def test_msa_module():
 
     assert pairwise.shape == pairwise_out.shape
 
+@pytest.mark.parametrize('checkpoint', (False, True))
+@pytest.mark.parametrize('serial', (False, True))
 @pytest.mark.parametrize('use_linear_attn', (False, True))
 @pytest.mark.parametrize('use_colt5_attn', (False, True))
-def test_diffusion_transformer(use_linear_attn, use_colt5_attn):
+def test_diffusion_transformer(
+    checkpoint,
+    serial,
+    use_linear_attn,
+    use_colt5_attn
+):
 
-    single = torch.randn(2, 16, 384)
-    pairwise = torch.randn(2, 16, 16, 128)
+    single = torch.randn(2, 16, 384).requires_grad_()
+    pairwise = torch.randn(2, 16, 16, 128).requires_grad_()
     mask = torch.randint(0, 2, (2, 16)).bool()
 
     diffusion_transformer = DiffusionTransformer(
         depth = 2,
         heads = 16,
+        serial = serial,
+        checkpoint = checkpoint,
         use_linear_attn = use_linear_attn,
         use_colt5_attn = use_colt5_attn
     )
@@ -252,6 +261,10 @@ def test_diffusion_transformer(use_linear_attn, use_colt5_attn):
 
     assert single.shape == single_out.shape
 
+    if checkpoint:
+        loss = single_out.sum()
+        loss.backward()
+
 def test_sequence_local_attn():
     atoms = torch.randn(2, 17, 32)
     attn_bias = torch.randn(2, 17, 17)