Maxtwin xm support.

This is to support g3 internal distillation runs.
Contains minor import and type hint fixes between github and g3 only environments.

PiperOrigin-RevId: 925513738

Author: Google-ML-Automation

src/maxtext/trainers/post_train/distillation/distillation_utils.py

@@ -31,7 +31,7 @@
 
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils
-# Reuse MaxText's native checkpointing logic
+# Reuse MaxText's native checkpointing logic.
 from maxtext.common.checkpointing import GrainCheckpointHandler, GrainCheckpointSave, GrainCheckpointRestore
 from tunix.sft import checkpoint_manager as tunix_checkpoint_manager
 from tunix.sft import peft_trainer

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -213,8 +213,6 @@ class MaxTextDistillationTrainer(peft_trainer.PeftTrainer):
   (positions, segment_ids) are passed to the model.
   """
 
-  checkpoint_manager: distillation_utils.MaxTextCheckpointManager | None
-
   def __init__(
       self,
       model,
@@ -234,6 +232,7 @@ def __init__(
     super().__init__(model=model, optimizer=dummy_optimizer, training_config=training_config, **kwargs)
 
     self.strategy = strategy
+    self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None
 
     # Per-step per-device TFLOPs (constants for the run): student fwd+bwd + teacher fwd-only.
     (