Move common modules to src/maxtext/common

Author: SurbhiJainUSC

src/MaxText/decode.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/elastic_train.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/experimental/rl/grpo_trainer.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,

src/MaxText/gcloud_stub.py

@@ -445,7 +445,7 @@ def workload_monitor():
     return _workload_monitor_stub()
 
   try:
-    from MaxText.gcp_workload_monitor import GCPWorkloadMonitor  # type: ignore  # pylint: disable=import-outside-toplevel
+    from maxtext.common.gcp_workload_monitor import GCPWorkloadMonitor  # type: ignore  # pylint: disable=import-outside-toplevel
 
     return GCPWorkloadMonitor, False
   except Exception:  # ModuleNotFoundError / ImportError  # pylint: disable=broad-exception-caught
@@ -484,7 +484,7 @@ def vertex_tensorboard_components():
     return _vertex_tb_stub()
 
   try:
-    from MaxText.vertex_tensorboard import VertexTensorboardManager  # type: ignore  # pylint: disable=import-outside-toplevel
+    from maxtext.common.vertex_tensorboard import VertexTensorboardManager  # type: ignore  # pylint: disable=import-outside-toplevel
 
     return VertexTensorboardManager, False
   except Exception:  # pylint: disable=broad-exception-caught

src/MaxText/generate_param_only_checkpoint.py

@@ -32,7 +32,6 @@
 from jax.sharding import Mesh
 from jax import random
 
-from MaxText import checkpointing
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import maxtext_utils
@@ -42,6 +41,7 @@
 from MaxText.layers import models, quantizations
 from MaxText.utils import gcs_utils
 from MaxText.utils import lora_utils
+from maxtext.common import checkpointing
 
 Transformer = models.transformer_as_linen
 

src/MaxText/inference_microbenchmark.py

@@ -26,9 +26,9 @@
 from MaxText import maxengine
 from MaxText import maxtext_utils
 from MaxText import prefill_packing
-from MaxText import profiler
 from MaxText import pyconfig
 from MaxText.utils import gcs_utils
+from maxtext.common import profiler
 
 import warnings
 

src/MaxText/layerwise_quantization.py

@@ -43,13 +43,13 @@
 from flax.linen import partitioning as nn_partitioning
 from flax import nnx
 
-from MaxText import checkpointing
 from MaxText import common_types
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import pyconfig
 from MaxText.layers import models, quantizations, deepseek
+from maxtext.common import checkpointing
 import orbax.checkpoint as ocp
 
 IGNORE = ocp.PLACEHOLDER

src/MaxText/maxtext_utils.py

@@ -35,14 +35,14 @@
 import orbax.checkpoint.experimental.emergency.checkpoint_manager as emergency_checkpoint_manager
 import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager
 
-from MaxText import checkpointing
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import multimodal_utils
 from MaxText import sharding
 from MaxText.configs import types
 from MaxText.common_types import DecoderBlockType, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
 from MaxText.inference.page_manager import PageState
+from maxtext.common import checkpointing
 
 OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient"
 

src/MaxText/sft_trainer.py

@@ -27,30 +27,29 @@
 
 from flax.linen import partitioning as nn_partitioning
 
-from MaxText import checkpointing
 from MaxText import exceptions
 from MaxText import max_utils
 from MaxText import max_logging
 from MaxText import maxtext_utils
-from MaxText import profiler
 from MaxText import pyconfig
 from MaxText import train_utils
 from MaxText import sharding
-from MaxText.data_loader import DataLoader
-from MaxText.metric_logger import MetricLogger
 from MaxText.train import (
     eval_step,
     get_first_step,
     train_step,
 )
 from MaxText.train_utils import setup_train_loop, validate_train_config
 from MaxText.utils import gcs_utils
-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.common.metric_logger import MetricLogger
 
 
 def train_loop(config, recorder, state=None):

src/MaxText/train.py

@@ -42,36 +42,35 @@
 from cloud_tpu_diagnostics.configuration import diagnostic_configuration
 from cloud_tpu_diagnostics.configuration import stack_trace_configuration
 
-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 train_utils
-from MaxText import profiler
 from MaxText import pyconfig
 from MaxText import sharding
 from MaxText.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
 from MaxText.common_types import ShardMode
 from MaxText.globals import EPS
-from MaxText.metric_logger import MetricLogger
 from MaxText.utils import gcs_utils
-from MaxText.utils.goodput_utils import (
-    GoodputEvent,
-    create_goodput_recorder,
-    maybe_monitor_goodput,
-    maybe_record_goodput,
-)
-from MaxText.vertex_tensorboard import VertexTensorboardManager
 # Placeholder: internal
 
 from MaxText.gradient_accumulation import gradient_accumulation_loss_and_grad
 from MaxText.vocabulary_tiling import vocab_tiling_linen_loss
 from MaxText.dpo_utils import _merge_dpo_state, _split_dpo_state, dpo_loss_fn
 from MaxText.train_utils import validate_train_config
-from MaxText.metric_logger import record_activation_metrics
 # pylint: disable=too-many-positional-arguments
 
+from maxtext.common import checkpointing, profiler
+from maxtext.common.goodput import (
+    GoodputEvent,
+    create_goodput_recorder,
+    maybe_monitor_goodput,
+    maybe_record_goodput,
+)
+from maxtext.common.metric_logger import MetricLogger, record_activation_metrics
+from maxtext.common.vertex_tensorboard import VertexTensorboardManager
+
 
 def get_first_step(state):
   return int(state.step)