Round fractional window size

mosaicml · Aug 25, 2024 · c59e5be · c59e5be
1 parent 57800e6
commit c59e5be
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/llmfoundry/callbacks/kill_loss_spike_callback.py b/llmfoundry/callbacks/kill_loss_spike_callback.py
@@ -102,17 +102,17 @@ def fit_start(self, state: State, logger: Logger) -> None:
             if state.max_duration.unit == TimeUnit.EPOCH:
                 self.window_size = max(
                     MIN_WINDOW_SIZE,
-                    (state.dataloader_len * state.max_duration.value / 20),
+                    round(state.dataloader_len * state.max_duration.value / 20),
                 )
             elif state.max_duration.unit == TimeUnit.BATCH:
                 self.window_size = max(
                     MIN_WINDOW_SIZE,
-                    state.max_duration.value / 20,
+                    round(state.max_duration.value / 20),
                 )
             elif state.max_duration.unit == TimeUnit.TOKEN:
                 self.window_size = max(
                     MIN_WINDOW_SIZE,
-                    state.max_duration.value / 20,
+                    round(state.max_duration.value / 20),
                 )
 
     def batch_end(self, state: State, logger: Logger) -> None: