diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py index 478062e8e..feb879d0b 100644 --- a/torchtitan/components/checkpoint.py +++ b/torchtitan/components/checkpoint.py @@ -292,7 +292,6 @@ def load_state_dict(state_dict): else: self.purge_thread = None - self.mp = None self.staging_future = None self.save_future = None if async_mode == AsyncMode.DISABLED: @@ -315,9 +314,6 @@ def __del__(self): def close(self): if hasattr(self, "enable_checkpoint") and self.enable_checkpoint: - if hasattr(self, "mp") and self.mp and self.mp.is_alive(): - self.mp_queue_send.put(Terminate()) - self.mp.join() if ( hasattr(self, "purge_thread") and self.purge_thread