clean up

allenai · Feb 23, 2025 · c858268 · c858268
1 parent 32a5894
commit c858268
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/scripts/train/OLMo2-7B-long-context.py b/src/scripts/train/OLMo2-7B-long-context.py
@@ -12,8 +12,7 @@
 from olmo_core.optim import AdamWConfig, CosWithWarmup, OptimGroupOverride
 from olmo_core.train import TrainerConfig
 from olmo_core.train.callbacks import CheckpointerCallback, CometCallback, WandBCallback
-from olmo_core.train.train_module import (
-    TransformerActivationCheckpointingConfig,
+from olmo_core.train.train_module import (  # TransformerActivationCheckpointingConfig,
     TransformerContextParallelConfig,
     TransformerDataParallelConfig,
     TransformerDataParallelWrappingStrategy,
@@ -24,6 +23,7 @@
 
 
 CONTEXT_LENGTH = 4 * 16_384
+# 64K length, 32 GPUs -> 2,750 TPS
 
 
 def build_model_config(common: CommonComponents) -> TransformerConfig:
@@ -55,7 +55,7 @@ def build_train_module_config(common: CommonComponents) -> TransformerTrainModul
             wrapping_strategy=TransformerDataParallelWrappingStrategy.fine_grained,
         ),
         cp_config=TransformerContextParallelConfig(degree=8),
-        ac_config=TransformerActivationCheckpointingConfig(),
+        #  ac_config=TransformerActivationCheckpointingConfig(),
         float8_config=Float8Config(enabled=True),
         max_grad_norm=1.0,
         scheduler=CosWithWarmup(warmup_steps=2000),