broccoli-ml 4.0.1__py3-none-any.whl → 5.0.0__py3-none-any.whl

broccoli/transformer.py

@@ -285,9 +285,11 @@ class FeedforwardBlock(nn.Module):
         normformer=False,
         post_norm=True,
         residual_path=True,
+        checkpoint=True,
     ):
         super().__init__()
 
+        self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
 
@@ -326,12 +328,18 @@ class FeedforwardBlock(nn.Module):
         )
 
     def forward(self, x):
+
+        if self.checkpoint:
+            processed = checkpoint(self.process, x, use_reentrant=False)
+        else:
+            processed = self.process(x)
+
         if self.residual_path and self.post_norm:
-            return self.layernorm(x + self.process(x))
+            return self.layernorm(x + processed)
         elif self.residual_path:
-            return x + self.process(x)
+            return x + processed
         else:
-            return self.process(x)
+            return processed
 
 
 class TransformerBlock(nn.Module):
@@ -365,6 +373,7 @@ class TransformerBlock(nn.Module):
         pre_norm=True,
         post_norm=False,
         normformer=False,
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -433,6 +442,7 @@ class TransformerBlock(nn.Module):
             normformer=normformer,
             post_norm=False,  # Handled outside the block
             residual_path=False,  # Handled outside the block
+            checkpoint=checkpoint_ff,
         )
 
     @property
@@ -445,17 +455,17 @@ class TransformerBlock(nn.Module):
             x = self.layer_norm_1(x)
             x = x + self.drop_path(self.attn(x, x, x))
             x = self.layer_norm_2(x)
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
             if self.post_norm:  # i.e. in addition! Pre and post.
                 x = self.layer_norm_3(x)
         elif self.post_norm:  # i.e. only, not prenorm, just post
             x = x + self.drop_path(self.attn(x, x, x))
             x = self.layer_norm_1(x)
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
             x = self.layer_norm_2(x)
         else:  # Not pre or post norm. Stand well back.
             x = x + self.drop_path(self.attn(x, x, x))
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
 
         return x
 
@@ -491,6 +501,7 @@ class TransformerEncoder(nn.Module):
         post_norm=False,
         normformer=False,
         msa_scaling="d",
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -567,6 +578,7 @@ class TransformerEncoder(nn.Module):
                     pre_norm=pre_norm,
                     post_norm=post_norm,
                     normformer=normformer,
+                    checkpoint_ff=checkpoint_ff,
                 )
                 for i in range(n_layers)
             ]

broccoli/vit.py

@@ -160,6 +160,7 @@ class ViTEncoder(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -321,6 +322,7 @@ class ViTEncoder(nn.Module):
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
+                checkpoint_ff=transformer_checkpoint_ff,
             )
         else:
             self.transformer = nn.Identity()
@@ -354,6 +356,7 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 residual_path=transformer_initial_ff_residual_path,
+                checkpoint=transformer_checkpoint_ff,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -426,6 +429,7 @@ class ViT(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -492,6 +496,7 @@ class ViT(nn.Module):
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
+            transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
         )
 

{broccoli_ml-4.0.1.dist-info → broccoli_ml-5.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 4.0.1
+Version: 5.0.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-4.0.1.dist-info → broccoli_ml-5.0.0.dist-info}/RECORD

@@ -4,10 +4,10 @@ broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
 broccoli/linear.py,sha256=Y7s-DzcwsOipRboNHc4HTScw4mJRalNoVFsNcxOB6a4,4872
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=uqSf8q30MF7Ds7LfqW8Pr206NXpSlf7o6770KZu81Ew,19883
+broccoli/transformer.py,sha256=eSBRF-HYJ-BxfisJCueUYCIYtHXgj1ewG5RxEfcmu-E,20128
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=_5uLcklmJ1Uoj7V7TkzF0UqroVnl8NCHun5B0mORmOg,18651
-broccoli_ml-4.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-4.0.1.dist-info/METADATA,sha256=vH_utDdo0-e2q8ReDrRHQ1d6fOzG4nzb9EWlqgyl4XY,1368
-broccoli_ml-4.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-4.0.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=BrNLOx4_gTY6xTwAn8xT-HOgUnSFtU6_m1CpJXuQiKY,18907
+broccoli_ml-5.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-5.0.0.dist-info/METADATA,sha256=NMkRLZfqhMZdBIn4BMHjF_-jyv3yYUlOrKeHRTt2rnE,1368
+broccoli_ml-5.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-5.0.0.dist-info/RECORD,,