Merge pull request #3054 from AI-Hypercomputer:move-maxtext-layers-2

PiperOrigin-RevId: 872641703

Author: Google-ML-Automation

src/MaxText/__init__.py

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

src/MaxText/generate_param_only_checkpoint.py

@@ -25,18 +25,16 @@
 from typing import Sequence
 
 from absl import app
-
 from etils import epath
-
 import jax
-from jax.sharding import Mesh
 from jax import random
-
+from jax.sharding import Mesh
 from MaxText import optimizers
 from MaxText import pyconfig
-from MaxText.common_types import DecoderBlockType, MODEL_MODE_TRAIN
-from MaxText.layers import models, quantizations
 from maxtext.common import checkpointing
+from MaxText.common_types import DecoderBlockType, MODEL_MODE_TRAIN
+from maxtext.layers import quantizations
+from maxtext.models import models
 from maxtext.utils import gcs_utils
 from maxtext.utils import lora_utils
 from maxtext.utils import max_logging

src/MaxText/integration/tunix/tunix_adapter.py

@@ -21,13 +21,13 @@
 
 from __future__ import annotations
 
-from typing import Optional, Tuple, Any
+from typing import Any, Optional, Tuple
 
-from jax import Array
 from flax import nnx
-from MaxText.layers.models import Transformer
-from MaxText.integration.tunix.utils import VllmWeightMapping
+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.models.models import Transformer
 
 
 class TunixMaxTextAdapter(nnx.Module):

src/MaxText/layerwise_quantization.py

@@ -33,24 +33,22 @@
 import os
 from typing import Any, Sequence
 
-from tqdm import tqdm
-
-import jax
-import jax.numpy as jnp
 from absl import app
 from aqt.jax.v2 import aqt_tensor
-
-from flax.linen import partitioning as nn_partitioning
 from flax import nnx
-
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
 from MaxText import common_types
 from MaxText import pyconfig
-from MaxText.layers import models, quantizations, deepseek
 from maxtext.common import checkpointing
+from maxtext.layers import quantizations
+from maxtext.models import deepseek, models
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 import orbax.checkpoint as ocp
+from tqdm import tqdm
 
 IGNORE = ocp.PLACEHOLDER
 PRNGKeyType = Any

src/MaxText/maxengine.py

@@ -39,7 +39,8 @@
 from MaxText import pyconfig
 from MaxText.common_types import MODEL_MODE_PREFILL, DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_AUTOREGRESSIVE
 from MaxText.globals import MAXTEXT_PKG_DIR
-from MaxText.layers import models, quantizations
+from maxtext.models import models
+from maxtext.layers import quantizations
 from maxtext.inference import inference_utils
 from maxtext.inference.page_manager import PageManager, PageState
 from maxtext.multimodal import processor as mm_processor

src/MaxText/train.py

@@ -40,15 +40,14 @@
 
 from MaxText import pyconfig
 from MaxText import sharding
-from MaxText.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
 from MaxText.common_types import ShardMode
 from MaxText.globals import EPS
 # Placeholder: internal
 
 from MaxText.gradient_accumulation import gradient_accumulation_loss_and_grad
 from MaxText.vocabulary_tiling import vocab_tiling_linen_loss
 # pylint: disable=too-many-positional-arguments
-
+from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
 from maxtext.common import checkpointing, profiler
 from maxtext.common.goodput import (
     GoodputEvent,

src/MaxText/train_compile.py

@@ -21,32 +21,29 @@
 as you would on the target hardware.
 """
 
-from typing import Sequence
+import functools
 import os
 import pickle
-import functools
+from typing import Sequence
 
 from absl import app
-
+from flax.linen import partitioning as nn_partitioning
 import jax
-from jax.experimental.topologies import get_topology_desc
-from jax.sharding import Mesh, AxisType
 from jax.experimental.serialize_executable import serialize
-
-from flax.linen import partitioning as nn_partitioning
-
+from jax.experimental.topologies import get_topology_desc
+from jax.sharding import AxisType, Mesh
 from MaxText import accelerator_to_spec_map
-from MaxText import train
 from MaxText import optimizers
 from MaxText import pyconfig
 from MaxText import sharding
+from MaxText import train
 from MaxText.common_types import MODEL_MODE_TRAIN, ShardMode
-from MaxText.layers import models
-from MaxText.layers import quantizations
+from maxtext.layers import quantizations
+from maxtext.models import models
+from maxtext.trainers.diloco import diloco
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
-from maxtext.trainers.diloco import diloco
 
 # pylint: disable=too-many-positional-arguments
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -39,26 +39,22 @@
 import os
 import sys
 
-from psutil import Process
-
-import numpy as np
-
 import jax
 from jax import random
 from jax.sharding import Mesh
-
-import tensorstore as ts
-
 from MaxText import optimizers
 from MaxText import pyconfig
+from maxtext.common import checkpointing
 from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.globals import MAXTEXT_PKG_DIR
-from MaxText.layers import quantizations
-from MaxText.layers.models import transformer_as_linen
-from maxtext.common import checkpointing
+from maxtext.layers import quantizations
+from maxtext.models.models import transformer_as_linen
 from maxtext.utils import max_logging
-from maxtext.utils import maxtext_utils
 from maxtext.utils import max_utils
+from maxtext.utils import maxtext_utils
+import numpy as np
+from psutil import Process
+import tensorstore as ts
 
 
 def fmt_size(num_bytes: int) -> str:

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -59,32 +59,32 @@
 """
 
 import argparse
+from functools import partial
+import json
 import os
-import time
 import sys
-import json
 import threading
-from functools import partial
-from typing import Sequence, List, Any, Callable
-import numpy as np
+import time
+from typing import Any, Callable, List, Sequence
 import absl
-
-from transformers import AutoConfig
+import flax.linen as nn
 from huggingface_hub import hf_hub_download, list_repo_files
-from safetensors import safe_open
 import jax
-import flax.linen as nn
-from orbax.checkpoint import type_handlers
 from MaxText import pyconfig
-from MaxText.common_types import MODEL_MODE_TRAIN
-from MaxText.layers import models, quantizations
 from maxtext.checkpoint_conversion.standalone_scripts.llama_or_mistral_ckpt import save_weights_to_checkpoint
 from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
-from maxtext.checkpoint_conversion.utils.utils import apply_hook_fns, HF_IDS, print_ram_usage, get_hf_model, MemoryMonitorTqdm, print_peak_memory, validate_and_filter_param_map_keys
+from maxtext.checkpoint_conversion.utils.utils import HF_IDS, MemoryMonitorTqdm, apply_hook_fns, get_hf_model, print_peak_memory, print_ram_usage, validate_and_filter_param_map_keys
+from MaxText.common_types import MODEL_MODE_TRAIN
 from maxtext.inference.inference_utils import str2bool
+from maxtext.layers import quantizations
+from maxtext.models import models
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
+import numpy as np
+from orbax.checkpoint import type_handlers
+from safetensors import safe_open
+from transformers import AutoConfig
 
 
 absl.logging.set_verbosity(absl.logging.INFO)  # for max_logging.log

src/maxtext/inference/kvcache.py

@@ -26,8 +26,8 @@
 from aqt.jax.v2.aqt_tensor import QTensor as KVTensor
 from aqt.jax.v2.flax import aqt_flax
 
-from MaxText.layers import nnx_wrappers
-from MaxText.layers.initializers import variable_to_logically_partitioned
+from maxtext.layers import nnx_wrappers
+from maxtext.layers.initializers import variable_to_logically_partitioned
 
 from MaxText.common_types import Array, AxisNames, AxisIdxes, Config, CACHE_BATCH_PREFILL, DType, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, MODEL_MODE_AUTOREGRESSIVE, CACHE_HEADS_NONE, DECODING_ACTIVE_SEQUENCE_INDICATOR
 from MaxText.common_types import CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV, CACHE_SCALE_BATCH, CACHE_SCALE_SEQUENCE, CACHE_SCALE_HEADS, CACHE_SCALE_KV