Dev global rank

[wiederm]

Pull Request Summary

This PR addresses an issue introduced in PR #308 where logging of the epoch time was intended to occur only on the main process (rank == 0). However, this led to a synchronization error that caused training to freeze after the first epoch. To resolve this, we adjusted the logging to ensure it only occurs on the main process without disrupting the training flow.

Key changes

• Adjusted Logging Conditions: Updated specific logging calls to include rank_zero_only=True where necessary. This change ensures that logging occurs only on the main process without causing synchronization issues across processes.
• Removed Redundant Logging Checks: Refactored the logging of certain metrics to avoid redundant if self.global_rank == 0 checks by using rank_zero_only directly in self.log() calls.
• Ensured Sync Consistency: For logging operations that require distribution synchronization (e.g., gradients, learning rate), we used sync_dist=True for safe distributed handling, while maintaining rank_zero_only=True for process-specific logging.

Associated Issue(s)

• issue log time per epoch #308

Pull Request Checklist

• Issue(s) raised/addressed and linked
• Includes appropriate unit test(s)
• Appropriate docstring(s) added/updated
• Appropriate .rst doc file(s) added/updated
• PR is ready for review

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.

Project coverage is 85.48%. Comparing base (c1430f2) to head (8b23f18).

Additional details and impacted files