From 3dd2e0af6ec75d3fbb02dedc191f094d75bd40ed Mon Sep 17 00:00:00 2001
From: root <23239305+b-chu@users.noreply.github.com>
Date: Tue, 20 Aug 2024 17:07:49 +0000
Subject: [PATCH] Change iteration timestamp for old checkpoints

---
 composer/trainer/trainer.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
index ffa0ff7225..2a65a40b43 100644
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -1943,6 +1943,18 @@ def __init__(
                     auto_microbatching,
                 )
 
+        # Set the iteration timestamp to the overall timestamp if loading from a checkpoint that was created before
+        # iteration was introduced in Composer v0.19.1. This is necessary to ensure that the iteration timestamp is
+        # accurate for callbacks and backwards compatibility for checkpoints.
+        if (
+            self.state.timestamp.iteration == 0 and self.state.timestamp.token_in_iteration == 0 and
+            self.state.timestamp.epoch_in_iteration == 0
+        ):
+            self.state.timestamp = self.state.timestamp.copy(
+                epoch_in_iteration=self.state.timestamp.epoch,
+                token_in_iteration=self.state.timestamp.token,
+            )
+
         self.engine.run_event(Event.AFTER_LOAD)
 
         # reseed here. This helps with a couple of issues: