Merge pull request #2684 from AI-Hypercomputer:sft_metrics

PiperOrigin-RevId: 832084340

Author: Google-ML-Automation

src/MaxText/sft/sft_trainer.py

@@ -46,7 +46,7 @@
 
 from orbax import checkpoint as ocp
 
-from tunix.sft import peft_trainer, profiler
+from tunix.sft import metrics_logger, peft_trainer, profiler
 
 from MaxText import max_utils
 from MaxText import max_logging
@@ -80,7 +80,7 @@ def get_tunix_config(mt_config):
   )
 
   # Metrics configurations
-  metrics_logging_options = peft_trainer.metrics_logger.MetricsLoggerOptions(log_dir=mt_config.tensorboard_dir)
+  metrics_logging_options = metrics_logger.MetricsLoggerOptions(log_dir=mt_config.tensorboard_dir)
 
   # Profiler configurations
   profiler_options = None

src/MaxText/utils/ckpt_conversion/to_maxtext.py

@@ -61,7 +61,6 @@
 import threading
 from functools import partial
 from typing import Sequence, List, Any, Callable
-from MaxText.inference_utils import str2bool
 import numpy as np
 import jax
 import psutil
@@ -80,6 +79,7 @@
 from MaxText import maxtext_utils
 from MaxText import pyconfig
 from MaxText.common_types import MODEL_MODE_TRAIN
+from MaxText.inference_utils import str2bool
 from MaxText.layers import models, quantizations
 from MaxText.checkpointing import save_checkpoint
 from MaxText.utils.ckpt_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING