clean up

allenai · Feb 24, 2025 · 2d84add · 2d84add
1 parent c858268
commit 2d84add
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/olmo_core/utils.py b/src/olmo_core/utils.py
@@ -392,6 +392,11 @@ def filter_warnings():
         message="Synchronization debug mode is a prototype feature.*",
         module="torch.cuda",
     )
+    warnings.filterwarnings(
+        action="ignore",
+        category=UserWarning,
+        message="TORCH_NCCL_AVOID_RECORD_STREAMS=1 has no effect .*",
+    )
     warnings.filterwarnings(
         action="ignore",
         category=FutureWarning,

diff --git a/src/scripts/train/OLMo2-7B-long-context.py b/src/scripts/train/OLMo2-7B-long-context.py
@@ -12,7 +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 (
     TransformerContextParallelConfig,
     TransformerDataParallelConfig,
     TransformerDataParallelWrappingStrategy,
@@ -55,7 +55,6 @@ def build_train_module_config(common: CommonComponents) -> TransformerTrainModul
             wrapping_strategy=TransformerDataParallelWrappingStrategy.fine_grained,
         ),
         cp_config=TransformerContextParallelConfig(degree=8),
-        #  ac_config=TransformerActivationCheckpointingConfig(),
         float8_config=Float8Config(enabled=True),
         max_grad_norm=1.0,
         scheduler=CosWithWarmup(warmup_steps=2000),