fused loss

allenai · Sep 9, 2024 · db448c8 · db448c8
1 parent a69f0bd
commit db448c8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/scripts/train/OLMo-13B.py b/src/scripts/train/OLMo-13B.py
@@ -60,6 +60,7 @@ def build_trainer_config(common: CommonComponents) -> TrainerConfig:
             metrics_collect_interval=10,
             cancel_check_interval=1,
             z_loss_multiplier=1e-5,
+            fused_loss=True,
         )
         .with_callback(
             "lr_scheduler", SchedulerCallback(scheduler=CosWithWarmup(warmup_steps=2000))