Merge pull request #3249 from AI-Hypercomputer:hengtaoguo-globals

PiperOrigin-RevId: 875405670

Author: Google-ML-Automation

src/MaxText/get_flops.py

@@ -15,7 +15,7 @@
 """ A wrapper file for easily calculating training TFLOPs. """
 
 from MaxText import pyconfig
-from MaxText.globals import MAXTEXT_PKG_DIR
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.utils.maxtext_utils import calculate_tflops_training_per_device
 import os
 from typing import Sequence, cast

src/MaxText/pyconfig.py

@@ -29,7 +29,7 @@
 import omegaconf
 
 from MaxText import pyconfig_deprecated
-from MaxText.globals import MAXTEXT_CONFIGS_DIR
+from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
 from maxtext.common.common_types import DecoderBlockType, ShardMode
 from maxtext.configs import types
 from maxtext.configs.types import MaxTextConfig

src/MaxText/pyconfig_deprecated.py

@@ -29,7 +29,7 @@
 import omegaconf
 
 from MaxText import accelerator_to_spec_map
-from MaxText.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR
+from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR
 from maxtext.common.common_types import AttentionType, DecoderBlockType, ShardMode
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -43,7 +43,7 @@
 from jax import random
 from jax.sharding import Mesh
 from MaxText import pyconfig
-from MaxText.globals import MAXTEXT_PKG_DIR
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.common import checkpointing
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.layers import quantizations

src/maxtext/common/checkpointing.py

@@ -22,7 +22,7 @@
 from etils import epath
 from flax.training import train_state
 import jax
-from MaxText.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
 from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator
 from maxtext.input_pipeline.multihost_dataloading import RemoteIterator
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator

src/maxtext/common/metric_logger.py

@@ -25,7 +25,7 @@
 
 import jax
 
-from MaxText.globals import EPS
+from maxtext.utils.globals import EPS
 from maxtext.common.gcloud_stub import mldiagnostics_modules
 from maxtext.common.gcloud_stub import workload_monitor
 from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics

src/maxtext/configs/types.py

@@ -30,8 +30,8 @@
 from maxtext.common.common_types import AttentionType, DecoderBlockType, ShardMode
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
+from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
 from MaxText import accelerator_to_spec_map
-from MaxText.globals import MAXTEXT_ASSETS_ROOT
 from pydantic.config import ConfigDict
 from pydantic.fields import Field
 from pydantic.functional_validators import field_validator, model_validator