Merge pull request #3005 from AI-Hypercomputer:common_modules

PiperOrigin-RevId: 861378591

Author: Google-ML-Automation

src/MaxText/checkpointing.py src/MaxText/common/checkpointing.pysrc/MaxText/checkpointing.py renamed to src/MaxText/common/checkpointing.py

@@ -220,16 +220,10 @@ def create_orbax_checkpoint_manager(
   p.mkdir(exist_ok=True, parents=True)
   if enable_continuous_checkpointing:
     save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy()
-    preservation_policy = preservation_policy_lib.LatestN(
-        max_num_checkpoints_to_keep
-    )
+    preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
   else:
-    save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy(
-        interval=save_interval_steps
-    )
-    preservation_policy = preservation_policy_lib.LatestN(
-        max_num_checkpoints_to_keep
-    )
+    save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy(interval=save_interval_steps)
+    preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
   manager = CheckpointManager(
       p,
       item_names=item_names,
@@ -239,7 +233,7 @@ def create_orbax_checkpoint_manager(
           enable_async_checkpointing=use_async,
           save_decision_policy=save_decision_policy,
           preservation_policy=preservation_policy,
-          ),
+      ),
       logger=orbax_logger,
   )
 
@@ -276,12 +270,8 @@ def create_orbax_emergency_checkpoint_manager(
       global_mesh=global_mesh,
       abstract_state=abstract_state,
       options=emergency_checkpoint_manager.CheckpointManagerOptions(
-          local=LocalCheckpointOptions(
-              save_interval_steps=local_save_interval_steps
-          ),
-          persistent=PersistentCheckpointOptions(
-              save_interval_steps=persistent_save_interval_steps
-          ),
+          local=LocalCheckpointOptions(save_interval_steps=local_save_interval_steps),
+          persistent=PersistentCheckpointOptions(save_interval_steps=persistent_save_interval_steps),
       ),
       logger=orbax_logger,
   )

src/MaxText/data_loader.py src/MaxText/common/data_loader.pysrc/MaxText/data_loader.py renamed to src/MaxText/common/data_loader.py

@@ -21,7 +21,7 @@
 
 from MaxText import exceptions
 from MaxText.sharding import get_input_data_sharding
-from MaxText.utils.goodput_utils import (
+from MaxText.common.goodput import (
     GoodputEvent,
     maybe_record_goodput,
 )

src/MaxText/gcp_workload_monitor.py …c/MaxText/common/gcp_workload_monitor.pysrc/MaxText/gcp_workload_monitor.py renamed to src/MaxText/common/gcp_workload_monitor.py src/MaxText/gcp_workload_monitor.py renamed to src/MaxText/common/gcp_workload_monitor.py

@@ -32,9 +32,8 @@
 from MaxText import maxtext_utils
 from MaxText.managed_mldiagnostics import ManagedMLDiagnostics
 from MaxText.utils import gcs_utils
-from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
 from MaxText.globals import EPS
-
+from MaxText.common.gcp_workload_monitor import GCPWorkloadMonitor
 from collections import defaultdict
 
 # Mapping MaxText metrics to managed profiler metrics

src/MaxText/utils/goodput_utils.py src/MaxText/common/goodput.pysrc/MaxText/utils/goodput_utils.py renamed to src/MaxText/common/goodput.py

@@ -26,9 +26,9 @@
 from MaxText import max_utils
 from MaxText import maxengine
 from MaxText import pyconfig
-from MaxText import profiler
 from MaxText import multimodal_utils
 from MaxText.multimodal import preprocessor
+from maxtext.common import profiler
 # Placeholder: internal
 
 # Number of text sequences to process in a single batch.

src/MaxText/metric_logger.py src/MaxText/common/metric_logger.pysrc/MaxText/metric_logger.py renamed to src/MaxText/common/metric_logger.py

@@ -60,27 +60,26 @@
 
 import tensorflow as tf
 
-from MaxText import checkpointing
 from MaxText import exceptions
 from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import train_utils
 from MaxText import max_logging
-from MaxText import profiler
 from MaxText import pyconfig
-from MaxText.data_loader import DataLoader
-from MaxText.metric_logger import MetricLogger
 from MaxText.train import get_first_step
 from MaxText.train_utils import setup_train_loop
 from MaxText.train import train_step
 from MaxText.train_utils import validate_train_config
-from MaxText.utils.goodput_utils import (
+from maxtext.common import checkpointing, profiler
+from MaxText.common.data_loader import DataLoader
+from MaxText.common.goodput import (
     GoodputEvent,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
 )
-from MaxText.vertex_tensorboard import VertexTensorboardManager
+from MaxText.common.metric_logger import MetricLogger
+from MaxText.common.vertex_tensorboard import VertexTensorboardManager
 
 logging.basicConfig()
 logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.INFO)

src/MaxText/profiler.py src/MaxText/common/profiler.pysrc/MaxText/profiler.py renamed to src/MaxText/common/profiler.py

@@ -67,33 +67,30 @@
 from ml_goodput_measurement.src.goodput import GoodputRecorder
 
 import MaxText as mt
-from MaxText import checkpointing
 from MaxText import exceptions
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import sharding
 from MaxText import train_utils
-from MaxText import profiler
 from MaxText import pyconfig
-from MaxText.checkpointing import CheckpointManager
 from MaxText.utils import gcs_utils
 from MaxText.inference import offline_engine
-from MaxText.data_loader import DataLoader
 from MaxText.experimental.rl import grpo_input_pipeline
 from MaxText.experimental.rl import grpo_utils
 from MaxText.globals import EPS
-from MaxText.metric_logger import MetricLogger
 from MaxText.train import get_first_step
 from MaxText.train_utils import validate_train_config
-from MaxText.utils.goodput_utils import (
+from maxtext.common import checkpointing, profiler
+from MaxText.common.data_loader import DataLoader
+from MaxText.common.goodput import (
     GoodputEvent,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
 )
-from MaxText.vertex_tensorboard import VertexTensorboardManager
-
+from MaxText.common.metric_logger import MetricLogger
+from MaxText.common.vertex_tensorboard import VertexTensorboardManager
 
 # pylint: disable=too-many-positional-arguments
 
@@ -505,7 +502,7 @@ def setup_train_loop(
     recorder: GoodputRecorder,
 ) -> tuple[
     jax.Array,
-    CheckpointManager,
+    checkpointing.CheckpointManager,
     TrainState,
     TrainState,
     mt.Transformer,