Merge pull request #3055 from AI-Hypercomputer:move-maxtext-integration

PiperOrigin-RevId: 873084919

Author: Google-ML-Automation

src/maxtext/examples/sft_qwen3_demo.ipynb

@@ -200,7 +200,7 @@
     "import MaxText\n",
     "from MaxText import pyconfig\n",
     "from maxtext.examples.sft_train_and_evaluate import evaluate_model, get_test_dataset\n",
-    "from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter\n",
+    "from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter\n",
     "from MaxText.globals import MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR\n",
     "from maxtext.trainers.post_train.sft import train_sft\n",
     "\n",

src/maxtext/examples/sft_train_and_evaluate.py

@@ -87,7 +87,7 @@
 
 from MaxText import pyconfig
 from MaxText.globals import MAXTEXT_REPO_ROOT
-from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter
+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
 from maxtext.utils import max_logging

src/MaxText/integration/__init__.py src/maxtext/integration/__init__.pysrc/MaxText/integration/__init__.py renamed to src/maxtext/integration/__init__.py

@@ -26,7 +26,7 @@
 from flax import nnx
 from jax import Array
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS  # pylint: disable=ungrouped-imports
-from MaxText.integration.tunix.utils import VllmWeightMapping
+from maxtext.integration.tunix.utils import VllmWeightMapping
 from maxtext.models.models import Transformer
 
 

src/MaxText/integration/tunix/__init__.py src/maxtext/integration/tunix/__init__.pysrc/MaxText/integration/tunix/__init__.py renamed to src/maxtext/integration/tunix/__init__.py

@@ -18,10 +18,10 @@
 dispatcher to retrieve the correct weight mapping configuration for a given
 model name. This allows for easy extension to support new models.
 """
-from MaxText.integration.tunix.weight_mapping.deepseek3 import DEEPSEEK_VLLM_MAPPING
-from MaxText.integration.tunix.weight_mapping.gpt_oss import GPT_OSS_VLLM_MAPPING
-from MaxText.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING
-from MaxText.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING
+from maxtext.integration.tunix.weight_mapping.deepseek3 import DEEPSEEK_VLLM_MAPPING
+from maxtext.integration.tunix.weight_mapping.gpt_oss import GPT_OSS_VLLM_MAPPING
+from maxtext.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING
+from maxtext.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING
 
 
 class StandaloneVllmWeightMapping: