Memory budget strategy for activation checkpointing

CHANGELOG.md

@@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added support for accessing Google on non-Google clusters via auth with service account keys.
 - Added support for revisions in `convert_checkpoint_from_hf.py` and the `load_hf_model` method of `olmo_core.nn.hf.checkpoint`.
 - `foreach` support in `SkipStepAdamW`.
+- Added `budget` mode for activation checkpointing configuration.
 
 ### Changed
 

src/olmo_core/nn/transformer/config.py

@@ -59,6 +59,8 @@ class TransformerActivationCheckpointingMode(StrEnum):
     """Checkpoint only selected modules."""
     selected_ops = "selected_ops"
     """Checkpoint only a specific set of operations."""
+    budget = "budget"
+    """Checkpoint based on a budget."""
 
 
 class TransformerType(StrEnum):

src/olmo_core/nn/transformer/model.py

@@ -542,6 +542,7 @@ def apply_activation_checkpointing(
         mode: TransformerActivationCheckpointingMode,
         block_interval: Optional[int] = None,
         modules: Optional[List[str]] = None,
+        activation_memory_budget: Optional[float] = None,
     ):
         """
         Apply activation checkpointing to the model.
@@ -551,7 +552,20 @@ def apply_activation_checkpointing(
             which blocks are wrapped.
         :param modules: Required when :data:`mode` is "selected_modules". A list of modules names
             to wrap for activation checkpointing. Globs are supported.
+        :param activation_memory_budget: The memory budget for activation checkpointing in the range
+            [0, 1]. 0 corresponds to the memory usage when recomputing all activations, and 1
+            corresponds to the memory usage when recomputing no activations (which is the default).
+            Requires compilation to be enabled.
         """
+
+        if mode == TransformerActivationCheckpointingMode.budget:
+            if activation_memory_budget is None:
+                raise ValueError("'activation_memory_budget' is required for 'budget' mode")
+            if activation_memory_budget < 0 or activation_memory_budget > 1:
+                raise ValueError("'activation_memory_budget' must be in the range [0, 1]")
+            torch._functorch.config.activation_memory_budget = activation_memory_budget
+            return
+
         from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
             checkpoint_wrapper as ptd_checkpoint_wrapper,
         )

src/olmo_core/train/train_module/transformer/common.py

@@ -98,6 +98,7 @@ def parallelize_model(
                 ac_config.mode,
                 block_interval=ac_config.block_interval,
                 modules=ac_config.modules,
+                activation_memory_budget=ac_config.activation_memory_budget,
             )
         log.info(f"Applied '{ac_config.mode}' activation checkpointing to the model")
 

src/olmo_core/train/train_module/transformer/config.py

@@ -221,6 +221,14 @@ class TransformerActivationCheckpointingConfig(Config):
     activation checkpointing. Globs are supported.
     """
 
+    activation_memory_budget: Optional[float] = None
+    """
+    Required when :data:`mode` is "budget". Memory budget for activation checkpointing in range [0, 1].
+    0 = recompute all activations, 1 = recompute none (default). Requires compilation to be enabled.
+
+    See https://pytorch.org/blog/activation-checkpointing-techniques/ for more details.
+    """
+
     def __post_init__(self):
         if (
             self.mode == TransformerActivationCheckpointingMode.selected_blocks

src/olmo_core/train/train_module/transformer/train_module.py

@@ -37,6 +37,7 @@
 from olmo_core.float8 import Float8Config
 from olmo_core.nn.lm_head import LMOutputWithLoss
 from olmo_core.nn.transformer import Transformer
+from olmo_core.nn.transformer.config import TransformerActivationCheckpointingMode
 from olmo_core.optim import OptimConfig, SkipStepOptimizer
 from olmo_core.optim.scheduler import Scheduler
 from olmo_core.utils import gc_cuda, get_default_device, log_once, move_to_device
@@ -142,6 +143,15 @@ def __init__(
                 "Training parallelism configs are only valid for distributed training"
             )
 
+        if (
+            ac_config is not None
+            and ac_config.mode == TransformerActivationCheckpointingMode.budget
+            and not compile_model
+        ):
+            raise OLMoConfigurationError(
+                "Activation checkpointing with 'budget' mode requires compilation to be enabled"
+            )
+
         # Parallelize model.
         self.model = parallelize_model(
             model,