"Adding block_until_ready to the _step_log call in _run_step to catch DATA_LOSS errors in the main thread instead of in the log emitting thread"

Author: lukebaumann

axlearn/common/trainer.py

@@ -1101,8 +1101,13 @@ def _run_step(
         if self.step % 100 == 0 or 0 <= self.step <= 5:
             self._step_log(
                 "loss=%s aux=%s",
-                outputs["loss"],
-                jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]),
+                jax.block_until_ready(outputs["loss"]),
+                jax.block_until_ready(
+                    jax.tree.map(
+                        lambda x: x.item() if x.ndim == 0 else f"T{x.shape}",
+                        outputs["aux"]
+                    ),
+                ),
             )
 
         self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]})