Move MaxText.globals

pylint

Author: hengtaoguo

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

src/maxtext/examples/sft_llama3_demo.ipynb

@@ -146,7 +146,7 @@
     "import datetime\n",
     "import os\n",
     "from MaxText import pyconfig\n",
-    "from MaxText.globals import MAXTEXT_PKG_DIR\n",
+    "from maxtext.utils.globals import MAXTEXT_PKG_DIR\n",
     "from maxtext.trainers.post_train.sft import train_sft\n",
     "import jax\n",
     "from huggingface_hub import login\n",

src/maxtext/examples/sft_train_and_evaluate.py

@@ -86,7 +86,7 @@
 from flax import nnx
 
 from MaxText import pyconfig
-from MaxText.globals import MAXTEXT_REPO_ROOT
+from maxtext.utils.globals import MAXTEXT_REPO_ROOT
 from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter
 from maxtext.input_pipeline import instruction_data_processing
 from maxtext.trainers.post_train.sft import train_sft

src/maxtext/experimental/agent/ckpt_conversion_agent/utils/save_param.py

@@ -28,7 +28,7 @@
 from transformers import AutoModelForCausalLM, AutoConfig
 
 from MaxText import pyconfig
-from MaxText.globals import MAXTEXT_PKG_DIR
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.inference.maxengine import maxengine
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils