[levanter] Wait for final checkpoint async save before process exit

The final checkpoint save in Levanter is async — the GCS write runs in a
background thread while the process continues.  In train_dpo.py, the
post-training code attempted to wait via `wait_until_finished()`, but
called it on a *newly created* Checkpointer instance (which had nothing
to wait for) instead of the original one that holds the pending save.
train_lm.py had no wait at all.

If preemption (or a normal exit) kills the process while the async write
is still in flight, the only checkpoint at the target training step is
silently lost.  This is especially damaging when the target step is not a
multiple of the permanent checkpoint interval (e.g. step 9917 with
`keep=[dict(every=10000)]`), since every checkpoint below 10000 is
temporary and subject to deletion on restart.

Fix:
- Store the Checkpointer as `self._checkpointer` on the Trainer so it is
  accessible after training.
- train_lm.py: add `trainer._checkpointer.wait_until_finished()` after
  training completes.
- train_dpo.py: replace the broken new-instance wait with
  `trainer._checkpointer.wait_until_finished()`.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Calvin-Xu

lib/levanter/src/levanter/main/train_dpo.py

@@ -488,12 +488,14 @@ def save_policy_hf_checkpoint(step):
         else:
             train_loader = train_loader.iter_from_step(0)
 
-        last_info = trainer.train(state, train_loader)
+        trainer.train(state, train_loader)
 
+        # Ensure the final checkpoint's async GCS write completes before the
+        # process exits.  Without this, preemption (or a normal exit race) can
+        # kill the process while the write is still in flight, silently losing
+        # the only checkpoint at the target training step.
         if trainer.config.checkpointer is not None:
-            trainer.run_hooks(last_info, force=True)
-            checkpointer = trainer.config.checkpointer.create(trainer.run_id)
-            checkpointer.wait_until_finished()
+            trainer._checkpointer.wait_until_finished()
 
     trainer.tracker.finish()
 

lib/levanter/src/levanter/main/train_lm.py

@@ -319,6 +319,13 @@ def compute_logits(model: LmHeadModel, example: LmExample):
         ## OK, actually run training!
         trainer.train(state, train_loader)
 
+        # Ensure the final checkpoint's async GCS write completes before the
+        # process exits.  Without this, preemption (or a normal exit race) can
+        # kill the process while the write is still in flight, silently losing
+        # the only checkpoint at the target training step.
+        if trainer.config.checkpointer is not None:
+            trainer._checkpointer.wait_until_finished()
+
     # This isn't necessary except when Levanter is run in a subprocess (as happens w/ ray)
     trainer.tracker.finish()
 

lib/levanter/src/levanter/trainer.py

@@ -579,10 +579,10 @@ def _add_default_hooks(self):
         self.add_hook(levanter.callbacks.pbar_logger(total=self.config.num_train_steps), every=1)
         self.add_hook(levanter.callbacks.log_step_info(self.config.num_train_steps), every=1)
         # engine.add_hook(callbacks.log_memory_usage(), every=1)
-        checkpointer = self.config.checkpointer.create(self.run_id)
+        self._checkpointer = self.config.checkpointer.create(self.run_id)
 
         def checkpoint_hook(info, force=False):
-            checkpointer.on_step(tree=info.state.saveable_state, step=info.step, force=force)
+            self._checkpointer.on_step(tree=info.state.saveable_state, step=info.step, force=force)
 
         self.add_hook(checkpoint_hook, every=1)  # checkpointer manages its own frequency