Merge pull request #3266 from AI-Hypercomputer:hengtaoguo-pyconfig

PiperOrigin-RevId: 875986137

Author: Google-ML-Automation

benchmarks/api_server/maxtext_generator.py

@@ -34,7 +34,7 @@
 
 from dataclasses import dataclass, field
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.inference.maxengine import maxengine
 from maxtext.multimodal import processor as mm_processor
 from maxtext.multimodal import utils as mm_utils

benchmarks/mmlu/mmlu_eval.py

@@ -56,7 +56,7 @@
 
 from tqdm import tqdm
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.inference.maxengine import maxengine
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils

src/MaxText/__init__.py

@@ -35,7 +35,7 @@
 
 from jax.sharding import Mesh
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.models import models
 from maxtext.trainers.post_train.dpo import dpo_utils
 from maxtext.utils import maxtext_utils

src/MaxText/maxengine_server.py

@@ -21,7 +21,7 @@
 import jax
 from absl import logging
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.utils import max_logging
 
 OLD_MODULE_PATH = "MaxText.maxengine_server"

src/MaxText/sft_trainer.py

@@ -27,7 +27,7 @@
 
 from flax.linen import partitioning as nn_partitioning
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.trainers.pre_train.train import (
     eval_step,
     get_first_step,

src/maxtext/checkpoint_conversion/compare_hf_ckpt.py

@@ -47,7 +47,7 @@
 from safetensors.torch import load as load_safetensors
 from safetensors import safe_open
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.checkpoint_conversion.utils.utils import HF_IDS, print_ram_usage, get_hf_model
 from maxtext.utils import max_logging
 

src/maxtext/checkpoint_conversion/load_and_quantize_checkpoint.py

@@ -21,7 +21,7 @@
 
 import jax
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.inference.maxengine import maxengine
 from maxtext.utils import max_utils
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

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

src/maxtext/checkpoint_conversion/standalone_scripts/llama_mistral_mixtral_orbax_to_hf.py

@@ -47,7 +47,7 @@
 
 from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.utils.generate_param_only_checkpoint import _read_train_checkpoint
 from maxtext.checkpoint_conversion.standalone_scripts import llama_or_mistral_ckpt
 from maxtext.common import checkpointing

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -66,7 +66,7 @@
 from absl import app
 from absl import flags
 
-from MaxText import pyconfig
+from maxtext.configs import pyconfig
 from maxtext.checkpoint_conversion.utils.param_mapping import (
     HOOK_FNS,
     PARAM_MAPPING,