Merge pull request #453 from hxdtest/only_rank0_log_metrics

Only rank0 log metrics to console
allenai · Mar 12, 2024 · ed47c29 · ed47c29
2 parents ad8198e + c699236
commit ed47c29
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,6 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Removed `AMDLayerNorm`, since the original layer norm bug has been fixed and we don't need this workaround anymore.
 
+### Fixed
+
+- Don't log garbage on nodes that aren't rank 0
+
 
 ## [v0.2.5](https://github.com/allenai/OLMo/releases/tag/v0.2.5) - 2024-03-06
 

diff --git a/olmo/train.py b/olmo/train.py
@@ -959,7 +959,10 @@ def on_trace_ready(p):
 
                     # Log metrics to console.
                     if self.global_step % self.cfg.console_log_interval == 0:
-                        self.log_metrics_to_console(f"[step={self.global_step}/{self.max_steps}]", metrics)
+                        if get_global_rank() == 0:
+                            self.log_metrics_to_console(f"[step={self.global_step}/{self.max_steps}]", metrics)
+                        else:
+                            log.info(f"[step={self.global_step}/{self.max_steps}]")
 
                     # Log metrics to W&B.
                     if (