Add support for activation checkpointing with DeepSpeed (#204)

* Update model_utils.py

* Update alphafold3.py

* Update model_utils.py

* Update tensor_typing.py

* Update tensor_typing.py

* Update test_af3.py

* Update alphafold3.py

* Update .env.sample

* Update test.yml

* Update alphafold3.py

* Update test_af3.py

* Update pyproject.toml

Author: amorehead

.env.sample

@@ -1,2 +1,3 @@
 TYPECHECK=True
 DEBUG=False
+DEEPSPEED_CHECKPOINTING=False

.github/workflows/test.yml

@@ -4,6 +4,7 @@ on: [push, pull_request]
 env:
   TYPECHECK: True
   DEBUG: True
+  DEEPSPEED_CHECKPOINTING: False
 
 jobs:
   build:

alphafold3_pytorch/alphafold3.py

@@ -13,7 +13,6 @@
 from torch import Tensor
 from torch.amp import autocast
 import torch.nn.functional as F
-from torch.utils.checkpoint import checkpoint, checkpoint_sequential
 
 from torch.nn import (
     Module,
@@ -73,6 +72,7 @@
     ExpressCoordinatesInFrame,
     RigidFrom3Points,
     calculate_weighted_rigid_align_weights,
+    package_available,
 )
 
 from frame_averaging_pytorch import FrameAverage
@@ -84,6 +84,7 @@
 import einx
 from einops import rearrange, repeat, reduce, einsum, pack, unpack
 from einops.layers.torch import Rearrange
+from environs import Env
 
 from tqdm import tqdm
 
@@ -169,10 +170,23 @@
 
 LinearNoBias = partial(Linear, bias = False)
 
+# environment
+
+env = Env()
+env.read_env()
+
 # always use non reentrant checkpointing
 
-checkpoint = partial(checkpoint, use_reentrant = False)
-checkpoint_sequential = partial(checkpoint_sequential, use_reentrant = False)
+DEEPSPEED_CHECKPOINTING = env.bool('DEEPSPEED_CHECKPOINTING', False)
+
+if DEEPSPEED_CHECKPOINTING:
+    assert package_available("deepspeed"), "DeepSpeed must be installed for checkpointing."
+
+    import deepspeed
+
+    checkpoint = deepspeed.checkpointing.checkpoint
+else:
+    checkpoint = partial(torch.utils.checkpoint.checkpoint, use_reentrant = False)
 
 # helper functions
 
@@ -1061,7 +1075,6 @@ def __init__(
         msa_pwa_heads = 8,
         msa_pwa_dim_head = 32,
         checkpoint = False,
-        checkpoint_segments = 1,
         pairwise_block_kwargs: dict = dict(),
         max_num_msa: int | None = None,
         layerscale_output: bool = True
@@ -1112,7 +1125,6 @@ def __init__(
             ]))
 
         self.checkpoint = checkpoint
-        self.checkpoint_segments = checkpoint_segments
 
         self.layers = layers
 
@@ -1182,19 +1194,19 @@ def inner(inputs):
                 return pairwise_repr, mask, msa, msa_mask
             return inner
 
-        def pairwise_block_wrapper(fn):
+        def msa_transition_wrapper(fn):
             @wraps(fn)
             def inner(inputs):
                 pairwise_repr, mask, msa, msa_mask = inputs
-                pairwise_repr = fn(pairwise_repr = pairwise_repr, mask = mask)
+                msa = fn(msa) + msa
                 return pairwise_repr, mask, msa, msa_mask
             return inner
 
-        def msa_transition_wrapper(fn):
+        def pairwise_block_wrapper(fn):
             @wraps(fn)
             def inner(inputs):
                 pairwise_repr, mask, msa, msa_mask = inputs
-                msa = fn(msa) + msa
+                pairwise_repr = fn(pairwise_repr = pairwise_repr, mask = mask)
                 return pairwise_repr, mask, msa, msa_mask
             return inner
 
@@ -1210,8 +1222,10 @@ def inner(inputs):
             wrapped_layers.append(msa_transition_wrapper(msa_transition))
             wrapped_layers.append(pairwise_block_wrapper(pairwise_block))
 
-        pairwise_repr, *_ = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs)
+        for layer in wrapped_layers:
+            inputs = checkpoint(layer, inputs)
 
+        pairwise_repr, *_ = inputs
         return pairwise_repr
 
     @typecheck
@@ -1318,7 +1332,6 @@ def __init__(
         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()
     ):
@@ -1357,7 +1370,6 @@ def __init__(
         # 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
@@ -1446,8 +1458,10 @@ 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)
+        for layer in wrapped_layers:
+            inputs = checkpoint(layer, inputs)
 
+        single_repr, pairwise_repr, _ = inputs
         return single_repr, pairwise_repr
 
     @typecheck
@@ -1590,7 +1604,6 @@ def __init__(
         pairwise_block_kwargs: dict = dict(),
         eps = 1e-5,
         checkpoint = False,
-        checkpoint_segments = 1,
         layerscale_output = True
     ):
         super().__init__()
@@ -1615,7 +1628,6 @@ def __init__(
         self.pairformer_stack = layers
 
         self.checkpoint = checkpoint
-        self.checkpoint_segments = checkpoint_segments
 
         self.final_norm = nn.LayerNorm(dim)
 
@@ -1666,8 +1678,10 @@ def inner(inputs):
         for block in self.pairformer_stack:
             wrapped_layers.append(block_wrapper(block))
 
-        templates, _ = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs)
+        for layer in wrapped_layers:
+            inputs = checkpoint(layer, inputs)
 
+        templates, _ = inputs
         return templates
 
     @typecheck
@@ -1877,7 +1891,6 @@ def __init__(
         add_residual = True,
         use_linear_attn = False,
         checkpoint = False,
-        checkpoint_segments = 1,
         linear_attn_kwargs = dict(
             heads = 8,
             dim_head = 16
@@ -1956,7 +1969,6 @@ def __init__(
         assert not (not serial and checkpoint), 'checkpointing can only be used for serial version of diffusion transformer'
 
         self.checkpoint = checkpoint
-        self.checkpoint_segments = checkpoint_segments
 
         self.layers = layers
 
@@ -2021,9 +2033,10 @@ def inner(inputs):
             wrapped_layers.append(attn_wrapper(attn))
             wrapped_layers.append(transition_wrapper(transition))
 
-        out = checkpoint_sequential(wrapped_layers, self.checkpoint_segments, inputs)
+        for layer in wrapped_layers:
+            inputs = checkpoint(layer, inputs)
 
-        noised_repr, *_ = out
+        noised_repr, *_ = inputs
         return noised_repr
 
     @typecheck
@@ -2314,10 +2327,6 @@ def __init__(
 
         self.attended_token_norm = nn.LayerNorm(dim_token)
 
-        # checkpointing
-
-        self.checkpoint_token_transformer = checkpoint_token_transformer
-
         # atom attention decoding related modules
 
         self.tokens_to_atom_decoder_input_cond = LinearNoBias(dim_token, dim_atom)
@@ -2332,6 +2341,7 @@ def __init__(
             serial = serial,
             use_linear_attn = use_linear_attn,
             linear_attn_kwargs = linear_attn_kwargs,
+            checkpoint = checkpoint_token_transformer,
             **atom_decoder_kwargs
         )
 
@@ -2484,18 +2494,11 @@ def forward(
             molecule_atom_lens = molecule_atom_lens
         )
 
-        # maybe checkpoint token transformer
-
-        token_transformer = self.token_transformer
-
-        if should_checkpoint(self, tokens, 'checkpoint_token_transformer'):
-            token_transformer = partial(checkpoint, token_transformer)
-
         # token transformer
 
         tokens = self.cond_tokens_with_cond_single(conditioned_single_repr) + tokens
 
-        tokens = token_transformer(
+        tokens = self.token_transformer(
             tokens,
             mask = mask,
             single_repr = conditioned_single_repr,
@@ -5991,6 +5994,7 @@ def __init__(
             dim_template_feats = dim_template_feats,
             dim = dim_template_model,
             dim_pairwise = dim_pairwise,
+            checkpoint=checkpoint_input_embedding,
             **template_embedder_kwargs
         )
 
@@ -6003,6 +6007,7 @@ def __init__(
             dim_pairwise = dim_pairwise,
             dim_msa_input = dim_msa_inputs,
             dim_additional_msa_feats = dim_additional_msa_feats,
+            checkpoint=checkpoint_input_embedding,
             **msa_module_kwargs,
         )
 
@@ -6011,6 +6016,7 @@ def __init__(
         self.pairformer = PairformerStack(
             dim_single = dim_single,
             dim_pairwise = dim_pairwise,
+            checkpoint=checkpoint_trunk_pairformer,
             **pairformer_stack
         )
 
@@ -6115,13 +6121,6 @@ def __init__(
 
         self.register_buffer('lddt_thresholds', torch.tensor([0.5, 1.0, 2.0, 4.0]))
 
-        # checkpointing related
-
-        self.checkpoint_trunk_pairformer = checkpoint_trunk_pairformer
-        self.checkpoint_diffusion_token_transformer = checkpoint_diffusion_token_transformer
-        self.checkpoint_distogram_head = checkpoint_distogram_head
-        self.checkpoint_confidence_head = checkpoint_confidence_head
-
         # loss related
 
         self.ignore_index = ignore_index
@@ -6510,16 +6509,9 @@ def forward(
 
                 pairwise = embedded_msa + pairwise
 
-            # maybe checkpoint trunk pairformer
-
-            pairformer = self.pairformer
-
-            if should_checkpoint(self, (single, pairwise), 'checkpoint_trunk_pairformer'):
-                pairformer = partial(checkpoint, pairformer)
-
             # main attention trunk (pairformer)
 
-            single, pairwise = pairformer(
+            single, pairwise = self.pairformer(
                 single_repr = single,
                 pairwise_repr = pairwise,
                 mask = mask
@@ -6650,12 +6642,7 @@ def forward(
 
             distance_labels = torch.where(distogram_mask, distance_labels, ignore)
 
-            distogram_head_fn = self.distogram_head
-
-            if should_checkpoint(self, pairwise, 'checkpoint_distogram_head'):
-                distogram_head_fn = partial(checkpoint, distogram_head_fn)
-
-            distogram_logits = distogram_head_fn(
+            distogram_logits = self.distogram_head(
                 pairwise,
                 molecule_atom_lens = molecule_atom_lens,
                 atom_feats = atom_feats

alphafold3_pytorch/tensor_typing.py

@@ -68,8 +68,13 @@ def __getitem__(self, shapes: str):
 
 if should_typecheck:
     logger.info("Type checking is enabled.")
+else:
+    logger.info("Type checking is disabled.")
+
 if IS_DEBUGGING:
     logger.info("Debugging is enabled.")
+else:
+    logger.info("Debugging is disabled.")
 
 __all__ = [
     Shaped,

alphafold3_pytorch/utils/model_utils.py

@@ -2,6 +2,7 @@
 from typing import Callable, List, Tuple, Union
 
 import einx
+import pkg_resources
 import torch
 import torch.nn.functional as F
 from einops import einsum, pack, rearrange, reduce, repeat, unpack
@@ -621,6 +622,19 @@ def should_checkpoint(
     )
 
 
+@typecheck
+def package_available(package_name: str) -> bool:
+    """Check if a package is available in your environment.
+
+    :param package_name: The name of the package to be checked.
+    :return: `True` if the package is available. `False` otherwise.
+    """
+    try:
+        return pkg_resources.require(package_name) is not None
+    except pkg_resources.DistributionNotFound:
+        return False
+
+
 # functions for deriving the frames for ligands
 # this follows the logic from Alphafold3 Supplementary section 4.3.2
 

pyproject.toml

@@ -41,7 +41,7 @@ dependencies = [
     "huggingface_hub>=0.21.4",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",
-    "numpy",
+    "numpy==1.23.5",
     "polars>=1.1.0",
     "pdbeccdutils>=0.8.5",
     "pydantic>=2.8.2",

tests/test_af3.py

@@ -818,22 +818,12 @@ def test_alphafold3_without_msa_and_templates():
             depth = 1
         ),
         pairformer_stack = dict(
-            checkpoint = True,
             depth = 2
         ),
         diffusion_module_kwargs = dict(
             atom_encoder_depth = 2,
-            atom_encoder_kwargs = dict(
-                checkpoint = True,
-            ),
             token_transformer_depth = 2,
-            token_transformer_kwargs = dict(
-                checkpoint = True,
-            ),
             atom_decoder_depth = 2,
-            atom_decoder_kwargs = dict(
-                checkpoint = True,
-            ),
         ),
     )