Add to other official scripts

allenai · Sep 4, 2024 · 26ed9b2 · 26ed9b2
1 parent e8fdbea
commit 26ed9b2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/scripts/train/OLMo-13B.py b/src/scripts/train/OLMo-13B.py
@@ -32,6 +32,7 @@
     ConfigSaverCallback,
     GPUMemoryMonitorCallback,
     GradClipperCallback,
+    ProfilerCallback,
     SchedulerCallback,
     SpeedMonitorCallback,
     WandBCallback,
@@ -176,6 +177,7 @@ def build_config(run_name: str, cluster: str, overrides: List[str]) -> Experimen
             ),
         )
         .with_callback("config_saver", ConfigSaverCallback())
+        .with_callback("profiler", ProfilerCallback(enabled=False))
     )
 
     return ExperimentConfig(

diff --git a/src/scripts/train/OLMo-7B.py b/src/scripts/train/OLMo-7B.py
@@ -32,6 +32,7 @@
     ConfigSaverCallback,
     GPUMemoryMonitorCallback,
     GradClipperCallback,
+    ProfilerCallback,
     SchedulerCallback,
     SpeedMonitorCallback,
     WandBCallback,
@@ -176,6 +177,7 @@ def build_config(run_name: str, cluster: str, overrides: List[str]) -> Experimen
             ),
         )
         .with_callback("config_saver", ConfigSaverCallback())
+        .with_callback("profiler", ProfilerCallback(enabled=False))
     )
 
     return ExperimentConfig(