Internal change

SurbhiJainUSC · Google-ML-Automation · commit d60808f5418b · 2026-01-26T16:46:38.000-08:00
PiperOrigin-RevId: 861412305
diff --git a/src/MaxText/elastic_train.py b/src/MaxText/elastic_train.py
@@ -71,15 +71,15 @@
 from MaxText.train import train_step
 from MaxText.train_utils import validate_train_config
 from maxtext.common import checkpointing, profiler
-from MaxText.common.data_loader import DataLoader
-from MaxText.common.goodput import (
+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
-from MaxText.common.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)
diff --git a/src/MaxText/experimental/rl/grpo_trainer.py b/src/MaxText/experimental/rl/grpo_trainer.py
@@ -82,15 +82,15 @@
 from MaxText.train import get_first_step
 from MaxText.train_utils import validate_train_config
 from maxtext.common import checkpointing, profiler
-from MaxText.common.data_loader import DataLoader
-from MaxText.common.goodput import (
+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
-from MaxText.common.vertex_tensorboard import VertexTensorboardManager
+from maxtext.common.metric_logger import MetricLogger
+from maxtext.common.vertex_tensorboard import VertexTensorboardManager
 
 # pylint: disable=too-many-positional-arguments
 
diff --git a/src/MaxText/gcloud_stub.py b/src/MaxText/gcloud_stub.py
@@ -445,7 +445,7 @@ def workload_monitor():
     return _workload_monitor_stub()
 
   try:
-    from MaxText.common.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.common.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
diff --git a/src/MaxText/sft_trainer.py b/src/MaxText/sft_trainer.py
@@ -42,14 +42,14 @@
 from MaxText.train_utils import setup_train_loop, validate_train_config
 from MaxText.utils import gcs_utils
 from maxtext.common import checkpointing, profiler
-from MaxText.common.data_loader import DataLoader
-from MaxText.common.goodput import (
+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
+from maxtext.common.metric_logger import MetricLogger
 
 
 def train_loop(config, recorder, state=None):
diff --git a/src/MaxText/train.py b/src/MaxText/train.py
@@ -62,14 +62,14 @@
 # pylint: disable=too-many-positional-arguments
 
 from maxtext.common import checkpointing, profiler
-from MaxText.common.goodput import (
+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
+from maxtext.common.metric_logger import MetricLogger, record_activation_metrics
+from maxtext.common.vertex_tensorboard import VertexTensorboardManager
 
 
 def get_first_step(state):
diff --git a/src/MaxText/train_utils.py b/src/MaxText/train_utils.py
@@ -26,8 +26,8 @@
 from MaxText.rampup_batch import create_rampup_manager
 from MaxText import model_creation_utils
 from maxtext.common import checkpointing
-from MaxText.common.data_loader import create_dataloader
-from MaxText.common.goodput import GoodputEvent, maybe_record_goodput
+from maxtext.common.data_loader import create_dataloader
+from maxtext.common.goodput import GoodputEvent, maybe_record_goodput
 
 
 def create_training_tools(config, model, mesh):
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py
@@ -21,7 +21,7 @@
 
 from MaxText import exceptions
 from MaxText.sharding import get_input_data_sharding
-from MaxText.common.goodput import (
+from maxtext.common.goodput import (
     GoodputEvent,
     maybe_record_goodput,
 )
diff --git a/src/maxtext/common/gcp_workload_monitor.py b/src/maxtext/common/gcp_workload_monitor.py
diff --git a/src/maxtext/common/goodput.py b/src/maxtext/common/goodput.py
diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
@@ -33,7 +33,7 @@
 from MaxText.managed_mldiagnostics import ManagedMLDiagnostics
 from MaxText.utils import gcs_utils
 from MaxText.globals import EPS
-from MaxText.common.gcp_workload_monitor import GCPWorkloadMonitor
+from maxtext.common.gcp_workload_monitor import GCPWorkloadMonitor
 from collections import defaultdict
 
 # Mapping MaxText metrics to managed profiler metrics
diff --git a/src/maxtext/common/profiler.py b/src/maxtext/common/profiler.py
diff --git a/src/maxtext/common/vertex_tensorboard.py b/src/maxtext/common/vertex_tensorboard.py
diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py
@@ -38,9 +38,9 @@
 from MaxText import sharding
 from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
 from MaxText.utils import gcs_utils
-from MaxText.common.data_loader import DataLoader
-from MaxText.common.goodput import GoodputEvent, record_goodput
-from MaxText.common.metric_logger import MetricLogger, MetadataKey
+from maxtext.common.data_loader import DataLoader
+from maxtext.common.goodput import GoodputEvent, record_goodput
+from maxtext.common.metric_logger import MetricLogger, MetadataKey
 
 
 class SFTTrainingHooks(TrainingHooks):
diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py
@@ -56,13 +56,13 @@
 from MaxText import pyconfig
 from MaxText import model_creation_utils
 from MaxText.train import loss_fn
-from MaxText.common.goodput import (
+from maxtext.common.goodput import (
     GoodputEvent,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
 )
-from MaxText.trainers.post_train.sft import hooks
+from maxtext.trainers.post_train.sft import hooks
 
 
 def get_tunix_config(mt_config):
diff --git a/tests/unit/data_loader_test.py b/tests/unit/data_loader_test.py
@@ -28,7 +28,7 @@
 from MaxText import exceptions
 from MaxText import pyconfig
 from MaxText.globals import MAXTEXT_PKG_DIR
-from MaxText.common.data_loader import DataLoader, RampUpDataLoader
+from maxtext.common.data_loader import DataLoader, RampUpDataLoader
 
 
 class DataLoaderTest(unittest.TestCase):
diff --git a/tests/unit/goodput_utils_test.py b/tests/unit/goodput_utils_test.py
@@ -19,7 +19,7 @@
 from MaxText import pyconfig
 from MaxText.globals import MAXTEXT_PKG_DIR
 from unittest import mock
-from MaxText.common.goodput import create_goodput_recorder, maybe_monitor_goodput, maybe_record_goodput, GoodputEvent
+from maxtext.common.goodput import create_goodput_recorder, maybe_monitor_goodput, maybe_record_goodput, GoodputEvent
 
 
 class GoodputUtilsTest(unittest.TestCase):
diff --git a/tools/gcs_benchmarks/standalone_dataloader.py b/tools/gcs_benchmarks/standalone_dataloader.py
@@ -31,7 +31,7 @@
 from MaxText import pyconfig
 from MaxText.train import get_first_step
 from MaxText.train_utils import validate_train_config, setup_train_loop
-from MaxText.common.data_loader import DataLoader
+from maxtext.common.data_loader import DataLoader
 
 
 def data_load_loop(config, state=None):

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@`
`21`	`21`
`22`	`22`	`from MaxText import exceptions`
`23`	`23`	`from MaxText.sharding import get_input_data_sharding`
`24`		`-from MaxText.common.goodput import (`
	`24`	`+from maxtext.common.goodput import (`
`25`	`25`	`GoodputEvent,`
`26`	`26`	`maybe_record_goodput,`
`27`	`27`	`)`